From d36a6f51e23c1929e3dc98784894592bd3b6763c Mon Sep 17 00:00:00 2001 From: Vasily Tsybenko Date: Wed, 28 May 2025 12:42:00 +0300 Subject: [PATCH] feat(mcptest): extend test server with prompt and resource support Add support for prompts and resources to the mcptest server, including: - New methods to add prompts and resources to test servers - Batch methods for adding multiple prompts/resources at once - Test cases demonstrating prompt and resource functionality - Helper methods for resource capabilities registration --- mcptest/mcptest.go | 35 ++++++++++- mcptest/mcptest_test.go | 112 ++++++++++++++++++++++++++++++++++- server/server.go | 128 ++++++++++++++++++++++------------------ 3 files changed, 214 insertions(+), 61 deletions(-) diff --git a/mcptest/mcptest.go b/mcptest/mcptest.go index 11dcad98..19f492cd 100644 --- a/mcptest/mcptest.go +++ b/mcptest/mcptest.go @@ -18,8 +18,11 @@ import ( // Server encapsulates an MCP server and manages resources like pipes and context. type Server struct { - name string - tools []server.ServerTool + name string + + tools []server.ServerTool + prompts []server.ServerPrompt + resources []server.ServerResource ctx context.Context cancel func() @@ -83,6 +86,32 @@ func (s *Server) AddTool(tool mcp.Tool, handler server.ToolHandlerFunc) { }) } +// AddPrompt adds a prompt to an unstarted server. +func (s *Server) AddPrompt(prompt mcp.Prompt, handler server.PromptHandlerFunc) { + s.prompts = append(s.prompts, server.ServerPrompt{ + Prompt: prompt, + Handler: handler, + }) +} + +// AddPrompts adds multiple prompts to an unstarted server. +func (s *Server) AddPrompts(prompts ...server.ServerPrompt) { + s.prompts = append(s.prompts, prompts...) +} + +// AddResource adds a resource to an unstarted server. +func (s *Server) AddResource(resource mcp.Resource, handler server.ResourceHandlerFunc) { + s.resources = append(s.resources, server.ServerResource{ + Resource: resource, + Handler: handler, + }) +} + +// AddResources adds multiple resources to an unstarted server. +func (s *Server) AddResources(resources ...server.ServerResource) { + s.resources = append(s.resources, resources...) +} + // Start starts the server in a goroutine. Make sure to defer Close() after Start(). // When using NewServer(), the returned server is already started. func (s *Server) Start() error { @@ -95,6 +124,8 @@ func (s *Server) Start() error { mcpServer := server.NewMCPServer(s.name, "1.0.0") mcpServer.AddTools(s.tools...) + mcpServer.AddPrompts(s.prompts...) + mcpServer.AddResources(s.resources...) logger := log.New(&s.logBuffer, "", 0) diff --git a/mcptest/mcptest_test.go b/mcptest/mcptest_test.go index b71b93ea..129ef39f 100644 --- a/mcptest/mcptest_test.go +++ b/mcptest/mcptest_test.go @@ -11,7 +11,7 @@ import ( "github.com/mark3labs/mcp-go/server" ) -func TestServer(t *testing.T) { +func TestServerWithTool(t *testing.T) { ctx := context.Background() srv, err := mcptest.NewServer(t, server.ServerTool{ @@ -77,3 +77,113 @@ func resultToString(result *mcp.CallToolResult) (string, error) { return b.String(), nil } + +func TestServerWithPrompt(t *testing.T) { + ctx := context.Background() + + srv := mcptest.NewUnstartedServer(t) + defer srv.Close() + + prompt := mcp.Prompt{ + Name: "greeting", + Description: "A greeting prompt", + Arguments: []mcp.PromptArgument{ + { + Name: "name", + Description: "The name to greet", + Required: true, + }, + }, + } + handler := func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{ + Description: "A greeting prompt", + Messages: []mcp.PromptMessage{ + { + Role: mcp.RoleUser, + Content: mcp.NewTextContent(fmt.Sprintf("Hello, %s!", request.Params.Arguments["name"])), + }, + }, + }, nil + } + + srv.AddPrompt(prompt, handler) + + err := srv.Start() + if err != nil { + t.Fatal(err) + } + + var getReq mcp.GetPromptRequest + getReq.Params.Name = "greeting" + getReq.Params.Arguments = map[string]string{"name": "John"} + getResult, err := srv.Client().GetPrompt(ctx, getReq) + if err != nil { + t.Fatal("GetPrompt:", err) + } + if getResult.Description != "A greeting prompt" { + t.Errorf("Expected prompt description 'A greeting prompt', got %q", getResult.Description) + } + if len(getResult.Messages) != 1 { + t.Fatalf("Expected 1 message, got %d", len(getResult.Messages)) + } + if getResult.Messages[0].Role != mcp.RoleUser { + t.Errorf("Expected message role 'user', got %q", getResult.Messages[0].Role) + } + content, ok := getResult.Messages[0].Content.(mcp.TextContent) + if !ok { + t.Fatalf("Expected TextContent, got %T", getResult.Messages[0].Content) + } + if content.Text != "Hello, John!" { + t.Errorf("Expected message content 'Hello, John!', got %q", content.Text) + } +} + +func TestServerWithResource(t *testing.T) { + ctx := context.Background() + + srv := mcptest.NewUnstartedServer(t) + defer srv.Close() + + resource := mcp.Resource{ + URI: "test://resource", + Name: "Test Resource", + Description: "A test resource", + MIMEType: "text/plain", + } + + handler := func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { + return []mcp.ResourceContents{ + mcp.TextResourceContents{ + URI: "test://resource", + MIMEType: "text/plain", + Text: "This is a test resource content.", + }, + }, nil + } + + srv.AddResource(resource, handler) + + err := srv.Start() + if err != nil { + t.Fatal(err) + } + + var readReq mcp.ReadResourceRequest + readReq.Params.URI = "test://resource" + readResult, err := srv.Client().ReadResource(ctx, readReq) + if err != nil { + t.Fatal("ReadResource:", err) + } + if len(readResult.Contents) != 1 { + t.Fatalf("Expected 1 content, got %d", len(readResult.Contents)) + } + textContent, ok := readResult.Contents[0].(mcp.TextResourceContents) + if !ok { + t.Fatalf("Expected TextResourceContents, got %T", readResult.Contents[0]) + } + want := "This is a test resource content." + if textContent.Text != want { + t.Errorf("Got %q, want %q", textContent.Text, want) + } +} diff --git a/server/server.go b/server/server.go index 7823d1f9..46e6d9c5 100644 --- a/server/server.go +++ b/server/server.go @@ -52,6 +52,18 @@ type ServerTool struct { Handler ToolHandlerFunc } +// ServerPrompt combines a Prompt with its handler function. +type ServerPrompt struct { + Prompt mcp.Prompt + Handler PromptHandlerFunc +} + +// ServerResource combines a Resource with its handler function. +type ServerResource struct { + Resource mcp.Resource + Handler ResourceHandlerFunc +} + // serverKey is the context key for storing the server instance type serverKey struct{} @@ -305,28 +317,16 @@ func NewMCPServer( return s } -// AddResource registers a new resource and its handler -func (s *MCPServer) AddResource( - resource mcp.Resource, - handler ResourceHandlerFunc, -) { - s.capabilitiesMu.RLock() - if s.capabilities.resources == nil { - s.capabilitiesMu.RUnlock() - - s.capabilitiesMu.Lock() - if s.capabilities.resources == nil { - s.capabilities.resources = &resourceCapabilities{} - } - s.capabilitiesMu.Unlock() - } else { - s.capabilitiesMu.RUnlock() - } +// AddResources registers multiple resources at once +func (s *MCPServer) AddResources(resources ...ServerResource) { + s.implicitlyRegisterResourceCapabilities() s.resourcesMu.Lock() - s.resources[resource.URI] = resourceEntry{ - resource: resource, - handler: handler, + for _, entry := range resources { + s.resources[entry.Resource.URI] = resourceEntry{ + resource: entry.Resource, + handler: entry.Handler, + } } s.resourcesMu.Unlock() @@ -337,6 +337,14 @@ func (s *MCPServer) AddResource( } } +// AddResource registers a new resource and its handler +func (s *MCPServer) AddResource( + resource mcp.Resource, + handler ResourceHandlerFunc, +) { + s.AddResources(ServerResource{Resource: resource, Handler: handler}) +} + // RemoveResource removes a resource from the server func (s *MCPServer) RemoveResource(uri string) { s.resourcesMu.Lock() @@ -357,18 +365,7 @@ func (s *MCPServer) AddResourceTemplate( template mcp.ResourceTemplate, handler ResourceTemplateHandlerFunc, ) { - s.capabilitiesMu.RLock() - if s.capabilities.resources == nil { - s.capabilitiesMu.RUnlock() - - s.capabilitiesMu.Lock() - if s.capabilities.resources == nil { - s.capabilities.resources = &resourceCapabilities{} - } - s.capabilitiesMu.Unlock() - } else { - s.capabilitiesMu.RUnlock() - } + s.implicitlyRegisterResourceCapabilities() s.resourcesMu.Lock() s.resourceTemplates[template.URITemplate.Raw()] = resourceTemplateEntry{ @@ -384,24 +381,15 @@ func (s *MCPServer) AddResourceTemplate( } } -// AddPrompt registers a new prompt handler with the given name -func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) { - s.capabilitiesMu.RLock() - if s.capabilities.prompts == nil { - s.capabilitiesMu.RUnlock() - - s.capabilitiesMu.Lock() - if s.capabilities.prompts == nil { - s.capabilities.prompts = &promptCapabilities{} - } - s.capabilitiesMu.Unlock() - } else { - s.capabilitiesMu.RUnlock() - } +// AddPrompts registers multiple prompts at once +func (s *MCPServer) AddPrompts(prompts ...ServerPrompt) { + s.implicitlyRegisterPromptCapabilities() s.promptsMu.Lock() - s.prompts[prompt.Name] = prompt - s.promptHandlers[prompt.Name] = handler + for _, entry := range prompts { + s.prompts[entry.Prompt.Name] = entry.Prompt + s.promptHandlers[entry.Prompt.Name] = entry.Handler + } s.promptsMu.Unlock() // When the list of available prompts changes, servers that declared the listChanged capability SHOULD send a notification. @@ -411,6 +399,11 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) { } } +// AddPrompt registers a new prompt handler with the given name +func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) { + s.AddPrompts(ServerPrompt{Prompt: prompt, Handler: handler}) +} + // DeletePrompts removes prompts from the server func (s *MCPServer) DeletePrompts(names ...string) { s.promptsMu.Lock() @@ -440,20 +433,39 @@ func (s *MCPServer) AddTool(tool mcp.Tool, handler ToolHandlerFunc) { // listChanged: true, but don't change the value if we've already explicitly // registered tools.listChanged false. func (s *MCPServer) implicitlyRegisterToolCapabilities() { + s.implicitlyRegisterCapabilities( + func() bool { return s.capabilities.tools != nil }, + func() { s.capabilities.tools = &toolCapabilities{listChanged: true} }, + ) +} + +func (s *MCPServer) implicitlyRegisterResourceCapabilities() { + s.implicitlyRegisterCapabilities( + func() bool { return s.capabilities.resources != nil }, + func() { s.capabilities.resources = &resourceCapabilities{} }, + ) +} + +func (s *MCPServer) implicitlyRegisterPromptCapabilities() { + s.implicitlyRegisterCapabilities( + func() bool { return s.capabilities.prompts != nil }, + func() { s.capabilities.prompts = &promptCapabilities{} }, + ) +} + +func (s *MCPServer) implicitlyRegisterCapabilities(check func() bool, register func()) { s.capabilitiesMu.RLock() - if s.capabilities.tools == nil { + if check() { s.capabilitiesMu.RUnlock() + return + } + s.capabilitiesMu.RUnlock() - s.capabilitiesMu.Lock() - if s.capabilities.tools == nil { - s.capabilities.tools = &toolCapabilities{ - listChanged: true, - } - } - s.capabilitiesMu.Unlock() - } else { - s.capabilitiesMu.RUnlock() + s.capabilitiesMu.Lock() + if !check() { + register() } + s.capabilitiesMu.Unlock() } // AddTools registers multiple tools at once