Skip to content

Commit f5cb4bd

Browse files
davidleitwdavidleitw
committed
feat(mcp, server): Implement 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 1eddde7 commit f5cb4bd

File tree

8 files changed

+1286
-14
lines changed

8 files changed

+1286
-14
lines changed

‎client/http_test.go

Lines changed: 64 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package client
22

33
import (
44
"context"
5+
"encoding/json"
56
"fmt"
67
"sync"
78
"testing"
@@ -10,9 +11,10 @@ import (
1011
"github.com/mark3labs/mcp-go/client/transport"
1112
"github.com/mark3labs/mcp-go/mcp"
1213
"github.com/mark3labs/mcp-go/server"
14+
"github.com/stretchr/testify/assert"
15+
"github.com/stretchr/testify/require"
1316
)
1417

15-
1618
func TestHTTPClient(t *testing.T) {
1719
hooks := &server.Hooks{}
1820
hooks.AddAfterCallTool(func(ctx context.Context, id any, message *mcp.CallToolRequest, result *mcp.CallToolResult) {
@@ -192,6 +194,67 @@ func TestHTTPClient(t *testing.T) {
192194
})
193195
}
194196

197+
func TestHTTPClient_ListTools_WithOutputSchema(t *testing.T) {
198+
// 1. Setup Server
199+
srv := server.NewMCPServer("test-server", "1.0.0")
200+
201+
// Define a tool with a structured output type, including descriptions.
202+
type WeatherData struct {
203+
Temperature float64 `json:"temperature" jsonschema:"description=The temperature in Celsius."`
204+
Conditions string `json:"conditions" jsonschema:"description=Weather conditions (e.g. Cloudy)."`
205+
}
206+
tool := mcp.NewTool("get_weather", mcp.WithOutputType[WeatherData]())
207+
srv.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
208+
return nil, nil // Handler not needed for this test
209+
})
210+
211+
// Use the dedicated test helper to create the server
212+
httpServer := server.NewTestStreamableHTTPServer(srv)
213+
defer httpServer.Close()
214+
215+
// 2. Setup Client
216+
// Use the correct client constructor
217+
client, err := NewStreamableHttpClient(httpServer.URL)
218+
require.NoError(t, err)
219+
220+
// Initialize the client session before making other requests.
221+
_, err = client.Initialize(context.Background(), mcp.InitializeRequest{
222+
Params: mcp.InitializeParams{
223+
ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
224+
ClientInfo: mcp.Implementation{
225+
Name: "test-client",
226+
Version: "1.0.0",
227+
},
228+
// Client does not need to declare tool capabilities,
229+
// it's a server-side declaration.
230+
Capabilities: mcp.ClientCapabilities{},
231+
},
232+
})
233+
require.NoError(t, err, "client not initialized")
234+
235+
// 3. Client calls ListTools
236+
result, err := client.ListTools(context.Background(), mcp.ListToolsRequest{})
237+
require.NoError(t, err)
238+
require.Len(t, result.Tools, 1, "Should retrieve one tool")
239+
240+
// 4. Assert on the received tool's OutputSchema
241+
retrievedTool := result.Tools[0]
242+
assert.Equal(t, "get_weather", retrievedTool.Name)
243+
require.NotNil(t, retrievedTool.OutputSchema, "OutputSchema should be present")
244+
245+
// Unmarshal and verify the content of the schema
246+
var schemaData map[string]interface{}
247+
err = json.Unmarshal(retrievedTool.OutputSchema, &schemaData)
248+
require.NoError(t, err)
249+
250+
properties := schemaData["properties"].(map[string]interface{})
251+
tempProp := properties["temperature"].(map[string]interface{})
252+
condProp := properties["conditions"].(map[string]interface{})
253+
254+
assert.Equal(t, "The temperature in Celsius.", tempProp["description"])
255+
assert.Equal(t, "Weather conditions (e.g. Cloudy).", condProp["description"])
256+
}
257+
195258
type SafeMap struct {
196259
mu sync.RWMutex
197260
data map[string]int

‎go.mod

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

55
require (
66
github.com/google/uuid v1.6.0
7+
github.com/invopop/jsonschema v0.13.0
8+
github.com/santhosh-tekuri/jsonschema v1.2.4
79
github.com/spf13/cast v1.7.1
810
github.com/stretchr/testify v1.9.0
911
github.com/yosida95/uritemplate/v3 v3.0.2
1012
)
1113

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

‎go.sum

Lines changed: 13 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,18 +10,27 @@ 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=
1625
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
26+
github.com/santhosh-tekuri/jsonschema v1.2.4 h1:hNhW8e7t+H1vgY+1QeEQpveR6D4+OwKPXCfD2aieJis=
27+
github.com/santhosh-tekuri/jsonschema v1.2.4/go.mod h1:TEAUOeZSmIxTTuHatJzrvARHiuO9LYd+cIxzgEHCQI4=
1728
github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
1829
github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
1930
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
2031
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=
2134
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
2235
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
2336
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=

0 commit comments

Comments
 (0)