Skip to content

Commit 3a777e4

Browse files
ezynda3ezynda3
committed
Update tool and prompt handlers
1 parent 459e921 commit 3a777e4

File tree

4 files changed

+24
-20
lines changed

4 files changed

+24
-20
lines changed

‎client/sse_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,11 @@ func TestSSEMCPClient(t *testing.T) {
2626
Type: "object",
2727
Properties: map[string]interface{}{},
2828
},
29-
}, func(ctx context.Context, arguments map[string]interface{}) (*mcp.CallToolResult, error) {
29+
}, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
3030
return &mcp.CallToolResult{}, nil
3131
})
3232

33-
// Create test server
33+
// Initialize
3434
testServer := server.NewTestServer(mcpServer)
3535
defer testServer.Close()
3636

‎examples/everything/main.go‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -216,7 +216,7 @@ func (s *MCPServer) handleResourceTemplate(
216216

217217
func (s *MCPServer) handleSimplePrompt(
218218
ctx context.Context,
219-
arguments map[string]string,
219+
request mcp.GetPromptRequest,
220220
) (*mcp.GetPromptResult, error) {
221221
return &mcp.GetPromptResult{
222222
Description: "A simple prompt without arguments",
@@ -234,8 +234,9 @@ func (s *MCPServer) handleSimplePrompt(
234234

235235
func (s *MCPServer) handleComplexPrompt(
236236
ctx context.Context,
237-
arguments map[string]string,
237+
request mcp.GetPromptRequest,
238238
) (*mcp.GetPromptResult, error) {
239+
arguments := request.Params.Arguments
239240
return &mcp.GetPromptResult{
240241
Description: "A complex prompt with arguments",
241242
Messages: []mcp.PromptMessage{
@@ -271,8 +272,9 @@ func (s *MCPServer) handleComplexPrompt(
271272

272273
func (s *MCPServer) handleEchoTool(
273274
ctx context.Context,
274-
arguments map[string]interface{},
275+
request mcp.CallToolRequest,
275276
) (*mcp.CallToolResult, error) {
277+
arguments := request.Params.Arguments
276278
message, ok := arguments["message"].(string)
277279
if !ok {
278280
return nil, fmt.Errorf("invalid message argument")
@@ -289,8 +291,9 @@ func (s *MCPServer) handleEchoTool(
289291

290292
func (s *MCPServer) handleAddTool(
291293
ctx context.Context,
292-
arguments map[string]interface{},
294+
request mcp.CallToolRequest,
293295
) (*mcp.CallToolResult, error) {
296+
arguments := request.Params.Arguments
294297
a, ok1 := arguments["a"].(float64)
295298
b, ok2 := arguments["b"].(float64)
296299
if !ok1 || !ok2 {
@@ -309,7 +312,7 @@ func (s *MCPServer) handleAddTool(
309312

310313
func (s *MCPServer) handleSendNotification(
311314
ctx context.Context,
312-
arguments map[string]interface{},
315+
request mcp.CallToolRequest,
313316
) (*mcp.CallToolResult, error) {
314317

315318
server := server.ServerFromContext(ctx)
@@ -346,8 +349,9 @@ func (s *MCPServer) ServeStdio() error {
346349

347350
func (s *MCPServer) handleLongRunningOperationTool(
348351
ctx context.Context,
349-
arguments map[string]interface{},
352+
request mcp.CallToolRequest,
350353
) (*mcp.CallToolResult, error) {
354+
arguments := request.Params.Arguments
351355
duration, _ := arguments["duration"].(float64)
352356
steps, _ := arguments["steps"].(float64)
353357
stepDuration := duration / steps
@@ -414,7 +418,7 @@ func (s *MCPServer) handleLongRunningOperationTool(
414418

415419
func (s *MCPServer) handleGetTinyImageTool(
416420
ctx context.Context,
417-
arguments map[string]interface{},
421+
request mcp.CallToolRequest,
418422
) (*mcp.CallToolResult, error) {
419423
return &mcp.CallToolResult{
420424
Content: []interface{}{

‎server/server.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,10 +32,10 @@ type ResourceHandlerFunc func(ctx context.Context, request mcp.ReadResourceReque
3232
type ResourceTemplateHandlerFunc func(ctx context.Context, request mcp.ReadResourceRequest) ([]interface{}, error)
3333

3434
// PromptHandlerFunc handles prompt requests with given arguments.
35-
type PromptHandlerFunc func(ctx context.Context, arguments map[string]string) (*mcp.GetPromptResult, error)
35+
type PromptHandlerFunc func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error)
3636

3737
// ToolHandlerFunc handles tool calls with given arguments.
38-
type ToolHandlerFunc func(ctx context.Context, arguments map[string]interface{}) (*mcp.CallToolResult, error)
38+
type ToolHandlerFunc func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error)
3939

4040
// NotificationContext provides client identification for notifications
4141
type NotificationContext struct {
@@ -613,7 +613,7 @@ func (s *MCPServer) handleGetPrompt(
613613
)
614614
}
615615

616-
result, err := handler(ctx, request.Params.Arguments)
616+
result, err := handler(ctx, request)
617617
if err != nil {
618618
return createErrorResponse(id, mcp.INTERNAL_ERROR, err.Error())
619619
}
@@ -654,7 +654,7 @@ func (s *MCPServer) handleToolCall(
654654
)
655655
}
656656

657-
result, err := handler(ctx, request.Params.Arguments)
657+
result, err := handler(ctx, request)
658658
if err != nil {
659659
return createErrorResponse(id, mcp.INTERNAL_ERROR, err.Error())
660660
}

‎server/server_test.go‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -67,11 +67,11 @@ func TestMCPServer_Capabilities(t *testing.T) {
6767
Type: "object",
6868
Properties: map[string]interface{}{},
6969
},
70-
}, func(ctx context.Context, arguments map[string]interface{}) (*mcp.CallToolResult, error) {
70+
}, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
7171
return &mcp.CallToolResult{}, nil
7272
})
7373

74-
// Create a new server with all capabilities and a tool
74+
// Create and handle the initialize message
7575
server = NewMCPServer(
7676
"test-server",
7777
"1.0.0",
@@ -88,7 +88,7 @@ func TestMCPServer_Capabilities(t *testing.T) {
8888
Type: "object",
8989
Properties: map[string]interface{}{},
9090
},
91-
}, func(ctx context.Context, arguments map[string]interface{}) (*mcp.CallToolResult, error) {
91+
}, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
9292
return &mcp.CallToolResult{}, nil
9393
})
9494

@@ -327,14 +327,14 @@ func TestMCPServer_PromptHandling(t *testing.T) {
327327

328328
server.AddPrompt(
329329
testPrompt,
330-
func(ctx context.Context, arguments map[string]string) (*mcp.GetPromptResult, error) {
330+
func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
331331
return &mcp.GetPromptResult{
332332
Messages: []mcp.PromptMessage{
333333
{
334334
Role: mcp.RoleAssistant,
335335
Content: mcp.TextContent{
336336
Type: "text",
337-
Text: "Test prompt with arg1: " + arguments["arg1"],
337+
Text: "Test prompt with arg1: " + request.Params.Arguments["arg1"],
338338
},
339339
},
340340
},
@@ -489,7 +489,7 @@ func TestMCPServer_HandleUndefinedHandlers(t *testing.T) {
489489
Type: "object",
490490
Properties: map[string]interface{}{},
491491
},
492-
}, func(ctx context.Context, arguments map[string]interface{}) (*mcp.CallToolResult, error) {
492+
}, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
493493
return &mcp.CallToolResult{}, nil
494494
})
495495

@@ -646,7 +646,7 @@ func createTestServer() *MCPServer {
646646
Name: "test-tool",
647647
Description: "Test tool",
648648
},
649-
func(ctx context.Context, arguments map[string]interface{}) (*mcp.CallToolResult, error) {
649+
func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
650650
return &mcp.CallToolResult{
651651
Content: []interface{}{
652652
mcp.TextContent{

0 commit comments

Comments
 (0)