Skip to content

Commit ae446b4

Browse files
davidleitwdavidleitw
committed
feat(mcp, server): Redesign tool output schema with generics
This commit replaces the builder-style `WithOutput*` functions with a single, type-safe generic function, `WithOutputType[T any]`. Developers can now define a tool's output schema using a standard Go struct with `json` and `jsonschema` tags, simplifying the API and improving developer experience. Key changes: - Adds `invopop/jsonschema` for schema generation from structs. - Improves server-side validation to correctly skip validation on tool errors. - Adds new generic helper functions (`NewToolResultStructured`, etc.) for creating structured results. - Updates and adds tests to cover the new API and validation logic.
1 parent d49b6b7 commit ae446b4

File tree

8 files changed

+652
-480
lines changed

8 files changed

+652
-480
lines changed

‎client/http_test.go

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,15 @@ package client
22

33
import (
44
"context"
5+
"encoding/json"
56
"fmt"
6-
"github.com/mark3labs/mcp-go/mcp"
7-
"github.com/mark3labs/mcp-go/server"
87
"testing"
98
"time"
9+
10+
"github.com/mark3labs/mcp-go/mcp"
11+
"github.com/mark3labs/mcp-go/server"
12+
"github.com/stretchr/testify/assert"
13+
"github.com/stretchr/testify/require"
1014
)
1115

1216
func TestHTTPClient(t *testing.T) {
@@ -109,3 +113,64 @@ func TestHTTPClient(t *testing.T) {
109113
}
110114
})
111115
}
116+
117+
func TestHTTPClient_ListTools_WithOutputSchema(t *testing.T) {
118+
// 1. Setup Server
119+
srv := server.NewMCPServer("test-server", "1.0.0")
120+
121+
// Define a tool with a structured output type, including descriptions.
122+
type WeatherData struct {
123+
Temperature float64 `json:"temperature" jsonschema:"description=The temperature in Celsius."`
124+
Conditions string `json:"conditions" jsonschema:"description=Weather conditions (e.g. Cloudy)."`
125+
}
126+
tool := mcp.NewTool("get_weather", mcp.WithOutputType[WeatherData]())
127+
srv.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
128+
return nil, nil // Handler not needed for this test
129+
})
130+
131+
// Use the dedicated test helper to create the server
132+
httpServer := server.NewTestStreamableHTTPServer(srv)
133+
defer httpServer.Close()
134+
135+
// 2. Setup Client
136+
// Use the correct client constructor
137+
client, err := NewStreamableHttpClient(httpServer.URL)
138+
require.NoError(t, err)
139+
140+
// Initialize the client session before making other requests.
141+
_, err = client.Initialize(context.Background(), mcp.InitializeRequest{
142+
Params: mcp.InitializeParams{
143+
ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
144+
ClientInfo: mcp.Implementation{
145+
Name: "test-client",
146+
Version: "1.0.0",
147+
},
148+
// Client does not need to declare tool capabilities,
149+
// it's a server-side declaration.
150+
Capabilities: mcp.ClientCapabilities{},
151+
},
152+
})
153+
require.NoError(t, err, "client not initialized")
154+
155+
// 3. Client calls ListTools
156+
result, err := client.ListTools(context.Background(), mcp.ListToolsRequest{})
157+
require.NoError(t, err)
158+
require.Len(t, result.Tools, 1, "Should retrieve one tool")
159+
160+
// 4. Assert on the received tool's OutputSchema
161+
retrievedTool := result.Tools[0]
162+
assert.Equal(t, "get_weather", retrievedTool.Name)
163+
require.NotNil(t, retrievedTool.OutputSchema, "OutputSchema should be present")
164+
165+
// Unmarshal and verify the content of the schema
166+
var schemaData map[string]interface{}
167+
err = json.Unmarshal(retrievedTool.OutputSchema, &schemaData)
168+
require.NoError(t, err)
169+
170+
properties := schemaData["properties"].(map[string]interface{})
171+
tempProp := properties["temperature"].(map[string]interface{})
172+
condProp := properties["conditions"].(map[string]interface{})
173+
174+
assert.Equal(t, "The temperature in Celsius.", tempProp["description"])
175+
assert.Equal(t, "Weather conditions (e.g. Cloudy).", condProp["description"])
176+
}

‎go.mod

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,19 @@ go 1.23
44

55
require (
66
github.com/google/uuid v1.6.0
7+
github.com/invopop/jsonschema v0.13.0
78
github.com/santhosh-tekuri/jsonschema v1.2.4
89
github.com/spf13/cast v1.7.1
910
github.com/stretchr/testify v1.9.0
1011
github.com/yosida95/uritemplate/v3 v3.0.2
1112
)
1213

1314
require (
15+
github.com/bahlo/generic-list-go v0.2.0 // indirect
16+
github.com/buger/jsonparser v1.1.1 // indirect
1417
github.com/davecgh/go-spew v1.1.1 // indirect
18+
github.com/mailru/easyjson v0.7.7 // indirect
1519
github.com/pmezard/go-difflib v1.0.0 // indirect
20+
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
1621
gopkg.in/yaml.v3 v3.0.1 // indirect
1722
)

‎go.sum

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
2+
github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
3+
github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs=
4+
github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
15
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
26
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
37
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
@@ -6,10 +10,15 @@ github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
610
github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
711
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
812
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
13+
github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E=
14+
github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0=
15+
github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
916
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
1017
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
1118
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
1219
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
20+
github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
21+
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
1322
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
1423
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
1524
github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8=
@@ -20,6 +29,8 @@ github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
2029
github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
2130
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
2231
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
32+
github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
33+
github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
2334
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
2435
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
2536
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=

0 commit comments

Comments
 (0)