Skip to content

Commit 8f8587f

Browse files
jbajba
committed
fixed more tests
1 parent bdf0233 commit 8f8587f

File tree

9 files changed

+74
-71
lines changed

9 files changed

+74
-71
lines changed

‎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: 22 additions & 23 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

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

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

450-
return &res, nil
450+
return &res, 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

456456
relations, err := k.createRelations(req.Params.Arguments.Relations)
457457
if err != nil {
@@ -469,8 +469,8 @@ func (k knowledgeBase) CreateRelations(ctx context.Context, req *mcp.ServerReque
469469
return &res, nil
470470
}
471471

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

475475
observations, err := k.addObservations(req.Params.Arguments.Observations)
476476
if err != nil {
@@ -488,8 +488,8 @@ func (k knowledgeBase) AddObservations(ctx context.Context, req *mcp.ServerReque
488488
return &res, nil
489489
}
490490

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

494494
err := k.deleteEntities(req.Params.Arguments.EntityNames)
495495
if err != nil {
@@ -503,8 +503,8 @@ func (k knowledgeBase) DeleteEntities(ctx context.Context, req *mcp.ServerReques
503503
return &res, nil
504504
}
505505

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

509509
err := k.deleteObservations(req.Params.Arguments.Deletions)
510510
if err != nil {
@@ -518,8 +518,8 @@ func (k knowledgeBase) DeleteObservations(ctx context.Context, req *mcp.ServerRe
518518
return &res, nil
519519
}
520520

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

524524
err := k.deleteRelations(req.Params.Arguments.Relations)
525525
if err != nil {
@@ -533,24 +533,23 @@ func (k knowledgeBase) DeleteRelations(ctx context.Context, req *mcp.ServerReque
533533
return &res, nil
534534
}
535535

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

539539
graph, err := k.loadGraph()
540540
if err != nil {
541-
return nil, err
541+
return nil, KnowledgeGraph{}, err
542542
}
543543

544544
res.Content = []mcp.Content{
545545
&mcp.TextContent{Text: "Graph read successfully"},
546546
}
547547

548-
res.StructuredContent = graph
549-
return &res, nil
548+
return &res, graph, nil
550549
}
551550

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

555554
graph, err := k.searchNodes(req.Params.Arguments.Query)
556555
if err != nil {
@@ -565,8 +564,8 @@ func (k knowledgeBase) SearchNodes(ctx context.Context, req *mcp.ServerRequest[*
565564
return &res, nil
566565
}
567566

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

571570
graph, err := k.openNodes(req.Params.Arguments.Names)
572571
if err != nil {

‎examples/server/sequentialthinking/main.go‎

Lines changed: 17 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -231,9 +231,7 @@ func deepCopyThoughts(thoughts []*Thought) []*Thought {
231231
}
232232

233233
// StartThinking begins a new sequential thinking session for a complex problem.
234-
func StartThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[StartThinkingArgs]]) (*mcp.CallToolResultFor[any], error) {
235-
args := req.Params.Arguments
236-
234+
func StartThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args StartThinkingArgs) (*mcp.CallToolResult, any, error) {
237235
sessionID := args.SessionID
238236
if sessionID == "" {
239237
sessionID = randText()
@@ -255,20 +253,18 @@ func StartThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolPara
255253

256254
store.SetSession(session)
257255

258-
return &mcp.CallToolResultFor[any]{
256+
return &mcp.CallToolResult{
259257
Content: []mcp.Content{
260258
&mcp.TextContent{
261259
Text: fmt.Sprintf("Started thinking session '%s' for problem: %s\nEstimated steps: %d\nReady for your first thought.",
262260
sessionID, args.Problem, estimatedSteps),
263261
},
264262
},
265-
}, nil
263+
}, nil, nil
266264
}
267265

268266
// ContinueThinking adds the next thought step, revises a previous step, or creates a branch in the thinking process.
269-
func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[ContinueThinkingArgs]]) (*mcp.CallToolResultFor[any], error) {
270-
args := req.Params.Arguments
271-
267+
func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args ContinueThinkingArgs) (*mcp.CallToolResult, any, error) {
272268
// Handle revision of existing thought
273269
if args.ReviseStep != nil {
274270
err := store.CompareAndSwap(args.SessionID, func(session *ThinkingSession) (*ThinkingSession, error) {
@@ -283,17 +279,17 @@ func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolP
283279
return session, nil
284280
})
285281
if err != nil {
286-
return nil, err
282+
return nil, nil, err
287283
}
288284

289-
return &mcp.CallToolResultFor[any]{
285+
return &mcp.CallToolResult{
290286
Content: []mcp.Content{
291287
&mcp.TextContent{
292288
Text: fmt.Sprintf("Revised step %d in session '%s':\n%s",
293289
*args.ReviseStep, args.SessionID, args.Thought),
294290
},
295291
},
296-
}, nil
292+
}, nil, nil
297293
}
298294

299295
// Handle branching
@@ -322,20 +318,20 @@ func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolP
322318
return session, nil
323319
})
324320
if err != nil {
325-
return nil, err
321+
return nil, nil, err
326322
}
327323

328324
// Save the branch session
329325
store.SetSession(branchSession)
330326

331-
return &mcp.CallToolResultFor[any]{
327+
return &mcp.CallToolResult{
332328
Content: []mcp.Content{
333329
&mcp.TextContent{
334330
Text: fmt.Sprintf("Created branch '%s' from session '%s'. You can now continue thinking in either session.",
335331
branchID, args.SessionID),
336332
},
337333
},
338-
}, nil
334+
}, nil, nil
339335
}
340336

341337
// Add new thought
@@ -381,27 +377,25 @@ func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolP
381377
return session, nil
382378
})
383379
if err != nil {
384-
return nil, err
380+
return nil, nil, err
385381
}
386382

