Skip to content

Commit de0304a

Browse files
vasayxtxvasayxtx
committed
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
1 parent c7c0e13 commit de0304a

File tree

3 files changed

+199
-39
lines changed

3 files changed

+199
-39
lines changed

‎mcptest/mcptest.go

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,11 @@ import (
1818

1919
// Server encapsulates an MCP server and manages resources like pipes and context.
2020
type Server struct {
21-
name string
22-
tools []server.ServerTool
21+
name string
22+
23+
tools []server.ServerTool
24+
prompts []server.ServerPrompt
25+
resources []server.ServerResource
2326

2427
ctx context.Context
2528
cancel func()
@@ -83,6 +86,32 @@ func (s *Server) AddTool(tool mcp.Tool, handler server.ToolHandlerFunc) {
8386
})
8487
}
8588

89+
// AddPrompt adds a prompt to an unstarted server.
90+
func (s *Server) AddPrompt(prompt mcp.Prompt, handler server.PromptHandlerFunc) {
91+
s.prompts = append(s.prompts, server.ServerPrompt{
92+
Prompt: prompt,
93+
Handler: handler,
94+
})
95+
}
96+
97+
// AddPrompts adds multiple prompts to an unstarted server.
98+
func (s *Server) AddPrompts(prompts ...server.ServerPrompt) {
99+
s.prompts = append(s.prompts, prompts...)
100+
}
101+
102+
// AddResource adds a resource to an unstarted server.
103+
func (s *Server) AddResource(resource mcp.Resource, handler server.ResourceHandlerFunc) {
104+
s.resources = append(s.resources, server.ServerResource{
105+
Resource: resource,
106+
Handler: handler,
107+
})
108+
}
109+
110+
// AddResources adds multiple resources to an unstarted server.
111+
func (s *Server) AddResources(resources ...server.ServerResource) {
112+
s.resources = append(s.resources, resources...)
113+
}
114+
86115
// Start starts the server in a goroutine. Make sure to defer Close() after Start().
87116
// When using NewServer(), the returned server is already started.
88117
func (s *Server) Start() error {
@@ -95,6 +124,8 @@ func (s *Server) Start() error {
95124
mcpServer := server.NewMCPServer(s.name, "1.0.0")
96125

97126
mcpServer.AddTools(s.tools...)
127+
mcpServer.AddPrompts(s.prompts...)
128+
mcpServer.AddResources(s.resources...)
98129

99130
logger := log.New(&s.logBuffer, "", 0)
100131

‎mcptest/mcptest_test.go

Lines changed: 111 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import (
1111
"github.com/mark3labs/mcp-go/server"
1212
)
1313

14-
func TestServer(t *testing.T) {
14+
func TestServerWithTool(t *testing.T) {
1515
ctx := context.Background()
1616

1717
srv, err := mcptest.NewServer(t, server.ServerTool{
@@ -77,3 +77,113 @@ func resultToString(result *mcp.CallToolResult) (string, error) {
7777

7878
return b.String(), nil
7979
}
80+
81+
func TestServerWithPrompt(t *testing.T) {
82+
ctx := context.Background()
83+
84+
srv := mcptest.NewUnstartedServer(t)
85+
defer srv.Close()
86+
87+
prompt := mcp.Prompt{
88+
Name: "greeting",
89+
Description: "A greeting prompt",
90+
Arguments: []mcp.PromptArgument{
91+
{
92+
Name: "name",
93+
Description: "The name to greet",
94+
Required: true,
95+
},
96+
},
97+
}
98+
handler := func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
99+
return &mcp.GetPromptResult{
100+
Description: "A greeting prompt",
101+
Messages: []mcp.PromptMessage{
102+
{
103+
Role: mcp.RoleUser,
104+
Content: mcp.NewTextContent(fmt.Sprintf("Hello, %s!", request.Params.Arguments["name"])),
105+
},
106+
},
107+
}, nil
108+
}
109+
110+
srv.AddPrompt(prompt, handler)
111+
112+
err := srv.Start()
113+
if err != nil {
114+
t.Fatal(err)
115+
}
116+
117+
var getReq mcp.GetPromptRequest
118+
getReq.Params.Name = "greeting"
119+
getReq.Params.Arguments = map[string]string{"name": "John"}
120+
getResult, err := srv.Client().GetPrompt(ctx, getReq)
121+
if err != nil {
122+
t.Fatal("GetPrompt:", err)
123+
}
124+
if getResult.Description != "A greeting prompt" {
125+
t.Errorf("Expected prompt description 'A greeting prompt', got %q", getResult.Description)
126+
}
127+
if len(getResult.Messages) != 1 {
128+
t.Fatalf("Expected 1 message, got %d", len(getResult.Messages))
129+
}
130+
if getResult.Messages[0].Role != mcp.RoleUser {
131+
t.Errorf("Expected message role 'user', got %q", getResult.Messages[0].Role)
132+
}
133+
content, ok := getResult.Messages[0].Content.(mcp.TextContent)
134+
if !ok {
135+
t.Fatalf("Expected TextContent, got %T", getResult.Messages[0].Content)
136+
}
137+
if content.Text != "Hello, John!" {
138+
t.Errorf("Expected message content 'Hello, John!', got %q", content.Text)
139+
}
140+
}
141+
142+
func TestServerWithResource(t *testing.T) {
143+
ctx := context.Background()
144+
145+
srv := mcptest.NewUnstartedServer(t)
146+
defer srv.Close()
147+
148+
resource := mcp.Resource{
149+
URI: "test://resource",
150+
Name: "Test Resource",
151+
Description: "A test resource",
152+
MIMEType: "text/plain",
153+
}
154+
155+
handler := func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
156+
return []mcp.ResourceContents{
157+
mcp.TextResourceContents{
158+
URI: "test://resource",
159+
MIMEType: "text/plain",
160+
Text: "This is a test resource content.",
161+
},
162+
}, nil
163+
}
164+
165+
srv.AddResource(resource, handler)
166+
167+
err := srv.Start()
168+
if err != nil {
169+
t.Fatal(err)
170+
}
171+
172+
var readReq mcp.ReadResourceRequest
173+
readReq.Params.URI = "test://resource"
174+
readResult, err := srv.Client().ReadResource(ctx, readReq)
175+
if err != nil {
176+
t.Fatal("ReadResource:", err)
177+
}
178+
if len(readResult.Contents) != 1 {
179+
t.Fatalf("Expected 1 content, got %d", len(readResult.Contents))
180+
}
181+
textContent, ok := readResult.Contents[0].(mcp.TextResourceContents)
182+
if !ok {
183+
t.Fatalf("Expected TextResourceContents, got %T", readResult.Contents[0])
184+
}
185+
want := "This is a test resource content."
186+
if textContent.Text != want {
187+
t.Errorf("Got %q, want %q", textContent.Text, want)
188+
}
189+
}

‎server/server.go

Lines changed: 55 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,18 @@ type ServerTool struct {
5252
Handler ToolHandlerFunc
5353
}
5454

55+
// ServerPrompt combines a Prompt with its handler function.
56+
type ServerPrompt struct {
57+
Prompt mcp.Prompt
58+
Handler PromptHandlerFunc
59+
}
60+
61+
// ServerResource combines a Resource with its handler function.
62+
type ServerResource struct {
63+
Resource mcp.Resource
64+
Handler ResourceHandlerFunc
65+
}
66+
5567
// serverKey is the context key for storing the server instance
5668
type serverKey struct{}
5769

@@ -305,28 +317,16 @@ func NewMCPServer(
305317
return s
306318
}
307319

308-
// AddResource registers a new resource and its handler
309-
func (s *MCPServer) AddResource(
310-
resource mcp.Resource,
311-
handler ResourceHandlerFunc,
312-
) {
313-
s.capabilitiesMu.RLock()
314-
if s.capabilities.resources == nil {
315-
s.capabilitiesMu.RUnlock()
316-
317-
s.capabilitiesMu.Lock()
318-
if s.capabilities.resources == nil {
319-
s.capabilities.resources = &resourceCapabilities{}
320-
}
321-
s.capabilitiesMu.Unlock()
322-
} else {
323-
s.capabilitiesMu.RUnlock()
324-
}
320+
// AddResources registers multiple resources at once
321+
func (s *MCPServer) AddResources(resources ...ServerResource) {
322+
s.implicitlyRegisterResourceCapabilities()
325323

326324
s.resourcesMu.Lock()
327-
s.resources[resource.URI] = resourceEntry{
328-
resource: resource,
329-
handler: handler,
325+
for _, entry := range resources {
326+
s.resources[entry.Resource.URI] = resourceEntry{
327+
resource: entry.Resource,
328+
handler: entry.Handler,
329+
}
330330
}
331331
s.resourcesMu.Unlock()
332332

@@ -337,6 +337,14 @@ func (s *MCPServer) AddResource(
337337
}
338338
}
339339

340+
// AddResource registers a new resource and its handler
341+
func (s *MCPServer) AddResource(
342+
resource mcp.Resource,
343+
handler ResourceHandlerFunc,
344+
) {
345+
s.AddResources(ServerResource{Resource: resource, Handler: handler})
346+
}
347+
340348
// RemoveResource removes a resource from the server
341349
func (s *MCPServer) RemoveResource(uri string) {
342350
s.resourcesMu.Lock()
@@ -357,18 +365,7 @@ func (s *MCPServer) AddResourceTemplate(
357365
template mcp.ResourceTemplate,
358366
handler ResourceTemplateHandlerFunc,
359367
) {
360-
s.capabilitiesMu.RLock()
361-
if s.capabilities.resources == nil {
362-
s.capabilitiesMu.RUnlock()
363-
364-
s.capabilitiesMu.Lock()
365-
if s.capabilities.resources == nil {
366-
s.capabilities.resources = &resourceCapabilities{}
367-
}
368-
s.capabilitiesMu.Unlock()
369-
} else {
370-
s.capabilitiesMu.RUnlock()
371-
}
368+
s.implicitlyRegisterResourceCapabilities()
372369

373370
s.resourcesMu.Lock()
374371
s.resourceTemplates[template.URITemplate.Raw()] = resourceTemplateEntry{
@@ -384,8 +381,8 @@ func (s *MCPServer) AddResourceTemplate(
384381
}
385382
}
386383

387-
// AddPrompt registers a new prompt handler with the given name
388-
func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
384+
// AddPrompts registers multiple prompts at once
385+
func (s *MCPServer) AddPrompts(prompts ...ServerPrompt) {
389386
s.capabilitiesMu.RLock()
390387
if s.capabilities.prompts == nil {
391388
s.capabilitiesMu.RUnlock()
@@ -400,8 +397,10 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
400397
}
401398

402399
s.promptsMu.Lock()
403-
s.prompts[prompt.Name] = prompt
404-
s.promptHandlers[prompt.Name] = handler
400+
for _, entry := range prompts {
401+
s.prompts[entry.Prompt.Name] = entry.Prompt
402+
s.promptHandlers[entry.Prompt.Name] = entry.Handler
403+
}
405404
s.promptsMu.Unlock()
406405

407406
// When the list of available prompts changes, servers that declared the listChanged capability SHOULD send a notification.
@@ -411,6 +410,11 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
411410
}
412411
}
413412

413+
// AddPrompt registers a new prompt handler with the given name
414+
func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
415+
s.AddPrompts(ServerPrompt{Prompt: prompt, Handler: handler})
416+
}
417+
414418
// DeletePrompts removes prompts from the server
415419
func (s *MCPServer) DeletePrompts(names ...string) {
416420
s.promptsMu.Lock()
@@ -456,6 +460,21 @@ func (s *MCPServer) implicitlyRegisterToolCapabilities() {
456460
}
457461
}
458462

463+
func (s *MCPServer) implicitlyRegisterResourceCapabilities() {
464+
s.capabilitiesMu.RLock()
465+
if s.capabilities.resources == nil {
466+
s.capabilitiesMu.RUnlock()
467+
468+
s.capabilitiesMu.Lock()
469+
if s.capabilities.resources == nil {
470+
s.capabilities.resources = &resourceCapabilities{}
471+
}
472+
s.capabilitiesMu.Unlock()
473+
} else {
474+
s.capabilitiesMu.RUnlock()
475+
}
476+
}
477+
459478
// AddTools registers multiple tools at once
460479
func (s *MCPServer) AddTools(tools ...ServerTool) {
461480
s.implicitlyRegisterToolCapabilities()

0 commit comments

Comments
 (0)