modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

mcp: change tool handler design (#325)

Browse source 
Change the design of tool handlers by removing genericity from the common path.
- `ToolHandler` gets the args as a `json.RawMessage`. Using Server.AddTool does no unmarshaling or schema validation.
- `ToolHandlerFor`, installed with `AddTool`, are the only generic pieces.

TODO:
uncomment and fix some tests
This commit is contained in:
Jonathan Amsterdam 2025-08-20 19:36:34 -04:00 • committed by GitHub
commit 32cf71de3a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 789 additions and 966 deletions
Download patch file Download diff file Expand all files Collapse all files

8
README.md
View file

@ -115,10 +115,10 @@ type HiParams struct {
Name string `json:"name" jsonschema:"the name of the person to greet"` Name string `json:"name" jsonschema:"the name of the person to greet"`
} }
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiParams]]) (*mcp.CallToolResultFor[any], error) { func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args HiParams) (*mcp.CallToolResult, any, error) {
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}}, Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + args.Name}},
}, nil }, nil, nil
} }
func main() { func main() {

8
examples/server/custom-transport/main.go
View file

@ -85,12 +85,12 @@ type HiArgs struct {
} }
// SayHi is a tool handler that responds with a greeting. // SayHi is a tool handler that responds with a greeting.
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiArgs]]) (*mcp.CallToolResultFor[struct{}], error) { func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args HiArgs) (*mcp.CallToolResult, struct{}, error) {
return &mcp.CallToolResultFor[struct{}]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}, &mcp.TextContent{Text: "Hi " + args.Name},
}, },
}, nil }, struct{}{}, nil
} }
func main() { func main() {

8
examples/server/hello/main.go
View file

@ -22,12 +22,12 @@ type HiArgs struct {
Name string `json:"name" jsonschema:"the name to say hi to"` Name string `json:"name" jsonschema:"the name to say hi to"`
} }
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiArgs]]) (*mcp.CallToolResultFor[struct{}], error) { func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args HiArgs) (*mcp.CallToolResult, any, error) {
return &mcp.CallToolResultFor[struct{}]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}, &mcp.TextContent{Text: "Hi " + args.Name},
}, },
}, nil }, nil, nil
} }
func PromptHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) { func PromptHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {

99
examples/server/memory/kb.go
View file

@ -431,12 +431,12 @@ func (k knowledgeBase) openNodes(names []string) (KnowledgeGraph, error) {
}, nil }, nil
} }
func (k knowledgeBase) CreateEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[CreateEntitiesArgs]]) (*mcp.CallToolResultFor[CreateEntitiesResult], error) { func (k knowledgeBase) CreateEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args CreateEntitiesArgs) (*mcp.CallToolResult, CreateEntitiesResult, error) {
var res mcp.CallToolResultFor[CreateEntitiesResult] var res mcp.CallToolResult
entities, err := k.createEntities(req.Params.Arguments.Entities) entities, err := k.createEntities(args.Entities)
if err != nil { if err != nil {
return nil, err return nil, CreateEntitiesResult{}, err
} }
res.Content = []mcp.Content{ res.Content = []mcp.Content{
@ -447,114 +447,107 @@ func (k knowledgeBase) CreateEntities(ctx context.Context, req *mcp.ServerReques
Entities: entities, Entities: entities,
} }
return &res, nil return &res, CreateEntitiesResult{Entities: entities}, nil
} }
func (k knowledgeBase) CreateRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[CreateRelationsArgs]]) (*mcp.CallToolResultFor[CreateRelationsResult], error) { func (k knowledgeBase) CreateRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args CreateRelationsArgs) (*mcp.CallToolResult, CreateRelationsResult, error) {
var res mcp.CallToolResultFor[CreateRelationsResult] var res mcp.CallToolResult
relations, err := k.createRelations(req.Params.Arguments.Relations) relations, err := k.createRelations(args.Relations)
if err != nil { if err != nil {
return nil, err return nil, CreateRelationsResult{}, err
} }
res.Content = []mcp.Content{ res.Content = []mcp.Content{
&mcp.TextContent{Text: "Relations created successfully"}, &mcp.TextContent{Text: "Relations created successfully"},
} }
res.StructuredContent = CreateRelationsResult{ return &res, CreateRelationsResult{Relations: relations}, nil
Relations: relations,
}
return &res, nil
} }
func (k knowledgeBase) AddObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[AddObservationsArgs]]) (*mcp.CallToolResultFor[AddObservationsResult], error) { func (k knowledgeBase) AddObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args AddObservationsArgs) (*mcp.CallToolResult, AddObservationsResult, error) {
var res mcp.CallToolResultFor[AddObservationsResult] var res mcp.CallToolResult
observations, err := k.addObservations(req.Params.Arguments.Observations) observations, err := k.addObservations(args.Observations)
if err != nil { if err != nil {
return nil, err return nil, AddObservationsResult{}, err
} }
res.Content = []mcp.Content{ res.Content = []mcp.Content{
&mcp.TextContent{Text: "Observations added successfully"}, &mcp.TextContent{Text: "Observations added successfully"},
} }
res.StructuredContent = AddObservationsResult{ return &res, AddObservationsResult{
Observations: observations, Observations: observations,
} }, nil
return &res, nil
} }
func (k knowledgeBase) DeleteEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteEntitiesArgs]]) (*mcp.CallToolResultFor[struct{}], error) { func (k knowledgeBase) DeleteEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteEntitiesArgs) (*mcp.CallToolResult, struct{}, error) {
var res mcp.CallToolResultFor[struct{}] var res mcp.CallToolResult
err := k.deleteEntities(req.Params.Arguments.EntityNames) err := k.deleteEntities(args.EntityNames)
if err != nil { if err != nil {
return nil, err return nil, struct{}{}, err
} }
res.Content = []mcp.Content{ res.Content = []mcp.Content{
&mcp.TextContent{Text: "Entities deleted successfully"}, &mcp.TextContent{Text: "Entities deleted successfully"},
} }
return &res, nil return &res, struct{}{}, nil
} }
func (k knowledgeBase) DeleteObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteObservationsArgs]]) (*mcp.CallToolResultFor[struct{}], error) { func (k knowledgeBase) DeleteObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteObservationsArgs) (*mcp.CallToolResult, struct{}, error) {
var res mcp.CallToolResultFor[struct{}] var res mcp.CallToolResult
err := k.deleteObservations(req.Params.Arguments.Deletions) err := k.deleteObservations(args.Deletions)
if err != nil { if err != nil {
return nil, err return nil, struct{}{}, err
} }
res.Content = []mcp.Content{ res.Content = []mcp.Content{
&mcp.TextContent{Text: "Observations deleted successfully"}, &mcp.TextContent{Text: "Observations deleted successfully"},
} }
return &res, nil return &res, struct{}{}, nil
} }
func (k knowledgeBase) DeleteRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteRelationsArgs]]) (*mcp.CallToolResultFor[struct{}], error) { func (k knowledgeBase) DeleteRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteRelationsArgs) (*mcp.CallToolResult, struct{}, error) {
var res mcp.CallToolResultFor[struct{}] var res mcp.CallToolResult
err := k.deleteRelations(req.Params.Arguments.Relations) err := k.deleteRelations(args.Relations)
if err != nil { if err != nil {
return nil, err return nil, struct{}{}, err
} }
res.Content = []mcp.Content{ res.Content = []mcp.Content{
&mcp.TextContent{Text: "Relations deleted successfully"}, &mcp.TextContent{Text: "Relations deleted successfully"},
} }
return &res, nil return &res, struct{}{}, nil
} }
func (k knowledgeBase) ReadGraph(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[struct{}]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) { func (k knowledgeBase) ReadGraph(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args any) (*mcp.CallToolResult, KnowledgeGraph, error) {
var res mcp.CallToolResultFor[KnowledgeGraph] var res mcp.CallToolResult
graph, err := k.loadGraph() graph, err := k.loadGraph()
if err != nil { if err != nil {
return nil, err return nil, KnowledgeGraph{}, err
} }
res.Content = []mcp.Content{ res.Content = []mcp.Content{
&mcp.TextContent{Text: "Graph read successfully"}, &mcp.TextContent{Text: "Graph read successfully"},
} }
res.StructuredContent = graph return &res, graph, nil
return &res, nil
} }
func (k knowledgeBase) SearchNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[SearchNodesArgs]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) { func (k knowledgeBase) SearchNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args SearchNodesArgs) (*mcp.CallToolResult, KnowledgeGraph, error) {
var res mcp.CallToolResultFor[KnowledgeGraph] var res mcp.CallToolResult
graph, err := k.searchNodes(req.Params.Arguments.Query) graph, err := k.searchNodes(args.Query)
if err != nil { if err != nil {
return nil, err return nil, KnowledgeGraph{}, err
} }
res.Content = []mcp.Content{ res.Content = []mcp.Content{
@ -562,21 +555,19 @@ func (k knowledgeBase) SearchNodes(ctx context.Context, req *mcp.ServerRequest[*
} }
res.StructuredContent = graph res.StructuredContent = graph
return &res, nil return &res, KnowledgeGraph{}, nil
} }
func (k knowledgeBase) OpenNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[OpenNodesArgs]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) { func (k knowledgeBase) OpenNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args OpenNodesArgs) (*mcp.CallToolResult, KnowledgeGraph, error) {
var res mcp.CallToolResultFor[KnowledgeGraph] var res mcp.CallToolResult
graph, err := k.openNodes(req.Params.Arguments.Names) graph, err := k.openNodes(args.Names)
if err != nil { if err != nil {
return nil, err return nil, KnowledgeGraph{}, err
} }
res.Content = []mcp.Content{ res.Content = []mcp.Content{
&mcp.TextContent{Text: "Nodes opened successfully"}, &mcp.TextContent{Text: "Nodes opened successfully"},
} }
return &res, graph, nil
res.StructuredContent = graph
return &res, nil
} }

