| 1 | package plugin |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "sync" |
| 8 | "testing" |
| 9 | |
| 10 | mcpjsonrpc "github.com/modelcontextprotocol/go-sdk/jsonrpc" |
| 11 | mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" |
| 12 | "reasonix/internal/mcpinteraction" |
| 13 | ) |
| 14 | |
| 15 | // scriptedBroker answers every elicitation with a canned result. |
| 16 | type scriptedBroker struct { |
| 17 | mu sync.Mutex |
| 18 | got []mcpinteraction.Request |
| 19 | answer mcpinteraction.Result |
| 20 | answerFn func(mcpinteraction.Request) mcpinteraction.Result |
| 21 | } |
| 22 | |
| 23 | func (b *scriptedBroker) Interact(_ context.Context, req mcpinteraction.Request) (mcpinteraction.Result, error) { |
| 24 | b.mu.Lock() |
| 25 | b.got = append(b.got, req) |
| 26 | answer := b.answer |
| 27 | if b.answerFn != nil { |
| 28 | answer = b.answerFn(req) |
| 29 | } |
| 30 | b.mu.Unlock() |
| 31 | return answer, nil |
| 32 | } |
| 33 | |
| 34 | func (b *scriptedBroker) requests() []mcpinteraction.Request { |
| 35 | b.mu.Lock() |
| 36 | defer b.mu.Unlock() |
| 37 | return append([]mcpinteraction.Request(nil), b.got...) |
| 38 | } |
| 39 | |
| 40 | // mrtrFixtureServer builds a server whose tool returns an input-required |
| 41 | // elicitation on the first call and completes on the second, recording the |
| 42 | // requestState of each attempt so the retry contract is observable. |
| 43 | type mrtrFixtureServer struct { |
| 44 | mu sync.Mutex |
| 45 | attempts int |
| 46 | states []string |
| 47 | numInputs int // how many parallel input requests to return (multi-input test) |
| 48 | } |
| 49 | |
| 50 | func (f *mrtrFixtureServer) handler(ctx context.Context, req *mcpsdk.CallToolRequest) (*mcpsdk.CallToolResult, error) { |
| 51 | f.mu.Lock() |
| 52 | f.attempts++ |
| 53 | f.mu.Unlock() |
| 54 | ss := req.GetSession().(*mcpsdk.ServerSession) |
| 55 | state := "" |
| 56 | if ip := ss.InitializeParams(); ip != nil { |
| 57 | state = ip.ProtocolVersion |
| 58 | } |
| 59 | f.mu.Lock() |
| 60 | f.states = append(f.states, state) |
| 61 | attempt := f.attempts |
| 62 | f.mu.Unlock() |
| 63 | if attempt == 1 { |
| 64 | f.mu.Lock() |
| 65 | n := f.numInputs |
| 66 | f.mu.Unlock() |
| 67 | elicitations := make([]*mcpsdk.ElicitParams, n) |
| 68 | for i := range elicitations { |
| 69 | elicitations[i] = &mcpsdk.ElicitParams{ |
| 70 | Message: fmt.Sprintf("question %d", i+1), |
| 71 | RequestedSchema: map[string]any{ |
| 72 | "type": "object", |
| 73 | "properties": map[string]any{ |
| 74 | "answer": map[string]any{"type": "string"}, |
| 75 | }, |
| 76 | "required": []any{"answer"}, |
| 77 | }, |
| 78 | } |
| 79 | } |
| 80 | requests := make(mcpsdk.InputRequestMap, len(elicitations)) |
| 81 | for i, e := range elicitations { |
| 82 | requests[fmt.Sprintf("input-%d", i+1)] = e |
| 83 | } |
| 84 | return &mcpsdk.CallToolResult{InputRequests: requests}, nil |
| 85 | } |
| 86 | return &mcpsdk.CallToolResult{ |
| 87 | Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: "done"}}, |
| 88 | }, nil |
| 89 | } |
| 90 | |
| 91 | func (f *mrtrFixtureServer) server() *mcpsdk.Server { |
| 92 | server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "mrtr-fixture", Version: "1"}, nil) |
| 93 | server.AddTool(&mcpsdk.Tool{ |
| 94 | Name: "ask_then_do", Description: "elicits then completes", |
| 95 | InputSchema: map[string]any{"type": "object", "properties": map[string]any{}}, |
| 96 | }, f.handler) |
| 97 | return server |
| 98 | } |
| 99 | |
| 100 | // runMRTTool drives one tools/call through the profile-aware SDK session with |
| 101 | // the broker attached to the call context, exactly as the agent does. |
| 102 | func runMRTTool(t *testing.T, broker mcpinteraction.Broker, fixture *mrtrFixtureServer) (*mcpsdk.CallToolResult, error) { |
| 103 | t.Helper() |
| 104 | lifeCtx, cancelLife := context.WithCancel(context.Background()) |
| 105 | transport := &sdkSessionTransport{ |
| 106 | name: "mrtr", spec: Spec{Name: "mrtr", Type: "http"}, profile: HostProfileInteractive, |
| 107 | lifeCtx: lifeCtx, cancel: cancelLife, state: SessionStateConnecting, |
| 108 | } |
| 109 | transport.endpointFactory = func(ctx context.Context) (sdkEndpoint, error) { |
| 110 | clientSide, serverSide := mcpsdk.NewInMemoryTransports() |
| 111 | go func() { _ = fixture.server().Run(ctx, serverSide) }() |
| 112 | return sdkEndpoint{transport: clientSide}, nil |
| 113 | } |
| 114 | t.Cleanup(transport.close) |
| 115 | |
| 116 | managed, err := transport.acquire(t.Context()) |
| 117 | if err != nil { |
| 118 | return nil, err |
| 119 | } |
| 120 | callCtx := t.Context() |
| 121 | if broker != nil { |
| 122 | callCtx = mcpinteraction.WithBroker(callCtx, broker) |
| 123 | } |
| 124 | raw, err := invokeSDKMethod(callCtx, managed.session, "tools/call", map[string]any{ |
| 125 | "name": "ask_then_do", "arguments": map[string]any{}, |
| 126 | }) |
| 127 | if err != nil { |
| 128 | return nil, err |
| 129 | } |
| 130 | var result mcpsdk.CallToolResult |
| 131 | if err := json.Unmarshal(raw, &result); err != nil { |
| 132 | return nil, err |
| 133 | } |
| 134 | return &result, nil |
| 135 | } |
| 136 | |
| 137 | func TestMRTRFormElicitationRetriesWithUserAnswer(t *testing.T) { |
| 138 | fixture := &mrtrFixtureServer{numInputs: 1} |
| 139 | broker := &scriptedBroker{answer: mcpinteraction.Result{ |
| 140 | Action: mcpinteraction.ActionAccept, |
| 141 | Content: map[string]any{"answer": "from the user"}, |
| 142 | }} |
| 143 | res, err := runMRTTool(t, broker, fixture) |
| 144 | if err != nil { |
| 145 | t.Fatalf("tools/call: %v", err) |
| 146 | } |
| 147 | if len(res.Content) == 0 { |
| 148 | t.Fatal("no content on completed result") |
| 149 | } |
| 150 | if got := fixture.attempts; got != 2 { |
| 151 | t.Fatalf("attempts = %d, want 2 (retry after elicitation)", got) |
| 152 | } |
| 153 | reqs := broker.requests() |
| 154 | if len(reqs) != 1 { |
| 155 | t.Fatalf("broker saw %d elicitations, want 1", len(reqs)) |
| 156 | } |
| 157 | if reqs[0].Server != "mrtr" || reqs[0].Mode != "form" || reqs[0].Message != "question 1" { |
| 158 | t.Fatalf("broker request = %+v", reqs[0]) |
| 159 | } |
| 160 | var schema map[string]any |
| 161 | if err := json.Unmarshal(reqs[0].RequestedSchema, &schema); err != nil { |
| 162 | t.Fatalf("requested schema not valid JSON: %v", err) |
| 163 | } |
| 164 | } |
| 165 | |
| 166 | func TestMRTRDeclineAndCancelReachServer(t *testing.T) { |
| 167 | for _, tc := range []struct { |
| 168 | name string |
| 169 | answer mcpinteraction.Result |
| 170 | }{ |
| 171 | {"decline", mcpinteraction.Result{Action: mcpinteraction.ActionDecline}}, |
| 172 | {"cancel", mcpinteraction.Result{Action: mcpinteraction.ActionCancel}}, |
| 173 | } { |
| 174 | t.Run(tc.name, func(t *testing.T) { |
| 175 | fixture := &mrtrFixtureServer{numInputs: 1} |
| 176 | res, err := runMRTTool(t, &scriptedBroker{answer: tc.answer}, fixture) |
| 177 | if err != nil { |
| 178 | t.Fatalf("tools/call: %v", err) |
| 179 | } |
| 180 | if fixture.attempts != 2 { |
| 181 | t.Fatalf("attempts = %d, want 2: decline/cancel answers must still complete the round trip", fixture.attempts) |
| 182 | } |
| 183 | if res == nil { |
| 184 | t.Fatal("nil result") |
| 185 | } |
| 186 | }) |
| 187 | } |
| 188 | } |
| 189 | |
| 190 | func TestMRTRMultipleInputRequestsOneRound(t *testing.T) { |
| 191 | fixture := &mrtrFixtureServer{numInputs: 3} |
| 192 | broker := &scriptedBroker{answer: mcpinteraction.Result{ |
| 193 | Action: mcpinteraction.ActionAccept, |
| 194 | Content: map[string]any{"answer": "ok"}, |
| 195 | }} |
| 196 | if _, err := runMRTTool(t, broker, fixture); err != nil { |
| 197 | t.Fatalf("tools/call: %v", err) |
| 198 | } |
| 199 | if got := len(broker.requests()); got != 3 { |
| 200 | t.Fatalf("broker saw %d elicitations, want 3", got) |
| 201 | } |
| 202 | if fixture.attempts != 2 { |
| 203 | t.Fatalf("attempts = %d, want 2 (one retry carries all three answers)", fixture.attempts) |
| 204 | } |
| 205 | } |
| 206 | |
| 207 | func TestMRTRBrokerTravelsWithCallContext(t *testing.T) { |
| 208 | fixture := &mrtrFixtureServer{numInputs: 1} |
| 209 | // No broker on the call ctx: the handler must cancel, not guess, and the |
| 210 | // call still completes through the SDK's input-response path. |
| 211 | res, err := runMRTTool(t, nil, fixture) |
| 212 | if err != nil { |
| 213 | t.Fatalf("tools/call: %v", err) |
| 214 | } |
| 215 | if res == nil { |
| 216 | t.Fatal("nil result") |
| 217 | } |
| 218 | } |
| 219 | |
| 220 | func TestLegacy20251125PushElicitationUsesOnlyUnambiguousCall(t *testing.T) { |
| 221 | transport := &sdkSessionTransport{name: "legacy-push"} |
| 222 | broker := &scriptedBroker{answer: mcpinteraction.Result{ |
| 223 | Action: mcpinteraction.ActionAccept, |
| 224 | Content: map[string]any{"answer": "legacy"}, |
| 225 | }} |
| 226 | callCtx := mcpinteraction.WithBroker(t.Context(), broker) |
| 227 | unregister := transport.registerLegacyElicitationCall(callCtx, "2025-11-25", "tools/call") |
| 228 | defer unregister() |
| 229 | |
| 230 | res, err := transport.handleElicitation(context.Background(), &mcpsdk.ElicitRequest{Params: &mcpsdk.ElicitParams{ |
| 231 | Message: "legacy form", |
| 232 | }}) |
| 233 | if err != nil { |
| 234 | t.Fatal(err) |
| 235 | } |
| 236 | if res.Action != mcpinteraction.ActionAccept || res.Content["answer"] != "legacy" { |
| 237 | t.Fatalf("legacy response = %+v", res) |
| 238 | } |
| 239 | if got := len(broker.requests()); got != 1 { |
| 240 | t.Fatalf("legacy broker requests = %d, want 1", got) |
| 241 | } |
| 242 | } |
| 243 | |
| 244 | func TestLegacy20251125PushElicitationEndToEnd(t *testing.T) { |
| 245 | lifeCtx, cancelLife := context.WithCancel(context.Background()) |
| 246 | transport := &sdkSessionTransport{ |
| 247 | name: "legacy-push", spec: Spec{Name: "legacy-push", Type: "http"}, profile: HostProfileInteractive, |
| 248 | lifeCtx: lifeCtx, cancel: cancelLife, state: SessionStateConnecting, |
| 249 | } |
| 250 | transport.endpointFactory = func(ctx context.Context) (sdkEndpoint, error) { |
| 251 | clientSide, serverSide := mcpsdk.NewInMemoryTransports() |
| 252 | go serveLegacyPushFixture(ctx, serverSide) |
| 253 | return sdkEndpoint{transport: clientSide}, nil |
| 254 | } |
| 255 | t.Cleanup(transport.close) |
| 256 | |
| 257 | broker := &scriptedBroker{answer: mcpinteraction.Result{ |
| 258 | Action: mcpinteraction.ActionAccept, Content: map[string]any{"answer": "legacy accepted"}, |
| 259 | }} |
| 260 | callCtx := mcpinteraction.WithBroker(t.Context(), broker) |
| 261 | raw, err := transport.call(callCtx, "tools/call", map[string]any{"name": "legacy_ask", "arguments": map[string]any{}}) |
| 262 | if err != nil { |
| 263 | t.Fatal(err) |
| 264 | } |
| 265 | if got := transport.sessionDiagnostics().ProtocolVersion; got != "2025-11-25" { |
| 266 | t.Fatalf("protocol = %q, want 2025-11-25", got) |
| 267 | } |
| 268 | var result struct { |
| 269 | Content []struct { |
| 270 | Text string `json:"text"` |
| 271 | } `json:"content"` |
| 272 | } |
| 273 | if err := json.Unmarshal(raw, &result); err != nil { |
| 274 | t.Fatal(err) |
| 275 | } |
| 276 | if len(result.Content) != 1 || result.Content[0].Text != "legacy accepted" || len(broker.requests()) != 1 { |
| 277 | t.Fatalf("legacy push result = %s, broker requests = %d", raw, len(broker.requests())) |
| 278 | } |
| 279 | } |
| 280 | |
| 281 | func serveLegacyPushFixture(ctx context.Context, transport mcpsdk.Transport) { |
| 282 | conn, err := transport.Connect(ctx) |
| 283 | if err != nil { |
| 284 | return |
| 285 | } |
| 286 | defer conn.Close() |
| 287 | for { |
| 288 | message, err := conn.Read(ctx) |
| 289 | if err != nil { |
| 290 | return |
| 291 | } |
| 292 | request, ok := message.(*mcpjsonrpc.Request) |
| 293 | if !ok { |
| 294 | continue |
| 295 | } |
| 296 | switch request.Method { |
| 297 | case "server/discover": |
| 298 | _ = writeLegacyFixtureMessage(ctx, conn, map[string]any{ |
| 299 | "jsonrpc": "2.0", "id": request.ID.Raw(), |
| 300 | "error": map[string]any{"code": mcpjsonrpc.CodeMethodNotFound, "message": "method not found"}, |
| 301 | }) |
| 302 | case "initialize": |
| 303 | _ = writeLegacyFixtureMessage(ctx, conn, map[string]any{ |
| 304 | "jsonrpc": "2.0", "id": request.ID.Raw(), |
| 305 | "result": map[string]any{ |
| 306 | "protocolVersion": "2025-11-25", |
| 307 | "capabilities": map[string]any{"tools": map[string]any{}}, |
| 308 | "serverInfo": map[string]any{"name": "legacy-push", "version": "1"}, |
| 309 | }, |
| 310 | }) |
| 311 | case "tools/call": |
| 312 | if err := writeLegacyFixtureMessage(ctx, conn, map[string]any{ |
| 313 | "jsonrpc": "2.0", "id": 700, "method": "elicitation/create", |
| 314 | "params": map[string]any{ |
| 315 | "message": "legacy question", |
| 316 | "requestedSchema": map[string]any{ |
| 317 | "type": "object", |
| 318 | "properties": map[string]any{"answer": map[string]any{"type": "string"}}, |
| 319 | "required": []any{"answer"}, |
| 320 | }, |
| 321 | }, |
| 322 | }); err != nil { |
| 323 | return |
| 324 | } |
| 325 | answerMessage, err := conn.Read(ctx) |
| 326 | if err != nil { |
| 327 | return |
| 328 | } |
| 329 | answer, ok := answerMessage.(*mcpjsonrpc.Response) |
| 330 | if !ok || answer.Error != nil { |
| 331 | return |
| 332 | } |
| 333 | var elicitation struct { |
| 334 | Action string `json:"action"` |
| 335 | Content map[string]any `json:"content"` |
| 336 | } |
| 337 | if json.Unmarshal(answer.Result, &elicitation) != nil { |
| 338 | return |
| 339 | } |
| 340 | text, _ := elicitation.Content["answer"].(string) |
| 341 | _ = writeLegacyFixtureMessage(ctx, conn, map[string]any{ |
| 342 | "jsonrpc": "2.0", "id": request.ID.Raw(), |
| 343 | "result": map[string]any{ |
| 344 | "content": []any{map[string]any{"type": "text", "text": text}}, |
| 345 | }, |
| 346 | }) |
| 347 | } |
| 348 | } |
| 349 | } |
| 350 | |
| 351 | func writeLegacyFixtureMessage(ctx context.Context, conn mcpsdk.Connection, wire map[string]any) error { |
| 352 | encoded, err := json.Marshal(wire) |
| 353 | if err != nil { |
| 354 | return err |
| 355 | } |
| 356 | message, err := mcpjsonrpc.DecodeMessage(encoded) |
| 357 | if err != nil { |
| 358 | return err |
| 359 | } |
| 360 | return conn.Write(ctx, message) |
| 361 | } |
| 362 | |
| 363 | func TestLegacyPushElicitationCancelsAmbiguousConcurrentCalls(t *testing.T) { |
| 364 | transport := &sdkSessionTransport{name: "legacy-concurrent"} |
| 365 | first := &scriptedBroker{answer: mcpinteraction.Result{Action: mcpinteraction.ActionAccept}} |
| 366 | second := &scriptedBroker{answer: mcpinteraction.Result{Action: mcpinteraction.ActionAccept}} |
| 367 | unregisterFirst := transport.registerLegacyElicitationCall(mcpinteraction.WithBroker(t.Context(), first), "2025-11-25", "tools/call") |
| 368 | defer unregisterFirst() |
| 369 | unregisterSecond := transport.registerLegacyElicitationCall(mcpinteraction.WithBroker(t.Context(), second), "2025-11-25", "tools/call") |
| 370 | defer unregisterSecond() |
| 371 | |
| 372 | res, err := transport.handleElicitation(context.Background(), &mcpsdk.ElicitRequest{Params: &mcpsdk.ElicitParams{Message: "ambiguous"}}) |
| 373 | if err != nil { |
| 374 | t.Fatal(err) |
| 375 | } |
| 376 | if res.Action != mcpinteraction.ActionCancel { |
| 377 | t.Fatalf("ambiguous action = %q, want cancel", res.Action) |
| 378 | } |
| 379 | if len(first.requests()) != 0 || len(second.requests()) != 0 { |
| 380 | t.Fatal("ambiguous legacy elicitation crossed into a call broker") |
| 381 | } |
| 382 | } |
| 383 | |
| 384 | func TestURLModeElicitationReachesBrokerSafely(t *testing.T) { |
| 385 | lifeCtx, cancelLife := context.WithCancel(context.Background()) |
| 386 | transport := &sdkSessionTransport{ |
| 387 | name: "url-mode", spec: Spec{Name: "url-mode", Type: "http"}, profile: HostProfileInteractive, |
| 388 | lifeCtx: lifeCtx, cancel: cancelLife, state: SessionStateConnecting, |
| 389 | } |
| 390 | t.Cleanup(transport.close) |
| 391 | broker := &scriptedBroker{answer: mcpinteraction.Result{Action: mcpinteraction.ActionAccept}} |
| 392 | ctx := mcpinteraction.WithBroker(t.Context(), broker) |
| 393 | |
| 394 | res, err := transport.handleElicitation(ctx, &mcpsdk.ElicitRequest{Params: &mcpsdk.ElicitParams{ |
| 395 | Mode: "url", Message: "finish signup", URL: "https://auth.example.com/consent?state=xyz", |
| 396 | }}) |
| 397 | if err != nil { |
| 398 | t.Fatalf("handleElicitation: %v", err) |
| 399 | } |
| 400 | if res.Action != mcpinteraction.ActionAccept { |
| 401 | t.Fatalf("action = %q", res.Action) |
| 402 | } |
| 403 | reqs := broker.requests() |
| 404 | if len(reqs) != 1 || reqs[0].URL != "https://auth.example.com/consent?state=xyz" || reqs[0].ElicitationID != "" { |
| 405 | t.Fatalf("broker request = %+v", reqs) |
| 406 | } |
| 407 | |
| 408 | // Dangerous URLs are cancelled before any UI sees them. |
| 409 | broker2 := &scriptedBroker{answer: mcpinteraction.Result{Action: mcpinteraction.ActionAccept}} |
| 410 | ctx2 := mcpinteraction.WithBroker(t.Context(), broker2) |
| 411 | res, err = transport.handleElicitation(ctx2, &mcpsdk.ElicitRequest{Params: &mcpsdk.ElicitParams{ |
| 412 | Mode: "url", Message: "bad", URL: "javascript:alert(1)", |
| 413 | }}) |
| 414 | if err != nil { |
| 415 | t.Fatalf("handleElicitation: %v", err) |
| 416 | } |
| 417 | if res.Action != mcpinteraction.ActionCancel { |
| 418 | t.Fatalf("dangerous URL action = %q, want cancel", res.Action) |
| 419 | } |
| 420 | if len(broker2.requests()) != 0 { |
| 421 | t.Fatal("dangerous URL reached the broker") |
| 422 | } |
| 423 | } |
| 424 |