Skip to content

Commit 72b1723

Browse files
committed
feat: implement structured tool output support
This commit introduces the core functionality for structured tool outputs. In `mcp/tools.go`: - Adds `OutputSchema` to `Tool` and `StructuredContent` to `CallToolResult`. - Renames `ToolInputSchema` to `ToolSchema` for unified handling. - Implements `ValidateStructuredOutput` with schema compilation and caching. - Adds a series of `WithOutput*` functions for programmatic schema creation. - Adds comments for new public APIs and internal logic. In `mcp/utils.go`: - Adds `NewStructuredToolResult` and `NewStructuredToolError` helpers to simplify creating structured results with backward-compatible text content.
1 parent be4ee49 commit 72b1723

File tree

3 files changed

+425
-14
lines changed

3 files changed

+425
-14
lines changed

mcp/tools.go

Lines changed: 237 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
package mcp
22

33
import (
4+
"bytes"
45
"encoding/json"
56
"errors"
67
"fmt"
78
"reflect"
89
"strconv"
10+
11+
"github.com/santhosh-tekuri/jsonschema"
912
)
1013

1114
var errToolSchemaConflict = errors.New("provide either InputSchema or RawInputSchema, not both")
@@ -36,6 +39,8 @@ type ListToolsResult struct {
3639
type CallToolResult struct {
3740
Result
3841
Content []Content `json:"content"` // Can be TextContent, ImageContent, AudioContent, or EmbeddedResource
42+
// Structured content that conforms to the tool's output schema
43+
StructuredContent any `json:"structuredContent,omitempty"`
3944
// Whether the tool call ended in an error.
4045
//
4146
// If not set, this is assumed to be false (the call was successful).
@@ -475,7 +480,11 @@ type Tool struct {
475480
// A human-readable description of the tool.
476481
Description string `json:"description,omitempty"`
477482
// A JSON Schema object defining the expected parameters for the tool.
478-
InputSchema ToolInputSchema `json:"inputSchema"`
483+
InputSchema ToolSchema `json:"inputSchema"`
484+
// A JSON Schema object defining the expected output for the tool.
485+
OutputSchema ToolSchema `json:"outputSchema,omitempty"`
486+
// Compiled JSON schema validator for output validation, cached for performance
487+
compiledOutputSchema *jsonschema.Schema `json:"-"`
479488
// Alternative to InputSchema - allows arbitrary JSON Schema to be provided
480489
RawInputSchema json.RawMessage `json:"-"` // Hide this from JSON marshaling
481490
// Optional properties describing tool behavior
@@ -487,11 +496,68 @@ func (t Tool) GetName() string {
487496
return t.Name
488497
}
489498

499+
// HasOutputSchema returns true if the tool has an output schema defined.
500+
// This indicates that the tool can return structured content.
501+
func (t Tool) HasOutputSchema() bool {
502+
return t.OutputSchema.Type != ""
503+
}
504+
505+
// validateStructuredOutput performs the actual validation using the compiled schema
506+
func (t Tool) validateStructuredOutput(result *CallToolResult) error {
507+
return t.compiledOutputSchema.ValidateInterface(result.StructuredContent)
508+
}
509+
510+
// ensureOutputSchemaValidator compiles and caches the JSON schema validator if not already done
511+
func (t *Tool) ensureOutputSchemaValidator() error {
512+
if t.compiledOutputSchema != nil {
513+
return nil
514+
}
515+
516+
schemaBytes, err := t.OutputSchema.MarshalJSON()
517+
if err != nil {
518+
return err
519+
}
520+
521+
compiler := jsonschema.NewCompiler()
522+
523+
const validatorKey = "output-schema-validator"
524+
if err := compiler.AddResource(validatorKey, bytes.NewReader(schemaBytes)); err != nil {
525+
return err
526+
}
527+
528+
compiledSchema, err := compiler.Compile(validatorKey)
529+
if err != nil {
530+
return err
531+
}
532+
533+
t.compiledOutputSchema = compiledSchema
534+
return nil
535+
}
536+
537+
// ValidateStructuredOutput validates the structured content against the tool's output schema.
538+
// Returns nil if the tool has no output schema or if validation passes.
539+
// Returns an error if the tool has an output schema but the structured content is invalid.
540+
func (t Tool) ValidateStructuredOutput(result *CallToolResult) error {
541+
if !t.HasOutputSchema() {
542+
return nil
543+
}
544+
545+
if result.StructuredContent == nil {
546+
return fmt.Errorf("tool %s has output schema but structuredContent is nil", t.Name)
547+
}
548+
549+
if err := t.ensureOutputSchemaValidator(); err != nil {
550+
return err
551+
}
552+
553+
return t.validateStructuredOutput(result)
554+
}
555+
490556
// MarshalJSON implements the json.Marshaler interface for Tool.
491557
// It handles marshaling either InputSchema or RawInputSchema based on which is set.
492558
func (t Tool) MarshalJSON() ([]byte, error) {
493559
// Create a map to build the JSON structure
494-
m := make(map[string]any, 3)
560+
m := make(map[string]any, 5)
495561

496562
// Add the name and description
497563
m["name"] = t.Name
@@ -510,29 +576,34 @@ func (t Tool) MarshalJSON() ([]byte, error) {
510576
m["inputSchema"] = t.InputSchema
511577
}
512578

579+
// Add output schema if defined
580+
if t.HasOutputSchema() {
581+
m["outputSchema"] = t.OutputSchema
582+
}
583+
513584
m["annotations"] = t.Annotations
514585

515586
return json.Marshal(m)
516587
}
517588

518-
type ToolInputSchema struct {
589+
type ToolSchema struct {
519590
Type string `json:"type"`
520591
Properties map[string]any `json:"properties,omitempty"`
521592
Required []string `json:"required,omitempty"`
522593
}
523594

524-
// MarshalJSON implements the json.Marshaler interface for ToolInputSchema.
525-
func (tis ToolInputSchema) MarshalJSON() ([]byte, error) {
595+
// MarshalJSON implements the json.Marshaler interface for ToolSchema.
596+
func (schema ToolSchema) MarshalJSON() ([]byte, error) {
526597
m := make(map[string]any)
527-
m["type"] = tis.Type
598+
m["type"] = schema.Type
528599

529600
// Marshal Properties to '{}' rather than `nil` when its length equals zero
530-
if tis.Properties != nil {
531-
m["properties"] = tis.Properties
601+
if schema.Properties != nil {
602+
m["properties"] = schema.Properties
532603
}
533604

534-
if len(tis.Required) > 0 {
535-
m["required"] = tis.Required
605+
if len(schema.Required) > 0 {
606+
m["required"] = schema.Required
536607
}
537608

538609
return json.Marshal(m)
@@ -569,11 +640,17 @@ type PropertyOption func(map[string]any)
569640
func NewTool(name string, opts ...ToolOption) Tool {
570641
tool := Tool{
571642
Name: name,
572-
InputSchema: ToolInputSchema{
643+
InputSchema: ToolSchema{
573644
Type: "object",
574645
Properties: make(map[string]any),
575646
Required: nil, // Will be omitted from JSON if empty
576647
},
648+
OutputSchema: ToolSchema{
649+
Type: "",
650+
Properties: make(map[string]any),
651+
Required: nil, // Will be omitted from JSON if empty
652+
},
653+
compiledOutputSchema: nil,
577654
Annotations: ToolAnnotation{
578655
Title: "",
579656
ReadOnlyHint: ToBoolPtr(false),
@@ -607,8 +684,7 @@ func NewToolWithRawSchema(name, description string, schema json.RawMessage) Tool
607684
return tool
608685
}
609686

610-
// WithDescription adds a description to the Tool.
611-
// The description should provide a clear, human-readable explanation of what the tool does.
687+
// WithDescription sets the description field of the Tool.
612688
func WithDescription(description string) ToolOption {
613689
return func(t *Tool) {
614690
t.Description = description
@@ -1076,3 +1152,151 @@ func WithBooleanItems(opts ...PropertyOption) PropertyOption {
10761152
schema["items"] = itemSchema
10771153
}
10781154
}
1155+
1156+
//
1157+
// Output Schema Configuration Functions
1158+
//
1159+
1160+
// WithOutputSchema sets the output schema for the Tool.
1161+
// This allows the tool to define the structure of its return data.
1162+
func WithOutputSchema(schema ToolSchema) ToolOption {
1163+
return func(t *Tool) {
1164+
t.OutputSchema = schema
1165+
}
1166+
}
1167+
1168+
// WithOutputBoolean adds a boolean property to the tool's output schema.
1169+
// It accepts property options to configure the boolean property's behavior and constraints.
1170+
func WithOutputBoolean(name string, opts ...PropertyOption) ToolOption {
1171+
return func(t *Tool) {
1172+
// Initialize output schema if not set
1173+
if t.OutputSchema.Type == "" {
1174+
t.OutputSchema.Type = "object"
1175+
}
1176+
1177+
schema := map[string]any{
1178+
"type": "boolean",
1179+
}
1180+
1181+
for _, opt := range opts {
1182+
opt(schema)
1183+
}
1184+
1185+
// Remove required from property schema and add to OutputSchema.required
1186+
if required, ok := schema["required"].(bool); ok && required {
1187+
delete(schema, "required")
1188+
t.OutputSchema.Required = append(t.OutputSchema.Required, name)
1189+
}
1190+
1191+
t.OutputSchema.Properties[name] = schema
1192+
}
1193+
}
1194+
1195+
// WithOutputNumber adds a number property to the tool's output schema.
1196+
// It accepts property options to configure the number property's behavior and constraints.
1197+
func WithOutputNumber(name string, opts ...PropertyOption) ToolOption {
1198+
return func(t *Tool) {
1199+
// Initialize output schema if not set
1200+
if t.OutputSchema.Type == "" {
1201+
t.OutputSchema.Type = "object"
1202+
}
1203+
1204+
schema := map[string]any{
1205+
"type": "number",
1206+
}
1207+
1208+
for _, opt := range opts {
1209+
opt(schema)
1210+
}
1211+
1212+
// Remove required from property schema and add to OutputSchema.required
1213+
if required, ok := schema["required"].(bool); ok && required {
1214+
delete(schema, "required")
1215+
t.OutputSchema.Required = append(t.OutputSchema.Required, name)
1216+
}
1217+
1218+
t.OutputSchema.Properties[name] = schema
1219+
}
1220+
}
1221+
1222+
// WithOutputString adds a string property to the tool's output schema.
1223+
// It accepts property options to configure the string property's behavior and constraints.
1224+
func WithOutputString(name string, opts ...PropertyOption) ToolOption {
1225+
return func(t *Tool) {
1226+
// Initialize output schema if not set
1227+
if t.OutputSchema.Type == "" {
1228+
t.OutputSchema.Type = "object"
1229+
}
1230+
1231+
schema := map[string]any{
1232+
"type": "string",
1233+
}
1234+
1235+
for _, opt := range opts {
1236+
opt(schema)
1237+
}
1238+
1239+
// Remove required from property schema and add to OutputSchema.required
1240+
if required, ok := schema["required"].(bool); ok && required {
1241+
delete(schema, "required")
1242+
t.OutputSchema.Required = append(t.OutputSchema.Required, name)
1243+
}
1244+
1245+
t.OutputSchema.Properties[name] = schema
1246+
}
1247+
}
1248+
1249+
// WithOutputObject adds an object property to the tool's output schema.
1250+
// It accepts property options to configure the object property's behavior and constraints.
1251+
func WithOutputObject(name string, opts ...PropertyOption) ToolOption {
1252+
return func(t *Tool) {
1253+
// Initialize output schema if not set
1254+
if t.OutputSchema.Type == "" {
1255+
t.OutputSchema.Type = "object"
1256+
}
1257+
1258+
schema := map[string]any{
1259+
"type": "object",
1260+
"properties": map[string]any{},
1261+
}
1262+
1263+
for _, opt := range opts {
1264+
opt(schema)
1265+
}
1266+
1267+
// Remove required from property schema and add to OutputSchema.required
1268+
if required, ok := schema["required"].(bool); ok && required {
1269+
delete(schema, "required")
1270+
t.OutputSchema.Required = append(t.OutputSchema.Required, name)
1271+
}
1272+
1273+
t.OutputSchema.Properties[name] = schema
1274+
}
1275+
}
1276+
1277+
// WithOutputArray adds an array property to the tool's output schema.
1278+
// It accepts property options to configure the array property's behavior and constraints.
1279+
func WithOutputArray(name string, opts ...PropertyOption) ToolOption {
1280+
return func(t *Tool) {
1281+
// Initialize output schema if not set
1282+
if t.OutputSchema.Type == "" {
1283+
t.OutputSchema.Type = "object"
1284+
}
1285+
1286+
schema := map[string]any{
1287+
"type": "array",
1288+
}
1289+
1290+
for _, opt := range opts {
1291+
opt(schema)
1292+
}
1293+
1294+
// Remove required from property schema and add to OutputSchema.required
1295+
if required, ok := schema["required"].(bool); ok && required {
1296+
delete(schema, "required")
1297+
t.OutputSchema.Required = append(t.OutputSchema.Required, name)
1298+
}
1299+
1300+
t.OutputSchema.Properties[name] = schema
1301+
}
1302+
}

mcp/utils.go

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -353,6 +353,58 @@ func NewToolResultErrorf(format string, a ...any) *CallToolResult {
353353
}
354354
}
355355

356+
// NewStructuredToolResult creates a CallToolResult with structured content and automatically
357+
// generates backwards-compatible JSON text content as recommended by the MCP protocol.
358+
// This follows the protocol guideline that tools returning structured content SHOULD also
359+
// return functionally equivalent unstructured content.
360+
func NewStructuredToolResult(structuredContent any) *CallToolResult {
361+
// Serialize structured content to JSON for backwards compatibility
362+
jsonBytes, err := json.Marshal(structuredContent)
363+
var jsonText string
364+
if err != nil {
365+
// Fallback to error message if serialization fails
366+
jsonText = fmt.Sprintf("Error serializing structured content: %v", err)
367+
} else {
368+
jsonText = string(jsonBytes)
369+
}
370+
371+
return &CallToolResult{
372+
Content: []Content{
373+
TextContent{
374+
Type: "text",
375+
Text: jsonText,
376+
},
377+
},
378+
StructuredContent: structuredContent,
379+
IsError: false,
380+
}
381+
}
382+
383+
// NewStructuredToolError creates a CallToolResult for an error case with structured content
384+
// and automatically generates JSON text content for backwards compatibility.
385+
func NewStructuredToolError(structuredContent any) *CallToolResult {
386+
// Serialize structured content to JSON for backwards compatibility
387+
jsonBytes, err := json.Marshal(structuredContent)
388+
var jsonText string
389+
if err != nil {
390+
// Fallback to error message if serialization fails
391+
jsonText = fmt.Sprintf("Error serializing structured content: %v", err)
392+
} else {
393+
jsonText = string(jsonBytes)
394+
}
395+
396+
return &CallToolResult{
397+
Content: []Content{
398+
TextContent{
399+
Type: "text",
400+
Text: jsonText,
401+
},
402+
},
403+
StructuredContent: structuredContent,
404+
IsError: true,
405+
}
406+
}
407+
356408
// NewListResourcesResult creates a new ListResourcesResult
357409
func NewListResourcesResult(
358410
resources []Resource,

0 commit comments

Comments
 (0)