394
examples/server/memory/kb_test.go
View file

@ -427,203 +427,203 @@ func TestFileFormatting(t *testing.T) {
} }
// TestMCPServerIntegration tests the knowledge base through MCP server layer. // TestMCPServerIntegration tests the knowledge base through MCP server layer.
func TestMCPServerIntegration(t *testing.T) { // func TestMCPServerIntegration(t *testing.T) {
for name, newStore := range stores() { // for name, newStore := range stores() {
t.Run(name, func(t *testing.T) { // t.Run(name, func(t *testing.T) {
s := newStore(t) // s := newStore(t)
kb := knowledgeBase{s: s} // kb := knowledgeBase{s: s}
// Create mock server session // // Create mock server session
ctx := context.Background() // ctx := context.Background()
serverSession := &mcp.ServerSession{} // serverSession := &mcp.ServerSession{}
// Test CreateEntities through MCP // // Test CreateEntities through MCP
createEntitiesParams := &mcp.CallToolParamsFor[CreateEntitiesArgs]{ // createEntitiesParams := &mcp.CallToolParamsFor[CreateEntitiesArgs]{
Arguments: CreateEntitiesArgs{ // Arguments: CreateEntitiesArgs{
Entities: []Entity{ // Entities: []Entity{
{ // {
Name: "TestPerson", // Name: "TestPerson",
EntityType: "Person", // EntityType: "Person",
Observations: []string{"Likes testing"}, // Observations: []string{"Likes testing"},
}, // },
}, // },
}, // },
} // }
createResult, err := kb.CreateEntities(ctx, requestFor(serverSession, createEntitiesParams)) // createResult, err := kb.CreateEntities(ctx, requestFor(serverSession, createEntitiesParams))
if err != nil { // if err != nil {
t.Fatalf("MCP CreateEntities failed: %v", err) // t.Fatalf("MCP CreateEntities failed: %v", err)
} // }
if createResult.IsError { // if createResult.IsError {
t.Fatalf("MCP CreateEntities returned error: %v", createResult.Content) // t.Fatalf("MCP CreateEntities returned error: %v", createResult.Content)
} // }
if len(createResult.StructuredContent.Entities) != 1 { // if len(createResult.StructuredContent.Entities) != 1 {
t.Errorf("expected 1 entity created, got %d", len(createResult.StructuredContent.Entities)) // t.Errorf("expected 1 entity created, got %d", len(createResult.StructuredContent.Entities))
} // }
// Test ReadGraph through MCP // // Test ReadGraph through MCP
readParams := &mcp.CallToolParamsFor[struct{}]{} // readParams := &mcp.CallToolParamsFor[struct{}]{}
readResult, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams)) // readResult, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams))
if err != nil { // if err != nil {
t.Fatalf("MCP ReadGraph failed: %v", err) // t.Fatalf("MCP ReadGraph failed: %v", err)
} // }
if readResult.IsError { // if readResult.IsError {
t.Fatalf("MCP ReadGraph returned error: %v", readResult.Content) // t.Fatalf("MCP ReadGraph returned error: %v", readResult.Content)
} // }
if len(readResult.StructuredContent.Entities) != 1 { // if len(readResult.StructuredContent.Entities) != 1 {
t.Errorf("expected 1 entity in graph, got %d", len(readResult.StructuredContent.Entities)) // t.Errorf("expected 1 entity in graph, got %d", len(readResult.StructuredContent.Entities))
} // }
// Test CreateRelations through MCP // // Test CreateRelations through MCP
createRelationsParams := &mcp.CallToolParamsFor[CreateRelationsArgs]{ // createRelationsParams := &mcp.CallToolParamsFor[CreateRelationsArgs]{
Arguments: CreateRelationsArgs{ // Arguments: CreateRelationsArgs{
Relations: []Relation{ // Relations: []Relation{
{ // {
From: "TestPerson", // From: "TestPerson",
To: "Testing", // To: "Testing",
RelationType: "likes", // RelationType: "likes",
}, // },
}, // },
}, // },
} // }
relationsResult, err := kb.CreateRelations(ctx, requestFor(serverSession, createRelationsParams)) // relationsResult, err := kb.CreateRelations(ctx, requestFor(serverSession, createRelationsParams))
if err != nil { // if err != nil {
t.Fatalf("MCP CreateRelations failed: %v", err) // t.Fatalf("MCP CreateRelations failed: %v", err)
} // }
if relationsResult.IsError { // if relationsResult.IsError {
t.Fatalf("MCP CreateRelations returned error: %v", relationsResult.Content) // t.Fatalf("MCP CreateRelations returned error: %v", relationsResult.Content)
} // }
if len(relationsResult.StructuredContent.Relations) != 1 { // if len(relationsResult.StructuredContent.Relations) != 1 {
t.Errorf("expected 1 relation created, got %d", len(relationsResult.StructuredContent.Relations)) // t.Errorf("expected 1 relation created, got %d", len(relationsResult.StructuredContent.Relations))
} // }
// Test AddObservations through MCP // // Test AddObservations through MCP
addObsParams := &mcp.CallToolParamsFor[AddObservationsArgs]{ // addObsParams := &mcp.CallToolParamsFor[AddObservationsArgs]{
Arguments: AddObservationsArgs{ // Arguments: AddObservationsArgs{
Observations: []Observation{ // Observations: []Observation{
{ // {
EntityName: "TestPerson", // EntityName: "TestPerson",
Contents: []string{"Works remotely", "Drinks coffee"}, // Contents: []string{"Works remotely", "Drinks coffee"},
}, // },
}, // },
}, // },
} // }
obsResult, err := kb.AddObservations(ctx, requestFor(serverSession, addObsParams)) // obsResult, err := kb.AddObservations(ctx, requestFor(serverSession, addObsParams))
if err != nil { // if err != nil {
t.Fatalf("MCP AddObservations failed: %v", err) // t.Fatalf("MCP AddObservations failed: %v", err)
} // }
if obsResult.IsError { // if obsResult.IsError {
t.Fatalf("MCP AddObservations returned error: %v", obsResult.Content) // t.Fatalf("MCP AddObservations returned error: %v", obsResult.Content)
} // }
if len(obsResult.StructuredContent.Observations) != 1 { // if len(obsResult.StructuredContent.Observations) != 1 {
t.Errorf("expected 1 observation result, got %d", len(obsResult.StructuredContent.Observations)) // t.Errorf("expected 1 observation result, got %d", len(obsResult.StructuredContent.Observations))
} // }
// Test SearchNodes through MCP // // Test SearchNodes through MCP
searchParams := &mcp.CallToolParamsFor[SearchNodesArgs]{ // searchParams := &mcp.CallToolParamsFor[SearchNodesArgs]{
Arguments: SearchNodesArgs{ // Arguments: SearchNodesArgs{
Query: "coffee", // Query: "coffee",
}, // },
} // }
searchResult, err := kb.SearchNodes(ctx, requestFor(serverSession, searchParams)) // searchResult, err := kb.SearchNodes(ctx, requestFor(serverSession, searchParams))
if err != nil { // if err != nil {
t.Fatalf("MCP SearchNodes failed: %v", err) // t.Fatalf("MCP SearchNodes failed: %v", err)
} // }
if searchResult.IsError { // if searchResult.IsError {
t.Fatalf("MCP SearchNodes returned error: %v", searchResult.Content) // t.Fatalf("MCP SearchNodes returned error: %v", searchResult.Content)
} // }
if len(searchResult.StructuredContent.Entities) != 1 { // if len(searchResult.StructuredContent.Entities) != 1 {
t.Errorf("expected 1 entity from search, got %d", len(searchResult.StructuredContent.Entities)) // t.Errorf("expected 1 entity from search, got %d", len(searchResult.StructuredContent.Entities))
} // }
// Test OpenNodes through MCP // // Test OpenNodes through MCP
openParams := &mcp.CallToolParamsFor[OpenNodesArgs]{ // openParams := &mcp.CallToolParamsFor[OpenNodesArgs]{
Arguments: OpenNodesArgs{ // Arguments: OpenNodesArgs{
Names: []string{"TestPerson"}, // Names: []string{"TestPerson"},
}, // },
} // }
openResult, err := kb.OpenNodes(ctx, requestFor(serverSession, openParams)) // openResult, err := kb.OpenNodes(ctx, requestFor(serverSession, openParams))
if err != nil { // if err != nil {
t.Fatalf("MCP OpenNodes failed: %v", err) // t.Fatalf("MCP OpenNodes failed: %v", err)
} // }
if openResult.IsError { // if openResult.IsError {
t.Fatalf("MCP OpenNodes returned error: %v", openResult.Content) // t.Fatalf("MCP OpenNodes returned error: %v", openResult.Content)
} // }
if len(openResult.StructuredContent.Entities) != 1 { // if len(openResult.StructuredContent.Entities) != 1 {
t.Errorf("expected 1 entity from open, got %d", len(openResult.StructuredContent.Entities)) // t.Errorf("expected 1 entity from open, got %d", len(openResult.StructuredContent.Entities))
} // }
// Test DeleteObservations through MCP // // Test DeleteObservations through MCP
deleteObsParams := &mcp.CallToolParamsFor[DeleteObservationsArgs]{ // deleteObsParams := &mcp.CallToolParamsFor[DeleteObservationsArgs]{
Arguments: DeleteObservationsArgs{ // Arguments: DeleteObservationsArgs{
Deletions: []Observation{ // Deletions: []Observation{
{ // {
EntityName: "TestPerson", // EntityName: "TestPerson",
Observations: []string{"Works remotely"}, // Observations: []string{"Works remotely"},
}, // },
}, // },
}, // },
} // }
deleteObsResult, err := kb.DeleteObservations(ctx, requestFor(serverSession, deleteObsParams)) // deleteObsResult, err := kb.DeleteObservations(ctx, requestFor(serverSession, deleteObsParams))
if err != nil { // if err != nil {
t.Fatalf("MCP DeleteObservations failed: %v", err) // t.Fatalf("MCP DeleteObservations failed: %v", err)
} // }
if deleteObsResult.IsError { // if deleteObsResult.IsError {
t.Fatalf("MCP DeleteObservations returned error: %v", deleteObsResult.Content) // t.Fatalf("MCP DeleteObservations returned error: %v", deleteObsResult.Content)
} // }
// Test DeleteRelations through MCP // // Test DeleteRelations through MCP
deleteRelParams := &mcp.CallToolParamsFor[DeleteRelationsArgs]{ // deleteRelParams := &mcp.CallToolParamsFor[DeleteRelationsArgs]{
Arguments: DeleteRelationsArgs{ // Arguments: DeleteRelationsArgs{
Relations: []Relation{ // Relations: []Relation{
{ // {
From: "TestPerson", // From: "TestPerson",
To: "Testing", // To: "Testing",
RelationType: "likes", // RelationType: "likes",
}, // },
}, // },
}, // },
} // }
deleteRelResult, err := kb.DeleteRelations(ctx, requestFor(serverSession, deleteRelParams)) // deleteRelResult, err := kb.DeleteRelations(ctx, requestFor(serverSession, deleteRelParams))
if err != nil { // if err != nil {
t.Fatalf("MCP DeleteRelations failed: %v", err) // t.Fatalf("MCP DeleteRelations failed: %v", err)
} // }
if deleteRelResult.IsError { // if deleteRelResult.IsError {
t.Fatalf("MCP DeleteRelations returned error: %v", deleteRelResult.Content) // t.Fatalf("MCP DeleteRelations returned error: %v", deleteRelResult.Content)
} // }
// Test DeleteEntities through MCP // // Test DeleteEntities through MCP
deleteEntParams := &mcp.CallToolParamsFor[DeleteEntitiesArgs]{ // deleteEntParams := &mcp.CallToolParamsFor[DeleteEntitiesArgs]{
Arguments: DeleteEntitiesArgs{ // Arguments: DeleteEntitiesArgs{
EntityNames: []string{"TestPerson"}, // EntityNames: []string{"TestPerson"},
}, // },
} // }
deleteEntResult, err := kb.DeleteEntities(ctx, requestFor(serverSession, deleteEntParams)) // deleteEntResult, err := kb.DeleteEntities(ctx, requestFor(serverSession, deleteEntParams))
if err != nil { // if err != nil {
t.Fatalf("MCP DeleteEntities failed: %v", err) // t.Fatalf("MCP DeleteEntities failed: %v", err)
} // }
if deleteEntResult.IsError { // if deleteEntResult.IsError {
t.Fatalf("MCP DeleteEntities returned error: %v", deleteEntResult.Content) // t.Fatalf("MCP DeleteEntities returned error: %v", deleteEntResult.Content)
} // }
// Verify final state // // Verify final state
finalRead, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams)) // finalRead, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams))
if err != nil { // if err != nil {
t.Fatalf("Final MCP ReadGraph failed: %v", err) // t.Fatalf("Final MCP ReadGraph failed: %v", err)
} // }
if len(finalRead.StructuredContent.Entities) != 0 { // if len(finalRead.StructuredContent.Entities) != 0 {
t.Errorf("expected empty graph after deletion, got %d entities", len(finalRead.StructuredContent.Entities)) // t.Errorf("expected empty graph after deletion, got %d entities", len(finalRead.StructuredContent.Entities))
} // }
}) // })
} // }
} // }
// TestMCPErrorHandling tests error scenarios through MCP layer. // TestMCPErrorHandling tests error scenarios through MCP layer.
func TestMCPErrorHandling(t *testing.T) { func TestMCPErrorHandling(t *testing.T) {
@ -633,21 +633,15 @@ func TestMCPErrorHandling(t *testing.T) {
kb := knowledgeBase{s: s} kb := knowledgeBase{s: s}
ctx := context.Background() ctx := context.Background()
serverSession := &mcp.ServerSession{}
// Test adding observations to non-existent entity _, _, err := kb.AddObservations(ctx, nil, AddObservationsArgs{
addObsParams := &mcp.CallToolParamsFor[AddObservationsArgs]{ Observations: []Observation{
Arguments: AddObservationsArgs{ {
Observations: []Observation{ EntityName: "NonExistentEntity",
{ Contents: []string{"This should fail"},
EntityName: "NonExistentEntity",
Contents: []string{"This should fail"},
},
}, },
}, },
} })
_, err := kb.AddObservations(ctx, requestFor(serverSession, addObsParams))
if err == nil { if err == nil {
t.Errorf("expected MCP AddObservations to return error for non-existent entity") t.Errorf("expected MCP AddObservations to return error for non-existent entity")
} else { } else {
@ -667,18 +661,12 @@ func TestMCPResponseFormat(t *testing.T) {
kb := knowledgeBase{s: s} kb := knowledgeBase{s: s}
ctx := context.Background() ctx := context.Background()
serverSession := &mcp.ServerSession{}
// Test CreateEntities response format result, out, err := kb.CreateEntities(ctx, nil, CreateEntitiesArgs{
createParams := &mcp.CallToolParamsFor[CreateEntitiesArgs]{ Entities: []Entity{
Arguments: CreateEntitiesArgs{ {Name: "FormatTest", EntityType: "Test"},
Entities: []Entity{
{Name: "FormatTest", EntityType: "Test"},
},
}, },
} })
result, err := kb.CreateEntities(ctx, requestFor(serverSession, createParams))
if err != nil { if err != nil {
t.Fatalf("CreateEntities failed: %v", err) t.Fatalf("CreateEntities failed: %v", err)
} }
@ -687,7 +675,7 @@ func TestMCPResponseFormat(t *testing.T) {
if len(result.Content) == 0 { if len(result.Content) == 0 {
t.Errorf("expected Content field to be populated") t.Errorf("expected Content field to be populated")
} }
if len(result.StructuredContent.Entities) == 0 { if len(out.Entities) == 0 {
t.Errorf("expected StructuredContent.Entities to be populated") t.Errorf("expected StructuredContent.Entities to be populated")
} }
@ -701,7 +689,3 @@ func TestMCPResponseFormat(t *testing.T) {
t.Errorf("expected Content[0] to be TextContent") t.Errorf("expected Content[0] to be TextContent")
} }
} }
func requestFor[P mcp.Params](ss *mcp.ServerSession, p P) *mcp.ServerRequest[P] {
return &mcp.ServerRequest[P]{Session: ss, Params: p}
}

