feat(mcptest): extend test server with prompt and resource support

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)
 

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)
+	}
+}

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