Skip to content

Commit 905a8ca

Browse files
bytethmbytethm
authored
feat: add server-initiated requests and Roots capability support (#45)
* feat: add server-initiated requests and Roots capability support * fix
1 parent dac3659 commit 905a8ca

28 files changed

Lines changed: 4271 additions & 179 deletions

‎client.go‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"net/http"
1313
"net/url"
1414
"reflect"
15+
"sync"
1516
"sync/atomic"
1617

1718
"trpc.group/trpc-go/trpc-mcp-go/internal/errors"
@@ -60,6 +61,10 @@ type Connector interface {
6061
RegisterNotificationHandler(method string, handler NotificationHandler)
6162
// UnregisterNotificationHandler removes a notification handler.
6263
UnregisterNotificationHandler(method string)
64+
// SetRootsProvider sets the provider for responding to server's roots/list requests.
65+
SetRootsProvider(provider RootsProvider)
66+
// SendRootsListChangedNotification notifies server that roots changed.
67+
SendRootsListChangedNotification(ctx context.Context) error
6368
}
6469

6570
// SessionClient extends Connector with session management capabilities.
@@ -115,6 +120,10 @@ type Client struct {
115120
transportConfig *transportConfig
116121

117122
logger Logger // Logger for client transport (optional).
123+
124+
// Roots support.
125+
rootsProvider RootsProvider // Provider for roots information.
126+
rootsMu sync.RWMutex // Mutex for protecting the rootsProvider.
118127
}
119128

120129
// ClientOption client option function
@@ -149,6 +158,11 @@ func NewClient(serverURL string, clientInfo Implementation, options ...ClientOpt
149158
// Create transport layer if not previously set via options.
150159
if client.transport == nil {
151160
client.transport = newStreamableHTTPClientTransport(client.transportConfig, client.transportOptions...)
161+
162+
// Set client reference in transport for roots handling.
163+
if streamableTransport, ok := client.transport.(*streamableHTTPClientTransport); ok {
164+
streamableTransport.client = client
165+
}
152166
}
153167

154168
return client, nil
@@ -608,6 +622,20 @@ func (c *Client) ReadResource(ctx context.Context, readResourceReq *ReadResource
608622
return parseReadResourceResultFromJSON(rawResp)
609623
}
610624

625+
// SetRootsProvider sets the provider for responding to server's roots/list requests.
626+
func (c *Client) SetRootsProvider(provider RootsProvider) {
627+
c.rootsMu.Lock()
628+
defer c.rootsMu.Unlock()
629+
c.rootsProvider = provider
630+
}
631+
632+
// SendRootsListChangedNotification notifies server that roots changed.
633+
func (c *Client) SendRootsListChangedNotification(ctx context.Context) error {
634+
// Create roots list changed notification.
635+
notification := NewJSONRPCNotificationFromMap(MethodNotificationsRootsListChanged, nil)
636+
return c.transport.sendNotification(ctx, notification)
637+
}
638+
611639
func isZeroStruct(x interface{}) bool {
612640
return reflect.ValueOf(x).IsZero()
613641
}

‎e2e/server_roots_test.go‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
// Tencent is pleased to support the open source community by making trpc-mcp-go available.
2+
//
3+
// Copyright (C) 2025 Tencent. All rights reserved.
4+
//
5+
// trpc-mcp-go is licensed under the Apache License Version 2.0.
6+
7+
package e2e
8+
9+
import (
10+
"context"
11+
"fmt"
12+
"testing"
13+
"time"
14+
15+
"github.com/stretchr/testify/assert"
16+
"github.com/stretchr/testify/require"
17+
mcp "trpc.group/trpc-go/trpc-mcp-go"
18+
)
19+
20+
// TestServerRootsProvider_Streamable tests server-to-client roots functionality over Streamable transport.
21+
func TestServerRootsProvider_Streamable(t *testing.T) {
22+
// Start a test server
23+
serverURL, cleanup := StartTestServer(t, WithTestTools())
24+
defer cleanup()
25+
26+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
27+
defer cancel()
28+
29+
// Create client with RootsProvider capability
30+
client, err := mcp.NewClient(
31+
serverURL,
32+
mcp.Implementation{
33+
Name: "roots-streamable-test-client",
34+
Version: "1.0.0",
35+
},
36+
mcp.WithClientLogger(mcp.GetDefaultLogger()),
37+
)
38+
require.NoError(t, err)
39+
defer client.Close()
40+
41+
// Set up roots provider with test directories
42+
rootsProvider := mcp.NewDefaultRootsProvider()
43+
rootsProvider.AddRoot("/tmp", "Temporary Directory")
44+
rootsProvider.AddRoot("/home", "Home Directory")
45+
client.SetRootsProvider(rootsProvider)
46+
47+
// Initialize with roots capability
48+
initResult, err := client.Initialize(ctx, &mcp.InitializeRequest{
49+
Params: mcp.InitializeParams{
50+
ProtocolVersion: mcp.ProtocolVersion_2025_03_26,
51+
ClientInfo: mcp.Implementation{
52+
Name: "roots-streamable-test-client",
53+
Version: "1.0.0",
54+
},
55+
Capabilities: mcp.ClientCapabilities{
56+
Roots: &mcp.RootsCapability{
57+
ListChanged: true,
58+
},
59+
},
60+
},
61+
})
62+
require.NoError(t, err)
63+
assert.Equal(t, mcp.ProtocolVersion_2025_03_26, initResult.ProtocolVersion)
64+
65+
// Test sending roots list changed notification
66+
err = client.SendRootsListChangedNotification(ctx)
67+
require.NoError(t, err)
68+
69+
// Give the server a moment to process the notification
70+
time.Sleep(100 * time.Millisecond)
71+
}
72+
73+
// TestServerRootsProvider_Notification tests roots notification handling.
74+
func TestServerRootsProvider_Notification(t *testing.T) {
75+
// Start a test server
76+
serverURL, cleanup := StartTestServer(t, WithTestTools())
77+
defer cleanup()
78+
79+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
80+
defer cancel()
81+
82+
// Create client with RootsProvider capability
83+
client, err := mcp.NewClient(
84+
serverURL,
85+
mcp.Implementation{
86+
Name: "roots-notification-test-client",
87+
Version: "1.0.0",
88+
},
89+
mcp.WithClientLogger(mcp.GetDefaultLogger()),
90+
)
91+
require.NoError(t, err)
92+
defer client.Close()
93+
94+
// Set up roots provider with test directories
95+
rootsProvider := mcp.NewDefaultRootsProvider()
96+
rootsProvider.AddRoot("/opt", "Optional Directory")
97+
client.SetRootsProvider(rootsProvider)
98+
99+
// Initialize with roots capability
100+
_, err = client.Initialize(ctx, &mcp.InitializeRequest{
101+
Params: mcp.InitializeParams{
102+
ProtocolVersion: mcp.ProtocolVersion_2025_03_26,
103+
ClientInfo: mcp.Implementation{
104+
Name: "roots-notification-test-client",
105+
Version: "1.0.0",
106+
},
107+
Capabilities: mcp.ClientCapabilities{
108+
Roots: &mcp.RootsCapability{
109+
ListChanged: true,
110+
},
111+
},
112+
},
113+
})
114+
require.NoError(t, err)
115+
116+
// Register a tool that uses ListRoots
117+
server := mcp.GetServerFromContext(ctx)
118+
if server != nil {
119+
if s, ok := server.(*mcp.Server); ok {
120+
listRootsTool := mcp.NewTool("list-streamable-roots",
121+
mcp.WithDescription("List client's root directories via Streamable transport"),
122+
)
123+
124+
s.RegisterTool(listRootsTool, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
125+
// Call ListRoots to get client roots
126+
roots, err := s.ListRoots(ctx)
127+
if err != nil {
128+
return mcp.NewErrorResult(fmt.Sprintf("Failed to list roots: %v", err)), nil
129+
}
130+
131+
// Format response
132+
message := fmt.Sprintf("Streamable client has %d root directories:\n", len(roots.Roots))
133+
for i, root := range roots.Roots {
134+
message += fmt.Sprintf("%d. %s (%s)\n", i+1, root.Name, root.URI)
135+
}
136+
137+
return mcp.NewTextResult(message), nil
138+
})
139+
140+
// Call the tool to test ListRoots functionality
141+
result, err := client.CallTool(ctx, &mcp.CallToolRequest{
142+
Params: mcp.CallToolParams{
143+
Name: "list-streamable-roots",
144+
},
145+
})
146+
147+
if err == nil {
148+
// Verify the response contains our root directories
149+
textContent, ok := result.Content[0].(mcp.TextContent)
150+
require.True(t, ok)
151+
assert.Contains(t, textContent.Text, "Optional Directory")
152+
}
153+
}
154+
}
155+
156+
// Send roots list changed notification
157+
err = client.SendRootsListChangedNotification(ctx)
158+
require.NoError(t, err)
159+
160+
// Give the server a moment to process the notification
161+
time.Sleep(100 * time.Millisecond)
162+
}

‎e2e/sse_integration_test.go‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,87 @@ func TestSSEClientServer_ConnectionManagement(t *testing.T) {
175175
})
176176
}
177177

