diff --git a/internal/runtime/executor/codex_executor_execute.go b/internal/runtime/executor/codex_executor_execute.go index b738890c3f..1cb4da22ff 100644 --- a/internal/runtime/executor/codex_executor_execute.go +++ b/internal/runtime/executor/codex_executor_execute.go @@ -137,11 +137,12 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re data, errRead := io.ReadAll(httpResp.Body) helps.AppendAPIResponseChunk(ctx, e.cfg, data) - lines := bytes.Split(data, []byte("\n")) + events := helps.NewCodexSSEScanner(ctx, bytes.NewReader(data), len(data)+1) outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte sawOutputDelta := false - for _, line := range lines { + for events.Scan() { + line := events.Bytes() if !bytes.HasPrefix(line, dataTag) { continue } diff --git a/internal/runtime/executor/codex_executor_sse_events_test.go b/internal/runtime/executor/codex_executor_sse_events_test.go new file mode 100644 index 0000000000..558bfdb56a --- /dev/null +++ b/internal/runtime/executor/codex_executor_sse_events_test.go @@ -0,0 +1,456 @@ +package executor + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" + log "github.com/sirupsen/logrus" + logtest "github.com/sirupsen/logrus/hooks/test" +) + +func codexStreamChunks(t *testing.T, stream string, buffered bool) []string { + t.Helper() + cfg := &config.Config{} + cfg.Codex.StreamBootstrapBuffering = buffered + executor := NewCodexExecutor(cfg) + ctx := context.WithValue(t.Context(), "cliproxy.roundtripper", codexAllocationTransport{stream: stream}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "test"}} + result, err := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","input":"hello"}`)}, + cliproxyexecutor.Options{Stream: true, SourceFormat: sdktranslator.FormatOpenAIResponse}) + if err != nil { + t.Fatal(err) + } + var chunks []string + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream error: %v", chunk.Err) + } + if len(chunk.Payload) > 0 { + chunks = append(chunks, string(chunk.Payload)) + } + } + return chunks +} + +func TestCodexExecuteStreamTreatsAnEventSplitAcrossDataLinesAsOneEvent(t *testing.T) { + completed := `{"type":"response.completed","response":{"id":"resp_split","model":"gpt-5.4","status":"completed","output":[{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"output_text","text":"hi"}]}],"usage":{"input_tokens":5,"output_tokens":1,"total_tokens":6}}}` + cut := strings.Index(completed, `"output":`) + // With bootstrap buffering, the first output releases the stream before the completed event. + created := "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_split\",\"model\":\"gpt-5.4\"}}\n\n" + + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_1\",\"output_index\":0,\"content_index\":0,\"delta\":\"hi\"}\n\n" + oneLine := created + "event: response.completed\ndata: " + completed + "\n\n" + split := created + "event: response.completed\ndata: " + completed[:cut] + "\ndata: " + completed[cut:] + "\n\n" + for _, buffered := range []bool{false, true} { + t.Run(fmt.Sprintf("buffered=%t", buffered), func(t *testing.T) { + want := codexStreamChunks(t, oneLine, buffered) + got := codexStreamChunks(t, split, buffered) + if fmt.Sprintf("%q", got) != fmt.Sprintf("%q", want) { + t.Fatalf("split event chunks = %q\nwant the one-line chunks %q", got, want) + } + if terminal := want[len(want)-1]; !strings.Contains(terminal, `"response.completed"`) || !json.Valid([]byte(strings.TrimPrefix(terminal, "data: "))) { + t.Fatalf("last chunk = %q, want a valid response.completed", terminal) + } + }) + } +} + +func cloneLogHooks(src log.LevelHooks) log.LevelHooks { + cloned := make(log.LevelHooks, len(src)) + for lvl, hooks := range src { + cloned[lvl] = append([]log.Hook(nil), hooks...) + } + return cloned +} + +func TestCodexExecuteStreamPassesADoneSentinelAfterTheCompletedEvent(t *testing.T) { + origHooks := cloneLogHooks(log.StandardLogger().Hooks) + hook := new(logtest.Hook) + log.StandardLogger().AddHook(hook) + t.Cleanup(func() { log.StandardLogger().ReplaceHooks(origHooks) }) + + stream := codexAllocationStream(2, 16) + "data: [DONE]\n\n" + chunks := codexStreamChunks(t, stream, false) + sawCompleted := false + for _, chunk := range chunks { + if strings.Contains(chunk, `"response.completed"`) { + sawCompleted = true + } + } + if !sawCompleted { + t.Fatalf("chunks = %q, want the completed event", chunks) + } + for _, entry := range hook.AllEntries() { + if strings.Contains(entry.Message, "SSE event data is not valid JSON") { + t.Fatalf("unexpected framing warning: %s", entry.Message) + } + } +} + +func TestCodexExecuteNonStreamReadsAnEventSplitAcrossDataLines(t *testing.T) { + completed := `{"type":"response.completed","response":{"id":"resp_split","object":"response","model":"gpt-5.4","status":"completed","output":[{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"output_text","text":"hi"}]}],"usage":{"input_tokens":5,"output_tokens":1,"total_tokens":6}}}` + cut := strings.Index(completed, `"output":`) + stream := "event: response.completed\r\ndata: " + completed[:cut] + "\r\ndata: " + completed[cut:] + "\r\n\r\n" + executor := NewCodexExecutor(&config.Config{}) + ctx := context.WithValue(t.Context(), "cliproxy.roundtripper", codexAllocationTransport{stream: stream}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "test"}} + resp, err := executor.Execute(ctx, auth, cliproxyexecutor.Request{Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","input":"hello"}`)}, + cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + var got struct { + ID string `json:"id"` + Output []json.RawMessage `json:"output"` + } + if errDecode := json.Unmarshal(resp.Payload, &got); errDecode != nil || got.ID != "resp_split" || len(got.Output) != 1 { + t.Fatalf("payload = %s (decode error %v), want the completed response", resp.Payload, errDecode) + } +} + +func TestCodexExecuteStreamKeepsAnInvalidSplitEventInvalid(t *testing.T) { + // Split inside a string, so the SSE-joined data has a raw newline in it and is not JSON. + stream := "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_bad\",\"model\":\"gpt-5.4\"}}\n\n" + + "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_bad\",\"status\":\"comp\ndata: leted\",\"output\":[]}}\n\n" + chunks := codexStreamChunks(t, stream, false) + last := chunks[len(chunks)-1] + var values []string + for _, line := range strings.Split(last, "\n") { + values = append(values, strings.TrimPrefix(strings.TrimPrefix(line, "data:"), " ")) + } + if !strings.HasPrefix(last, "data: ") || strings.Count(last, "\ndata: ") != 1 || json.Valid([]byte(strings.Join(values, "\n"))) { + t.Fatalf("last chunk = %q, want the invalid event's two data lines unchanged", last) + } +} + +func TestCodexExecuteStreamAssemblesSplitEventsWithAndWithoutBootstrapBuffering(t *testing.T) { + // A handshake event split across three data lines, and a terminal event split across two data lines. + stream := "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_split\",\"model\":\"gpt-5.4\"}}\n\n" + + "data: {\"type\":\"response.in_progress\",\"response\":{\"id\":\"resp_split\",\"output\":[\ndata: {\"x\":1},\ndata: {\"x\":1}]}}\n\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_split\",\"status\":\"completed\",\ndata: \"output\":[]}}\n\n" + fragments := func(chunks []string) int { + count := 0 + for _, chunk := range chunks { + if chunk == `data: {"x":1},` { + count++ + } + } + return count + } + for _, buffered := range []bool{false, true} { + t.Run(fmt.Sprintf("buffered=%t", buffered), func(t *testing.T) { + chunks := codexStreamChunks(t, stream, buffered) + if count := fragments(chunks); count != 0 { + t.Fatalf("buffered=%t: %d fragments emitted, want the event assembled", buffered, count) + } + last := chunks[len(chunks)-1] + if !strings.Contains(last, `"response.completed"`) || !json.Valid([]byte(strings.TrimPrefix(last, "data: "))) { + t.Fatalf("last chunk = %q, want a valid response.completed", last) + } + }) + } +} + +func TestCodexExecuteStreamMinimalReproFromIssue6575(t *testing.T) { + // Minimal reproduction from issue 6575: response.completed immediately follows response.created and is split. + stream := "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"r\",\"status\":\"in_progress\",\"output\":[]}}\n\n" + + "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\ndata: \"output\":[]}}\n\n" + for _, buffered := range []bool{false, true} { + t.Run(fmt.Sprintf("buffered=%t", buffered), func(t *testing.T) { + chunks := codexStreamChunks(t, stream, buffered) + if len(chunks) == 0 { + t.Fatal("expected chunks, got none") + } + last := chunks[len(chunks)-1] + if !strings.Contains(last, `"response.completed"`) || !json.Valid([]byte(strings.TrimPrefix(last, "data: "))) { + t.Fatalf("last chunk = %q, want valid response.completed", last) + } + }) + } +} + +func TestCodexExecuteStreamSplitHandshakeExceedingLineBudgetReleasesBootstrapWithoutFailover(t *testing.T) { + // A split in_progress event that spans 50 data lines (> codexBootstrapMaxBufferedFrames of 48), followed by overload. + var stream strings.Builder + stream.WriteString("event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_split\",\"model\":\"gpt-5.4\"}}\n\n") + stream.WriteString("data: {\"type\":\"response.in_progress\",\"response\":{\"id\":\"resp_split\",\"items\":[\n") + for i := 0; i < 50; i++ { + stream.WriteString("data: \n") + } + stream.WriteString("data: ]}}\n\n") + // Followed by overload error + stream.WriteString("event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"server_error\",\"code\":\"server_overloaded\",\"message\":\"The server had an error while processing your request.\",\"param\":\"input\"}}\n\n") + + // When bootstrap buffering is enabled, the 50 lines exceeded the line budget (48 frames), + // so the stream is released and the overload error is delivered in-stream rather than failing over. + cfg := &config.Config{} + cfg.Codex.StreamBootstrapBuffering = true + executor := NewCodexExecutor(cfg) + ctx := context.WithValue(t.Context(), "cliproxy.roundtripper", codexAllocationTransport{stream: stream.String()}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "test"}} + result, err := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","input":"hello"}`)}, + cliproxyexecutor.Options{Stream: true, SourceFormat: sdktranslator.FormatOpenAIResponse}) + if err != nil { + t.Fatalf("expected stream to be released without failover error: %v", err) + } + if result == nil { + t.Fatal("expected non-nil stream result") + } + sawError := false + for chunk := range result.Chunks { + if chunk.Err != nil { + sawError = true + } + } + if !sawError { + t.Fatal("expected in-stream delivery of overload error after budget exhausted") + } +} + +func TestCodexExecuteStreamEventExceedsBudgetWhileAssemblingReleasesStream(t *testing.T) { + // Line 1: event: response.created + // Line 2: data: {"type":"response.created",...} + // Line 3: blank line + // Line 4: data: {"type":"response.in_progress",... + // Lines 5-47: 43 empty data lines + // Line 48: {"key_at_boundary": 123} - exact 48th physical line triggering budget limit + // Lines 49-50: {"key_after_boundary": 456} and ]}} closing event + var stream strings.Builder + stream.WriteString("event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_split\",\"model\":\"gpt-5.4\"}}\n\n") + stream.WriteString("data: {\"type\":\"response.in_progress\",\"response\":{\"id\":\"resp_split\",\"items\":[\n") + for i := 0; i < 43; i++ { + stream.WriteString("data: \n") + } + // Line 48: triggers the 48-line frame budget limit during assembly. + stream.WriteString("data: {\"key_at_boundary\":123},\n") + // Lines 49-50: streamed in goroutine after bootstrap release. + stream.WriteString("data: {\"key_after_boundary\":456}\n") + stream.WriteString("data: ]}}\n") // without blank separator, next line is error + stream.WriteString("event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"server_error\",\"code\":\"server_overloaded\",\"message\":\"overloaded\"}}\n\n") + + cfg := &config.Config{} + cfg.Codex.StreamBootstrapBuffering = true + executor := NewCodexExecutor(cfg) + ctx := context.WithValue(t.Context(), "cliproxy.roundtripper", codexAllocationTransport{stream: stream.String()}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "test"}} + result, err := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","input":"hello"}`)}, + cliproxyexecutor.Options{Stream: true, SourceFormat: sdktranslator.FormatOpenAIResponse}) + if err != nil { + t.Fatalf("expected stream to be released without failover error: %v", err) + } + if result == nil { + t.Fatal("expected non-nil stream result") + } + sawError := false + var inProgressChunk string + for chunk := range result.Chunks { + if chunk.Err != nil { + sawError = true + } else if strings.Contains(string(chunk.Payload), "response.in_progress") { + inProgressChunk = string(chunk.Payload) + } + } + if !sawError { + t.Fatal("expected in-stream delivery of overload error after budget exhausted") + } + if !strings.Contains(inProgressChunk, `"key_at_boundary":123`) || !strings.Contains(inProgressChunk, `"key_after_boundary":456`) { + t.Fatalf("inProgressChunk = %q, want both key_at_boundary and key_after_boundary preserved", inProgressChunk) + } + if !json.Valid([]byte(strings.TrimPrefix(inProgressChunk, "data: "))) { + t.Fatalf("inProgressChunk = %q, want valid JSON without dropped lines", inProgressChunk) + } +} + +func TestCodexExecuteStreamEventExceedsTimeBudgetWhileAssemblingReleasesStream(t *testing.T) { + // Time budget check while assembling a multi-line event. + t0 := time.Date(2026, 9, 14, 12, 0, 0, 0, time.UTC) + clock := withMockClock(t, t0) + + bootstrapStarted := make(chan struct{}) + var once sync.Once + cleanup := setCodexBootstrapNowForTest(func() time.Time { + now := clock.now() + once.Do(func() { + close(bootstrapStarted) + }) + return now + }) + t.Cleanup(cleanup) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_t\"}}\n\n")) + _, _ = w.Write([]byte("data: {\"type\":\"response.in_progress\",\"response\":{\"id\":\"resp_t\",\"items\":[\n")) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + + <-bootstrapStarted + clock.advance(11 * time.Second) + + _, _ = w.Write([]byte("data: {\"t_item\":1}\ndata: ]}}\n\n")) + _, _ = w.Write([]byte("event: error\ndata: " + codexOverloadEvent + "\n\n")) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + })) + defer server.Close() + + req, opts := codexTestRequest() + result, err := NewCodexExecutor(codexBufferingConfigWithTimeout(true, "10s")).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + if err != nil { + t.Fatalf("the time budget must release the stream: %v", err) + } + if result == nil { + t.Fatal("expected a stream result once the time budget released the stream") + } + body, streamErr := drainChunks(result) + if streamErr == nil { + t.Fatal("expected the overload to arrive in-stream after the time budget released") + } + if !strings.Contains(body, "response.in_progress") || !strings.Contains(body, `"t_item":1`) { + t.Fatalf("body = %q, want valid assembled in_progress with t_item:1", body) + } +} + +func TestCodexExecuteStreamSplitCompletedAfterOutputAndClockAdvance(t *testing.T) { + t0 := time.Date(2026, 9, 14, 12, 0, 0, 0, time.UTC) + clock := withMockClock(t, t0) + + bootstrapStarted := make(chan struct{}) + var once sync.Once + cleanup := setCodexBootstrapNowForTest(func() time.Time { + now := clock.now() + once.Do(func() { + close(bootstrapStarted) + }) + return now + }) + t.Cleanup(cleanup) + + completed := `{"type":"response.completed","response":{"id":"resp_clock","model":"gpt-5.4","status":"completed","output":[{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7}}}` + cut := strings.Index(completed, `"output":`) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_clock\",\"model\":\"gpt-5.4\"}}\n\n")) + _, _ = w.Write([]byte("event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_1\",\"output_index\":0,\"content_index\":0,\"delta\":\"ok\"}\n\n")) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + + <-bootstrapStarted + // Advance clock past the 10s bootstrap timeout. + clock.advance(20 * time.Second) + + // Send split response.completed. + _, _ = w.Write([]byte("event: response.completed\ndata: " + completed[:cut] + "\ndata: " + completed[cut:] + "\n\n")) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + })) + defer server.Close() + + req, opts := codexTestRequest() + result, err := NewCodexExecutor(codexBufferingConfigWithTimeout(true, "10s")).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result == nil { + t.Fatal("expected non-nil stream result") + } + body, streamErr := drainChunks(result) + if streamErr != nil { + t.Fatalf("unexpected stream error: %v", streamErr) + } + if !strings.Contains(body, "response.completed") || !strings.Contains(body, `"id":"resp_clock"`) { + t.Fatalf("body = %q, want valid assembled response.completed after clock advance", body) + } +} + +func TestCodexExecuteStreamMixedFramingCommentInsideSplitEventCountsBudget(t *testing.T) { + // 46 data lines held in an in_progress event, then 2 comment lines (total 48 lines), then overload. + var stream strings.Builder + stream.WriteString("data: {\"type\":\"response.in_progress\",\"items\":[\n") + for i := 0; i < 46; i++ { + stream.WriteString("data: \n") + } + // 47th line: comment inside the split event + stream.WriteString(": keepalive-1\n") + // 48th line: comment inside the split event (hits budget limit) + stream.WriteString(": keepalive-2\n") + // 49th line: overload without failover + stream.WriteString("event: error\ndata: " + codexOverloadEvent + "\n\n") + // finishes event + stream.WriteString("data: ]}}\n\n") + + cfg := &config.Config{} + cfg.Codex.StreamBootstrapBuffering = true + executor := NewCodexExecutor(cfg) + ctx := context.WithValue(t.Context(), "cliproxy.roundtripper", codexAllocationTransport{stream: stream.String()}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "test"}} + result, err := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","input":"hello"}`)}, + cliproxyexecutor.Options{Stream: true, SourceFormat: sdktranslator.FormatOpenAIResponse}) + if err != nil { + t.Fatalf("expected stream to be released without failover error: %v", err) + } + if result == nil { + t.Fatal("expected non-nil stream result") + } + sawError := false + for chunk := range result.Chunks { + if chunk.Err != nil { + sawError = true + } + } + if !sawError { + t.Fatal("expected in-stream delivery of overload error after mixed framing budget exhausted") + } +} + +func TestCodexExecuteStreamEventExceedsByteBudgetWhileAssemblingReleasesStream(t *testing.T) { + // A split event whose data lines exceed codexBootstrapMaxBufferedBytes (1 MiB) while assembling, + // followed by overload. + var stream strings.Builder + stream.WriteString("data: {\"type\":\"response.in_progress\",\"items\":[\n") + // Write 5 data lines of 250KB each = 1.25MB > 1MB limit + largeChunk := strings.Repeat("x", 250*1024) + for i := 0; i < 5; i++ { + fmt.Fprintf(&stream, "data: %q,\n", largeChunk) + } + stream.WriteString("data: \"last\"]}\n\n") + stream.WriteString("event: error\ndata: " + codexOverloadEvent + "\n\n") + + cfg := &config.Config{} + cfg.Codex.StreamBootstrapBuffering = true + executor := NewCodexExecutor(cfg) + ctx := context.WithValue(t.Context(), "cliproxy.roundtripper", codexAllocationTransport{stream: stream.String()}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "test"}} + result, err := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","input":"hello"}`)}, + cliproxyexecutor.Options{Stream: true, SourceFormat: sdktranslator.FormatOpenAIResponse}) + if err != nil { + t.Fatalf("expected stream to be released without failover error: %v", err) + } + if result == nil { + t.Fatal("expected non-nil stream result") + } + sawError := false + for chunk := range result.Chunks { + if chunk.Err != nil { + sawError = true + } + } + if !sawError { + t.Fatal("expected in-stream delivery of overload error after byte budget exhausted") + } +} diff --git a/internal/runtime/executor/codex_executor_stream.go b/internal/runtime/executor/codex_executor_stream.go index f59c58d1cb..f468586ca8 100644 --- a/internal/runtime/executor/codex_executor_stream.go +++ b/internal/runtime/executor/codex_executor_stream.go @@ -1,7 +1,6 @@ package executor import ( - "bufio" "bytes" "context" "io" @@ -154,27 +153,42 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au bootstrapStart = nowCodexBootstrap() } - scanner := bufio.NewScanner(httpResp.Body) - scanner.Buffer(nil, 52_428_800) // 50MB + bufferedFrames := 0 + bufferedBytes := 0 + budgetExhausted := false + exhaustedReason := "" + + scanner := helps.NewCodexSSEScanner(ctx, httpResp.Body, 52_428_800) // 50MB + scanner.SetOnLine(func(rawLine []byte) { + helps.AppendAPIResponseChunk(ctx, e.cfg, rawLine) + }) + if buffering { + scanner.SetBudgetChecker(func(line []byte, heldLines, heldBytes int) bool { + timeSinceStart := nowCodexBootstrap().Sub(bootstrapStart) + if bootstrapTimeout > 0 && timeSinceStart >= bootstrapTimeout { + budgetExhausted = true + exhaustedReason = "time budget" + return true + } + if bufferedFrames+heldLines >= codexBootstrapMaxBufferedFrames { + budgetExhausted = true + exhaustedReason = "frame budget" + return true + } + if bufferedBytes+heldBytes >= codexBootstrapMaxBufferedBytes { + budgetExhausted = true + exhaustedReason = "byte budget" + return true + } + return false + }) + } claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) var param any outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte var bufferedChunks [][]byte - // bufferedFrames counts the scanned lines this loop holds and bufferedBytes sums each line - // together with the chunks it translates into. Every iteration either holds the line or leaves - // the loop, so this is one unit per line read. Counting only the chunks would bound nothing for - // a downstream format that renders a frame as zero chunks, and counting only some line kinds - // would let the upstream's choice of framing decide whether the bound advances at all. The cost - // is that a verbose framing spends the budget faster: the three-line event:/data:/blank shape - // protects roughly a third as many events as a stream of bare ": keepalive" comments does. - // - // In addition to the frame and byte budgets, bootstrapTimeout bounds how long trickled - // frames may hold the downstream headers. A peer that never terminates a line is bounded - // by the caller's request context. - bufferedFrames := 0 - bufferedBytes := 0 var initialChunks [][]byte streamStarted := false immediateTerminal := false @@ -193,7 +207,12 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au if buffering { for scanner.Scan() { line := scanner.Bytes() - helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if budgetExhausted { + helps.LogWithRequestID(ctx).Debugf("codex executor: bootstrap %s exhausted after %d lines / %d bytes / %v, releasing stream without overload probing", exhaustedReason, bufferedFrames, bufferedBytes, nowCodexBootstrap().Sub(bootstrapStart)) + initialChunks = nil + streamStarted = true + break + } var translatedLine []byte isHandshake := false terminalSuccess := false @@ -282,14 +301,17 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, body, translatedLine, ¶m, claudeInputTokens) if isHandshake && !terminalSuccess { - frameBytes := len(line) + rawLineCount := scanner.RawLines() + frameBytes := scanner.RawBytes() for i := range chunks { frameBytes += len(chunks[i]) } timeSinceStart := nowCodexBootstrap().Sub(bootstrapStart) timeoutReached := bootstrapTimeout > 0 && timeSinceStart >= bootstrapTimeout - if !timeoutReached && bufferedFrames < codexBootstrapMaxBufferedFrames && bufferedBytes+frameBytes <= codexBootstrapMaxBufferedBytes { - bufferedFrames++ + totalFrames := bufferedFrames + scanner.HeldLines() + rawLineCount + totalBytes := bufferedBytes + scanner.HeldBytes() + frameBytes + if !timeoutReached && totalFrames <= codexBootstrapMaxBufferedFrames && totalBytes <= codexBootstrapMaxBufferedBytes { + bufferedFrames += rawLineCount bufferedBytes += frameBytes bufferedChunks = append(bufferedChunks, chunks...) continue @@ -297,7 +319,9 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au exhausted := "frame budget" if timeoutReached { exhausted = "time budget" - } else if bufferedFrames < codexBootstrapMaxBufferedFrames { + } else if totalFrames > codexBootstrapMaxBufferedFrames { + exhausted = "frame budget" + } else { exhausted = "byte budget" } helps.LogWithRequestID(ctx).Debugf("codex executor: bootstrap %s exhausted after %d lines / %d bytes / %v, releasing stream without overload probing", exhausted, bufferedFrames, bufferedBytes, timeSinceStart) @@ -310,6 +334,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au } break } + scanner.SetBudgetChecker(nil) if !streamStarted && bootstrapTerminalErr == nil { closeBootstrapBody() @@ -385,7 +410,6 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au }() for scanner.Scan() { line := scanner.Bytes() - helps.AppendAPIResponseChunk(ctx, e.cfg, line) var translatedLine []byte terminalSuccess := false diff --git a/internal/runtime/executor/helps/codex_sse_scanner.go b/internal/runtime/executor/helps/codex_sse_scanner.go new file mode 100644 index 0000000000..35f553290d --- /dev/null +++ b/internal/runtime/executor/helps/codex_sse_scanner.go @@ -0,0 +1,330 @@ +package helps + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "io" + "strconv" + "strings" + + log "github.com/sirupsen/logrus" +) + +// codexSSEDiagnosticLines caps how many data lines one invalid-event diagnostic describes. +const codexSSEDiagnosticLines = 8 + +var codexSSEDataPrefix = []byte("data:") + +type codexSSEQueuedItem struct { + payload []byte + rawLines int + rawBytes int +} + +// CodexSSEScanner reads an upstream Responses SSE stream line by line, as bufio.Scanner does, but +// delivers an event whose data the upstream split across several "data:" lines as one "data:" line. +// In SSE an event's data may span several data lines, joined with "\n". Taking each data line as a +// whole event instead cuts such an event: on a split response.completed the stream ends at the +// first fragment and the rest is never read. +// +// Every event on this path is one JSON value. A data line that is complete JSON on its own, or +// cannot start JSON at all (an empty value, [DONE]), is delivered byte-for-byte as read with no +// lookahead, which keeps upstreams that omit the blank line between events working as before. A +// data line that opens a JSON object or array without completing it starts an event that collects +// the following data lines until one of these: +// - their joined value closes its outermost object or array; +// - a blank line ends the event; +// - the stream ends; +// - the collected bytes would exceed the line limit. +// +// A completed event is delivered compacted onto one line. One that never becomes valid JSON is +// delivered as the same multi-line data field, so it still fails downstream. The first such event +// in a stream gets one warning that describes its framing without any of its bytes. +// +// The scanner merges data lines and nothing else. It holds only the data lines of an event being +// collected and delivers every other line the moment it reads it, in the order read. A comment, +// id: or retry: line inside such an event therefore comes before the event's data instead of +// between its pieces, which SSE treats alike: an event's fields take effect when it is dispatched. +// The blank line or over-limit line that ends a collection follows the event's data. +// +// SetAssembly(false) makes the scanner read exactly as bufio.Scanner does, for a caller that +// budgets every upstream line while it holds a stream. +type CodexSSEScanner struct { + ctx context.Context + lines *bufio.Scanner + maxBytes int + // stock is set while assembly is off. + stock bool + // held is the event being collected: its data lines as read, joined with "\n" in one buffer, so + // its memory follows its bytes, not its line count. count is how many lines it holds and size + // their bytes. joined is their joined value, and nesting follows it. + held []byte + count int + size int + joined []byte + nesting codexJSONNesting + queue []codexSSEQueuedItem + current codexSSEQueuedItem + // warned is set once the stream has had its invalid-event warning. + warned bool + onLine func(line []byte) + budgetCheck func(line []byte, heldLines int, heldBytes int) bool +} + +// NewCodexSSEScanner returns a CodexSSEScanner over r, with assembly on. maxBytes bounds one line, +// as bufio.Scanner.Buffer does, and also the data lines collected for one event. +func NewCodexSSEScanner(ctx context.Context, r io.Reader, maxBytes int) *CodexSSEScanner { + lines := bufio.NewScanner(r) + lines.Buffer(nil, maxBytes) + return &CodexSSEScanner{ctx: ctx, lines: lines, maxBytes: maxBytes} +} + +// SetOnLine sets a callback invoked for every physical line read from upstream. +func (s *CodexSSEScanner) SetOnLine(fn func(line []byte)) { + s.onLine = fn +} + +// HeldLines returns how many data lines are currently held in an incomplete event. +func (s *CodexSSEScanner) HeldLines() int { + return s.count +} + +// HeldBytes returns how many bytes are currently held in an incomplete event. +func (s *CodexSSEScanner) HeldBytes() int { + return s.size +} + +// SetBudgetChecker sets an optional callback invoked on each raw line read while an event is being held. +// If it returns true, the scanner yields from Scan() with an empty payload to signal that the +// caller should release buffering, while preserving the held assembly state for subsequent Scan() calls. +func (s *CodexSSEScanner) SetBudgetChecker(fn func(line []byte, heldLines int, heldBytes int) bool) { + s.budgetCheck = fn +} + +// SetAssembly turns event assembly on or off. While it is off, no data line starts a collection and +// every line is delivered as read, exactly as bufio.Scanner delivers it. An event already being +// collected is still completed. +func (s *CodexSSEScanner) SetAssembly(on bool) { s.stock = !on } + +// Scan advances to the next line, or to the next assembled event data line. +func (s *CodexSSEScanner) Scan() bool { + if len(s.queue) > 0 { + s.current, s.queue = s.queue[0], s.queue[1:] + return true + } + for s.lines.Scan() { + line := s.lines.Bytes() + if s.onLine != nil { + s.onLine(line) + } + if !bytes.HasPrefix(line, codexSSEDataPrefix) { + if len(line) == 0 && s.count > 0 { + s.release([]byte{}) + return s.Scan() + } + s.current = codexSSEQueuedItem{payload: line, rawLines: 1, rawBytes: len(line)} + return true + } + if s.count == 0 { + if s.stock || !codexSSEIncompleteJSON(bytes.TrimSpace(line[len(codexSSEDataPrefix):])) { + s.current = codexSSEQueuedItem{payload: line, rawLines: 1, rawBytes: len(line)} + return true + } + } else if s.size+len(line) > s.maxBytes { + // Bounded like a single line: stop collecting and let this line be read on its own. + s.release(bytes.Clone(line)) + return s.Scan() + } + + closed := s.hold(line) + if !closed && s.budgetCheck != nil && s.budgetCheck(line, s.count, s.size) { + s.budgetCheck = nil + s.current = codexSSEQueuedItem{payload: nil, rawLines: 0, rawBytes: 0} + return true + } + if closed { + s.release() + return s.Scan() + } + } + if s.count > 0 { + s.release() + return s.Scan() + } + return false +} + +// hold adds a data line to the event being collected and reports whether its data has closed. +func (s *CodexSSEScanner) hold(line []byte) bool { + value := codexSSEValue(line) + if s.count > 0 { + s.held = append(s.held, '\n') + s.joined = append(s.joined, '\n') + s.nesting.feed([]byte{'\n'}) + } + s.held = append(s.held, line...) + s.joined = append(s.joined, value...) + s.count++ + s.size += len(line) + return s.nesting.feed(value) +} + +// release ends the collection and queues the event as one line, followed by the line that ended +// it, if any. +func (s *CodexSSEScanner) release(end ...[]byte) { + assembled := s.assemble(s.held, s.count, s.joined) + s.queue = append(s.queue, codexSSEQueuedItem{ + payload: assembled, + rawLines: s.count, + rawBytes: s.size, + }) + for _, e := range end { + s.queue = append(s.queue, codexSSEQueuedItem{ + payload: e, + rawLines: 1, + rawBytes: len(e), + }) + } + s.held, s.count, s.size, s.joined, s.nesting = nil, 0, 0, nil, codexJSONNesting{} +} + +// codexSSEValue is one data line's value: the text after "data:" without the one leading space. +func codexSSEValue(line []byte) []byte { + return bytes.TrimPrefix(line[len(codexSSEDataPrefix):], []byte(" ")) +} + +// codexSSEClosed reports whether joined event data is complete JSON. Only a value that now ends +// with a closing bracket can have become complete, so other values are not re-validated. +func codexSSEClosed(joined []byte) bool { + trimmed := bytes.TrimRight(joined, " \t\r\n") + if len(trimmed) == 0 || (trimmed[len(trimmed)-1] != '}' && trimmed[len(trimmed)-1] != ']') { + return false + } + return json.Valid(joined) +} + +// codexSSEIncompleteJSON reports whether one data line's value starts a JSON object or array that +// it does not complete. Values that cannot start JSON, such as an empty value or [DONE], are not. +func codexSSEIncompleteJSON(value []byte) bool { + if len(value) == 0 || (value[0] != '{' && value[0] != '[') || bytes.Equal(value, []byte("[DONE]")) { + return false + } + return !codexSSEClosed(value) +} + +// codexJSONNesting follows the nesting of a JSON text fed in pieces, so an event split across many +// data lines is validated once, when its outermost object or array closes, not after every piece. +type codexJSONNesting struct { + depth int + inString bool + escaped bool +} + +// feed consumes p and reports whether the outermost object or array closed within it. +func (n *codexJSONNesting) feed(p []byte) bool { + for _, c := range p { + switch { + case n.inString: + switch { + case n.escaped: + n.escaped = false + case c == '\\': + n.escaped = true + case c == '"': + n.inString = false + } + case c == '"': + n.inString = true + case c == '{' || c == '[': + n.depth++ + case c == '}' || c == ']': + n.depth-- + if n.depth <= 0 { + return true + } + } + } + return false +} + +// Bytes returns the current line. It stays valid until the next call to Scan. +func (s *CodexSSEScanner) Bytes() []byte { return s.current.payload } + +// RawLines returns the count of physical upstream lines consumed to produce the current event. +func (s *CodexSSEScanner) RawLines() int { + if s.current.rawLines == 0 { + return 1 + } + return s.current.rawLines +} + +// RawBytes returns the total bytes of physical upstream lines consumed to produce the current event. +func (s *CodexSSEScanner) RawBytes() int { + if s.current.rawBytes == 0 { + return len(s.current.payload) + } + return s.current.rawBytes +} + +// Err returns the error that stopped the scan, as bufio.Scanner does. +func (s *CodexSSEScanner) Err() error { return s.lines.Err() } + +// assemble returns the event held as lines data lines: compacted onto one line when its data is +// valid JSON, otherwise as held. +func (s *CodexSSEScanner) assemble(held []byte, lines int, joined []byte) []byte { + if lines == 1 { + s.diagnose(held, lines, bytes.TrimSpace(joined)) + return held + } + if json.Valid(joined) { + var compacted bytes.Buffer + compacted.Grow(len(codexSSEDataPrefix) + 1 + len(joined)) + compacted.Write(codexSSEDataPrefix) + compacted.WriteByte(' ') + if errCompact := json.Compact(&compacted, joined); errCompact == nil { + return compacted.Bytes() + } + } + s.diagnose(held, lines, joined) + return held +} + +// diagnose logs one bounded warning for the first event in the stream whose data is not valid +// JSON, so a broken upstream cannot turn every event into a log line. It describes how many data +// lines carried it, each line's length and whether it is valid on its own, and where the joined +// data fails to parse. It never logs data bytes. +func (s *CodexSSEScanner) diagnose(held []byte, lines int, data []byte) { + if s.warned || len(data) == 0 || bytes.Equal(data, []byte("[DONE]")) || json.Valid(data) { + return + } + s.warned = true + offset := int64(-1) + var syntaxErr *json.SyntaxError + if errors.As(json.Unmarshal(data, new(json.RawMessage)), &syntaxErr) { + offset = syntaxErr.Offset + } + shown := min(lines, codexSSEDiagnosticLines) + lengths := make([]string, shown) + valid := make([]string, shown) + for i := range shown { + var line []byte + line, held, _ = bytes.Cut(held, []byte("\n")) + value := bytes.TrimSpace(line[len(codexSSEDataPrefix):]) + lengths[i] = strconv.Itoa(len(value)) + valid[i] = strconv.FormatBool(json.Valid(value)) + } + more := "" + if lines > shown { + more = ",+" + strconv.Itoa(lines-shown) + } + LogWithRequestID(s.ctx).WithFields(log.Fields{ + "data_lines": lines, + "line_bytes": strings.Join(lengths, ",") + more, + "lines_valid": strings.Join(valid, ",") + more, + "data_bytes": len(data), + "json_error_offset": offset, + }).Warn("codex executor: upstream SSE event data is not valid JSON") +} diff --git a/internal/runtime/executor/helps/codex_sse_scanner_test.go b/internal/runtime/executor/helps/codex_sse_scanner_test.go new file mode 100644 index 0000000000..289dd2e933 --- /dev/null +++ b/internal/runtime/executor/helps/codex_sse_scanner_test.go @@ -0,0 +1,341 @@ +package helps + +import ( + "bufio" + "context" + "fmt" + "io" + "strings" + "testing" + "time" + + log "github.com/sirupsen/logrus" + logtest "github.com/sirupsen/logrus/hooks/test" +) + +// wellFormedCodexSSEStream is a Responses stream with one event per data line, comments, event +// lines, events without a blank separator, and the [DONE] sentinel. +func wellFormedCodexSSEStream() string { + var stream strings.Builder + stream.WriteString("event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"model\":\"gpt-5.4\"}}\n\n") + for i := range 6 { + fmt.Fprintf(&stream, ": comment-%04d\nevent: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":%q}\n\n", i, strings.Repeat("x", 64)) + } + stream.WriteString("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"status\":\"completed\",\"output\":[]}}\n\n") + return stream.String() +} + +func cloneLogHooks(src log.LevelHooks) log.LevelHooks { + cloned := make(log.LevelHooks, len(src)) + for lvl, hooks := range src { + cloned[lvl] = append([]log.Hook(nil), hooks...) + } + return cloned +} + +func captureCodexSSEWarnings(t *testing.T) *logtest.Hook { + t.Helper() + origHooks := cloneLogHooks(log.StandardLogger().Hooks) + hook := new(logtest.Hook) + log.StandardLogger().AddHook(hook) + t.Cleanup(func() { log.StandardLogger().ReplaceHooks(origHooks) }) + return hook +} + +func scanCodexSSE(t *testing.T, stream string) []string { + t.Helper() + scanner := NewCodexSSEScanner(context.Background(), strings.NewReader(stream), 1<<20) + var lines []string + for scanner.Scan() { + lines = append(lines, string(scanner.Bytes())) + } + if errScan := scanner.Err(); errScan != nil { + t.Fatalf("scan error: %v", errScan) + } + return lines +} + +func TestCodexSSEScannerAssemblesEventData(t *testing.T) { + const marker = "body-bytes-must-not-be-logged" + cases := []struct { + name string + stream string + want []string + warnings int + }{ + { + name: "one-line events pass through as read", + stream: "event: response.created\ndata: {\"type\":\"response.created\"}\n\n: keep-alive\n\ndata: {\"type\":\"x\",\"v\":\"a b\"}\n\n", + want: []string{"event: response.created", `data: {"type":"response.created"}`, "", ": keep-alive", "", `data: {"type":"x","v":"a b"}`, ""}, + }, + { + name: "CRLF line endings", + stream: "event: a\r\ndata: {\"type\":\"a\"}\r\n\r\n", + want: []string{"event: a", `data: {"type":"a"}`, ""}, + }, + { + name: "data lines of one event are joined and compacted", + stream: "event: response.completed\ndata: {\"type\":\"response.completed\",\ndata: \"response\":{\"id\":\"r\",\"output\":[]}}\n\n", + want: []string{"event: response.completed", `data: {"type":"response.completed","response":{"id":"r","output":[]}}`, ""}, + }, + { + name: "joined CRLF event", + stream: "data: {\"a\":\r\ndata: [1,\r\ndata: 2]}\r\n\r\n", + want: []string{`data: {"a":[1,2]}`, ""}, + }, + { + name: "comment and field lines inside an event come before its data", + stream: "event: e\ndata: {\"a\":\n: keep-alive\nid: 7\nretry: 3000\ndata: 1}\n\n", + want: []string{"event: e", ": keep-alive", "id: 7", "retry: 3000", `data: {"a":1}`, ""}, + }, + { + name: "unterminated one-line event at end of stream", + stream: "data: {\"type\":\"response.completed\"}", + want: []string{`data: {"type":"response.completed"}`}, + }, + { + name: "unterminated multi-line event at end of stream", + stream: "data: {\"type\":\ndata: \"response.completed\"}\n", + want: []string{`data: {"type":"response.completed"}`}, + }, + { + name: "complete events without blank separators stay separate", + stream: "data: {\"a\":1}\ndata: {\"b\":2}\n", + want: []string{`data: {"a":1}`, `data: {"b":2}`}, + }, + { + name: "split event followed by an event without a blank separator", + stream: "data: {\"a\":\ndata: 1}\ndata: {\"b\":2}\n", + want: []string{`data: {"a":1}`, `data: {"b":2}`}, + }, + { + name: "an invalid fragment absorbs the following data lines up to the blank line", + stream: "data: {\"a\":\"" + marker + "\ndata: {\"b\":2}\n\n", + want: []string{"data: {\"a\":\"" + marker + "\ndata: {\"b\":2}", ""}, + warnings: 1, + }, + { + name: "only the first invalid event in a stream is described", + stream: "data: {\"a\":\n\ndata: {\"b\":\n\ndata: {\"c\":3}\n\n", + want: []string{"data: {\"a\":", "", "data: {\"b\":", "", `data: {"c":3}`, ""}, + warnings: 1, + }, + { + name: "a complete object on its own line inside an open array stays in the event", + stream: "data: {\"type\":\"response.completed\",\"response\":{\"output\":[\ndata: {\"id\":\"m\",\"type\":\"message\"}\ndata: ]}}\n\n", + want: []string{`data: {"type":"response.completed","response":{"output":[{"id":"m","type":"message"}]}}`, ""}, + }, + { + name: "unterminated incomplete event at end of stream is still delivered", + stream: "data: {\"type\":\"response.completed\",\"t\":\"" + marker + "\",", + want: []string{`data: {"type":"response.completed","t":"` + marker + `",`}, + warnings: 1, + }, + { + name: "empty and non-JSON data lines take the stock path", + stream: "data:\n\ndata: \n\ndata: ping\n\n", + want: []string{"data:", "", "data: ", "", "data: ping", ""}, + }, + { + name: "done sentinel", + stream: "data: [DONE]\n\n", + want: []string{"data: [DONE]", ""}, + }, + { + name: "invalid one-line event passes through with one warning", + stream: "data: {\"type\":\"response.completed\",\"t\":\"" + marker + "\",\n\n", + want: []string{`data: {"type":"response.completed","t":"` + marker + `",`, ""}, + warnings: 1, + }, + { + name: "event split inside a string stays invalid and keeps its data lines", + stream: "data: {\"t\":\"" + marker + "\ndata: " + marker + "\"}\n\n", + want: []string{`data: {"t":"` + marker + "\ndata: " + marker + `"}`, ""}, + warnings: 1, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + hook := captureCodexSSEWarnings(t) + got := scanCodexSSE(t, tc.stream) + if fmt.Sprintf("%q", got) != fmt.Sprintf("%q", tc.want) { + t.Fatalf("lines = %q\nwant %q", got, tc.want) + } + if len(hook.AllEntries()) != tc.warnings { + t.Fatalf("warnings = %d, want %d", len(hook.AllEntries()), tc.warnings) + } + for _, entry := range hook.AllEntries() { + rendered, errRender := entry.String() + if errRender != nil { + t.Fatal(errRender) + } + _, hasLines := entry.Data["data_lines"] + _, hasOffset := entry.Data["json_error_offset"] + if entry.Level != log.WarnLevel || strings.Contains(rendered, marker) || !hasLines || !hasOffset { + t.Fatalf("diagnostic = %q, want a bounded warning without data bytes", rendered) + } + } + }) + } +} + +func TestCodexSSEScannerBoundsAnEventToTheLineLimit(t *testing.T) { + hook := captureCodexSSEWarnings(t) + // Each fragment fits the 64-byte limit; together they do not. + stream := "data: {\"a\":\"" + strings.Repeat("x", 40) + "\ndata: " + strings.Repeat("y", 40) + "\ndata: " + strings.Repeat("z", 40) + "\"}\n\ndata: {\"b\":2}\n\n" + scanner := NewCodexSSEScanner(context.Background(), strings.NewReader(stream), 64) + var got []string + for scanner.Scan() { + got = append(got, string(scanner.Bytes())) + } + if errScan := scanner.Err(); errScan != nil { + t.Fatalf("scan error: %v", errScan) + } + // The collected fragment is flushed as an invalid event; the lines after it are read as stock reads them. + want := []string{ + "data: {\"a\":\"" + strings.Repeat("x", 40), + "data: " + strings.Repeat("y", 40), + "data: " + strings.Repeat("z", 40) + "\"}", + "", + `data: {"b":2}`, + "", + } + if fmt.Sprintf("%q", got) != fmt.Sprintf("%q", want) { + t.Fatalf("lines = %q\nwant %q", got, want) + } + if len(hook.AllEntries()) != 1 { + t.Fatalf("warnings = %d, want 1 for the flushed invalid event", len(hook.AllEntries())) + } +} + +func TestCodexSSEScannerAssemblesAnEventOfThousandsOfLinesQuickly(t *testing.T) { + // A 2 MB response.completed whose upstream wrapped it into 2,000 data lines between array + // elements. Every line ends with "}", so only nesting, not the last byte, tells the event apart + // from a finished one. + item := func(i int) string { + return fmt.Sprintf(`{"type":"message","content":[{"type":"output_text","text":"%s"}]}`, strings.Repeat(string(rune('a'+i%26)), 960)) + } + head := `{"type":"response.completed","response":{"id":"resp_big","output":[` + var payload, stream strings.Builder + payload.WriteString(head + item(0)) + stream.WriteString("data: " + head + item(0) + "\n") + for i := 1; i < 2000; i++ { + payload.WriteString("," + item(i)) + stream.WriteString("data: ," + item(i) + "\n") + } + payload.WriteString("]}}") + stream.WriteString("data: ]}}\n\n") + started := time.Now() + scanner := NewCodexSSEScanner(context.Background(), strings.NewReader(stream.String()), 52_428_800) + var got []string + for scanner.Scan() { + got = append(got, string(scanner.Bytes())) + } + elapsed := time.Since(started) + if len(got) != 2 || got[0] != "data: "+payload.String() { + t.Fatalf("assembled %d lines, want the event on one line and its blank line", len(got)) + } + if elapsed > 2*time.Second { + t.Fatalf("assembling 2,001 data lines took %v; validation must not grow with every line", elapsed) + } +} + +func TestCodexSSEScannerHoldsAnEventInMemoryThatFollowsItsBytes(t *testing.T) { + // One valid event of 100,000 empty data lines: about 600 KB, well within the line limit. + stream := "data: [\n" + strings.Repeat("data:\n", 100_000) + "data: ]\n\n" + var got []string + allocs := testing.AllocsPerRun(1, func() { + got = got[:0] + scanner := NewCodexSSEScanner(context.Background(), strings.NewReader(stream), 1<<20) + for scanner.Scan() { + got = append(got, string(scanner.Bytes())) + } + }) + if fmt.Sprintf("%q", got) != fmt.Sprintf("%q", []string{"data: []", ""}) { + t.Fatalf("lines = %q", got) + } + // Holding each line separately costs at least one allocation per line. + if allocs > 1000 { + t.Fatalf("assembling 100,002 data lines took %.0f allocations; holding an event must not allocate per line", allocs) + } +} + +func TestCodexSSEScannerReadsLineByLineWhileAssemblyIsOff(t *testing.T) { + stream := "data: {\"a\":\ndata: 1}\n\ndata: {\"b\":\ndata: 2}\n\n" + scanner := NewCodexSSEScanner(context.Background(), strings.NewReader(stream), 1<<20) + scanner.SetAssembly(false) + var got []string + for scanner.Scan() { + got = append(got, string(scanner.Bytes())) + if len(got) == 3 { + // The first event has been read as bufio.Scanner reads it; assemble from here on. + scanner.SetAssembly(true) + } + } + want := []string{"data: {\"a\":", "data: 1}", "", `data: {"b":2}`, ""} + if fmt.Sprintf("%q", got) != fmt.Sprintf("%q", want) { + t.Fatalf("lines = %q\nwant %q", got, want) + } +} + +// codexSSEChunkReader returns one chunk per Read and counts the reads. +type codexSSEChunkReader struct { + chunks []string + reads int +} + +func (r *codexSSEChunkReader) Read(p []byte) (int, error) { + if r.reads == len(r.chunks) { + return 0, io.EOF + } + r.reads++ + return copy(p, r.chunks[r.reads-1]), nil +} + +func TestCodexSSEScannerDeliversNonDataLinesAsTheyAreRead(t *testing.T) { + reader := &codexSSEChunkReader{chunks: []string{"data: {\"a\":\n: keep-alive\nid: 7\n", "data: 1}\n\n"}} + scanner := NewCodexSSEScanner(context.Background(), reader, 1<<20) + var got []string + for scanner.Scan() { + got = append(got, fmt.Sprintf("%q@%d", scanner.Bytes(), reader.reads)) + } + // Only the data lines wait for the rest of their event; the other lines are not held for it. + want := []string{`": keep-alive"@1`, `"id: 7"@1`, `"data: {\"a\":1}"@2`, `""@2`} + if fmt.Sprintf("%q", got) != fmt.Sprintf("%q", want) { + t.Fatalf("lines = %q\nwant %q", got, want) + } +} + +func TestCodexSSEScannerDiagnosticDescribesFraming(t *testing.T) { + hook := captureCodexSSEWarnings(t) + // A split inside a string: each fragment is invalid on its own, and so is their SSE join. + scanCodexSSE(t, "data: {\"type\":\"response.completed\",\"t\":\"abc\ndata: def\"}\n\n") + entries := hook.AllEntries() + if len(entries) != 1 { + t.Fatalf("warnings = %d, want 1", len(entries)) + } + want := log.Fields{"data_lines": 2, "line_bytes": "37,5", "lines_valid": "false,false", "data_bytes": 43, "json_error_offset": int64(38)} + for key, value := range want { + if entries[0].Data[key] != value { + t.Fatalf("diagnostic field %s = %v, want %v (fields %v)", key, entries[0].Data[key], value, entries[0].Data) + } + } +} + +func TestCodexSSEScannerMatchesBufioScannerOnWellFormedStreams(t *testing.T) { + stream := wellFormedCodexSSEStream() + "data: {\"type\":\"x\"}\ndata: {\"type\":\"y\"}\n: trailing comment\n\nevent: done\ndata: [DONE]\n\n" + for _, newline := range []string{"\n", "\r\n"} { + input := strings.ReplaceAll(stream, "\n", newline) + t.Run(fmt.Sprintf("newline=%q", newline), func(t *testing.T) { + stock := bufio.NewScanner(strings.NewReader(input)) + var want []string + for stock.Scan() { + want = append(want, stock.Text()) + } + got := scanCodexSSE(t, input) + if fmt.Sprintf("%q", got) != fmt.Sprintf("%q", want) { + t.Fatalf("lines = %q\nwant bufio.Scanner's %q", got, want) + } + }) + } +}