| 1 | package providerext |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "fmt" |
| 7 | "strings" |
| 8 | "sync" |
| 9 | "testing" |
| 10 | "time" |
| 11 | |
| 12 | "reasonix/internal/extension/protocol" |
| 13 | "reasonix/internal/provider" |
| 14 | ) |
| 15 | |
| 16 | func TestStreamDeliversOutOfOrderChunksInOrder(t *testing.T) { |
| 17 | fc := newFakeClient("demo", demoDescriptor()) |
| 18 | r := testResolver(t, baseCatalog(), nil, fc) |
| 19 | out, id := openTestStream(t, r, fc, nil) |
| 20 | |
| 21 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 2, Chunk: textChunk("b")}) |
| 22 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 3, Chunk: textChunk("c")}) |
| 23 | select { |
| 24 | case chunk := <-out: |
| 25 | t.Fatalf("received chunk %q before the missing seq 1 arrived", chunk.Text) |
| 26 | case <-time.After(50 * time.Millisecond): |
| 27 | } |
| 28 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: textChunk("a")}) |
| 29 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 3}) |
| 30 | |
| 31 | chunks := collectChunks(t, out) |
| 32 | if got := texts(chunks); fmt.Sprint(got) != "[a b c]" { |
| 33 | t.Fatalf("delivered texts = %v, want in-order [a b c]", got) |
| 34 | } |
| 35 | for _, chunk := range chunks { |
| 36 | if chunk.Type != provider.ChunkText { |
| 37 | t.Fatalf("chunk type = %v", chunk.Type) |
| 38 | } |
| 39 | } |
| 40 | } |
| 41 | |
| 42 | func TestStreamDropsDuplicateAndStaleChunks(t *testing.T) { |
| 43 | fc := newFakeClient("demo", demoDescriptor()) |
| 44 | r := testResolver(t, baseCatalog(), nil, fc) |
| 45 | out, id := openTestStream(t, r, fc, nil) |
| 46 | |
| 47 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: textChunk("first")}) |
| 48 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: textChunk("duplicate")}) |
| 49 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 2}) |
| 50 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 2, Chunk: textChunk("second")}) |
| 51 | // A stale replay of seq 1 after delivery must not resurrect it. |
| 52 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: textChunk("stale")}) |
| 53 | |
| 54 | chunks := collectChunks(t, out) |
| 55 | if got := texts(chunks); fmt.Sprint(got) != "[first second]" { |
| 56 | t.Fatalf("delivered texts = %v, want [first second]", got) |
| 57 | } |
| 58 | } |
| 59 | |
| 60 | func TestStreamCleanEndClosesChannel(t *testing.T) { |
| 61 | fc := newFakeClient("demo", demoDescriptor()) |
| 62 | r := testResolver(t, baseCatalog(), nil, fc) |
| 63 | out, id := openTestStream(t, r, fc, nil) |
| 64 | |
| 65 | // The zero-chunk sentinel: end with LastSeq 0 closes immediately. |
| 66 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 0}) |
| 67 | chunks := collectChunks(t, out) |
| 68 | if len(chunks) != 0 { |
| 69 | t.Fatalf("chunks = %v, want none", chunks) |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | func TestStreamIdleWatchdogRefreshesOnProviderChunk(t *testing.T) { |
| 74 | fc := newFakeClient("demo", demoDescriptor()) |
| 75 | r := testResolver(t, baseCatalog(), nil, fc) |
| 76 | r.idleTimeout = 80 * time.Millisecond |
| 77 | out, id := openTestStream(t, r, fc, nil) |
| 78 | |
| 79 | time.Sleep(50 * time.Millisecond) |
| 80 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: textChunk("progress")}) |
| 81 | time.Sleep(50 * time.Millisecond) |
| 82 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 1}) |
| 83 | |
| 84 | chunks := collectChunks(t, out) |
| 85 | if got := texts(chunks); fmt.Sprint(got) != "[progress]" { |
| 86 | t.Fatalf("chunks = %v, want progress without idle cancellation", got) |
| 87 | } |
| 88 | } |
| 89 | |
| 90 | func TestStreamIdleWatchdogCancelsSilentExtension(t *testing.T) { |
| 91 | fc := newFakeClient("demo", demoDescriptor()) |
| 92 | r := testResolver(t, baseCatalog(), nil, fc) |
| 93 | r.idleTimeout = 30 * time.Millisecond |
| 94 | out, _ := openTestStream(t, r, fc, nil) |
| 95 | |
| 96 | chunks := collectChunks(t, out) |
| 97 | if len(chunks) != 1 || chunks[0].Type != provider.ChunkError || chunks[0].Err == nil || !strings.Contains(chunks[0].Err.Error(), "stalled") { |
| 98 | t.Fatalf("silent stream chunks = %+v, want stalled interruption", chunks) |
| 99 | } |
| 100 | } |
| 101 | |
| 102 | func TestStreamMissingChunkAtEndInterrupts(t *testing.T) { |
| 103 | fc := newFakeClient("demo", demoDescriptor()) |
| 104 | r := testResolver(t, baseCatalog(), nil, fc) |
| 105 | out, id := openTestStream(t, r, fc, nil) |
| 106 | |
| 107 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: textChunk("a")}) |
| 108 | // seq 2 never arrives; the frozen boundary demands it. |
| 109 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 3}) |
| 110 | |
| 111 | chunks := collectChunks(t, out) |
| 112 | if len(chunks) != 2 { |
| 113 | t.Fatalf("chunks = %v, want the delivered text plus the gap error", texts(chunks)) |
| 114 | } |
| 115 | terminal := chunks[1] |
| 116 | if terminal.Type != provider.ChunkError || !provider.IsStreamInterrupted(terminal.Err) { |
| 117 | t.Fatalf("terminal = %+v, want interrupted ChunkError", terminal) |
| 118 | } |
| 119 | if !strings.Contains(terminal.Err.Error(), "missing chunk 2 of 3") { |
| 120 | t.Fatalf("gap error = %q, want the missing seq named", terminal.Err) |
| 121 | } |
| 122 | } |
| 123 | |
| 124 | func TestStreamLateChunksAfterEndDropped(t *testing.T) { |
| 125 | fc := newFakeClient("demo", demoDescriptor()) |
| 126 | r := testResolver(t, baseCatalog(), nil, fc) |
| 127 | out, id := openTestStream(t, r, fc, nil) |
| 128 | |
| 129 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: textChunk("a")}) |
| 130 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 1}) |
| 131 | chunks := collectChunks(t, out) |
| 132 | if got := texts(chunks); fmt.Sprint(got) != "[a]" { |
| 133 | t.Fatalf("chunks = %v", got) |
| 134 | } |
| 135 | |
| 136 | // Late traffic for a completed stream is dropped, never resurrected: the |
| 137 | // channel stays closed and nothing new arrives. |
| 138 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 2, Chunk: textChunk("late")}) |
| 139 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 2}) |
| 140 | select { |
| 141 | case chunk, ok := <-out: |
| 142 | if ok { |
| 143 | t.Fatalf("late delivery %q after the stream closed", chunk.Text) |
| 144 | } |
| 145 | case <-time.After(50 * time.Millisecond): |
| 146 | t.Fatal("stream channel should already be closed") |
| 147 | } |
| 148 | } |
| 149 | |
| 150 | func TestStreamRejectsBufferedChunkBeyondFrozenEnd(t *testing.T) { |
| 151 | fc := newFakeClient("demo", demoDescriptor()) |
| 152 | r := testResolver(t, baseCatalog(), nil, fc) |
| 153 | out, id := openTestStream(t, r, fc, nil) |
| 154 | |
| 155 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 2, Chunk: textChunk("beyond")}) |
| 156 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 1}) |
| 157 | |
| 158 | chunks := collectChunks(t, out) |
| 159 | if len(chunks) != 1 || chunks[0].Type != provider.ChunkError || !provider.IsStreamInterrupted(chunks[0].Err) { |
| 160 | t.Fatalf("chunks = %+v, want interrupted protocol error", chunks) |
| 161 | } |
| 162 | if !strings.Contains(chunks[0].Err.Error(), "exceeds frozen LastSeq 1") { |
| 163 | t.Fatalf("error = %q, want frozen boundary detail", chunks[0].Err) |
| 164 | } |
| 165 | } |
| 166 | |
| 167 | func TestStreamRejectsLateChunkBeyondFrozenEnd(t *testing.T) { |
| 168 | fc := newFakeClient("demo", demoDescriptor()) |
| 169 | r := testResolver(t, baseCatalog(), nil, fc) |
| 170 | out, id := openTestStream(t, r, fc, nil) |
| 171 | |
| 172 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 2}) |
| 173 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 3, Chunk: textChunk("late")}) |
| 174 | |
| 175 | chunks := collectChunks(t, out) |
| 176 | if len(chunks) != 1 || chunks[0].Type != provider.ChunkError || !provider.IsStreamInterrupted(chunks[0].Err) { |
| 177 | t.Fatalf("chunks = %+v, want interrupted protocol error", chunks) |
| 178 | } |
| 179 | } |
| 180 | |
| 181 | func TestStreamRejectsConflictingDuplicateEnd(t *testing.T) { |
| 182 | fc := newFakeClient("demo", demoDescriptor()) |
| 183 | r := testResolver(t, baseCatalog(), nil, fc) |
| 184 | out, id := openTestStream(t, r, fc, nil) |
| 185 | |
| 186 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 2}) |
| 187 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 3}) |
| 188 | |
| 189 | chunks := collectChunks(t, out) |
| 190 | if len(chunks) != 1 || chunks[0].Type != provider.ChunkError || !provider.IsStreamInterrupted(chunks[0].Err) { |
| 191 | t.Fatalf("chunks = %+v, want interrupted protocol error", chunks) |
| 192 | } |
| 193 | } |
| 194 | |
| 195 | func TestStreamCancelSendsCancelAndCloses(t *testing.T) { |
| 196 | fc := newFakeClient("demo", demoDescriptor()) |
| 197 | r := testResolver(t, baseCatalog(), nil, fc) |
| 198 | |
| 199 | p, err := r.Resolve(provider.Selection{Ref: "plugin/demo/fake/x"}) |
| 200 | if err != nil { |
| 201 | t.Fatalf("Resolve: %v", err) |
| 202 | } |
| 203 | ctx, cancel := context.WithCancel(context.Background()) |
| 204 | out, err := p.Stream(ctx, provider.Request{Messages: []provider.Message{{Role: provider.RoleUser}}}) |
| 205 | if err != nil { |
| 206 | t.Fatalf("Stream: %v", err) |
| 207 | } |
| 208 | id := fc.openedParams(t).StreamID |
| 209 | |
| 210 | cancel() |
| 211 | fc.waitCancel(t, id) |
| 212 | chunks := collectChunks(t, out) |
| 213 | // Cancellation aborts delivery (the consumer is gone): any error chunk |
| 214 | // that does beat the abort must be the interruption, never a hard failure. |
| 215 | for _, chunk := range chunks { |
| 216 | if chunk.Type == provider.ChunkError && !provider.IsStreamInterrupted(chunk.Err) { |
| 217 | t.Fatalf("post-cancel chunk = %+v, want interruption only", chunk) |
| 218 | } |
| 219 | } |
| 220 | } |
| 221 | |
| 222 | func TestStreamErrorChunkIsDefensivelyRedacted(t *testing.T) { |
| 223 | fc := newFakeClient("demo", demoDescriptor()) |
| 224 | r := testResolver(t, baseCatalog(), nil, fc) |
| 225 | out, id := openTestStream(t, r, fc, nil) |
| 226 | const secret = "sk-abcdef1234567890SECRETKEY" |
| 227 | |
| 228 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: protocol.ProviderChunk{ |
| 229 | Type: protocol.ChunkError, |
| 230 | Error: &protocol.ProviderError{Code: protocol.ProviderFailed, Message: "provider rejected api_key=" + secret}, |
| 231 | }}) |
| 232 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 1}) |
| 233 | |
| 234 | chunks := collectChunks(t, out) |
| 235 | if len(chunks) != 1 || chunks[0].Type != provider.ChunkError { |
| 236 | t.Fatalf("chunks = %+v", chunks) |
| 237 | } |
| 238 | if chunks[0].Err == nil || strings.Contains(chunks[0].Err.Error(), secret) { |
| 239 | t.Fatalf("error leaked credential: %v", chunks[0].Err) |
| 240 | } |
| 241 | if !strings.Contains(chunks[0].Err.Error(), "provider rejected api_key=") { |
| 242 | t.Fatalf("error lost diagnostic context: %v", chunks[0].Err) |
| 243 | } |
| 244 | if provider.IsStreamInterrupted(chunks[0].Err) { |
| 245 | t.Fatal("provider_failed mapped to an interruption") |
| 246 | } |
| 247 | } |
| 248 | |
| 249 | func TestStreamInterruptedErrorChunkMapsToStreamInterrupted(t *testing.T) { |
| 250 | fc := newFakeClient("demo", demoDescriptor()) |
| 251 | r := testResolver(t, baseCatalog(), nil, fc) |
| 252 | out, id := openTestStream(t, r, fc, nil) |
| 253 | |
| 254 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: protocol.ProviderChunk{ |
| 255 | Type: protocol.ChunkError, |
| 256 | Error: &protocol.ProviderError{Code: protocol.ProviderInterrupted, Message: "The extension provider stream was interrupted."}, |
| 257 | }}) |
| 258 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 1}) |
| 259 | |
| 260 | chunks := collectChunks(t, out) |
| 261 | if len(chunks) != 1 || !provider.IsStreamInterrupted(chunks[0].Err) { |
| 262 | t.Fatalf("chunks = %+v, want StreamInterruptedError", chunks) |
| 263 | } |
| 264 | } |
| 265 | |
| 266 | func TestStreamEndErrorBecomesTerminalChunkError(t *testing.T) { |
| 267 | fc := newFakeClient("demo", demoDescriptor()) |
| 268 | r := testResolver(t, baseCatalog(), nil, fc) |
| 269 | out, id := openTestStream(t, r, fc, nil) |
| 270 | const secret = "sk-abcdef1234567890SECRETKEY" |
| 271 | |
| 272 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: textChunk("partial")}) |
| 273 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 1, Error: "provider rejected token=" + secret}) |
| 274 | |
| 275 | chunks := collectChunks(t, out) |
| 276 | if len(chunks) != 2 { |
| 277 | t.Fatalf("chunks = %v", texts(chunks)) |
| 278 | } |
| 279 | terminal := chunks[1] |
| 280 | if terminal.Type != provider.ChunkError || terminal.Err == nil || strings.Contains(terminal.Err.Error(), secret) { |
| 281 | t.Fatalf("terminal = %+v, want the host-redacted end error", terminal) |
| 282 | } |
| 283 | if !strings.Contains(terminal.Err.Error(), "provider rejected token=") { |
| 284 | t.Fatalf("terminal error lost diagnostic context: %q", terminal.Err) |
| 285 | } |
| 286 | if provider.IsStreamInterrupted(terminal.Err) { |
| 287 | t.Fatal("a clean failure must not read as an interruption") |
| 288 | } |
| 289 | } |
| 290 | |
| 291 | func TestStreamEndInterruptedBecomesStreamInterrupted(t *testing.T) { |
| 292 | fc := newFakeClient("demo", demoDescriptor()) |
| 293 | r := testResolver(t, baseCatalog(), nil, fc) |
| 294 | out, id := openTestStream(t, r, fc, nil) |
| 295 | |
| 296 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 0, Interrupted: true}) |
| 297 | chunks := collectChunks(t, out) |
| 298 | if len(chunks) != 1 || !provider.IsStreamInterrupted(chunks[0].Err) { |
| 299 | t.Fatalf("chunks = %+v, want StreamInterruptedError", chunks) |
| 300 | } |
| 301 | } |
| 302 | |
| 303 | func TestStreamChunkTypesRoundTripThroughDTO(t *testing.T) { |
| 304 | fc := newFakeClient("demo", demoDescriptor()) |
| 305 | r := testResolver(t, baseCatalog(), nil, fc) |
| 306 | out, id := openTestStream(t, r, fc, nil) |
| 307 | |
| 308 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: protocol.ProviderChunk{ |
| 309 | Type: protocol.ChunkReasoning, Text: "thinking", Signature: "sig-123", |
| 310 | }}) |
| 311 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 2, Chunk: protocol.ProviderChunk{ |
| 312 | Type: protocol.ChunkToolCallStart, ToolCall: &protocol.ProviderToolCall{ID: "call-1", Name: "bash"}, |
| 313 | }}) |
| 314 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 3, Chunk: protocol.ProviderChunk{ |
| 315 | Type: protocol.ChunkToolCallDelta, ToolCall: &protocol.ProviderToolCall{ID: "call-1", Name: "bash"}, ArgChars: 42, |
| 316 | }}) |
| 317 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 4, Chunk: protocol.ProviderChunk{ |
| 318 | Type: protocol.ChunkToolCall, |
| 319 | ToolCall: &protocol.ProviderToolCall{ |
| 320 | ID: "call-1", Name: "bash", Arguments: `{"cmd":"ls"}`, ThoughtSignature: "gemini-sig", |
| 321 | }, |
| 322 | }}) |
| 323 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 5, Chunk: protocol.ProviderChunk{ |
| 324 | Type: protocol.ChunkUsage, |
| 325 | Usage: &protocol.ProviderUsage{ |
| 326 | PromptTokens: 10, CompletionTokens: 20, TotalTokens: 30, |
| 327 | CacheHitTokens: 4, CacheMissTokens: 6, ReasoningTokens: 8, FinishReason: "tool_calls", |
| 328 | }, |
| 329 | }}) |
| 330 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: id, LastSeq: 5}) |
| 331 | |
| 332 | chunks := collectChunks(t, out) |
| 333 | if len(chunks) != 5 { |
| 334 | t.Fatalf("chunks = %d, want 5", len(chunks)) |
| 335 | } |
| 336 | if chunks[0].Type != provider.ChunkReasoning || chunks[0].Text != "thinking" || chunks[0].Signature != "sig-123" { |
| 337 | t.Fatalf("reasoning chunk = %+v", chunks[0]) |
| 338 | } |
| 339 | if chunks[1].Type != provider.ChunkToolCallStart || chunks[1].ToolCall == nil || chunks[1].ToolCall.ID != "call-1" { |
| 340 | t.Fatalf("tool-call-start chunk = %+v", chunks[1]) |
| 341 | } |
| 342 | if chunks[2].Type != provider.ChunkToolCallArgsDelta || chunks[2].ArgChars != 42 { |
| 343 | t.Fatalf("args-delta chunk = %+v", chunks[2]) |
| 344 | } |
| 345 | if chunks[3].Type != provider.ChunkToolCall || chunks[3].ToolCall.Arguments != `{"cmd":"ls"}` || chunks[3].ToolCall.ThoughtSignature != "gemini-sig" { |
| 346 | t.Fatalf("tool-call chunk = %+v", chunks[3]) |
| 347 | } |
| 348 | usage := chunks[4].Usage |
| 349 | if chunks[4].Type != provider.ChunkUsage || usage == nil || |
| 350 | usage.PromptTokens != 10 || usage.CompletionTokens != 20 || usage.TotalTokens != 30 || |
| 351 | usage.CacheHitTokens != 4 || usage.CacheMissTokens != 6 || usage.ReasoningTokens != 8 || |
| 352 | usage.FinishReason != "tool_calls" { |
| 353 | t.Fatalf("usage chunk = %+v", chunks[4]) |
| 354 | } |
| 355 | } |
| 356 | |
| 357 | func TestStreamDisconnectMidStreamInterrupts(t *testing.T) { |
| 358 | fc := newFakeClient("demo", demoDescriptor()) |
| 359 | r := testResolver(t, baseCatalog(), nil, fc) |
| 360 | out, id := openTestStream(t, r, fc, nil) |
| 361 | |
| 362 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: textChunk("a")}) |
| 363 | fc.kill() // mid-stream crash: no end, no more chunks, ever |
| 364 | |
| 365 | chunks := collectChunks(t, out) |
| 366 | if len(chunks) != 2 { |
| 367 | t.Fatalf("chunks = %v, want delivered text plus the interruption", texts(chunks)) |
| 368 | } |
| 369 | terminal := chunks[1] |
| 370 | if terminal.Type != provider.ChunkError || !provider.IsStreamInterrupted(terminal.Err) { |
| 371 | t.Fatalf("terminal = %+v, want StreamInterruptedError", terminal) |
| 372 | } |
| 373 | if !strings.Contains(terminal.Err.Error(), "demo") { |
| 374 | t.Fatalf("interruption = %q, want the plugin named", terminal.Err) |
| 375 | } |
| 376 | } |
| 377 | |
| 378 | func TestStreamFailsFastAfterCrash(t *testing.T) { |
| 379 | fc := newFakeClient("demo", demoDescriptor()) |
| 380 | r := testResolver(t, baseCatalog(), nil, fc) |
| 381 | fc.kill() |
| 382 | |
| 383 | p, err := r.Resolve(provider.Selection{Ref: "plugin/demo/fake/x"}) |
| 384 | if err != nil { |
| 385 | t.Fatalf("Resolve: %v", err) |
| 386 | } |
| 387 | _, err = p.Stream(context.Background(), provider.Request{}) |
| 388 | if !provider.IsStreamInterrupted(err) { |
| 389 | t.Fatalf("Stream error = %v, want fail-fast StreamInterruptedError", err) |
| 390 | } |
| 391 | if opens := len(fc.opened); opens != 0 { |
| 392 | t.Fatalf("stream opens = %d, want none after the crash", opens) |
| 393 | } |
| 394 | } |
| 395 | |
| 396 | func TestStreamOpenDeclined(t *testing.T) { |
| 397 | fc := newFakeClient("demo", demoDescriptor()) |
| 398 | fc.accept = false |
| 399 | r := testResolver(t, baseCatalog(), nil, fc) |
| 400 | |
| 401 | p, err := r.Resolve(provider.Selection{Ref: "plugin/demo/fake/x"}) |
| 402 | if err != nil { |
| 403 | t.Fatalf("Resolve: %v", err) |
| 404 | } |
| 405 | _, err = p.Stream(context.Background(), provider.Request{}) |
| 406 | if err == nil || !strings.Contains(err.Error(), "declined") { |
| 407 | t.Fatalf("Stream error = %v, want declined", err) |
| 408 | } |
| 409 | } |
| 410 | |
| 411 | func TestStreamOpenInterruptedErrorMapsToStreamInterrupted(t *testing.T) { |
| 412 | fc := newFakeClient("demo", demoDescriptor()) |
| 413 | fc.openErr = &protocol.ProtocolError{Reason: protocol.ErrProviderInterrupted, Message: "extension sidecar demo crashed"} |
| 414 | r := testResolver(t, baseCatalog(), nil, fc) |
| 415 | |
| 416 | p, err := r.Resolve(provider.Selection{Ref: "plugin/demo/fake/x"}) |
| 417 | if err != nil { |
| 418 | t.Fatalf("Resolve: %v", err) |
| 419 | } |
| 420 | _, err = p.Stream(context.Background(), provider.Request{}) |
| 421 | if !provider.IsStreamInterrupted(err) { |
| 422 | t.Fatalf("Stream error = %v, want StreamInterruptedError", err) |
| 423 | } |
| 424 | } |
| 425 | |
| 426 | func TestStreamOpenGenericErrorPassesThrough(t *testing.T) { |
| 427 | fc := newFakeClient("demo", demoDescriptor()) |
| 428 | fc.openErr = errors.New("transport wedged") |
| 429 | r := testResolver(t, baseCatalog(), nil, fc) |
| 430 | |
| 431 | p, err := r.Resolve(provider.Selection{Ref: "plugin/demo/fake/x"}) |
| 432 | if err != nil { |
| 433 | t.Fatalf("Resolve: %v", err) |
| 434 | } |
| 435 | _, err = p.Stream(context.Background(), provider.Request{}) |
| 436 | if err == nil || !strings.Contains(err.Error(), "transport wedged") { |
| 437 | t.Fatalf("Stream error = %v", err) |
| 438 | } |
| 439 | if provider.IsStreamInterrupted(err) { |
| 440 | t.Fatal("generic open failure mapped to an interruption") |
| 441 | } |
| 442 | } |
| 443 | |
| 444 | func TestStreamOpenCarriesRequestEffortAndSeqBase(t *testing.T) { |
| 445 | descriptor := demoDescriptor() |
| 446 | descriptor.Efforts = []string{"low", "high"} |
| 447 | fc := newFakeClient("demo", descriptor) |
| 448 | r := testResolver(t, baseCatalog(), nil, fc) |
| 449 | |
| 450 | effort := "high" |
| 451 | p, err := r.Resolve(provider.Selection{Ref: "plugin/demo/fake/x", Effort: &effort}) |
| 452 | if err != nil { |
| 453 | t.Fatalf("Resolve: %v", err) |
| 454 | } |
| 455 | temperature := 0.5 |
| 456 | out, err := p.Stream(context.Background(), provider.Request{ |
| 457 | Messages: []provider.Message{ |
| 458 | {Role: provider.RoleSystem, Content: "sys"}, |
| 459 | {Role: provider.RoleUser, Content: "hi", Images: []string{"data:image/png;base64,AA=="}}, |
| 460 | {Role: provider.RoleAssistant, Content: "prev", ReasoningContent: "because", ReasoningSignature: "rs"}, |
| 461 | }, |
| 462 | Tools: []provider.ToolSchema{{Name: "bash", Description: "run", Parameters: []byte(`{"type":"object"}`)}}, |
| 463 | Temperature: &temperature, |
| 464 | MaxTokens: 128, |
| 465 | }) |
| 466 | if err != nil { |
| 467 | t.Fatalf("Stream: %v", err) |
| 468 | } |
| 469 | params := fc.openedParams(t) |
| 470 | if params.ProviderRef != "plugin/demo/fake/x" || params.Model != "x" || params.Effort != "high" { |
| 471 | t.Fatalf("open params = %+v", params) |
| 472 | } |
| 473 | if params.SeqBase != 1 { |
| 474 | t.Fatalf("SeqBase = %d, want 1-based chunk numbering", params.SeqBase) |
| 475 | } |
| 476 | if !strings.HasPrefix(params.StreamID, "es_") { |
| 477 | t.Fatalf("StreamID = %q, want the es_ prefix", params.StreamID) |
| 478 | } |
| 479 | req := params.Request |
| 480 | if len(req.Messages) != 3 || len(req.Tools) != 1 { |
| 481 | t.Fatalf("request = %+v", req) |
| 482 | } |
| 483 | if req.Messages[1].Images[0] != "data:image/png;base64,AA==" || req.Messages[2].ReasoningSignature != "rs" { |
| 484 | t.Fatalf("request messages did not convert: %+v", req.Messages) |
| 485 | } |
| 486 | if req.Tools[0].Name != "bash" || string(req.Tools[0].Parameters) != `{"type":"object"}` { |
| 487 | t.Fatalf("request tools did not convert: %+v", req.Tools) |
| 488 | } |
| 489 | if req.Temperature == nil || *req.Temperature != 0.5 || req.MaxTokens != 128 { |
| 490 | t.Fatalf("request scalars = %+v", req) |
| 491 | } |
| 492 | |
| 493 | // Finish the stream cleanly so its watcher cannot outlive the test. |
| 494 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: params.StreamID, LastSeq: 0}) |
| 495 | collectChunks(t, out) |
| 496 | } |
| 497 | |
| 498 | func TestProviderReasoningPoliciesComeFromDescriptor(t *testing.T) { |
| 499 | descriptor := demoDescriptor() |
| 500 | descriptor.ToolCallReasoning = true |
| 501 | descriptor.ReasoningRoundTrip = true |
| 502 | descriptor.WarnOnMissingToolCallReasoning = true |
| 503 | fc := newFakeClient("demo", descriptor) |
| 504 | r := testResolver(t, baseCatalog(), nil, fc) |
| 505 | |
| 506 | p, err := r.Resolve(provider.Selection{Ref: "plugin/demo/fake/x"}) |
| 507 | if err != nil { |
| 508 | t.Fatalf("Resolve: %v", err) |
| 509 | } |
| 510 | if !provider.RequiresToolCallReasoning(p) || !provider.RequiresReasoningRoundTrip(p) || !provider.WarnOnMissingToolCallReasoning(p) { |
| 511 | t.Fatal("descriptor reasoning policies did not propagate") |
| 512 | } |
| 513 | if identity := p.(interface{ MissingToolCallReasoningWarningIdentity() string }).MissingToolCallReasoningWarningIdentity(); !strings.Contains(identity, "demo") || !strings.Contains(identity, "plugin/demo/fake/x") { |
| 514 | t.Fatalf("warning identity = %q", identity) |
| 515 | } |
| 516 | } |
| 517 | |
| 518 | func TestRouteUnknownStreamDropped(t *testing.T) { |
| 519 | r := testResolver(t, baseCatalog(), nil) |
| 520 | // No stream registered: routing must not panic or create state. |
| 521 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: "es_nope", Seq: 1, Chunk: textChunk("x")}) |
| 522 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: "es_nope", LastSeq: 1}) |
| 523 | r.mu.Lock() |
| 524 | registered := len(r.streams) |
| 525 | r.mu.Unlock() |
| 526 | if registered != 0 { |
| 527 | t.Fatalf("unknown routing created %d streams", registered) |
| 528 | } |
| 529 | } |
| 530 | |
| 531 | func TestStreamDeliveryOverflowTerminates(t *testing.T) { |
| 532 | r := testResolver(t, baseCatalog(), nil) |
| 533 | stream := &extensionStream{ |
| 534 | out: make(chan provider.Chunk, 1), |
| 535 | done: make(chan struct{}), |
| 536 | deliveryWake: make(chan struct{}, 1), |
| 537 | nextSeq: 1, |
| 538 | pending: map[int64]provider.Chunk{}, |
| 539 | delivery: make([]provider.Chunk, deliveryQueueLimit-1), |
| 540 | } |
| 541 | r.mu.Lock() |
| 542 | r.streams["overflow"] = stream |
| 543 | stream.pending[1] = provider.Chunk{Type: provider.ChunkText, Text: "overflow"} |
| 544 | r.flushLocked("overflow", stream) |
| 545 | _, stillRegistered := r.streams["overflow"] |
| 546 | final := stream.deliveryFinal |
| 547 | queued := append([]provider.Chunk(nil), stream.delivery...) |
| 548 | r.mu.Unlock() |
| 549 | |
| 550 | if stillRegistered || !final { |
| 551 | t.Fatal("overflowing stream was not terminated") |
| 552 | } |
| 553 | if len(queued) != deliveryQueueLimit || queued[len(queued)-1].Err == nil || |
| 554 | !provider.IsStreamInterrupted(queued[len(queued)-1].Err) { |
| 555 | t.Fatalf("overflow queue = %d chunks, terminal %v", len(queued), queued[len(queued)-1].Err) |
| 556 | } |
| 557 | } |
| 558 | |
| 559 | func TestStreamDisconnectDoesNotBlockOnBackpressure(t *testing.T) { |
| 560 | fc := newFakeClient("demo", demoDescriptor()) |
| 561 | r := testResolver(t, baseCatalog(), nil, fc) |
| 562 | stream := &extensionStream{ |
| 563 | client: fc, |
| 564 | out: make(chan provider.Chunk, 1), |
| 565 | done: make(chan struct{}), |
| 566 | abortDelivery: make(chan struct{}), |
| 567 | deliveryWake: make(chan struct{}, 1), |
| 568 | nextSeq: 1, |
| 569 | pending: map[int64]provider.Chunk{ |
| 570 | 1: {Type: provider.ChunkText, Text: "one"}, |
| 571 | 2: {Type: provider.ChunkText, Text: "two"}, |
| 572 | }, |
| 573 | } |
| 574 | r.mu.Lock() |
| 575 | r.streams["backpressure"] = stream |
| 576 | r.mu.Unlock() |
| 577 | go r.deliverStream(stream) |
| 578 | |
| 579 | r.mu.Lock() |
| 580 | r.flushLocked("backpressure", stream) |
| 581 | r.mu.Unlock() |
| 582 | |
| 583 | deadline := time.Now().Add(time.Second) |
| 584 | for len(stream.out) != 1 && time.Now().Before(deadline) { |
| 585 | time.Sleep(time.Millisecond) |
| 586 | } |
| 587 | if len(stream.out) != 1 { |
| 588 | t.Fatal("stream never filled its output buffer") |
| 589 | } |
| 590 | |
| 591 | fc.kill() |
| 592 | // The watchStream goroutine only exists for streams opened through |
| 593 | // Resolver.open; this hand-built stream finishes the way the broker's |
| 594 | // Detach does, directly. |
| 595 | r.mu.Lock() |
| 596 | r.finishLocked("backpressure", stream, provider.Chunk{Type: provider.ChunkError, Err: &provider.StreamInterruptedError{ |
| 597 | Err: errors.New("extension sidecar demo disconnected"), |
| 598 | }}) |
| 599 | r.mu.Unlock() |
| 600 | var chunks []provider.Chunk |
| 601 | for chunk := range stream.out { |
| 602 | chunks = append(chunks, chunk) |
| 603 | } |
| 604 | // The disconnect finishes the stream without aborting delivery: buffered |
| 605 | // chunks drain ahead of the terminal interruption. |
| 606 | if len(chunks) != 3 || chunks[0].Text != "one" || chunks[1].Text != "two" || |
| 607 | !provider.IsStreamInterrupted(chunks[2].Err) { |
| 608 | t.Fatalf("delivered chunks = %#v, want ordered text followed by interruption", chunks) |
| 609 | } |
| 610 | } |
| 611 | |
| 612 | func TestStreamAbandonedConsumerDoesNotLeakDelivery(t *testing.T) { |
| 613 | r := testResolver(t, baseCatalog(), nil) |
| 614 | stream := &extensionStream{ |
| 615 | out: make(chan provider.Chunk, 1), |
| 616 | abortDelivery: make(chan struct{}), |
| 617 | deliveryWake: make(chan struct{}, 1), |
| 618 | delivery: []provider.Chunk{ |
| 619 | {Type: provider.ChunkText, Text: "one"}, |
| 620 | {Type: provider.ChunkText, Text: "two"}, |
| 621 | }, |
| 622 | } |
| 623 | exited := make(chan struct{}) |
| 624 | go func() { |
| 625 | r.deliverStream(stream) |
| 626 | close(exited) |
| 627 | }() |
| 628 | deadline := time.Now().Add(time.Second) |
| 629 | for len(stream.out) != 1 && time.Now().Before(deadline) { |
| 630 | time.Sleep(time.Millisecond) |
| 631 | } |
| 632 | if len(stream.out) != 1 { |
| 633 | t.Fatal("delivery did not fill the abandoned consumer buffer") |
| 634 | } |
| 635 | r.mu.Lock() |
| 636 | r.abortDeliveryLocked(stream) |
| 637 | r.mu.Unlock() |
| 638 | select { |
| 639 | case <-exited: |
| 640 | case <-time.After(time.Second): |
| 641 | t.Fatal("delivery goroutine remained blocked after abort") |
| 642 | } |
| 643 | } |
| 644 | |
| 645 | func TestConcurrentStreamsOnOneSidecar(t *testing.T) { |
| 646 | fc := newFakeClient("demo", demoDescriptor()) |
| 647 | r := testResolver(t, baseCatalog(), nil, fc) |
| 648 | |
| 649 | const streamCount = 8 |
| 650 | const chunkCount = 20 |
| 651 | type handle struct { |
| 652 | out <-chan provider.Chunk |
| 653 | id string |
| 654 | } |
| 655 | handles := make([]handle, 0, streamCount) |
| 656 | for i := range streamCount { |
| 657 | p, err := r.Resolve(provider.Selection{Ref: "plugin/demo/fake/x"}) |
| 658 | if err != nil { |
| 659 | t.Fatalf("Resolve: %v", err) |
| 660 | } |
| 661 | out, err := p.Stream(context.Background(), provider.Request{Messages: []provider.Message{{Role: provider.RoleUser}}}) |
| 662 | if err != nil { |
| 663 | t.Fatalf("Stream %d: %v", i, err) |
| 664 | } |
| 665 | fc.mu.Lock() |
| 666 | id := fc.opened[len(fc.opened)-1].StreamID |
| 667 | fc.mu.Unlock() |
| 668 | handles = append(handles, handle{out: out, id: id}) |
| 669 | } |
| 670 | |
| 671 | // Interleave chunk routing for every stream from separate goroutines. |
| 672 | var wg sync.WaitGroup |
| 673 | for i, h := range handles { |
| 674 | wg.Add(1) |
| 675 | go func(i int, h handle) { |
| 676 | defer wg.Done() |
| 677 | for seq := int64(1); seq <= chunkCount; seq++ { |
| 678 | r.RouteStreamChunk(protocol.StreamChunkParams{ |
| 679 | StreamID: h.id, Seq: seq, |
| 680 | Chunk: textChunk(fmt.Sprintf("s%d-c%d", i, seq)), |
| 681 | }) |
| 682 | } |
| 683 | r.RouteStreamEnd(protocol.StreamEndParams{StreamID: h.id, LastSeq: chunkCount}) |
| 684 | }(i, h) |
| 685 | } |
| 686 | wg.Wait() |
| 687 | |
| 688 | for i, h := range handles { |
| 689 | chunks := collectChunks(t, h.out) |
| 690 | if len(chunks) != chunkCount { |
| 691 | t.Fatalf("stream %d delivered %d chunks, want %d", i, len(chunks), chunkCount) |
| 692 | } |
| 693 | for seq := 1; seq <= chunkCount; seq++ { |
| 694 | want := fmt.Sprintf("s%d-c%d", i, seq) |
| 695 | if chunks[seq-1].Text != want { |
| 696 | t.Fatalf("stream %d chunk %d = %q, want %q", i, seq, chunks[seq-1].Text, want) |
| 697 | } |
| 698 | } |
| 699 | } |
| 700 | } |
| 701 | |
| 702 | // TestStreamPendingWindowOverflowInterrupts: a sidecar emitting ever-higher |
| 703 | // sequences without the missing next chunk must not grow the pending buffer |
| 704 | // without bound — the stream fails interrupted once the sequence window is |
| 705 | // exceeded. |
| 706 | func TestStreamPendingWindowOverflowInterrupts(t *testing.T) { |
| 707 | fc := newFakeClient("demo", demoDescriptor()) |
| 708 | r := testResolver(t, baseCatalog(), nil, fc) |
| 709 | out, id := openTestStream(t, r, fc, nil) |
| 710 | |
| 711 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: 1, Chunk: textChunk("first")}) |
| 712 | // Seqs 2..256 sit inside the pending window; none is delivered while seq |
| 713 | // 2 is missing... feed a gap first: seq 3 skips 2, so nextSeq stalls. |
| 714 | for seq := int64(3); seq <= pendingWindowLimit+1; seq++ { |
| 715 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: seq, Chunk: textChunk("gap")}) |
| 716 | } |
| 717 | select { |
| 718 | case chunk := <-out: |
| 719 | if chunk.Type != provider.ChunkText { |
| 720 | t.Fatalf("unexpected early terminal chunk: %+v", chunk) |
| 721 | } |
| 722 | case <-time.After(50 * time.Millisecond): |
| 723 | t.Fatal("seq 1 should have been delivered immediately") |
| 724 | } |
| 725 | // The first chunk beyond the window terminates the stream. |
| 726 | r.RouteStreamChunk(protocol.StreamChunkParams{StreamID: id, Seq: pendingWindowLimit + 2, Chunk: textChunk("overflow")}) |
| 727 | |
| 728 | chunks := collectChunks(t, out) |
| 729 | last := chunks[len(chunks)-1] |
| 730 | if last.Type != provider.ChunkError || !provider.IsStreamInterrupted(last.Err) { |
| 731 | t.Fatalf("terminal chunk = %+v, want interrupted error", last) |
| 732 | } |
| 733 | } |
| 734 | |
| 735 | func TestReasoningSelectionRejectsUndeclaredBeforeSidecarIO(t *testing.T) { |
| 736 | descriptor := demoDescriptor() |
| 737 | descriptor.Efforts = []string{"low", "high"} |
| 738 | fc := newFakeClient("demo", descriptor) |
| 739 | r := testResolver(t, baseCatalog(), nil, fc) |
| 740 | bad := "medium" |
| 741 | _, err := r.Resolve(provider.Selection{Ref: "plugin/demo/fake/x", Effort: &bad}) |
| 742 | var unsupported *provider.UnsupportedReasoningEffort |
| 743 | if !errors.As(err, &unsupported) { |
| 744 | t.Fatalf("selection error=%v", err) |
| 745 | } |
| 746 | p, err := r.Resolve(provider.Selection{Ref: "plugin/demo/fake/x"}) |
| 747 | if err != nil { |
| 748 | t.Fatal(err) |
| 749 | } |
| 750 | _, err = p.Stream(context.Background(), provider.Request{EffortOverride: bad}) |
| 751 | if !errors.As(err, &unsupported) { |
| 752 | t.Fatalf("override error=%v", err) |
| 753 | } |
| 754 | fc.mu.Lock() |
| 755 | defer fc.mu.Unlock() |
| 756 | if len(fc.opened) != 0 { |
| 757 | t.Fatal("invalid effort reached sidecar") |
| 758 | } |
| 759 | } |
| 760 |