mcp: validate tool output (#352)

Validate the handler's returned output against the output schema of a tool.

Fixes #301.

Author: jba

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")

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{}

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