diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index a16ac83..28a8b97 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -257,6 +257,26 @@ func TestEndToEnd(t *testing.T) { t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff) } + // Check output schema validation. + badout := &Tool{ + Name: "badout", + OutputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "x": {Type: "string"}, + }, + }, + } + AddTool(s, badout, func(_ context.Context, _ *CallToolRequest, arg map[string]any) (*CallToolResult, map[string]any, error) { + return nil, map[string]any{"x": 1}, nil + }) + _, err = cs.CallTool(ctx, &CallToolParams{Name: "badout"}) + wantMsg := `has type "integer", want "string"` + if err == nil || !strings.Contains(err.Error(), wantMsg) { + t.Errorf("\ngot %q\nwant error message containing %q", err, wantMsg) + } + + // Check tools-changed notifications. s.AddTool(&Tool{Name: "T", InputSchema: &jsonschema.Schema{Type: "object"}}, nopHandler) waitForNotification(t, "tools") s.RemoveTools("T") diff --git a/mcp/server.go b/mcp/server.go index 114489a..0c1acc2 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -247,8 +247,12 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan }, nil } - // TODO(v0.3.0): Validate out. - _ = outputResolved + // Validate output schema, if any. + // Skip if out is nil: we've removed "null" from the output schema, so nil won't validate. + if v := reflect.ValueOf(out); v.Kind() == reflect.Pointer && v.IsNil() { + } else if err := validateSchema(outputResolved, &out); err != nil { + return nil, fmt.Errorf("tool output: %w", err) + } if res == nil { res = &CallToolResult{} diff --git a/mcp/tool.go b/mcp/tool.go index bd10a07..53a3c7a 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -9,6 +9,7 @@ import ( "context" "encoding/json" "fmt" + // "log" "github.com/google/jsonschema-go/jsonschema" ) @@ -42,13 +43,16 @@ func unmarshalSchema(data json.RawMessage, resolved *jsonschema.Resolved, v any) if err := dec.Decode(v); err != nil { return fmt.Errorf("unmarshaling: %w", err) } - // TODO: test with nil args. + return validateSchema(resolved, v) +} + +func validateSchema(resolved *jsonschema.Resolved, value any) error { if resolved != nil { - if err := resolved.ApplyDefaults(v); err != nil { - return fmt.Errorf("applying defaults from \n\t%s\nto\n\t%s:\n%w", schemaJSON(resolved.Schema()), data, err) + if err := resolved.ApplyDefaults(value); err != nil { + return fmt.Errorf("applying defaults from \n\t%s\nto\n\t%v:\n%w", schemaJSON(resolved.Schema()), value, err) } - if err := resolved.Validate(v); err != nil { - return fmt.Errorf("validating\n\t%s\nagainst\n\t %s:\n %w", data, schemaJSON(resolved.Schema()), err) + if err := resolved.Validate(value); err != nil { + return fmt.Errorf("validating\n\t%v\nagainst\n\t %s:\n %w", value, schemaJSON(resolved.Schema()), err) } } return nil