fix(tool): Return the OutputSchema from the tool definition #571
Conversation
This commit returns the `OutputSchema` from the tool definition as per the MCP spec: https://modelcontextprotocol.io/specification/2025-06-18/server/tools#output-schema The behaviour of `WithOutputSchema()` also changed as `RawOutputSchema` is not longer populated, but `OutputSchema` is from the `T` generic type. The only way now to set `RawOutputSchema` is through the `WithRawOutputSchema()` method.
…-output-schema fix(tool): Return the `OutputSchema` from the tool definition
WalkthroughIntroduces typed ToolArgumentsSchema with aliases for input/output schemas, adds a public Tool.OutputSchema, updates marshaling to handle OutputSchema vs RawOutputSchema with conflict checks, and adjusts WithOutputSchema to populate OutputSchema. Tests updated to assert OutputSchema instead of RawOutputSchema. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Assessment against linked issues
Assessment against linked issues: Out-of-scope changes
Possibly related issues
Possibly related PRs
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
|
As mentioned in the PR description, I've performed the change on the WithOutputSchema that basically set the OutputSchema from the passed type T, and RawOutputSchema is not set anymore. This is quite different from what has been done on WithInputSchema(), so let me know your thought about this implementation, and if I need to perform a change in this method too 🤓 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
mcp/tools.go (1)
631-649: Include empty “properties” and explicit “additionalProperties”: false for object schemas
Per MCP spec, “properties” is optional; to satisfy OpenAI function-calling you must emit an empty object and disallow additional properties when type is “object.”
Apply:if tis.Properties != nil { m["properties"] = tis.Properties } else if tis.Type == "object" { - m["properties"] = map[string]any{} + m["properties"] = map[string]any{} + m["additionalProperties"] = false }
🧹 Nitpick comments (6)
mcp/tools.go (4)
568-574: Typed OutputSchema added—consider Unmarshal symmetry and raw retentionDefault unmarshaling will populate OutputSchema, but we lose lossless round‑trip and backward access to RawOutputSchema. Suggest adding Tool.UnmarshalJSON to populate both OutputSchema (typed) and RawOutputSchema (raw) for symmetry with input and to fully address #563.
Example addition outside this hunk:
+// UnmarshalJSON implements custom JSON unmarshaling for Tool to retain raw and typed schemas. +func (t *Tool) UnmarshalJSON(data []byte) error { + type wire struct { + Meta *Meta `json:"_meta,omitempty"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + InputSchema json.RawMessage `json:"inputSchema"` + OutputSchema json.RawMessage `json:"outputSchema,omitempty"` + Annotations ToolAnnotation `json:"annotations"` + } + var w wire + if err := json.Unmarshal(data, &w); err != nil { + return err + } + t.Meta = w.Meta + t.Name = w.Name + t.Description = w.Description + t.Annotations = w.Annotations + + // Input schema: prefer typed, fall back to raw + if len(w.InputSchema) > 0 { + var typedIn ToolInputSchema + if err := json.Unmarshal(w.InputSchema, &typedIn); err == nil && typedIn.Type != "" { + t.InputSchema = typedIn + t.RawInputSchema = nil + } else { + t.InputSchema = ToolInputSchema{} + t.RawInputSchema = append(json.RawMessage(nil), w.InputSchema...) + } + } + + // Output schema: populate typed and retain raw for lossless round-trip + if len(w.OutputSchema) > 0 { + _ = json.Unmarshal(w.OutputSchema, &t.OutputSchema) // best-effort + t.RawOutputSchema = append(json.RawMessage(nil), w.OutputSchema...) + } + return nil +}
604-612: Conflict error mentions InputSchema for an OutputSchema conflictThe wrapped error message is misleading when both OutputSchema and RawOutputSchema are set.
Apply:
- if t.OutputSchema.Type != "" { - return nil, fmt.Errorf("tool %s has both OutputSchema and RawOutputSchema set: %w", t.Name, errToolSchemaConflict) - } + if t.OutputSchema.Type != "" { + return nil, fmt.Errorf("tool %s has both OutputSchema and RawOutputSchema set: provide either OutputSchema or RawOutputSchema, not both", t.Name) + }(Optional) Add a dedicated var:
+var errToolOutputSchemaConflict = errors.New("provide either OutputSchema or RawOutputSchema, not both")And use it here.
619-629: Avoid lossy output schemas: include root-level additionalProperties (optional)ToolArgumentsSchema drops root additionalProperties and similar keywords on unmarshal/remarshal.
Apply:
type ToolArgumentsSchema struct { Defs map[string]any `json:"$defs,omitempty"` Type string `json:"type"` Properties map[string]any `json:"properties,omitempty"` Required []string `json:"required,omitempty"` + AdditionalProperties any `json:"additionalProperties,omitempty"` }
794-803: Harden WithOutputSchema: ensure non-nil map and validate typeMinor robustness: initialize Properties if nil; optionally validate non-object types.
if err := json.Unmarshal(mcpSchema, &t.OutputSchema); err != nil { // Skip and maintain backward compatibility return } // Always set the type to "object" as of the current MCP spec t.OutputSchema.Type = "object" +if t.OutputSchema.Properties == nil { + t.OutputSchema.Properties = map[string]any{} +}mcp/tools_test.go (2)
595-597: assert.NotNil on a struct is ineffective; assert concrete fieldsOutputSchema is a struct and never nil. Assert Type/properties and JSON presence instead.
Apply:
- // Check that RawOutputSchema was set - assert.NotNil(t, tool.OutputSchema) + // Check that OutputSchema was set as an object with properties + assert.Equal(t, "object", tool.OutputSchema.Type) + // Marshal and verify structure is present in JSON below
598-610: Add round-trip and conflict tests for output schema (covers #563 semantics)Ensure unmarshaling populates OutputSchema and that dual schema conflicts error.
Proposed additions outside this hunk:
func TestUnmarshalToolWithOutputSchema(t *testing.T) { type TestOutput struct { Name string `json:"name"` } tool := NewTool("t", WithOutputSchema[TestOutput](), WithString("input", Required())) data, err := json.Marshal(tool) assert.NoError(t, err) var out Tool assert.NoError(t, json.Unmarshal(data, &out)) assert.Equal(t, "t", out.Name) assert.Equal(t, "object", out.OutputSchema.Type) } func TestToolWithBothOutputSchemasError(t *testing.T) { type TestOutput struct{ X string `json:"x"` } tool := NewTool("dual-output", WithOutputSchema[TestOutput]()) tool.RawOutputSchema = json.RawMessage(`{"type":"object","properties":{"y":{"type":"string"}}}`) _, err := json.Marshal(tool) assert.Error(t, err) }Want me to push these test updates?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
mcp/tools.go(3 hunks)mcp/tools_test.go(1 hunks)
🧰 Additional context used
🧠 Learnings (4)
📓 Common learnings
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:0-0
Timestamp: 2025-03-04T07:00:57.111Z
Learning: The Tool struct in the mark3labs/mcp-go project should handle both InputSchema and RawInputSchema consistently between MarshalJSON and UnmarshalJSON methods, even though the tools response from MCP server typically doesn't contain rawInputSchema.
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:0-0
Timestamp: 2025-03-04T07:00:57.111Z
Learning: The Tool struct in mark3labs/mcp-go handles both InputSchema and RawInputSchema formats. When unmarshaling JSON, it first tries to parse into a structured ToolInputSchema format, and if that fails, it falls back to using the raw schema format, providing symmetry with the MarshalJSON method.
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:107-137
Timestamp: 2025-03-04T06:59:43.882Z
Learning: Tool responses from the MCP server shouldn't contain RawInputSchema, which is why the UnmarshalJSON method for the Tool struct is implemented to handle only the structured InputSchema format.
📚 Learning: 2025-03-04T07:00:57.111Z
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:0-0
Timestamp: 2025-03-04T07:00:57.111Z
Learning: The Tool struct in the mark3labs/mcp-go project should handle both InputSchema and RawInputSchema consistently between MarshalJSON and UnmarshalJSON methods, even though the tools response from MCP server typically doesn't contain rawInputSchema.
Applied to files:
mcp/tools_test.gomcp/tools.go
📚 Learning: 2025-03-04T06:59:43.882Z
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:107-137
Timestamp: 2025-03-04T06:59:43.882Z
Learning: Tool responses from the MCP server shouldn't contain RawInputSchema, which is why the UnmarshalJSON method for the Tool struct is implemented to handle only the structured InputSchema format.
Applied to files:
mcp/tools_test.gomcp/tools.go
📚 Learning: 2025-03-04T07:00:57.111Z
Learnt from: xinwo
PR: mark3labs/mcp-go#35
File: mcp/tools.go:0-0
Timestamp: 2025-03-04T07:00:57.111Z
Learning: The Tool struct in mark3labs/mcp-go handles both InputSchema and RawInputSchema formats. When unmarshaling JSON, it first tries to parse into a structured ToolInputSchema format, and if that fails, it falls back to using the raw schema format, providing symmetry with the MarshalJSON method.
Applied to files:
mcp/tools_test.gomcp/tools.go
|
@ezynda3 @alex210501 The commit b924391 is breaking the entire library, please see #572 |
|
@matheuscscp Taking a look |
3 participants
Description
This PR returns the
OutputSchemafrom the tool definition as per the MCP spec: https://modelcontextprotocol.io/specification/2025-06-18/server/tools#output-schemaThe behaviour of
WithOutputSchema()also changed asRawOutputSchemais not longer populated, butOutputSchemais from theTgeneric type. The only way now to setRawOutputSchemais through theWithRawOutputSchema()method.Fixes #563
Type of Change
Checklist
MCP Spec Compliance
Additional Information
Summary by CodeRabbit