185 lines
5.7 KiB
Go
185 lines
5.7 KiB
Go
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
|
// Use of this source code is governed by an MIT-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
// The toolschemas example demonstrates how to create tools using both the
|
|
// low-level [ToolHandler] and high level [ToolHandlerFor], as well as how to
|
|
// customize schemas in both cases.
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
|
|
"github.com/google/jsonschema-go/jsonschema"
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
)
|
|
|
|
// Input is the input into all the tools handlers below.
|
|
type Input struct {
|
|
Name string `json:"name" jsonschema:"the person to greet"`
|
|
}
|
|
|
|
// Output is the structured output of the tool.
|
|
//
|
|
// Not every tool needs to have structured output.
|
|
type Output struct {
|
|
Greeting string `json:"greeting" jsonschema:"the greeting to send to the user"`
|
|
}
|
|
|
|
// simpleGreeting is an [mcp.ToolHandlerFor] that only cares about input and output.
|
|
func simpleGreeting(_ context.Context, _ *mcp.CallToolRequest, input Input) (*mcp.CallToolResult, Output, error) {
|
|
return nil, Output{"Hi " + input.Name}, nil
|
|
}
|
|
|
|
// manualGreeter handles the parsing and validation of input and output manually.
|
|
//
|
|
// Therefore, it needs to close over its resolved schemas, to use them in
|
|
// validation.
|
|
type manualGreeter struct {
|
|
inputSchema *jsonschema.Resolved
|
|
outputSchema *jsonschema.Resolved
|
|
}
|
|
|
|
func (t *manualGreeter) greet(_ context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
// errf produces a 'tool error', embedding the error in a CallToolResult.
|
|
errf := func(format string, args ...any) *mcp.CallToolResult {
|
|
return &mcp.CallToolResult{
|
|
Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf(format, args...)}},
|
|
IsError: true,
|
|
}
|
|
}
|
|
// Handle the parsing and validation of input and output.
|
|
//
|
|
// Note that errors here are treated as tool errors, not protocol errors.
|
|
|
|
// First, unmarshal to a map[string]any and validate.
|
|
if err := unmarshalAndValidate(req.Params.Arguments, t.inputSchema); err != nil {
|
|
return errf("invalid input: %v", err), nil
|
|
}
|
|
|
|
// Now unmarshal again to input.
|
|
var input Input
|
|
if err := json.Unmarshal(req.Params.Arguments, &input); err != nil {
|
|
return errf("failed to unmarshal arguments: %v", err), nil
|
|
}
|
|
output := Output{Greeting: "Hi " + input.Name}
|
|
outputJSON, err := json.Marshal(output)
|
|
if err != nil {
|
|
return errf("output failed to marshal: %v", err), nil
|
|
}
|
|
//
|
|
if err := unmarshalAndValidate(outputJSON, t.outputSchema); err != nil {
|
|
return errf("invalid output: %v", err), nil
|
|
}
|
|
|
|
return &mcp.CallToolResult{
|
|
Content: []mcp.Content{&mcp.TextContent{Text: string(outputJSON)}},
|
|
StructuredContent: output,
|
|
}, nil
|
|
}
|
|
|
|
// unmarshalAndValidate unmarshals data to a map[string]any, then validates that against res.
|
|
func unmarshalAndValidate(data []byte, res *jsonschema.Resolved) error {
|
|
var m map[string]any
|
|
if err := json.Unmarshal(data, &m); err != nil {
|
|
return err
|
|
}
|
|
return res.Validate(m)
|
|
}
|
|
|
|
var (
|
|
inputSchema = &jsonschema.Schema{
|
|
Type: "object",
|
|
Properties: map[string]*jsonschema.Schema{
|
|
"name": {Type: "string", MaxLength: jsonschema.Ptr(10)},
|
|
},
|
|
}
|
|
outputSchema = &jsonschema.Schema{
|
|
Type: "object",
|
|
Properties: map[string]*jsonschema.Schema{
|
|
"greeting": {Type: "string"},
|
|
},
|
|
}
|
|
)
|
|
|
|
func newManualGreeter() (*manualGreeter, error) {
|
|
resIn, err := inputSchema.Resolve(nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
resOut, err := outputSchema.Resolve(nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &manualGreeter{
|
|
inputSchema: resIn,
|
|
outputSchema: resOut,
|
|
}, nil
|
|
}
|
|
|
|
func main() {
|
|
server := mcp.NewServer(&mcp.Implementation{Name: "greeter"}, nil)
|
|
|
|
// Add the 'greeting' tool in a few different ways.
|
|
|
|
// First, we can just use [mcp.AddTool], and get the out-of-the-box handling
|
|
// it provides for schema inference, validation, parsing, and packing the
|
|
// result.
|
|
mcp.AddTool(server, &mcp.Tool{Name: "simple greeting"}, simpleGreeting)
|
|
|
|
// Alternatively, we can create our schemas entirely manually, and add them
|
|
// using [mcp.Server.AddTool]. Since we're using the 'raw' API, we have to do
|
|
// the parsing and validation ourselves
|
|
manual, err := newManualGreeter()
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
server.AddTool(&mcp.Tool{
|
|
Name: "manual greeting",
|
|
InputSchema: inputSchema,
|
|
OutputSchema: outputSchema,
|
|
}, manual.greet)
|
|
|
|
// We can even use raw schema values. In this case, note that we're not
|
|
// validating the input at all.
|
|
server.AddTool(&mcp.Tool{
|
|
Name: "unvalidated greeting",
|
|
InputSchema: json.RawMessage(`{"type":"object","properties":{"user":{"type":"string"}}}`),
|
|
}, func(_ context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
// Note: no validation!
|
|
var args struct{ User string }
|
|
if err := json.Unmarshal(req.Params.Arguments, &args); err != nil {
|
|
return nil, err
|
|
}
|
|
return &mcp.CallToolResult{
|
|
Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + args.User}},
|
|
}, nil
|
|
})
|
|
|
|
// Finally, note that we can also use custom schemas with a ToolHandlerFor.
|
|
// We can do this in two ways: by using one of the schema values constructed
|
|
// above, or by using jsonschema.For and adjusting the resulting schema.
|
|
mcp.AddTool(server, &mcp.Tool{
|
|
Name: "customized greeting 1",
|
|
InputSchema: inputSchema,
|
|
// OutputSchema will still be derived from Output.
|
|
}, simpleGreeting)
|
|
|
|
customSchema, err := jsonschema.For[Input](nil)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
customSchema.Properties["name"].MaxLength = jsonschema.Ptr(10)
|
|
mcp.AddTool(server, &mcp.Tool{
|
|
Name: "customized greeting 2",
|
|
InputSchema: customSchema,
|
|
}, simpleGreeting)
|
|
|
|
// Now run the server.
|
|
if err := server.Run(context.Background(), &mcp.StdioTransport{}); err != nil {
|
|
log.Printf("Server failed: %v", err)
|
|
}
|
|
}
|