387-
return &mcp.CallToolResultFor[any]{
383+
return &mcp.CallToolResult{
388384
Content: []mcp.Content{
389385
&mcp.TextContent{
390386
Text: fmt.Sprintf("Session '%s' - %s:\n%s%s",
391387
args.SessionID, progress, args.Thought, statusMsg),
392388
},
393389
},
394-
}, nil
390+
}, nil, nil
395391
}
396392

397393
// ReviewThinking provides a complete review of the thinking process for a session.
398-
func ReviewThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[ReviewThinkingArgs]]) (*mcp.CallToolResultFor[any], error) {
399-
args := req.Params.Arguments
400-
394+
func ReviewThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args ReviewThinkingArgs) (*mcp.CallToolResult, any, error) {
401395
// Get a snapshot of the session to avoid race conditions
402396
sessionSnapshot, exists := store.SessionSnapshot(args.SessionID)
403397
if !exists {
404-
return nil, fmt.Errorf("session %s not found", args.SessionID)
398+
return nil, nil, fmt.Errorf("session %s not found", args.SessionID)
405399
}
406400

407401
var review strings.Builder
@@ -424,13 +418,13 @@ func ReviewThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolPar
424418
fmt.Fprintf(&review, "%d. %s%s\n", i+1, thought.Content, status)
425419
}
426420

427-
return &mcp.CallToolResultFor[any]{
421+
return &mcp.CallToolResult{
428422
Content: []mcp.Content{
429423
&mcp.TextContent{
430424
Text: review.String(),
431425
},
432426
},
433-
}, nil
427+
}, nil, nil
434428
}
435429

436430
// ThinkingHistory handles resource requests for thinking session data and history.

‎examples/server/sse/main.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,12 @@ type SayHiParams struct {
2424
Name string `json:"name"`
2525
}
2626

27-
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[SayHiParams]]) (*mcp.CallToolResultFor[any], error) {
28-
return &mcp.CallToolResultFor[any]{
27+
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args SayHiParams) (*mcp.CallToolResult, any, error) {
28+
return &mcp.CallToolResult{
2929
Content: []mcp.Content{
30-
&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name},
30+
&mcp.TextContent{Text: "Hi " + args.Name},
3131
},
32-
}, nil
32+
}, nil, nil
3333
}
3434

3535
func main() {

‎internal/readme/server/server.go‎

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

19-
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiParams]]) (*mcp.CallToolResultFor[any], error) {
20-
return &mcp.CallToolResultFor[any]{
21-
Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}},
22-
}, nil
19+
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args HiParams) (*mcp.CallToolResult, any, error) {
20+
return &mcp.CallToolResult{
21+
Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + args.Name}},
22+
}, nil, nil
2323
}
2424

2525
func main() {

‎mcp/client_list_test.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,11 @@ func TestList(t *testing.T) {
2525
var wantTools []*mcp.Tool
2626
for _, name := range []string{"apple", "banana", "cherry"} {
2727
t := &mcp.Tool{Name: name, Description: name + " tool"}
28-
wantTools = append(wantTools, t)
29-
mcp.AddTool(server, t, SayHi)
28+
// Previously, we could add t to wantTools, because AddTool modified
29+
// it. Now, AddTool gets a copy.
30+
tt, h := mcp.ToolFor(t, SayHi)
31+
wantTools = append(wantTools, tt)
32+
server.AddTool(tt, h)
3033
}
3134
t.Run("list", func(t *testing.T) {
3235
res, err := clientSession.ListTools(ctx, nil)

‎mcp/server.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,8 +225,13 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan
225225

226226
// Call typed handler.
227227
res, out, err := h(ctx, req, in)
228+
// TODO(rfindley): investigate why server errors are embedded in this strange way,
229+
// rather than returned as jsonrpc2 server errors.
228230
if err != nil {
229-
return nil, err
231+
return &CallToolResult{
232+
Content: []Content{&TextContent{Text: err.Error()}},
233+
IsError: true,
234+
}, nil
230235
}
231236

232237
// TODO(v0.3.0): Validate out.

‎mcp/streamable_test.go‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -625,12 +625,14 @@ func TestStreamableServerTransport(t *testing.T) {
625625
// Create a server containing a single tool, which runs the test tool
626626
// behavior, if any.
627627
server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
628-
server.AddTool(&Tool{Name: "tool"}, func(ctx context.Context, req *ServerRequest[*CallToolParams]) (*CallToolResult, error) {
629-
if test.tool != nil {
630-
test.tool(t, ctx, req.Session)
631-
}
632-
return &CallToolResult{}, nil
633-
})
628+
server.AddTool(
629+
&Tool{Name: "tool", InputSchema: &jsonschema.Schema{}},
630+
func(ctx context.Context, req *ServerRequest[*CallToolParams]) (*CallToolResult, error) {
631+
if test.tool != nil {
632+
test.tool(t, ctx, req.Session)
633+
}
634+
return &CallToolResult{}, nil
635+
})
634636

635637
// Start the streamable handler.
636638
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)

0 commit comments

Comments
 (0)