178+
// TestSSEClientServer_RootsProvider tests server-to-client roots functionality over SSE transport.
179+
func TestSSEClientServer_RootsProvider(t *testing.T) {
180+
server, sseEndpoint, cleanup := startSSEServer(t)
181+
defer cleanup()
182+
183+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
184+
defer cancel()
185+
186+
// Create SSE Client with RootsProvider capability
187+
client, err := mcp.NewSSEClient(
188+
sseEndpoint,
189+
mcp.Implementation{
190+
Name: "roots-sse-test-client",
191+
Version: "1.0.0",
192+
},
193+
mcp.WithClientLogger(mcp.GetDefaultLogger()),
194+
)
195+
require.NoError(t, err)
196+
defer client.Close()
197+
198+
// Set up roots provider with test directories
199+
rootsProvider := mcp.NewDefaultRootsProvider()
200+
rootsProvider.AddRoot("/tmp", "Temporary Directory")
201+
rootsProvider.AddRoot("/var", "Variable Directory")
202+
client.SetRootsProvider(rootsProvider)
203+
204+
// Initialize with roots capability
205+
initResult, err := client.Initialize(ctx, &mcp.InitializeRequest{
206+
Params: mcp.InitializeParams{
207+
ProtocolVersion: mcp.ProtocolVersion_2025_03_26,
208+
ClientInfo: mcp.Implementation{
209+
Name: "roots-sse-test-client",
210+
Version: "1.0.0",
211+
},
212+
Capabilities: mcp.ClientCapabilities{
213+
Roots: &mcp.RootsCapability{
214+
ListChanged: true,
215+
},
216+
},
217+
},
218+
})
219+
require.NoError(t, err)
220+
assert.Equal(t, mcp.ProtocolVersion_2025_03_26, initResult.ProtocolVersion)
221+
222+
// Register a tool that uses ListRoots
223+
listRootsTool := mcp.NewTool("list-sse-roots",
224+
mcp.WithDescription("List client's root directories via SSE"),
225+
)
226+
227+
server.RegisterTool(listRootsTool, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
228+
// Call ListRoots to get client roots
229+
roots, err := server.ListRoots(ctx)
230+
if err != nil {
231+
return mcp.NewErrorResult(fmt.Sprintf("Failed to list roots: %v", err)), nil
232+
}
233+
234+
// Format response
235+
message := fmt.Sprintf("SSE client has %d root directories:\n", len(roots.Roots))
236+
for i, root := range roots.Roots {
237+
message += fmt.Sprintf("%d. %s (%s)\n", i+1, root.Name, root.URI)
238+
}
239+
240+
return mcp.NewTextResult(message), nil
241+
})
242+
243+
// Call the tool to test ListRoots functionality
244+
result, err := client.CallTool(ctx, &mcp.CallToolRequest{
245+
Params: mcp.CallToolParams{
246+
Name: "list-sse-roots",
247+
},
248+
})
249+
require.NoError(t, err)
250+
require.NotEmpty(t, result.Content)
251+
252+
// Verify the response contains our root directories
253+
textContent, ok := result.Content[0].(mcp.TextContent)
254+
require.True(t, ok)
255+
assert.Contains(t, textContent.Text, "Temporary Directory")
256+
assert.Contains(t, textContent.Text, "Variable Directory")
257+
}
258+
178259
// startSSEServer starts an SSE server for testing.
179260
func startSSEServer(t *testing.T) (*mcp.SSEServer, string, func()) {
180261
t.Helper()

‎e2e/stdio_integration_test.go‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -327,3 +327,81 @@ func TestStdioClientServer_ConcurrentOperations(t *testing.T) {
327327
}
328328
})
329329
}
330+
331+
// TestStdioClientServer_RootsProvider tests server-to-client roots functionality over STDIO transport.
332+
func TestStdioClientServer_RootsProvider(t *testing.T) {
333+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
334+
defer cancel()
335+
336+
// Create STDIO server configuration
337+
serverConfig := mcp.StdioTransportConfig{
338+
ServerParams: mcp.StdioServerParameters{
339+
Command: "go",
340+
Args: []string{"run", "./test_server/main.go"},
341+
},
342+
Timeout: 10 * time.Second,
343+
}
344+
345+
// Create STDIO client with RootsProvider capability
346+
client, err := mcp.NewStdioClient(
347+
serverConfig,
348+
mcp.Implementation{
349+
Name: "roots-stdio-test-client",
350+
Version: "1.0.0",
351+
},
352+
mcp.WithStdioLogger(mcp.GetDefaultLogger()),
353+
)
354+
require.NoError(t, err)
355+
defer client.Close()
356+
357+
// Set up roots provider with test directories
358+
rootsProvider := mcp.NewDefaultRootsProvider()
359+
rootsProvider.AddRoot("/etc", "Configuration Directory")
360+
rootsProvider.AddRoot("/usr", "User Directory")
361+
client.SetRootsProvider(rootsProvider)
362+
363+
// Initialize with roots capability
364+
initResult, err := client.Initialize(ctx, &mcp.InitializeRequest{
365+
Params: mcp.InitializeParams{
366+
ProtocolVersion: mcp.ProtocolVersion_2025_03_26,
367+
ClientInfo: mcp.Implementation{
368+
Name: "roots-stdio-test-client",
369+
Version: "1.0.0",
370+
},
371+
Capabilities: mcp.ClientCapabilities{
372+
Roots: &mcp.RootsCapability{
373+
ListChanged: true,
374+
},
375+
},
376+
},
377+
})
378+
require.NoError(t, err)
379+
assert.Equal(t, mcp.ProtocolVersion_2025_03_26, initResult.ProtocolVersion)
380+
381+
// Test sending roots list changed notification
382+
err = client.SendRootsListChangedNotification(ctx)
383+
require.NoError(t, err)
384+
385+
// Give the server a moment to process the notification
386+
time.Sleep(100 * time.Millisecond)
387+
388+
// Test calling the list-roots tool
389+
result, err := client.CallTool(ctx, &mcp.CallToolRequest{
390+
Params: mcp.CallToolParams{
391+
Name: "list-roots",
392+
},
393+
})
394+
395+
if err == nil {
396+
// If the tool exists, verify the response
397+
require.NotEmpty(t, result.Content)
398+
textContent, ok := result.Content[0].(mcp.TextContent)
399+
require.True(t, ok)
400+
assert.Contains(t, textContent.Text, "Configuration Directory")
401+
assert.Contains(t, textContent.Text, "User Directory")
402+
} else {
403+
// The tool might not exist in the test server, which is acceptable
404+
// We're primarily testing that the notification was sent successfully
405+
t.Logf("Tool 'list-roots' not available: %v", err)
406+
}
407+
}

0 commit comments

Comments
 (0)