From 18919425f9d2999c2c1c6bc35b8328866d11e6d4 Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Mon, 8 Jun 2026 20:13:13 +0100 Subject: [PATCH] fix(aibridge): initiate SSE stream before agentic continuation to avoid IsStreaming race (#26139) The agentic loop has a race between the main goroutine and the `Start` goroutine on the shared `ResponseWriter`. When an iteration's response contains only injected-tool events (no text to relay), `Start` may not have called `InitiateStream` by the time main reaches the `IsStreaming` check on the next iteration. The `IsStreaming` check then returns false, main writes a JSON error via `writeUpstreamError`, and `Start` later writes SSE headers and events on top, producing a malformed JSON+SSE response: ``` {\"error\":{\"message\":\"all configured keys are rate-limited\",\"type\":\"rate_limit_error\"},\"request_id\":\"\",\"type\":\"error\"}event: message_start\n..." ``` Fix: explicitly call `events.InitiateStream(w)` at the agentic continuation point so the SSE stream is committed before the next iteration runs. Keeps `messages` consistent with the pattern already used in `chatcompletions/streaming.go`. `sync.Once` makes the double call safe. Related: coder/internal#1524 Related: coder/coder#25654 Closes: https://linear.app/codercom/issue/AIGOV-336/flake-teststreaminginterception-agenticloopfailoveragentic-all-keys > [!NOTE] > Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira --- .../single_injected_tool_no_preamble.txtar | 42 ++++++++++++++++++ aibridge/fixtures/fixtures.go | 3 ++ aibridge/intercept/messages/streaming.go | 5 +++ .../messages/streaming_internal_test.go | 4 -- .../integrationtest/bridge_internal_test.go | 44 +++++++++++++++++++ .../internal/integrationtest/mockupstream.go | 13 ++++++ 6 files changed, 107 insertions(+), 4 deletions(-) create mode 100644 aibridge/fixtures/anthropic/single_injected_tool_no_preamble.txtar diff --git a/aibridge/fixtures/anthropic/single_injected_tool_no_preamble.txtar b/aibridge/fixtures/anthropic/single_injected_tool_no_preamble.txtar new file mode 100644 index 0000000000..5ab09da55d --- /dev/null +++ b/aibridge/fixtures/anthropic/single_injected_tool_no_preamble.txtar @@ -0,0 +1,42 @@ +Coder MCP tools automatically injected, with the model responding with only a tool call and no text preamble. + +-- request -- +{ + "model": "claude-sonnet-4-20250514", + "max_tokens": 1024, + "messages": [ + { + "role": "user", + "content": "list coder workspace IDs for admin" + } + ] +} + +-- streaming -- +event: message_start +data: {"type":"message_start","message":{"id":"msg_01JWGa2JHsKBHL28Cjr2dvPK","type":"message","role":"assistant","model":"claude-sonnet-4-20250514","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":7545,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":1,"service_tier":"standard"}} } + +event: content_block_start +data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_01TSQLR6R6wBUqoxGPjQKDAj","name":"bmcp_coder_coder_list_workspaces","input":{}} } + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":""} } + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"owner\""} } + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":": \"ad"} } + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"min\"}"} } + +event: content_block_stop +data: {"type":"content_block_stop","index":0 } + +event: message_delta +data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":74}} + +event: message_stop +data: {"type":"message_stop" } + diff --git a/aibridge/fixtures/fixtures.go b/aibridge/fixtures/fixtures.go index c731e0fb9c..835b7d2715 100644 --- a/aibridge/fixtures/fixtures.go +++ b/aibridge/fixtures/fixtures.go @@ -24,6 +24,9 @@ var ( //go:embed anthropic/single_injected_tool.txtar AntSingleInjectedTool []byte + //go:embed anthropic/single_injected_tool_no_preamble.txtar + AntSingleInjectedToolNoPreamble []byte + //go:embed anthropic/fallthrough.txtar AntFallthrough []byte diff --git a/aibridge/intercept/messages/streaming.go b/aibridge/intercept/messages/streaming.go index 47c49528a9..badea17b9f 100644 --- a/aibridge/intercept/messages/streaming.go +++ b/aibridge/intercept/messages/streaming.go @@ -483,6 +483,11 @@ newStream: // Causes a new stream to be run with updated messages. isFirst = false + // Commit the SSE stream before the next iteration so a + // later IsStreaming check always takes the SSE branch + // instead of racing with the Start goroutine. + // sync.Once makes this safe. + events.InitiateStream(w) continue newStream } diff --git a/aibridge/intercept/messages/streaming_internal_test.go b/aibridge/intercept/messages/streaming_internal_test.go index 5fc7da00df..40b8344777 100644 --- a/aibridge/intercept/messages/streaming_internal_test.go +++ b/aibridge/intercept/messages/streaming_internal_test.go @@ -454,10 +454,6 @@ func TestStreamingInterception_AgenticLoopFailover(t *testing.T) { // keys 429 during the agentic continuation. // Then: 3 requests, error injected as SSE event, both // keys temporary. - // - // Known flake: race in eventstream.IsStreaming() can - // produce a malformed response on the all-keys-exhausted - // path. See https://github.com/coder/internal/issues/1524. name: "agentic_all_keys_fail", responses: []upstreamResponse{ {statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody}, diff --git a/aibridge/internal/integrationtest/bridge_internal_test.go b/aibridge/internal/integrationtest/bridge_internal_test.go index 8c6804c247..dcf96134aa 100644 --- a/aibridge/internal/integrationtest/bridge_internal_test.go +++ b/aibridge/internal/integrationtest/bridge_internal_test.go @@ -149,6 +149,50 @@ func TestAnthropicMessages(t *testing.T) { }) } }) + + // When the upstream's first response is an injected tool call with no + // text preamble and the next upstream call fails, the response must + // remain a well-formed SSE stream. The upstream error is relayed as a + // well-formed SSE event. + t.Run("streaming injected tool call no preamble with upstream 500", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong) + t.Cleanup(cancel) + + fix := fixtures.Parse(t, fixtures.AntSingleInjectedToolNoPreamble) + upstream := newMockUpstream(ctx, t, + newFixtureResponse(fix), + newErrorResponse(http.StatusInternalServerError), + ) + + mockMCP := setupMCPForTest(t, defaultTracer) + bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, withMCP(mockMCP)) + + reqBody, err := sjson.SetBytes(fix.Request(), "stream", true) + require.NoError(t, err) + resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody) + require.NoError(t, err) + defer resp.Body.Close() + + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, "text/event-stream", resp.Header.Get("Content-Type")) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + bodyStr := string(body) + + // Once iteration 1 succeeded the response is committed as SSE, + // so the iteration-2 error MUST be an SSE event and not a raw JSON body. + require.Contains(t, bodyStr, "event: error", + "iteration-2 error must be relayed as an SSE event") + + // Tool was invoked despite the iteration-2 failure. + require.Len(t, mockMCP.getCallsByTool(mockToolName), 1, + "expected MCP tool to be invoked exactly once") + + bridgeServer.Recorder.VerifyAllInterceptionsEnded(t) + }) } func TestAnthropicMessagesModelThoughts(t *testing.T) { diff --git a/aibridge/internal/integrationtest/mockupstream.go b/aibridge/internal/integrationtest/mockupstream.go index ea493a7639..cbef82047f 100644 --- a/aibridge/internal/integrationtest/mockupstream.go +++ b/aibridge/internal/integrationtest/mockupstream.go @@ -65,6 +65,19 @@ func newFixtureToolResponse(fix fixtures.Fixture) upstreamResponse { return resp } +// newErrorResponse returns an upstreamResponse that replays a raw HTTP error +// response with the given status code. Used to drive iteration-N error paths +// from inside a multi-call mockUpstream scripted-response list. +func newErrorResponse(status int) upstreamResponse { + body := fmt.Sprintf(`{"error":{"message":%q}}`, http.StatusText(status)) + raw := fmt.Sprintf("HTTP/1.1 %d %s\r\n", status, http.StatusText(status)) + raw += "x-should-retry: false\r\n" + raw += "Content-Type: application/json\r\n" + raw += fmt.Sprintf("Content-Length: %d\r\n\r\n%s", len(body), body) + rawBytes := []byte(raw) + return upstreamResponse{Streaming: rawBytes, Blocking: rawBytes} +} + // receivedRequest captures the details of a single request handled by mockUpstream. type receivedRequest struct { Method string