Skip to content

Commit 32cf71d

Browse files
jbajba
authored
mcp: change tool handler design (#325)
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
1 parent 79f063b commit 32cf71d

23 files changed

+789
-966
lines changed

‎README.md

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -115,10 +115,10 @@ type HiParams struct {
115115
Name string `json:"name" jsonschema:"the name of the person to greet"`
116116
}
117117

118-
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiParams]]) (*mcp.CallToolResultFor[any], error) {
119-
return &mcp.CallToolResultFor[any]{
120-
Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}},
121-
}, nil
118+
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args HiParams) (*mcp.CallToolResult, any, error) {
119+
return &mcp.CallToolResult{
120+
Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + args.Name}},
121+
}, nil, nil
122122
}
123123

124124
func main() {

‎examples/server/custom-transport/main.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -85,12 +85,12 @@ type HiArgs struct {
8585
}
8686

8787
// SayHi is a tool handler that responds with a greeting.
88-
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
89-
return &mcp.CallToolResultFor[struct{}]{
88+
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args HiArgs) (*mcp.CallToolResult, struct{}, error) {
89+
return &mcp.CallToolResult{
9090
Content: []mcp.Content{
91-
&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name},
91+
&mcp.TextContent{Text: "Hi " + args.Name},
9292
},
93-
}, nil
93+
}, struct{}{}, nil
9494
}
9595

9696
func main() {

‎examples/server/hello/main.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,12 @@ type HiArgs struct {
2222
Name string `json:"name" jsonschema:"the name to say hi to"`
2323
}
2424

25-
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
26-
return &mcp.CallToolResultFor[struct{}]{
25+
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args HiArgs) (*mcp.CallToolResult, any, error) {
26+
return &mcp.CallToolResult{
2727
Content: []mcp.Content{
28-
&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name},
28+
&mcp.TextContent{Text: "Hi " + args.Name},
2929
},
30-
}, nil
30+
}, nil, nil
3131
}
3232

3333
func PromptHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {

‎examples/server/memory/kb.go

Lines changed: 45 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -431,12 +431,12 @@ func (k knowledgeBase) openNodes(names []string) (KnowledgeGraph, error) {
431431
}, nil
432432
}
433433

434-
func (k knowledgeBase) CreateEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[CreateEntitiesArgs]]) (*mcp.CallToolResultFor[CreateEntitiesResult], error) {
435-
var res mcp.CallToolResultFor[CreateEntitiesResult]
434+
func (k knowledgeBase) CreateEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args CreateEntitiesArgs) (*mcp.CallToolResult, CreateEntitiesResult, error) {
435+
var res mcp.CallToolResult
436436

437-
entities, err := k.createEntities(req.Params.Arguments.Entities)
437+
entities, err := k.createEntities(args.Entities)
438438
if err != nil {
439-
return nil, err
439+
return nil, CreateEntitiesResult{}, err
440440
}
441441

442442
res.Content = []mcp.Content{
@@ -447,136 +447,127 @@ func (k knowledgeBase) CreateEntities(ctx context.Context, req *mcp.ServerReques
447447
Entities: entities,
448448
}
449449

450-
return &res, nil
450+
return &res, CreateEntitiesResult{Entities: entities}, nil
451451
}
452452

453-
func (k knowledgeBase) CreateRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[CreateRelationsArgs]]) (*mcp.CallToolResultFor[CreateRelationsResult], error) {
454-
var res mcp.CallToolResultFor[CreateRelationsResult]
453+
func (k knowledgeBase) CreateRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args CreateRelationsArgs) (*mcp.CallToolResult, CreateRelationsResult, error) {
454+
var res mcp.CallToolResult
455455

456-
relations, err := k.createRelations(req.Params.Arguments.Relations)
456+
relations, err := k.createRelations(args.Relations)
457457
if err != nil {
458-
return nil, err
458+
return nil, CreateRelationsResult{}, err
459459
}
460460

461461
res.Content = []mcp.Content{
462462
&mcp.TextContent{Text: "Relations created successfully"},
463463
}
464464

465-
res.StructuredContent = CreateRelationsResult{
466-
Relations: relations,
467-
}
468-
469-
return &res, nil
465+
return &res, CreateRelationsResult{Relations: relations}, nil
470466
}
471467

472-
func (k knowledgeBase) AddObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[AddObservationsArgs]]) (*mcp.CallToolResultFor[AddObservationsResult], error) {
473-
var res mcp.CallToolResultFor[AddObservationsResult]
468+
func (k knowledgeBase) AddObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args AddObservationsArgs) (*mcp.CallToolResult, AddObservationsResult, error) {
469+
var res mcp.CallToolResult
474470

475-
observations, err := k.addObservations(req.Params.Arguments.Observations)
471+
observations, err := k.addObservations(args.Observations)
476472
if err != nil {
477-
return nil, err
473+
return nil, AddObservationsResult{}, err
478474
}
479475

480476
res.Content = []mcp.Content{
481477
&mcp.TextContent{Text: "Observations added successfully"},
482478
}
483479

484-
res.StructuredContent = AddObservationsResult{
480+
return &res, AddObservationsResult{
485481
Observations: observations,
486-
}
487-
488-
return &res, nil
482+
}, nil
489483
}
490484

