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