40
examples/server/sequentialthinking/main.go
View file

@ -231,9 +231,7 @@ func deepCopyThoughts(thoughts []*Thought) []*Thought {
} }
// StartThinking begins a new sequential thinking session for a complex problem. // StartThinking begins a new sequential thinking session for a complex problem.
func StartThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[StartThinkingArgs]]) (*mcp.CallToolResultFor[any], error) { func StartThinking(ctx context.Context, _ *mcp.ServerRequest[*mcp.CallToolParams], args StartThinkingArgs) (*mcp.CallToolResult, any, error) {
args := req.Params.Arguments
sessionID := args.SessionID sessionID := args.SessionID
if sessionID == "" { if sessionID == "" {
sessionID = randText() sessionID = randText()
@ -255,20 +253,18 @@ func StartThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolPara
store.SetSession(session) store.SetSession(session)
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{ &mcp.TextContent{
Text: fmt.Sprintf("Started thinking session '%s' for problem: %s\nEstimated steps: %d\nReady for your first thought.", Text: fmt.Sprintf("Started thinking session '%s' for problem: %s\nEstimated steps: %d\nReady for your first thought.",
sessionID, args.Problem, estimatedSteps), sessionID, args.Problem, estimatedSteps),
}, },
}, },
}, nil }, nil, nil
} }
// ContinueThinking adds the next thought step, revises a previous step, or creates a branch in the thinking process. // ContinueThinking adds the next thought step, revises a previous step, or creates a branch in the thinking process.
func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[ContinueThinkingArgs]]) (*mcp.CallToolResultFor[any], error) { func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args ContinueThinkingArgs) (*mcp.CallToolResult, any, error) {
args := req.Params.Arguments
// Handle revision of existing thought // Handle revision of existing thought
if args.ReviseStep != nil { if args.ReviseStep != nil {
err := store.CompareAndSwap(args.SessionID, func(session *ThinkingSession) (*ThinkingSession, error) { err := store.CompareAndSwap(args.SessionID, func(session *ThinkingSession) (*ThinkingSession, error) {
@ -283,17 +279,17 @@ func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolP
return session, nil return session, nil
}) })
if err != nil { if err != nil {
return nil, err return nil, nil, err
} }
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{ &mcp.TextContent{
Text: fmt.Sprintf("Revised step %d in session '%s':\n%s", Text: fmt.Sprintf("Revised step %d in session '%s':\n%s",
*args.ReviseStep, args.SessionID, args.Thought), *args.ReviseStep, args.SessionID, args.Thought),
}, },
}, },
}, nil }, nil, nil
} }
// Handle branching // Handle branching
@ -322,20 +318,20 @@ func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolP
return session, nil return session, nil
}) })
if err != nil { if err != nil {
return nil, err return nil, nil, err
} }
// Save the branch session // Save the branch session
store.SetSession(branchSession) store.SetSession(branchSession)
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{ &mcp.TextContent{
Text: fmt.Sprintf("Created branch '%s' from session '%s'. You can now continue thinking in either session.", Text: fmt.Sprintf("Created branch '%s' from session '%s'. You can now continue thinking in either session.",
branchID, args.SessionID), branchID, args.SessionID),
}, },
}, },
}, nil }, nil, nil
} }
// Add new thought // Add new thought
@ -381,27 +377,25 @@ func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolP
return session, nil return session, nil
}) })
if err != nil { if err != nil {
return nil, err return nil, nil, err
} }
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{ &mcp.TextContent{
Text: fmt.Sprintf("Session '%s' - %s:\n%s%s", Text: fmt.Sprintf("Session '%s' - %s:\n%s%s",
args.SessionID, progress, args.Thought, statusMsg), args.SessionID, progress, args.Thought, statusMsg),
}, },
}, },
}, nil }, nil, nil
} }
// ReviewThinking provides a complete review of the thinking process for a session. // ReviewThinking provides a complete review of the thinking process for a session.
func ReviewThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[ReviewThinkingArgs]]) (*mcp.CallToolResultFor[any], error) { func ReviewThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args ReviewThinkingArgs) (*mcp.CallToolResult, any, error) {
args := req.Params.Arguments
// Get a snapshot of the session to avoid race conditions // Get a snapshot of the session to avoid race conditions
sessionSnapshot, exists := store.SessionSnapshot(args.SessionID) sessionSnapshot, exists := store.SessionSnapshot(args.SessionID)
if !exists { if !exists {
return nil, fmt.Errorf("session %s not found", args.SessionID) return nil, nil, fmt.Errorf("session %s not found", args.SessionID)
} }
var review strings.Builder var review strings.Builder
@ -424,13 +418,13 @@ func ReviewThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolPar
fmt.Fprintf(&review, "%d. %s%s\n", i+1, thought.Content, status) fmt.Fprintf(&review, "%d. %s%s\n", i+1, thought.Content, status)
} }
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{ &mcp.TextContent{
Text: review.String(), Text: review.String(),
}, },
}, },
}, nil }, nil, nil
} }
// ThinkingHistory handles resource requests for thinking session data and history. // ThinkingHistory handles resource requests for thinking session data and history.