491-
func (k knowledgeBase) DeleteEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteEntitiesArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
492-
var res mcp.CallToolResultFor[struct{}]
485+
func (k knowledgeBase) DeleteEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteEntitiesArgs) (*mcp.CallToolResult, struct{}, error) {
486+
var res mcp.CallToolResult
493487

494-
err := k.deleteEntities(req.Params.Arguments.EntityNames)
488+
err := k.deleteEntities(args.EntityNames)
495489
if err != nil {
496-
return nil, err
490+
return nil, struct{}{}, err
497491
}
498492

499493
res.Content = []mcp.Content{
500494
&mcp.TextContent{Text: "Entities deleted successfully"},
501495
}
502496

503-
return &res, nil
497+
return &res, struct{}{}, nil
504498
}
505499

506-
func (k knowledgeBase) DeleteObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteObservationsArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
507-
var res mcp.CallToolResultFor[struct{}]
500+
func (k knowledgeBase) DeleteObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteObservationsArgs) (*mcp.CallToolResult, struct{}, error) {
501+
var res mcp.CallToolResult
508502

509-
err := k.deleteObservations(req.Params.Arguments.Deletions)
503+
err := k.deleteObservations(args.Deletions)
510504
if err != nil {
511-
return nil, err
505+
return nil, struct{}{}, err
512506
}
513507

514508
res.Content = []mcp.Content{
515509
&mcp.TextContent{Text: "Observations deleted successfully"},
516510
}
517511

518-
return &res, nil
512+
return &res, struct{}{}, nil
519513
}
520514

521-
func (k knowledgeBase) DeleteRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteRelationsArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
522-
var res mcp.CallToolResultFor[struct{}]
515+
func (k knowledgeBase) DeleteRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteRelationsArgs) (*mcp.CallToolResult, struct{}, error) {
516+
var res mcp.CallToolResult
523517

524-
err := k.deleteRelations(req.Params.Arguments.Relations)
518+
err := k.deleteRelations(args.Relations)
525519
if err != nil {
526-
return nil, err
520+
return nil, struct{}{}, err
527521
}
528522

529523
res.Content = []mcp.Content{
530524
&mcp.TextContent{Text: "Relations deleted successfully"},
531525
}
532526

533-
return &res, nil
527+
return &res, struct{}{}, nil
534528
}
535529

536-
func (k knowledgeBase) ReadGraph(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[struct{}]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
537-
var res mcp.CallToolResultFor[KnowledgeGraph]
530+
func (k knowledgeBase) ReadGraph(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args any) (*mcp.CallToolResult, KnowledgeGraph, error) {
531+
var res mcp.CallToolResult
538532

539533
graph, err := k.loadGraph()
540534
if err != nil {
541-
return nil, err
535+
return nil, KnowledgeGraph{}, err
542536
}
543537

544538
res.Content = []mcp.Content{
545539
&mcp.TextContent{Text: "Graph read successfully"},
546540
}
547541

548-
res.StructuredContent = graph
549-
return &res, nil
542+
return &res, graph, nil
550543
}
551544

552-
func (k knowledgeBase) SearchNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[SearchNodesArgs]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
553-
var res mcp.CallToolResultFor[KnowledgeGraph]
545+
func (k knowledgeBase) SearchNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args SearchNodesArgs) (*mcp.CallToolResult, KnowledgeGraph, error) {
546+
var res mcp.CallToolResult
554547

555-
graph, err := k.searchNodes(req.Params.Arguments.Query)
548+
graph, err := k.searchNodes(args.Query)
556549
if err != nil {
557-
return nil, err
550+
return nil, KnowledgeGraph{}, err
558551
}
559552

560553
res.Content = []mcp.Content{
561554
&mcp.TextContent{Text: "Nodes searched successfully"},
562555
}
563556

564557
res.StructuredContent = graph
565-
return &res, nil
558+
return &res, KnowledgeGraph{}, nil
566559
}
567560

568-
func (k knowledgeBase) OpenNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[OpenNodesArgs]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
569-
var res mcp.CallToolResultFor[KnowledgeGraph]
561+
func (k knowledgeBase) OpenNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args OpenNodesArgs) (*mcp.CallToolResult, KnowledgeGraph, error) {
562+
var res mcp.CallToolResult
570563

571-
graph, err := k.openNodes(req.Params.Arguments.Names)
564+
graph, err := k.openNodes(args.Names)
572565
if err != nil {
573-
return nil, err
566+
return nil, KnowledgeGraph{}, err
574567
}
575568

576569
res.Content = []mcp.Content{
577570
&mcp.TextContent{Text: "Nodes opened successfully"},
578571
}
579-
580-
res.StructuredContent = graph
581-
return &res, nil
572+
return &res, graph, nil
582573
}

0 commit comments

Comments
 (0)