@@ -15,10 +15,12 @@ import (
1515 "maps"
1616 "net/url"
1717 "path/filepath"
18+ "reflect"
1819 "slices"
1920 "sync"
2021 "time"
2122
23+ "github.com/google/jsonschema-go/jsonschema"
2224 "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
2325 "github.com/modelcontextprotocol/go-sdk/internal/util"
2426 "github.com/modelcontextprotocol/go-sdk/jsonrpc"
@@ -138,59 +140,100 @@ func (s *Server) RemovePrompts(names ...string) {
138140 func () bool { return s .prompts .remove (names ... ) })
139141}
140142
141- // AddTool adds a [Tool] to the server, or replaces one with the same name.
143+ // AddRawTool adds a [Tool] to the server, or replaces one with the same name.
142144// The Tool argument must not be modified after this call.
143145//
144146// The tool's input schema must be non-nil. For a tool that takes no input,
145147// or one where any input is valid, set [Tool.InputSchema] to the empty schema,
146148// &jsonschema.Schema{}.
147- func (s * Server ) AddTool (t * Tool , h ToolHandler ) {
148- if t .InputSchema == nil {
149- // This prevents the tool author from forgetting to write a schema where
150- // one should be provided. If we papered over this by supplying the empty
151- // schema, then every input would be validated and the problem wouldn't be
152- // discovered until runtime, when the LLM sent bad data.
153- panic (fmt .Sprintf ("adding tool %q: nil input schema" , t .Name ))
154- }
155- if err := addToolErr (s , t , h ); err != nil {
156- panic (err )
157- }
149+ //
150+ // When the handler is invoked as part of a CallTool request, req.Params.Arguments
151+ // will be a json.RawMessage. Unmarshaling the arguments and validating them against the
152+ // input schema are the handler author's responsibility.
153+ func (s * Server ) AddRawTool (t * Tool , h RawToolHandler ) {
154+ st := & serverTool {tool : t , handler : h }
155+ // Assume there was a change, since add replaces existing tools.
156+ // (It's possible a tool was replaced with an identical one, but not worth checking.)
157+ // TODO: Batch these changes by size and time? The typescript SDK doesn't.
158+ // TODO: Surface notify error here? best not, in case we need to batch.
159+ s .changeAndNotify (notificationToolListChanged , & ToolListChangedParams {},
160+ func () bool { s .tools .add (st ); return true })
158161}
159162
160- // AddTool adds a [Tool] to the server, or replaces one with the same name.
161163// If the tool's input schema is nil, it is set to the schema inferred from the In
162164// type parameter, using [jsonschema.For].
163165// If the tool's output schema is nil and the Out type parameter is not the empty
164166// interface, then the output schema is set to the schema inferred from Out.
165- // The Tool argument must not be modified after this call.
166- func AddTool [ In , Out any ]( s * Server , t * Tool , h ToolHandlerFor [ In , Out ]) {
167- if err := addToolErr ( s , t , h ); err != nil {
168- panic (err )
167+ func RawToolHandlerFor [ In , Out any ]( t * Tool , h ToolHandlerFor [ In , Out ]) RawToolHandler {
168+ hh , err := toolForErr ( t , h )
169+ if err != nil {
170+ panic (fmt . Sprintf ( "ToolFor: tool %q: %v" , t . Name , err ) )
169171 }
172+ return hh
170173}
171174
172- func addToolErr [ In , Out any ]( s * Server , t * Tool , h ToolHandlerFor [ In , Out ]) ( err error ) {
173- defer util . Wrapf ( & err , "adding tool %q" , t . Name )
174- // If the exact same Tool pointer has already been registered under this name,
175- // avoid rebuilding schemas and re-registering. This prevents duplicate
176- // registration from causing errors (and unnecessary work).
177- s . mu . Lock ( )
178- if existing , ok := s . tools . get ( t . Name ); ok && existing . tool == t {
179- s . mu . Unlock ( )
180- return nil
175+ // TODO(v0.3.0): test
176+ func toolForErr [ In , Out any ]( t * Tool , h ToolHandlerFor [ In , Out ]) ( RawToolHandler , error ) {
177+ var err error
178+ inputSchema := t . InputSchema
179+ if inputSchema == nil {
180+ inputSchema , err = jsonschema. For [ In ]( nil )
181+ if err != nil {
182+ return nil , fmt . Errorf ( "input schema: %w" , err )
183+ }
181184 }
182- s .mu .Unlock ()
183- st , err := newServerTool (t , h )
185+ inputResolved , err := inputSchema .Resolve (& jsonschema.ResolveOptions {ValidateDefaults : true })
184186 if err != nil {
185- return err
187+ return nil , fmt . Errorf ( "resolving input schema: %w" , err )
186188 }
187- // Assume there was a change, since add replaces existing tools.
188- // (It's possible a tool was replaced with an identical one, but not worth checking.)
189- // TODO: Batch these changes by size and time? The typescript SDK doesn't.
190- // TODO: Surface notify error here? best not, in case we need to batch.
191- s .changeAndNotify (notificationToolListChanged , & ToolListChangedParams {},
192- func () bool { s .tools .add (st ); return true })
193- return nil
189+
190+ outputSchema := t .OutputSchema
191+ if outputSchema == nil && reflect .TypeFor [Out ]() != reflect .TypeFor [any ]() {
192+ outputSchema , err = jsonschema.For [Out ](nil )
193+ }
194+ if err != nil {
195+ return nil , fmt .Errorf ("output schema: %w" , err )
196+ }
197+ outputResolved , err := outputSchema .Resolve (& jsonschema.ResolveOptions {ValidateDefaults : true })
198+ if err != nil {
199+ return nil , fmt .Errorf ("resolving output schema: %w" , err )
200+ }
201+
202+ th := func (ctx context.Context , req * ServerRequest [* CallToolParams ]) (* CallToolResult , error ) {
203+ // Unmarshal and validate args.
204+ rawArgs := req .Params .Arguments .(json.RawMessage )
205+ var in In
206+ if rawArgs != nil {
207+ if err := unmarshalSchema (rawArgs , inputResolved , & in ); err != nil {
208+ return nil , err
209+ }
210+ }
211+
212+ // Call typed handler.
213+ res , out , err := h (ctx , req , in )
214+ if err != nil {
215+ return nil , err
216+ }
217+
218+ // TODO(v0.3.0): Validate out.
219+ _ = outputResolved
220+
221+ // TODO: return the serialized JSON in a TextContent block, as per spec?
222+ // https://modelcontextprotocol.io/specification/2025-06-18/server/tools#structured-content
223+ // But people may use res.Content for other things.
224+ if res == nil {
225+ res = & CallToolResult {}
226+ }
227+ res .StructuredContent = out
228+ return res , nil
229+ }
230+
231+ return th , nil
232+ }
233+
234+ // AddTool is a convenience for s.AddRawTool(t, RawToolHandler(t, h)).
235+ func AddTool [In , Out any ](s * Server , t * Tool , h ToolHandlerFor [In , Out ]) {
236+ s .AddRawTool (t , RawToolHandlerFor (t , h ))
194237}
195238
196239// RemoveTools removes the tools with the given names.
@@ -335,7 +378,7 @@ func (s *Server) listTools(_ context.Context, req *ServerRequest[*ListToolsParam
335378 })
336379}
337380
338- func (s * Server ) callTool (ctx context.Context , req * ServerRequest [* CallToolParamsFor [json. RawMessage ] ]) (* CallToolResult , error ) {
381+ func (s * Server ) callTool (ctx context.Context , req * ServerRequest [* CallToolParams ]) (* CallToolResult , error ) {
339382 s .mu .Lock ()
340383 st , ok := s .tools .get (req .Params .Name )
341384 s .mu .Unlock ()
0 commit comments