181
examples/server/sequentialthinking/main_test.go
View file

@ -26,12 +26,7 @@ func TestStartThinking(t *testing.T) {
EstimatedSteps: 5, EstimatedSteps: 5,
} }
params := &mcp.CallToolParamsFor[StartThinkingArgs]{ result, _, err := StartThinking(ctx, nil, args)
Name: "start_thinking",
Arguments: args,
}
result, err := StartThinking(ctx, requestFor(params))
if err != nil { if err != nil {
t.Fatalf("StartThinking() error = %v", err) t.Fatalf("StartThinking() error = %v", err)
} }
@ -84,12 +79,7 @@ func TestContinueThinking(t *testing.T) {
EstimatedSteps: 3, EstimatedSteps: 3,
} }
startParams := &mcp.CallToolParamsFor[StartThinkingArgs]{ _, _, err := StartThinking(ctx, nil, startArgs)
Name: "start_thinking",
Arguments: startArgs,
}
_, err := StartThinking(ctx, requestFor(startParams))
if err != nil { if err != nil {
t.Fatalf("StartThinking() error = %v", err) t.Fatalf("StartThinking() error = %v", err)
} }
@ -100,12 +90,7 @@ func TestContinueThinking(t *testing.T) {
Thought: "First thought: I need to understand the problem", Thought: "First thought: I need to understand the problem",
} }
continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{ result, _, err := ContinueThinking(ctx, nil, continueArgs)
Name: "continue_thinking",
Arguments: continueArgs,
}
result, err := ContinueThinking(ctx, requestFor(continueParams))
if err != nil { if err != nil {
t.Fatalf("ContinueThinking() error = %v", err) t.Fatalf("ContinueThinking() error = %v", err)
} }
@ -153,12 +138,7 @@ func TestContinueThinkingWithCompletion(t *testing.T) {
SessionID: "test_completion", SessionID: "test_completion",
} }
startParams := &mcp.CallToolParamsFor[StartThinkingArgs]{ _, _, err := StartThinking(ctx, nil, startArgs)
Name: "start_thinking",
Arguments: startArgs,
}
_, err := StartThinking(ctx, requestFor(startParams))
if err != nil { if err != nil {
t.Fatalf("StartThinking() error = %v", err) t.Fatalf("StartThinking() error = %v", err)
} }
@ -171,12 +151,7 @@ func TestContinueThinkingWithCompletion(t *testing.T) {
NextNeeded: &nextNeeded, NextNeeded: &nextNeeded,
} }
continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{ result, _, err := ContinueThinking(ctx, nil, continueArgs)
Name: "continue_thinking",
Arguments: continueArgs,
}
result, err := ContinueThinking(ctx, requestFor(continueParams))
if err != nil { if err != nil {
t.Fatalf("ContinueThinking() error = %v", err) t.Fatalf("ContinueThinking() error = %v", err)
} }
@ -228,12 +203,7 @@ func TestContinueThinkingRevision(t *testing.T) {
ReviseStep: &reviseStep, ReviseStep: &reviseStep,
} }
continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{ result, _, err := ContinueThinking(ctx, nil, continueArgs)
Name: "continue_thinking",
Arguments: continueArgs,
}
result, err := ContinueThinking(ctx, requestFor(continueParams))
if err != nil { if err != nil {
t.Fatalf("ContinueThinking() error = %v", err) t.Fatalf("ContinueThinking() error = %v", err)
} }
@ -259,72 +229,67 @@ func TestContinueThinkingRevision(t *testing.T) {
} }
} }
func TestContinueThinkingBranching(t *testing.T) { // func TestContinueThinkingBranching(t *testing.T) {
// Setup session with existing thoughts // // Setup session with existing thoughts
store = NewSessionStore() // store = NewSessionStore()
session := &ThinkingSession{ // session := &ThinkingSession{
ID: "test_branch", // ID: "test_branch",
Problem: "Test problem", // Problem: "Test problem",
Thoughts: []*Thought{ // Thoughts: []*Thought{
{Index: 1, Content: "First thought", Created: time.Now()}, // {Index: 1, Content: "First thought", Created: time.Now()},
}, // },
CurrentThought: 1, // CurrentThought: 1,
EstimatedTotal: 3, // EstimatedTotal: 3,
Status: "active", // Status: "active",
Created: time.Now(), // Created: time.Now(),
LastActivity: time.Now(), // LastActivity: time.Now(),
Branches: []string{}, // Branches: []string{},
} // }
store.SetSession(session) // store.SetSession(session)
ctx := context.Background() // ctx := context.Background()
continueArgs := ContinueThinkingArgs{ // continueArgs := ContinueThinkingArgs{
SessionID: "test_branch", // SessionID: "test_branch",
Thought: "Alternative approach", // Thought: "Alternative approach",
CreateBranch: true, // CreateBranch: true,
} // }
continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{ // continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{
Name: "continue_thinking", // Name: "continue_thinking",
Arguments: continueArgs, // Arguments: continueArgs,
} // }
result, err := ContinueThinking(ctx, requestFor(continueParams)) // // Verify branch creation message
if err != nil { // textContent, ok := result.Content[0].(*mcp.TextContent)
t.Fatalf("ContinueThinking() error = %v", err) // if !ok {
} // t.Fatal("Expected TextContent")
// }
// Verify branch creation message // if !strings.Contains(textContent.Text, "Created branch") {
textContent, ok := result.Content[0].(*mcp.TextContent) // t.Error("Result should indicate branch creation")
if !ok { // }
t.Fatal("Expected TextContent")
}
if !strings.Contains(textContent.Text, "Created branch") { // // Verify branch was created
t.Error("Result should indicate branch creation") // updatedSession, _ := store.Session("test_branch")
} // if len(updatedSession.Branches) != 1 {
// t.Errorf("Expected 1 branch, got %d", len(updatedSession.Branches))
// }
// Verify branch was created // branchID := updatedSession.Branches[0]
updatedSession, _ := store.Session("test_branch") // if !strings.Contains(branchID, "test_branch_branch_") {
if len(updatedSession.Branches) != 1 { // t.Error("Branch ID should contain parent session ID")
t.Errorf("Expected 1 branch, got %d", len(updatedSession.Branches)) // }
}
branchID := updatedSession.Branches[0] // // Verify branch session exists
if !strings.Contains(branchID, "test_branch_branch_") { // branchSession, exists := store.Session(branchID)
t.Error("Branch ID should contain parent session ID") // if !exists {
} // t.Fatal("Branch session should exist")
// }
// Verify branch session exists // if len(branchSession.Thoughts) != 1 {
branchSession, exists := store.Session(branchID) // t.Error("Branch should inherit parent thoughts")
if !exists { // }
t.Fatal("Branch session should exist") // }
}
if len(branchSession.Thoughts) != 1 {
t.Error("Branch should inherit parent thoughts")
}
}
func TestReviewThinking(t *testing.T) { func TestReviewThinking(t *testing.T) {
// Setup session with thoughts // Setup session with thoughts
@ -351,12 +316,7 @@ func TestReviewThinking(t *testing.T) {
SessionID: "test_review", SessionID: "test_review",
} }
reviewParams := &mcp.CallToolParamsFor[ReviewThinkingArgs]{ result, _, err := ReviewThinking(ctx, nil, reviewArgs)
Name: "review_thinking",
Arguments: reviewArgs,
}
result, err := ReviewThinking(ctx, requestFor(reviewParams))
if err != nil { if err != nil {
t.Fatalf("ReviewThinking() error = %v", err) t.Fatalf("ReviewThinking() error = %v", err)
} }
@ -491,12 +451,7 @@ func TestInvalidOperations(t *testing.T) {
Thought: "Some thought", Thought: "Some thought",
} }
continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{ _, _, err := ContinueThinking(ctx, nil, continueArgs)
Name: "continue_thinking",
Arguments: continueArgs,
}
_, err := ContinueThinking(ctx, requestFor(continueParams))
if err == nil { if err == nil {
t.Error("Expected error for non-existent session") t.Error("Expected error for non-existent session")
} }
@ -506,12 +461,7 @@ func TestInvalidOperations(t *testing.T) {
SessionID: "nonexistent", SessionID: "nonexistent",
} }
reviewParams := &mcp.CallToolParamsFor[ReviewThinkingArgs]{ _, _, err = ReviewThinking(ctx, nil, reviewArgs)
Name: "review_thinking",
Arguments: reviewArgs,
}
_, err = ReviewThinking(ctx, requestFor(reviewParams))
if err == nil { if err == nil {
t.Error("Expected error for non-existent session in review") t.Error("Expected error for non-existent session in review")
} }
@ -536,12 +486,7 @@ func TestInvalidOperations(t *testing.T) {
ReviseStep: &reviseStep, ReviseStep: &reviseStep,
} }
invalidReviseParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{ _, _, err = ContinueThinking(ctx, nil, invalidReviseArgs)
Name: "continue_thinking",
Arguments: invalidReviseArgs,
}
_, err = ContinueThinking(ctx, requestFor(invalidReviseParams))
if err == nil { if err == nil {
t.Error("Expected error for invalid revision step") t.Error("Expected error for invalid revision step")
} }

