| 1 | package plugin |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "context" |
| 6 | "encoding/json" |
| 7 | "fmt" |
| 8 | "io" |
| 9 | "testing" |
| 10 | "time" |
| 11 | |
| 12 | "reasonix/internal/tool" |
| 13 | ) |
| 14 | |
| 15 | type discardWriteCloser struct{} |
| 16 | |
| 17 | func (discardWriteCloser) Write(p []byte) (int, error) { return len(p), nil } |
| 18 | func (discardWriteCloser) Close() error { return nil } |
| 19 | |
| 20 | // TestStdioCallReturnsOnContextCancel pins that a stdio call unblocks when its |
| 21 | // context is cancelled even though the server never replies. The stdio child is |
| 22 | // bound to the session, not the turn, so without this a hung server would hang a |
| 23 | // cancelled turn forever. No reader goroutine runs here, so the reply never |
| 24 | // arrives — only ctx cancellation can return the call. |
| 25 | func TestStdioCallReturnsOnContextCancel(t *testing.T) { |
| 26 | tr := &stdioTransport{ |
| 27 | name: "hung", |
| 28 | stdin: discardWriteCloser{}, |
| 29 | pending: map[int]chan rpcResponse{}, |
| 30 | } |
| 31 | |
| 32 | ctx, cancel := context.WithCancel(context.Background()) |
| 33 | done := make(chan error, 1) |
| 34 | go func() { |
| 35 | _, err := tr.call(ctx, "tools/call", map[string]any{}) |
| 36 | done <- err |
| 37 | }() |
| 38 | |
| 39 | time.Sleep(100 * time.Millisecond) // let the call park in its select |
| 40 | cancel() |
| 41 | select { |
| 42 | case err := <-done: |
| 43 | if err == nil { |
| 44 | t.Fatal("cancelled call returned nil error") |
| 45 | } |
| 46 | case <-time.After(2 * time.Second): |
| 47 | t.Fatal("stdio call did not return within 2s of ctx cancel — a hung server hangs the turn") |
| 48 | } |
| 49 | } |
| 50 | |
| 51 | func TestStdioCallRespectsExistingDeadline(t *testing.T) { |
| 52 | tr := &stdioTransport{ |
| 53 | name: "server", |
| 54 | stdin: discardWriteCloser{}, |
| 55 | pending: map[int]chan rpcResponse{}, |
| 56 | } |
| 57 | |
| 58 | ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) |
| 59 | defer cancel() |
| 60 | done := make(chan error, 1) |
| 61 | go func() { |
| 62 | _, err := tr.call(ctx, "tools/call", map[string]any{}) |
| 63 | done <- err |
| 64 | }() |
| 65 | |
| 66 | select { |
| 67 | case err := <-done: |
| 68 | if err == nil { |
| 69 | t.Fatal("timed-out call returned nil error") |
| 70 | } |
| 71 | case <-time.After(1 * time.Second): |
| 72 | t.Fatal("stdio call did not return within caller deadline") |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | func TestStdioCallCancelReturnsContextCanceled(t *testing.T) { |
| 77 | tr := &stdioTransport{ |
| 78 | name: "slow-server", |
| 79 | stdin: discardWriteCloser{}, |
| 80 | pending: map[int]chan rpcResponse{}, |
| 81 | } |
| 82 | |
| 83 | ctx, cancel := context.WithCancel(context.Background()) |
| 84 | done := make(chan error, 1) |
| 85 | go func() { |
| 86 | _, err := tr.call(ctx, "tools/call", map[string]any{}) |
| 87 | done <- err |
| 88 | }() |
| 89 | |
| 90 | time.Sleep(200 * time.Millisecond) |
| 91 | cancel() |
| 92 | |
| 93 | select { |
| 94 | case err := <-done: |
| 95 | if err == nil { |
| 96 | t.Fatal("cancelled call returned nil error") |
| 97 | } |
| 98 | if err != context.Canceled { |
| 99 | t.Fatalf("expected context.Canceled, got: %v", err) |
| 100 | } |
| 101 | case <-time.After(2 * time.Second): |
| 102 | t.Fatal("stdio call did not return within 2s of cancel") |
| 103 | } |
| 104 | } |
| 105 | |
| 106 | // Some MCP servers send capability-change notifications and a ping while the |
| 107 | // initialize call is in flight. The server must receive its ping response |
| 108 | // before it can finish the handshake; dropping server requests deadlocks both |
| 109 | // sides even though notifications themselves are harmless. |
| 110 | func TestStdioInitializeHandlesNotificationsAndServerPing(t *testing.T) { |
| 111 | workspaceRoot := t.TempDir() |
| 112 | serverReads, clientWrites := io.Pipe() |
| 113 | clientReads, serverWrites := io.Pipe() |
| 114 | t.Cleanup(func() { |
| 115 | _ = clientWrites.Close() |
| 116 | _ = serverReads.Close() |
| 117 | _ = serverWrites.Close() |
| 118 | _ = clientReads.Close() |
| 119 | }) |
| 120 | |
| 121 | tr := &stdioTransport{ |
| 122 | name: "matlab", |
| 123 | roots: mcpRoots(workspaceRoot), |
| 124 | stdin: clientWrites, |
| 125 | stdout: bufio.NewReader(clientReads), |
| 126 | stderr: &tailBuffer{limit: 1024}, |
| 127 | pending: map[int]chan rpcResponse{}, |
| 128 | } |
| 129 | go tr.readLoop() |
| 130 | |
| 131 | serverDone := make(chan error, 1) |
| 132 | go func() { |
| 133 | dec := json.NewDecoder(serverReads) |
| 134 | enc := json.NewEncoder(serverWrites) |
| 135 | var initialize struct { |
| 136 | ID int `json:"id"` |
| 137 | Method string `json:"method"` |
| 138 | Params struct { |
| 139 | Capabilities map[string]json.RawMessage `json:"capabilities"` |
| 140 | } `json:"params"` |
| 141 | } |
| 142 | if err := dec.Decode(&initialize); err != nil { |
| 143 | serverDone <- fmt.Errorf("decode initialize: %w", err) |
| 144 | return |
| 145 | } |
| 146 | if initialize.Method != "initialize" { |
| 147 | serverDone <- fmt.Errorf("first method = %q, want initialize", initialize.Method) |
| 148 | return |
| 149 | } |
| 150 | if _, ok := initialize.Params.Capabilities["roots"]; !ok { |
| 151 | serverDone <- fmt.Errorf("initialize capabilities = %v, want roots", initialize.Params.Capabilities) |
| 152 | return |
| 153 | } |
| 154 | for _, method := range []string{"notifications/tools/list_changed", "notifications/resources/list_changed"} { |
| 155 | if err := enc.Encode(map[string]any{"jsonrpc": "2.0", "method": method}); err != nil { |
| 156 | serverDone <- fmt.Errorf("encode %s: %w", method, err) |
| 157 | return |
| 158 | } |
| 159 | } |
| 160 | if err := enc.Encode(map[string]any{"jsonrpc": "2.0", "id": "server-roots", "method": "roots/list"}); err != nil { |
| 161 | serverDone <- fmt.Errorf("encode roots/list: %w", err) |
| 162 | return |
| 163 | } |
| 164 | var rootsResponse struct { |
| 165 | ID string `json:"id"` |
| 166 | Result struct { |
| 167 | Roots []mcpRoot `json:"roots"` |
| 168 | } `json:"result"` |
| 169 | } |
| 170 | if err := dec.Decode(&rootsResponse); err != nil { |
| 171 | serverDone <- fmt.Errorf("decode roots/list response: %w", err) |
| 172 | return |
| 173 | } |
| 174 | wantRoots := mcpRoots(workspaceRoot) |
| 175 | if rootsResponse.ID != "server-roots" || len(rootsResponse.Result.Roots) != 1 || rootsResponse.Result.Roots[0] != wantRoots[0] { |
| 176 | serverDone <- fmt.Errorf("roots/list response = %+v, want %+v", rootsResponse, wantRoots) |
| 177 | return |
| 178 | } |
| 179 | if err := enc.Encode(map[string]any{"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}); err != nil { |
| 180 | serverDone <- fmt.Errorf("encode ping: %w", err) |
| 181 | return |
| 182 | } |
| 183 | var pingResponse struct { |
| 184 | ID string `json:"id"` |
| 185 | Result map[string]any `json:"result"` |
| 186 | } |
| 187 | if err := dec.Decode(&pingResponse); err != nil { |
| 188 | serverDone <- fmt.Errorf("decode ping response: %w", err) |
| 189 | return |
| 190 | } |
| 191 | if pingResponse.ID != "server-ping" || pingResponse.Result == nil { |
| 192 | serverDone <- fmt.Errorf("ping response = %+v", pingResponse) |
| 193 | return |
| 194 | } |
| 195 | if err := enc.Encode(map[string]any{ |
| 196 | "jsonrpc": "2.0", |
| 197 | "id": initialize.ID, |
| 198 | "result": map[string]any{ |
| 199 | "protocolVersion": protocolVersion, |
| 200 | "serverInfo": map[string]any{"name": "matlab", "version": "0.11.2"}, |
| 201 | "capabilities": map[string]any{}, |
| 202 | }, |
| 203 | }); err != nil { |
| 204 | serverDone <- fmt.Errorf("encode initialize response: %w", err) |
| 205 | return |
| 206 | } |
| 207 | var initialized struct { |
| 208 | Method string `json:"method"` |
| 209 | } |
| 210 | if err := dec.Decode(&initialized); err != nil { |
| 211 | serverDone <- fmt.Errorf("decode initialized notification: %w", err) |
| 212 | return |
| 213 | } |
| 214 | if initialized.Method != "notifications/initialized" { |
| 215 | serverDone <- fmt.Errorf("final method = %q, want notifications/initialized", initialized.Method) |
| 216 | return |
| 217 | } |
| 218 | serverDone <- nil |
| 219 | }() |
| 220 | |
| 221 | ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) |
| 222 | defer cancel() |
| 223 | client := &Client{name: "matlab", t: tr, spec: Spec{WorkspaceRoot: workspaceRoot}} |
| 224 | if err := client.initialize(ctx); err != nil { |
| 225 | t.Fatalf("initialize with server notifications and ping: %v", err) |
| 226 | } |
| 227 | select { |
| 228 | case err := <-serverDone: |
| 229 | if err != nil { |
| 230 | t.Fatal(err) |
| 231 | } |
| 232 | case <-ctx.Done(): |
| 233 | t.Fatal("server did not complete the MCP initialization handshake") |
| 234 | } |
| 235 | } |
| 236 | |
| 237 | func TestStdioToolCallRoutesProgressNotification(t *testing.T) { |
| 238 | serverReads, clientWrites := io.Pipe() |
| 239 | clientReads, serverWrites := io.Pipe() |
| 240 | t.Cleanup(func() { |
| 241 | _ = clientWrites.Close() |
| 242 | _ = serverReads.Close() |
| 243 | _ = serverWrites.Close() |
| 244 | _ = clientReads.Close() |
| 245 | }) |
| 246 | |
| 247 | tr := &stdioTransport{ |
| 248 | name: "worker", |
| 249 | stdin: clientWrites, |
| 250 | stdout: bufio.NewReader(clientReads), |
| 251 | stderr: &tailBuffer{limit: 1024}, |
| 252 | pending: map[int]chan rpcResponse{}, |
| 253 | } |
| 254 | go tr.readLoop() |
| 255 | |
| 256 | serverDone := make(chan error, 1) |
| 257 | go func() { |
| 258 | dec := json.NewDecoder(serverReads) |
| 259 | enc := json.NewEncoder(serverWrites) |
| 260 | var request struct { |
| 261 | ID int `json:"id"` |
| 262 | Method string `json:"method"` |
| 263 | Params struct { |
| 264 | Meta map[string]any `json:"_meta"` |
| 265 | } `json:"params"` |
| 266 | } |
| 267 | if err := dec.Decode(&request); err != nil { |
| 268 | serverDone <- err |
| 269 | return |
| 270 | } |
| 271 | token, _ := request.Params.Meta["progressToken"].(string) |
| 272 | if request.Method != "tools/call" || token == "" { |
| 273 | serverDone <- fmt.Errorf("tools/call request = %+v, want progressToken", request) |
| 274 | return |
| 275 | } |
| 276 | if err := enc.Encode(map[string]any{ |
| 277 | "jsonrpc": "2.0", |
| 278 | "method": "notifications/progress", |
| 279 | "params": map[string]any{ |
| 280 | "progressToken": token, |
| 281 | "progress": 2, |
| 282 | "total": 5, |
| 283 | "message": "Indexing", |
| 284 | }, |
| 285 | }); err != nil { |
| 286 | serverDone <- err |
| 287 | return |
| 288 | } |
| 289 | if err := enc.Encode(map[string]any{"jsonrpc": "2.0", "id": request.ID, "result": map[string]any{"content": []any{}}}); err != nil { |
| 290 | serverDone <- err |
| 291 | return |
| 292 | } |
| 293 | serverDone <- nil |
| 294 | }() |
| 295 | |
| 296 | progress := make(chan string, 1) |
| 297 | ctx := tool.WithProgress(context.Background(), func(chunk string) { progress <- chunk }) |
| 298 | client := &Client{name: "worker", t: tr} |
| 299 | if _, err := client.call(ctx, "tools/call", map[string]any{"name": "index", "arguments": map[string]any{}}); err != nil { |
| 300 | t.Fatalf("tools/call: %v", err) |
| 301 | } |
| 302 | select { |
| 303 | case got := <-progress: |
| 304 | if got != "Indexing (2/5)\n" { |
| 305 | t.Fatalf("progress = %q", got) |
| 306 | } |
| 307 | case <-time.After(time.Second): |
| 308 | t.Fatal("progress notification was not routed") |
| 309 | } |
| 310 | if err := <-serverDone; err != nil { |
| 311 | t.Fatal(err) |
| 312 | } |
| 313 | } |
| 314 | |
| 315 | // readLoop is the only goroutine draining stdout, so it must never block on |
| 316 | // the shared stdin pipe: with both pipe buffers full, waiting on writeMu would |
| 317 | // deadlock against a client call whose own stdin write is jammed. Replies to |
| 318 | // server requests therefore go through a bounded queue that drops on overflow. |
| 319 | func TestStdioReadLoopStaysLiveWhenReplyWriterIsBlocked(t *testing.T) { |
| 320 | stdinReads, stdinWrites := io.Pipe() // nobody reads: reply writes block forever |
| 321 | stdoutReads, stdoutWrites := io.Pipe() |
| 322 | t.Cleanup(func() { |
| 323 | _ = stdinReads.Close() |
| 324 | _ = stdinWrites.Close() |
| 325 | _ = stdoutReads.Close() |
| 326 | _ = stdoutWrites.Close() |
| 327 | }) |
| 328 | |
| 329 | tr := &stdioTransport{ |
| 330 | name: "jammed", |
| 331 | stdin: stdinWrites, |
| 332 | stdout: bufio.NewReader(stdoutReads), |
| 333 | stderr: &tailBuffer{limit: 1024}, |
| 334 | pending: map[int]chan rpcResponse{}, |
| 335 | } |
| 336 | waiting := make(chan rpcResponse, 1) |
| 337 | tr.pending[7] = waiting |
| 338 | go tr.readLoop() |
| 339 | |
| 340 | // Flood well past the reply queue bound while the reply writer is stuck in |
| 341 | // its first stdin write; overflow must drop, not block readLoop. The writes |
| 342 | // run off the test goroutine so a deadlocked readLoop fails the timeout |
| 343 | // below instead of hanging the whole package; Cleanup unblocks the writer. |
| 344 | go func() { |
| 345 | for i := 0; i < 2*stdioReplyQueueBound; i++ { |
| 346 | line := fmt.Sprintf(`{"jsonrpc":"2.0","id":"srv-%d","method":"ping"}`+"\n", i) |
| 347 | if _, err := io.WriteString(stdoutWrites, line); err != nil { |
| 348 | return |
| 349 | } |
| 350 | } |
| 351 | _, _ = io.WriteString(stdoutWrites, `{"jsonrpc":"2.0","id":7,"result":{}}`+"\n") |
| 352 | }() |
| 353 | |
| 354 | select { |
| 355 | case resp := <-waiting: |
| 356 | if resp.ID != 7 { |
| 357 | t.Fatalf("routed response id = %d, want 7", resp.ID) |
| 358 | } |
| 359 | case <-time.After(2 * time.Second): |
| 360 | t.Fatal("readLoop stopped routing responses while the reply writer was blocked") |
| 361 | } |
| 362 | } |
| 363 |