fix: enforce server-to-client request scope by restricting client request ID context access

Author: guglielmo-san

mcp/client_example_test.go

@@ -78,24 +78,36 @@ func Example_sampling() {
 		},
 	})
 
-	// Connect the server and client...
+	// Create a server with a tool that uses sampling.
+	// Server-to-client requests like CreateMessage must be made within a
+	// client request handler (e.g. a tool call).
 	ct, st := mcp.NewInMemoryTransports()
 	s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
+	s.AddTool(&mcp.Tool{Name: "sample", InputSchema: map[string]any{"type": "object"}}, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		msg, err := req.Session.CreateMessage(ctx, &mcp.CreateMessageParams{})
+		if err != nil {
+			return nil, err
+		}
+		return &mcp.CallToolResult{
+			Content: []mcp.Content{&mcp.TextContent{Text: msg.Content.(*mcp.TextContent).Text}},
+		}, nil
+	})
 	session, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		log.Fatal(err)
 	}
 	defer session.Close()
 
-	if _, err := c.Connect(ctx, ct, nil); err != nil {
+	cs, err := c.Connect(ctx, ct, nil)
+	if err != nil {
 		log.Fatal(err)
 	}
 
-	msg, err := session.CreateMessage(ctx, &mcp.CreateMessageParams{})
+	res, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: "sample"})
 	if err != nil {
 		log.Fatal(err)
 	}
-	fmt.Println(msg.Content.(*mcp.TextContent).Text)
+	fmt.Println(res.Content[0].(*mcp.TextContent).Text)
 	// Output: would have created a message
 }
 
@@ -107,7 +119,27 @@ func Example_elicitation() {
 	ctx := context.Background()
 	ct, st := mcp.NewInMemoryTransports()
 
+	// Create a server with a tool that uses elicitation.
+	// Server-to-client requests like Elicit must be made within a client
+	// request handler (e.g. a tool call).
 	s := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
+	s.AddTool(&mcp.Tool{Name: "ask", InputSchema: map[string]any{"type": "object"}}, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		res, err := req.Session.Elicit(ctx, &mcp.ElicitParams{
+			Message: "Please provide information",
+			RequestedSchema: &jsonschema.Schema{
+				Type: "object",
+				Properties: map[string]*jsonschema.Schema{
+					"test": {Type: "string"},
+				},
+			},
+		})
+		if err != nil {
+			return nil, err
+		}
+		return &mcp.CallToolResult{
+			Content: []mcp.Content{&mcp.TextContent{Text: res.Content["test"].(string)}},
+		}, nil
+	})
 	ss, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		log.Fatal(err)
@@ -119,22 +151,16 @@ func Example_elicitation() {
 			return &mcp.ElicitResult{Action: "accept", Content: map[string]any{"test": "value"}}, nil
 		},
 	})
-	if _, err := c.Connect(ctx, ct, nil); err != nil {
+	cs, err := c.Connect(ctx, ct, nil)
+	if err != nil {
 		log.Fatal(err)
 	}
-	res, err := ss.Elicit(ctx, &mcp.ElicitParams{
-		Message: "This should fail",
-		RequestedSchema: &jsonschema.Schema{
-			Type: "object",
-			Properties: map[string]*jsonschema.Schema{
-				"test": {Type: "string"},
-			},
-		},
-	})
+
+	toolRes, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: "ask"})
 	if err != nil {
 		log.Fatal(err)
 	}
-	fmt.Println(res.Content["test"])
+	fmt.Println(toolRes.Content[0].(*mcp.TextContent).Text)
 	// Output: value
 }
 

mcp/elicitation_test.go

@@ -120,7 +120,7 @@ func TestElicitationURLMode(t *testing.T) {
 			}
 			defer cs.Close()
 
-			result, err := ss.Elicit(ctx, tc.params)
+			result, err := ss.Elicit(clientRequestCtx(ctx), tc.params)
 
 			if tc.wantErrMsg != "" {
 				if err == nil || !strings.Contains(err.Error(), tc.wantErrMsg) {
@@ -170,7 +170,7 @@ func TestElicitationCompleteNotification(t *testing.T) {
 
 		// 1. Server initiates a URL elicitation
 		elicitID := "testElicitationID-123"
-		resp, err := ss.Elicit(ctx, &ElicitParams{
+		resp, err := ss.Elicit(clientRequestCtx(ctx), &ElicitParams{
 			Mode:          "url",
 			Message:       "Please complete this form: ",
 			URL:           "https://example.com/form?id=" + elicitID,
@@ -251,7 +251,7 @@ func TestElicitationNoValidationWithoutAccept(t *testing.T) {
 			}
 			defer cs.Close()
 
-			res, err := ss.Elicit(ctx, &ElicitParams{
+			res, err := ss.Elicit(clientRequestCtx(ctx), &ElicitParams{
 				Message:         "Test bug",
 				RequestedSchema: schema,
 			})