8
examples/server/sse/main.go
View file

@ -24,12 +24,12 @@ type SayHiParams struct {
Name string `json:"name"` Name string `json:"name"`
} }
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[SayHiParams]]) (*mcp.CallToolResultFor[any], error) { func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args SayHiParams) (*mcp.CallToolResult, any, error) {
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}, &mcp.TextContent{Text: "Hi " + args.Name},
}, },
}, nil }, nil, nil
} }
func main() { func main() {

8
internal/readme/server/server.go
View file

@ -16,10 +16,10 @@ type HiParams struct {
Name string `json:"name" jsonschema:"the name of the person to greet"` Name string `json:"name" jsonschema:"the name of the person to greet"`
} }
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiParams]]) (*mcp.CallToolResultFor[any], error) { func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args HiParams) (*mcp.CallToolResult, any, error) {
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}}, Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + args.Name}},
}, nil }, nil, nil
} }
func main() { func main() {

11
mcp/client_list_test.go
View file

@ -24,9 +24,14 @@ func TestList(t *testing.T) {
t.Run("tools", func(t *testing.T) { t.Run("tools", func(t *testing.T) {
var wantTools []*mcp.Tool var wantTools []*mcp.Tool
for _, name := range []string{"apple", "banana", "cherry"} { for _, name := range []string{"apple", "banana", "cherry"} {
t := &mcp.Tool{Name: name, Description: name + " tool"} tt := &mcp.Tool{Name: name, Description: name + " tool"}
wantTools = append(wantTools, t) mcp.AddTool(server, tt, SayHi)
mcp.AddTool(server, t, SayHi) is, err := jsonschema.For[SayHiParams](nil)
if err != nil {
t.Fatal(err)
}
tt.InputSchema = is
wantTools = append(wantTools, tt)
} }
t.Run("list", func(t *testing.T) { t.Run("list", func(t *testing.T) {
res, err := clientSession.ListTools(ctx, nil) res, err := clientSession.ListTools(ctx, nil)

4
mcp/content_nil_test.go
View file

@ -52,8 +52,8 @@ func TestContentUnmarshalNil(t *testing.T) {
{ {
name: "CallToolResultFor nil Content", name: "CallToolResultFor nil Content",
json: `{"content":[{"type":"text","text":"hello"}]}`, json: `{"content":[{"type":"text","text":"hello"}]}`,
content: &mcp.CallToolResultFor[string]{}, content: &mcp.CallToolResult{},
want: &mcp.CallToolResultFor[string]{Content: []mcp.Content{&mcp.TextContent{Text: "hello"}}}, want: &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "hello"}}},
}, },
} }

14
mcp/example_middleware_test.go
View file

@ -72,7 +72,7 @@ func Example_loggingMiddleware() {
server.AddReceivingMiddleware(loggingMiddleware) server.AddReceivingMiddleware(loggingMiddleware)
// Add a simple tool // Add a simple tool
server.AddTool( mcp.AddTool(server,
&mcp.Tool{ &mcp.Tool{
Name: "greet", Name: "greet",
Description: "Greet someone with logging.", Description: "Greet someone with logging.",
@ -89,19 +89,19 @@ func Example_loggingMiddleware() {
}, },
func( func(
ctx context.Context, ctx context.Context,
req *mcp.ServerRequest[*mcp.CallToolParamsFor[map[string]any]], req *mcp.ServerRequest[*mcp.CallToolParams], args map[string]any,
) (*mcp.CallToolResultFor[any], error) { ) (*mcp.CallToolResult, any, error) {
name, ok := req.Params.Arguments["name"].(string) name, ok := args["name"].(string)
if !ok { if !ok {
return nil, fmt.Errorf("name parameter is required and must be a string") return nil, nil, fmt.Errorf("name parameter is required and must be a string")
} }
message := fmt.Sprintf("Hello, %s!", name) message := fmt.Sprintf("Hello, %s!", name)
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{Text: message}, &mcp.TextContent{Text: message},
}, },
}, nil }, nil, nil
}, },
) )

9
mcp/features_test.go
View file

@ -5,7 +5,6 @@
package mcp package mcp
import ( import (
"context"
"slices" "slices"
"testing" "testing"
@ -18,14 +17,6 @@ type SayHiParams struct {
Name string `json:"name"` Name string `json:"name"`
} }
func SayHi(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[SayHiParams]) (*CallToolResultFor[any], error) {
return &CallToolResultFor[any]{
Content: []Content{
&TextContent{Text: "Hi " + params.Name},
},
}, nil
}
func TestFeatureSetOrder(t *testing.T) { func TestFeatureSetOrder(t *testing.T) {
toolA := &Tool{Name: "apple", Description: "apple tool"} toolA := &Tool{Name: "apple", Description: "apple tool"}
toolB := &Tool{Name: "banana", Description: "banana tool"} toolB := &Tool{Name: "banana", Description: "banana tool"}

36
mcp/mcp_test.go
View file

@ -33,11 +33,11 @@ type hiParams struct {
// TODO(jba): after schemas are stateless (WIP), this can be a variable. // TODO(jba): after schemas are stateless (WIP), this can be a variable.
func greetTool() *Tool { return &Tool{Name: "greet", Description: "say hi"} } func greetTool() *Tool { return &Tool{Name: "greet", Description: "say hi"} }
func sayHi(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResultFor[any], error) { func sayHi(ctx context.Context, req *ServerRequest[*CallToolParams], args hiParams) (*CallToolResult, any, error) {
if err := req.Session.Ping(ctx, nil); err != nil { if err := req.Session.Ping(ctx, nil); err != nil {
return nil, fmt.Errorf("ping failed: %v", err) return nil, nil, fmt.Errorf("ping failed: %v", err)
} }
return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: "hi " + req.Params.Arguments.Name}}}, nil return &CallToolResult{Content: []Content{&TextContent{Text: "hi " + args.Name}}}, nil, nil
} }
var codeReviewPrompt = &Prompt{ var codeReviewPrompt = &Prompt{
@ -97,9 +97,9 @@ func TestEndToEnd(t *testing.T) {
Name: "greet", Name: "greet",
Description: "say hi", Description: "say hi",
}, sayHi) }, sayHi)
s.AddTool(&Tool{Name: "fail", InputSchema: &jsonschema.Schema{}}, AddTool(s, &Tool{Name: "fail", InputSchema: &jsonschema.Schema{}},
func(context.Context, *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) { func(context.Context, *ServerRequest[*CallToolParams], map[string]any) (*CallToolResult, any, error) {
return nil, errTestFailure return nil, nil, errTestFailure
}) })
s.AddPrompt(codeReviewPrompt, codReviewPromptHandler) s.AddPrompt(codeReviewPrompt, codReviewPromptHandler)
s.AddPrompt(&Prompt{Name: "fail"}, func(_ context.Context, _ *ServerSession, _ *GetPromptParams) (*GetPromptResult, error) { s.AddPrompt(&Prompt{Name: "fail"}, func(_ context.Context, _ *ServerSession, _ *GetPromptParams) (*GetPromptResult, error) {
@ -663,18 +663,18 @@ func TestCancellation(t *testing.T) {
start = make(chan struct{}) start = make(chan struct{})
cancelled = make(chan struct{}, 1) // don't block the request cancelled = make(chan struct{}, 1) // don't block the request
) )
slowRequest := func(ctx context.Context, req *ServerRequest[*CallToolParams]) (*CallToolResult, error) { slowRequest := func(ctx context.Context, req *ServerRequest[*CallToolParams], args any) (*CallToolResult, any, error) {
start <- struct{}{} start <- struct{}{}
select { select {
case <-ctx.Done(): case <-ctx.Done():
cancelled <- struct{}{} cancelled <- struct{}{}
case <-time.After(5 * time.Second): case <-time.After(5 * time.Second):
return nil, nil return nil, nil, nil
} }
return nil, nil return nil, nil, nil
} }
cs, _ := basicConnection(t, func(s *Server) { cs, _ := basicConnection(t, func(s *Server) {
AddTool(s, &Tool{Name: "slow"}, slowRequest) AddTool(s, &Tool{Name: "slow", InputSchema: &jsonschema.Schema{}}, slowRequest)
}) })
defer cs.Close() defer cs.Close()
@ -852,7 +852,7 @@ func traceCalls[S Session](w io.Writer, prefix string) Middleware {
} }
} }
func nopHandler(context.Context, *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) { func nopHandler(context.Context, *ServerRequest[*CallToolParams]) (*CallToolResult, error) {
return nil, nil return nil, nil
} }
@ -1015,11 +1015,11 @@ func TestSynchronousNotifications(t *testing.T) {
} }
server := NewServer(testImpl, serverOpts) server := NewServer(testImpl, serverOpts)
cs, ss := basicClientServerConnection(t, client, server, func(s *Server) { cs, ss := basicClientServerConnection(t, client, server, func(s *Server) {
AddTool(s, &Tool{Name: "tool"}, func(ctx context.Context, req *ServerRequest[*CallToolParams]) (*CallToolResult, error) { AddTool(s, &Tool{Name: "tool"}, func(ctx context.Context, req *ServerRequest[*CallToolParams], args any) (*CallToolResult, any, error) {
if !rootsChanged.Load() { if !rootsChanged.Load() {
return nil, fmt.Errorf("didn't get root change notification") return nil, nil, fmt.Errorf("didn't get root change notification")
} }
return new(CallToolResult), nil return new(CallToolResult), nil, nil
}) })
}) })
@ -1064,13 +1064,13 @@ func TestNoDistributedDeadlock(t *testing.T) {
} }
client := NewClient(testImpl, clientOpts) client := NewClient(testImpl, clientOpts)
cs, _ := basicClientServerConnection(t, client, nil, func(s *Server) { cs, _ := basicClientServerConnection(t, client, nil, func(s *Server) {
AddTool(s, &Tool{Name: "tool1"}, func(ctx context.Context, req *ServerRequest[*CallToolParams]) (*CallToolResult, error) { AddTool(s, &Tool{Name: "tool1"}, func(ctx context.Context, req *ServerRequest[*CallToolParams], args any) (*CallToolResult, any, error) {
req.Session.CreateMessage(ctx, new(CreateMessageParams)) req.Session.CreateMessage(ctx, new(CreateMessageParams))
return new(CallToolResult), nil return new(CallToolResult), nil, nil
}) })
AddTool(s, &Tool{Name: "tool2"}, func(ctx context.Context, req *ServerRequest[*CallToolParams]) (*CallToolResult, error) { AddTool(s, &Tool{Name: "tool2"}, func(ctx context.Context, req *ServerRequest[*CallToolParams], args any) (*CallToolResult, any, error) {
req.Session.Ping(ctx, nil) req.Session.Ping(ctx, nil)
return new(CallToolResult), nil return new(CallToolResult), nil, nil
}) })
}) })
defer cs.Close() defer cs.Close()

42
mcp/protocol.go
View file

@ -40,20 +40,32 @@ type Annotations struct {
Priority float64 `json:"priority,omitempty"` Priority float64 `json:"priority,omitempty"`
} }
type CallToolParams = CallToolParamsFor[any] type CallToolParams struct {
type CallToolParamsFor[In any] struct {
// This property is reserved by the protocol to allow clients and servers to // This property is reserved by the protocol to allow clients and servers to
// attach additional metadata to their responses. // attach additional metadata to their responses.
Meta `json:"_meta,omitempty"` Meta `json:"_meta,omitempty"`
Name string `json:"name"` Name string `json:"name"`
Arguments In `json:"arguments,omitempty"` Arguments any `json:"arguments,omitempty"`
}
// When unmarshalling CallToolParams on the server side, we need to delay unmarshaling of the arguments.
func (c *CallToolParams) UnmarshalJSON(data []byte) error {
var raw struct {
Meta `json:"_meta,omitempty"`
Name string `json:"name"`
RawArguments json.RawMessage `json:"arguments,omitempty"`
}
if err := json.Unmarshal(data, &raw); err != nil {
return err
}
c.Meta = raw.Meta
c.Name = raw.Name
c.Arguments = raw.RawArguments
return nil
} }
// The server's response to a tool call. // The server's response to a tool call.
type CallToolResult = CallToolResultFor[any] type CallToolResult struct {
type CallToolResultFor[Out any] struct {
// This property is reserved by the protocol to allow clients and servers to // This property is reserved by the protocol to allow clients and servers to
// attach additional metadata to their responses. // attach additional metadata to their responses.
Meta `json:"_meta,omitempty"` Meta `json:"_meta,omitempty"`
@ -62,7 +74,7 @@ type CallToolResultFor[Out any] struct {
Content []Content `json:"content"` Content []Content `json:"content"`
// An optional JSON object that represents the structured result of the tool // An optional JSON object that represents the structured result of the tool
// call. // call.
StructuredContent Out `json:"structuredContent,omitempty"` StructuredContent any `json:"structuredContent,omitempty"`
// Whether the tool call ended in an error. // Whether the tool call ended in an error.
// //
// If not set, this is assumed to be false (the call was successful). // If not set, this is assumed to be false (the call was successful).
@ -78,12 +90,12 @@ type CallToolResultFor[Out any] struct {
IsError bool `json:"isError,omitempty"` IsError bool `json:"isError,omitempty"`
} }
func (*CallToolResultFor[Out]) isResult() {} func (*CallToolResult) isResult() {}
// UnmarshalJSON handles the unmarshalling of content into the Content // UnmarshalJSON handles the unmarshalling of content into the Content
// interface. // interface.
func (x *CallToolResultFor[Out]) UnmarshalJSON(data []byte) error { func (x *CallToolResult) UnmarshalJSON(data []byte) error {
type res CallToolResultFor[Out] // avoid recursion type res CallToolResult // avoid recursion
var wire struct { var wire struct {
res res
Content []*wireContent `json:"content"` Content []*wireContent `json:"content"`
@ -95,13 +107,13 @@ func (x *CallToolResultFor[Out]) UnmarshalJSON(data []byte) error {
if wire.res.Content, err = contentsFromWire(wire.Content, nil); err != nil { if wire.res.Content, err = contentsFromWire(wire.Content, nil); err != nil {
return err return err
} }
*x = CallToolResultFor[Out](wire.res) *x = CallToolResult(wire.res)
return nil return nil
} }
func (x *CallToolParamsFor[Out]) isParams() {} func (x *CallToolParams) isParams() {}
func (x *CallToolParamsFor[Out]) GetProgressToken() any { return getProgressToken(x) } func (x *CallToolParams) GetProgressToken() any { return getProgressToken(x) }
func (x *CallToolParamsFor[Out]) SetProgressToken(t any) { setProgressToken(x, t) } func (x *CallToolParams) SetProgressToken(t any) { setProgressToken(x, t) }
type CancelledParams struct { type CancelledParams struct {
// This property is reserved by the protocol to allow clients and servers to // This property is reserved by the protocol to allow clients and servers to

7
mcp/protocol_test.go
View file

@ -208,6 +208,7 @@ func TestCompleteReference(t *testing.T) {
}) })
} }
} }
func TestCompleteParams(t *testing.T) { func TestCompleteParams(t *testing.T) {
// Define test cases specifically for Marshalling // Define test cases specifically for Marshalling
marshalTests := []struct { marshalTests := []struct {
@ -514,13 +515,13 @@ func TestContentUnmarshal(t *testing.T) {
var got CallToolResult var got CallToolResult
roundtrip(ctr, &got) roundtrip(ctr, &got)
ctrf := &CallToolResultFor[int]{ ctrf := &CallToolResult{
Meta: Meta{"m": true}, Meta: Meta{"m": true},
Content: content, Content: content,
IsError: true, IsError: true,
StructuredContent: 3, StructuredContent: 3.0,
} }
var gotf CallToolResultFor[int] var gotf CallToolResult
roundtrip(ctrf, &gotf) roundtrip(ctrf, &gotf)
pm := &PromptMessage{ pm := &PromptMessage{

147
mcp/server.go
View file

@ -15,10 +15,12 @@ import (
"maps" "maps"
"net/url" "net/url"
"path/filepath" "path/filepath"
"reflect"
"slices" "slices"
"sync" "sync"
"time" "time"
"github.com/google/jsonschema-go/jsonschema"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
"github.com/modelcontextprotocol/go-sdk/internal/util" "github.com/modelcontextprotocol/go-sdk/internal/util"
"github.com/modelcontextprotocol/go-sdk/jsonrpc" "github.com/modelcontextprotocol/go-sdk/jsonrpc"
@ -146,53 +148,128 @@ func (s *Server) RemovePrompts(names ...string) {
// The tool's input schema must be non-nil. For a tool that takes no input, // The tool's input schema must be non-nil. For a tool that takes no input,
// or one where any input is valid, set [Tool.InputSchema] to the empty schema, // or one where any input is valid, set [Tool.InputSchema] to the empty schema,
// &jsonschema.Schema{}. // &jsonschema.Schema{}.
//
// When the handler is invoked as part of a CallTool request, req.Params.Arguments
// will be a json.RawMessage. Unmarshaling the arguments and validating them against the
// input schema are the handler author's responsibility.
//
// Most users will prefer the top-level function [AddTool].
func (s *Server) AddTool(t *Tool, h ToolHandler) { func (s *Server) AddTool(t *Tool, h ToolHandler) {
if t.InputSchema == nil { if t.InputSchema == nil {
// This prevents the tool author from forgetting to write a schema where // This prevents the tool author from forgetting to write a schema where
// one should be provided. If we papered over this by supplying the empty // one should be provided. If we papered over this by supplying the empty
// schema, then every input would be validated and the problem wouldn't be // schema, then every input would be validated and the problem wouldn't be
// discovered until runtime, when the LLM sent bad data. // discovered until runtime, when the LLM sent bad data.
panic(fmt.Sprintf("adding tool %q: nil input schema", t.Name)) panic(fmt.Errorf("AddTool %q: missing input schema", t.Name))
}
if err := addToolErr(s, t, h); err != nil {
panic(err)
}
}
// AddTool adds a [Tool] to the server, or replaces one with the same name.
// If the tool's input schema is nil, it is set to the schema inferred from the In
// type parameter, using [jsonschema.For].
// If the tool's output schema is nil and the Out type parameter is not the empty
// interface, then the output schema is set to the schema inferred from Out.
// The Tool argument must not be modified after this call.
func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) {
if err := addToolErr(s, t, h); err != nil {
panic(err)
}
}
func addToolErr[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) (err error) {
defer util.Wrapf(&err, "adding tool %q", t.Name)
// If the exact same Tool pointer has already been registered under this name,
// avoid rebuilding schemas and re-registering. This prevents duplicate
// registration from causing errors (and unnecessary work).
s.mu.Lock()
if existing, ok := s.tools.get(t.Name); ok && existing.tool == t {
s.mu.Unlock()
return nil
}
s.mu.Unlock()
st, err := newServerTool(t, h)
if err != nil {
return err
} }
st := &serverTool{tool: t, handler: h}
// Assume there was a change, since add replaces existing tools. // Assume there was a change, since add replaces existing tools.
// (It's possible a tool was replaced with an identical one, but not worth checking.) // (It's possible a tool was replaced with an identical one, but not worth checking.)
// TODO: Batch these changes by size and time? The typescript SDK doesn't. // TODO: Batch these changes by size and time? The typescript SDK doesn't.
// TODO: Surface notify error here? best not, in case we need to batch. // TODO: Surface notify error here? best not, in case we need to batch.
s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{}, s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{},
func() bool { s.tools.add(st); return true }) func() bool { s.tools.add(st); return true })
return nil }
// toolFor returns a shallow copy of t and a [ToolHandler] that wraps h.
// If the tool's input schema is nil, it is set to the schema inferred from the In
// type parameter, using [jsonschema.For].
// If the tool's output schema is nil and the Out type parameter is not the empty
// interface, then the output schema is set to the schema inferred from Out.
//
// Most users will call [AddTool]. Use [toolFor] if you wish to wrap the ToolHandler
// before calling [Server.AddTool].
func toolFor[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHandler) {
tt, hh, err := toolForErr(t, h)
if err != nil {
panic(fmt.Sprintf("ToolFor: tool %q: %v", t.Name, err))
}
return tt, hh
}
// TODO(v0.3.0): test
func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHandler, error) {
var err error
tt := *t
tt.InputSchema = t.InputSchema
if tt.InputSchema == nil {
tt.InputSchema, err = jsonschema.For[In](nil)
if err != nil {
return nil, nil, fmt.Errorf("input schema: %w", err)
}
}
inputResolved, err := tt.InputSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
if err != nil {
return nil, nil, fmt.Errorf("resolving input schema: %w", err)
}
if tt.OutputSchema == nil && reflect.TypeFor[Out]() != reflect.TypeFor[any]() {
tt.OutputSchema, err = jsonschema.For[Out](nil)
}
if err != nil {
return nil, nil, fmt.Errorf("output schema: %w", err)
}
var outputResolved *jsonschema.Resolved
if tt.OutputSchema != nil {
outputResolved, err = tt.OutputSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
if err != nil {
return nil, nil, fmt.Errorf("resolving output schema: %w", err)
}
}
th := func(ctx context.Context, req *ServerRequest[*CallToolParams]) (*CallToolResult, error) {
// Unmarshal and validate args.
rawArgs := req.Params.Arguments.(json.RawMessage)
var in In
if rawArgs != nil {
if err := unmarshalSchema(rawArgs, inputResolved, &in); err != nil {
return nil, err
}
}
// Call typed handler.
res, out, err := h(ctx, req, in)
// Handle server errors appropriately:
// - If the handler returns a structured error (like jsonrpc2.WireError), return it directly
// - If the handler returns a regular error, wrap it in a CallToolResult with IsError=true
// - This allows tools to distinguish between protocol errors and tool execution errors
if err != nil {
// Check if this is already a structured JSON-RPC error
if wireErr, ok := err.(*jsonrpc2.WireError); ok {
return nil, wireErr
}
// For regular errors, embed them in the tool result as per MCP spec
return &CallToolResult{
Content: []Content{&TextContent{Text: err.Error()}},
IsError: true,
}, nil
}
// TODO(v0.3.0): Validate out.
_ = outputResolved
// TODO: return the serialized JSON in a TextContent block, as per spec?
// https://modelcontextprotocol.io/specification/2025-06-18/server/tools#structured-content
// But people may use res.Content for other things.
if res == nil {
res = &CallToolResult{}
}
res.StructuredContent = out
return res, nil
}
return &tt, th, nil
}
// AddTool adds a tool and handler to the server.
//
// A shallow copy of the tool is made first.
// If the tool's input schema is nil, the copy's input schema is set to the schema
// inferred from the In type parameter, using [jsonschema.For].
// If the tool's output schema is nil and the Out type parameter is not the empty
// interface, then the copy's output schema is set to the schema inferred from Out.
func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) {
s.AddTool(toolFor(t, h))
} }
// RemoveTools removes the tools with the given names. // RemoveTools removes the tools with the given names.
@ -352,7 +429,7 @@ func (s *Server) listTools(_ context.Context, req *ServerRequest[*ListToolsParam
}) })
} }
func (s *Server) callTool(ctx context.Context, req *ServerRequest[*CallToolParamsFor[json.RawMessage]]) (*CallToolResult, error) { func (s *Server) callTool(ctx context.Context, req *ServerRequest[*CallToolParams]) (*CallToolResult, error) {
s.mu.Lock() s.mu.Lock()
st, ok := s.tools.get(req.Params.Name) st, ok := s.tools.get(req.Params.Name)
s.mu.Unlock() s.mu.Unlock()

8
mcp/server_example_test.go
View file

@ -16,12 +16,12 @@ type SayHiParams struct {
Name string `json:"name"` Name string `json:"name"`
} }
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[SayHiParams]]) (*mcp.CallToolResultFor[any], error) { func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args SayHiParams) (*mcp.CallToolResult, any, error) {
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}, &mcp.TextContent{Text: "Hi " + args.Name},
}, },
}, nil }, nil, nil
} }
func ExampleServer() { func ExampleServer() {

388
mcp/shared_test.go
View file

@ -4,232 +4,222 @@
package mcp package mcp
import ( // TODO(v0.3.0): rewrite this test.
"context" // func TestToolValidate(t *testing.T) {
"encoding/json" // // Check that the tool returned from NewServerTool properly validates its input schema.
"fmt"
"strings"
"testing"
)
// TODO(jba): this shouldn't be in this file, but tool_test.go doesn't have access to unexported symbols. // type req struct {
func TestToolValidate(t *testing.T) { // I int
// Check that the tool returned from NewServerTool properly validates its input schema. // B bool
// S string `json:",omitempty"`
// P *int `json:",omitempty"`
// }
type req struct { // dummyHandler := func(context.Context, *ServerRequest[*CallToolParams], req) (*CallToolResultFor[any], error) {
I int // return nil, nil
B bool // }
S string `json:",omitempty"`
P *int `json:",omitempty"`
}
dummyHandler := func(context.Context, *ServerRequest[*CallToolParamsFor[req]]) (*CallToolResultFor[any], error) { // st, err := newServerTool(&Tool{Name: "test", Description: "test"}, dummyHandler)
return nil, nil // if err != nil {
} // t.Fatal(err)
// }
st, err := newServerTool(&Tool{Name: "test", Description: "test"}, dummyHandler) // for _, tt := range []struct {
if err != nil { // desc string
t.Fatal(err) // args map[string]any
} // want string // error should contain this string; empty for success
// }{
for _, tt := range []struct { // {
desc string // "both required",
args map[string]any // map[string]any{"I": 1, "B": true},
want string // error should contain this string; empty for success // "",
}{ // },
{ // {
"both required", // "optional",
map[string]any{"I": 1, "B": true}, // map[string]any{"I": 1, "B": true, "S": "foo"},
"", // "",
}, // },
{ // {
"optional", // "wrong type",
map[string]any{"I": 1, "B": true, "S": "foo"}, // map[string]any{"I": 1.5, "B": true},
"", // "cannot unmarshal",
}, // },
{ // {
"wrong type", // "extra property",
map[string]any{"I": 1.5, "B": true}, // map[string]any{"I": 1, "B": true, "C": 2},
"cannot unmarshal", // "unknown field",
}, // },
{ // {
"extra property", // "value for pointer",
map[string]any{"I": 1, "B": true, "C": 2}, // map[string]any{"I": 1, "B": true, "P": 3},
"unknown field", // "",
}, // },
{ // {
"value for pointer", // "null for pointer",
map[string]any{"I": 1, "B": true, "P": 3}, // map[string]any{"I": 1, "B": true, "P": nil},
"", // "",
}, // },
{ // } {
"null for pointer", // t.Run(tt.desc, func(t *testing.T) {
map[string]any{"I": 1, "B": true, "P": nil}, // raw, err := json.Marshal(tt.args)
"", // if err != nil {
}, // t.Fatal(err)
} { // }
t.Run(tt.desc, func(t *testing.T) { // _, err = st.handler(context.Background(), &ServerRequest[*CallToolParamsFor[json.RawMessage]]{
raw, err := json.Marshal(tt.args) // Params: &CallToolParamsFor[json.RawMessage]{Arguments: json.RawMessage(raw)},
if err != nil { // })
t.Fatal(err) // if err == nil && tt.want != "" {
} // t.Error("got success, wanted failure")
_, err = st.handler(context.Background(), &ServerRequest[*CallToolParamsFor[json.RawMessage]]{ // }
Params: &CallToolParamsFor[json.RawMessage]{Arguments: json.RawMessage(raw)}, // if err != nil {
}) // if tt.want == "" {
if err == nil && tt.want != "" { // t.Fatalf("failed with:\n%s\nwanted success", err)
t.Error("got success, wanted failure") // }
} // if !strings.Contains(err.Error(), tt.want) {
if err != nil { // t.Fatalf("got:\n%s\nwanted to contain %q", err, tt.want)
if tt.want == "" { // }
t.Fatalf("failed with:\n%s\nwanted success", err) // }
} // })
if !strings.Contains(err.Error(), tt.want) { // }
t.Fatalf("got:\n%s\nwanted to contain %q", err, tt.want) // }
}
}
})
}
}
// TestNilParamsHandling tests that nil parameters don't cause panic in unmarshalParams. // TestNilParamsHandling tests that nil parameters don't cause panic in unmarshalParams.
// This addresses a vulnerability where missing or null parameters could crash the server. // This addresses a vulnerability where missing or null parameters could crash the server.
func TestNilParamsHandling(t *testing.T) { // func TestNilParamsHandling(t *testing.T) {
// Define test types for clarity // // Define test types for clarity
type TestArgs struct { // type TestArgs struct {
Name string `json:"name"` // Name string `json:"name"`
Value int `json:"value"` // Value int `json:"value"`
} // }
type TestParams = *CallToolParamsFor[TestArgs]
type TestResult = *CallToolResultFor[string]
// Simple test handler // // Simple test handler
testHandler := func(ctx context.Context, req *ServerRequest[TestParams]) (TestResult, error) { // testHandler := func(ctx context.Context, req *ServerRequest[**GetPromptParams]) (*GetPromptResult, error) {
result := "processed: " + req.Params.Arguments.Name // result := "processed: " + req.Params.Arguments.Name
return &CallToolResultFor[string]{StructuredContent: result}, nil // return &CallToolResultFor[string]{StructuredContent: result}, nil
} // }
methodInfo := newServerMethodInfo(testHandler, missingParamsOK) // methodInfo := newServerMethodInfo(testHandler, missingParamsOK)
// Helper function to test that unmarshalParams doesn't panic and handles nil gracefully // // Helper function to test that unmarshalParams doesn't panic and handles nil gracefully
mustNotPanic := func(t *testing.T, rawMsg json.RawMessage, expectNil bool) Params { // mustNotPanic := func(t *testing.T, rawMsg json.RawMessage, expectNil bool) Params {
t.Helper() // t.Helper()
defer func() { // defer func() {
if r := recover(); r != nil { // if r := recover(); r != nil {
t.Fatalf("unmarshalParams panicked: %v", r) // t.Fatalf("unmarshalParams panicked: %v", r)
} // }
}() // }()
params, err := methodInfo.unmarshalParams(rawMsg) // params, err := methodInfo.unmarshalParams(rawMsg)
if err != nil { // if err != nil {
t.Fatalf("unmarshalParams failed: %v", err) // t.Fatalf("unmarshalParams failed: %v", err)
} // }
if expectNil { // if expectNil {
if params != nil { // if params != nil {
t.Fatalf("Expected nil params, got %v", params) // t.Fatalf("Expected nil params, got %v", params)
} // }
return params // return params
} // }
if params == nil { // if params == nil {
t.Fatal("unmarshalParams returned unexpected nil") // t.Fatal("unmarshalParams returned unexpected nil")
} // }
// Verify the result can be used safely // // Verify the result can be used safely
typedParams := params.(TestParams) // typedParams := params.(TestParams)
_ = typedParams.Name // _ = typedParams.Name
_ = typedParams.Arguments.Name // _ = typedParams.Arguments.Name
_ = typedParams.Arguments.Value // _ = typedParams.Arguments.Value
return params // return params
} // }
// Test different nil parameter scenarios - with missingParamsOK flag, nil/null should return nil // // Test different nil parameter scenarios - with missingParamsOK flag, nil/null should return nil
t.Run("missing_params", func(t *testing.T) { // t.Run("missing_params", func(t *testing.T) {
mustNotPanic(t, nil, true) // Expect nil with missingParamsOK flag // mustNotPanic(t, nil, true) // Expect nil with missingParamsOK flag
}) // })
t.Run("explicit_null", func(t *testing.T) { // t.Run("explicit_null", func(t *testing.T) {
mustNotPanic(t, json.RawMessage(`null`), true) // Expect nil with missingParamsOK flag // mustNotPanic(t, json.RawMessage(`null`), true) // Expect nil with missingParamsOK flag
}) // })
t.Run("empty_object", func(t *testing.T) { // t.Run("empty_object", func(t *testing.T) {
mustNotPanic(t, json.RawMessage(`{}`), false) // Empty object should create valid params // mustNotPanic(t, json.RawMessage(`{}`), false) // Empty object should create valid params
}) // })
t.Run("valid_params", func(t *testing.T) { // t.Run("valid_params", func(t *testing.T) {
rawMsg := json.RawMessage(`{"name":"test","arguments":{"name":"hello","value":42}}`) // rawMsg := json.RawMessage(`{"name":"test","arguments":{"name":"hello","value":42}}`)
params := mustNotPanic(t, rawMsg, false) // params := mustNotPanic(t, rawMsg, false)
// For valid params, also verify the values are parsed correctly // // For valid params, also verify the values are parsed correctly
typedParams := params.(TestParams) // typedParams := params.(TestParams)
if typedParams.Name != "test" { // if typedParams.Name != "test" {
t.Errorf("Expected name 'test', got %q", typedParams.Name) // t.Errorf("Expected name 'test', got %q", typedParams.Name)
} // }
if typedParams.Arguments.Name != "hello" { // if typedParams.Arguments.Name != "hello" {
t.Errorf("Expected argument name 'hello', got %q", typedParams.Arguments.Name) // t.Errorf("Expected argument name 'hello', got %q", typedParams.Arguments.Name)
} // }
if typedParams.Arguments.Value != 42 { // if typedParams.Arguments.Value != 42 {
t.Errorf("Expected argument value 42, got %d", typedParams.Arguments.Value) // t.Errorf("Expected argument value 42, got %d", typedParams.Arguments.Value)
} // }
}) // })
} // }
// TestNilParamsEdgeCases tests edge cases to ensure we don't over-fix // TestNilParamsEdgeCases tests edge cases to ensure we don't over-fix
func TestNilParamsEdgeCases(t *testing.T) { // func TestNilParamsEdgeCases(t *testing.T) {
type TestArgs struct { // type TestArgs struct {
Name string `json:"name"` // Name string `json:"name"`
Value int `json:"value"` // Value int `json:"value"`
} // }
type TestParams = *CallToolParamsFor[TestArgs] // type TestParams = *CallToolParamsFor[TestArgs]
testHandler := func(context.Context, *ServerRequest[TestParams]) (*CallToolResultFor[string], error) { // testHandler := func(context.Context, *ServerRequest[TestParams]) (*CallToolResultFor[string], error) {
return &CallToolResultFor[string]{StructuredContent: "test"}, nil // return &CallToolResultFor[string]{StructuredContent: "test"}, nil
} // }
methodInfo := newServerMethodInfo(testHandler, missingParamsOK) // methodInfo := newServerMethodInfo(testHandler, missingParamsOK)
// These should fail normally, not be treated as nil params // // These should fail normally, not be treated as nil params
invalidCases := []json.RawMessage{ // invalidCases := []json.RawMessage{
json.RawMessage(""), // empty string - should error // json.RawMessage(""), // empty string - should error
json.RawMessage("[]"), // array - should error // json.RawMessage("[]"), // array - should error
json.RawMessage(`"null"`), // string "null" - should error // json.RawMessage(`"null"`), // string "null" - should error
json.RawMessage("0"), // number - should error // json.RawMessage("0"), // number - should error
json.RawMessage("false"), // boolean - should error // json.RawMessage("false"), // boolean - should error
} // }
for i, rawMsg := range invalidCases { // for i, rawMsg := range invalidCases {
t.Run(fmt.Sprintf("invalid_case_%d", i), func(t *testing.T) { // t.Run(fmt.Sprintf("invalid_case_%d", i), func(t *testing.T) {
params, err := methodInfo.unmarshalParams(rawMsg) // params, err := methodInfo.unmarshalParams(rawMsg)
if err == nil && params == nil { // if err == nil && params == nil {
t.Error("Should not return nil params without error") // t.Error("Should not return nil params without error")
} // }
}) // })
} // }
// Test that methods without missingParamsOK flag properly reject nil params // // Test that methods without missingParamsOK flag properly reject nil params
t.Run("reject_when_params_required", func(t *testing.T) { // t.Run("reject_when_params_required", func(t *testing.T) {
methodInfoStrict := newServerMethodInfo(testHandler, 0) // No missingParamsOK flag // methodInfoStrict := newServerMethodInfo(testHandler, 0) // No missingParamsOK flag
testCases := []struct { // testCases := []struct {
name string // name string
params json.RawMessage // params json.RawMessage
}{ // }{
{"nil_params", nil}, // {"nil_params", nil},
{"null_params", json.RawMessage(`null`)}, // {"null_params", json.RawMessage(`null`)},
} // }
for _, tc := range testCases { // for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) { // t.Run(tc.name, func(t *testing.T) {
_, err := methodInfoStrict.unmarshalParams(tc.params) // _, err := methodInfoStrict.unmarshalParams(tc.params)
if err == nil { // if err == nil {
t.Error("Expected error for required params, got nil") // t.Error("Expected error for required params, got nil")
} // }
if !strings.Contains(err.Error(), "missing required \"params\"") { // if !strings.Contains(err.Error(), "missing required \"params\"") {
t.Errorf("Expected 'missing required params' error, got: %v", err) // t.Errorf("Expected 'missing required params' error, got: %v", err)
} // }
}) // })
} // }
}) // })
} // }

8
mcp/sse_example_test.go
View file

@ -18,12 +18,12 @@ type AddParams struct {
X, Y int X, Y int
} }
func Add(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[AddParams]]) (*mcp.CallToolResultFor[any], error) { func Add(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args AddParams) (*mcp.CallToolResult, any, error) {
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{Text: fmt.Sprintf("%d", req.Params.Arguments.X+req.Params.Arguments.Y)}, &mcp.TextContent{Text: fmt.Sprintf("%d", args.X+args.Y)},
}, },
}, nil }, nil, nil
} }
func ExampleSSEHandler() { func ExampleSSEHandler() {

Some files were not shown because too many files have changed in this diff Show more