fix(providers): recover Codex text from responses completion events (#631)

* fix(providers): recover Codex text from responses completion events

* fix(providers): honor Codex phase-aware response assembly

* refactor(providers): extract Codex stream state to separate file

Split codexMessageStreamState and related types/methods from codex.go
(457 lines) into codex_stream_state.go to comply with 200-line file
size guideline. Fix redundant nil check in registerMessageItem. Add
clarifying comment on reconcileCompleted edge case.

---------

Co-authored-by: viettranx <viettranx@gmail.com>
This commit is contained in:
Kai (Tam Nhu) Tranandviettranx authored and GitHub committed 2026-04-02 08:47:26 +07:00
1 parent c84f4ac905
commit 9fd9e1f47b
4 files changed
+645 -34

No files matched your search

+24 -10
View File
@@ -93,6 +93,7 @@ func (p *CodexProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk
result := &ChatResponse{FinishReason: "stop"}
toolCalls := make(map[string]*codexToolCallAcc) // keyed by item_id
streamState := newCodexMessageStreamState()
scanner := bufio.NewScanner(respBody)
scanner.Buffer(make([]byte, 0, SSEScanBufInit), SSEScanBufMax)
@@ -112,7 +113,7 @@ func (p *CodexProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk
continue
}
p.processSSEEvent(&event, result, toolCalls, onChunk)
p.processSSEEvent(&event, result, toolCalls, streamState, onChunk)
}
if err := scanner.Err(); err != nil {
@@ -151,14 +152,22 @@ func (p *CodexProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk
}
// processSSEEvent handles a single SSE event during streaming.
func (p *CodexProvider) processSSEEvent(event *codexSSEEvent, result *ChatResponse, toolCalls map[string]*codexToolCallAcc, onChunk func(StreamChunk)) {
func (p *CodexProvider) processSSEEvent(event *codexSSEEvent, result *ChatResponse, toolCalls map[string]*codexToolCallAcc, streamState *codexMessageStreamState, onChunk func(StreamChunk)) {
switch event.Type {
case "response.output_item.added":
if event.Item != nil {
streamState.registerMessageItem(event.ItemID, event.OutputIndex, event.Item)
}
case "response.output_text.delta":
if event.Delta != "" {
result.Content += event.Delta
if onChunk != nil {
onChunk(StreamChunk{Content: event.Delta})
}
streamState.recordTextDelta(event.ItemID, event.OutputIndex, event.ContentIndex, event.Delta, result, onChunk)
case "response.output_text.done":
streamState.recordFinalText(event.ItemID, event.OutputIndex, event.ContentIndex, event.Text, result, onChunk)
case "response.content_part.done":
if event.Part != nil && event.Part.Type == "output_text" {
streamState.recordFinalText(event.ItemID, event.OutputIndex, event.ContentIndex, event.Part.Text, result, onChunk)
}
case "response.function_call_arguments.delta":
@@ -175,9 +184,9 @@ func (p *CodexProvider) processSSEEvent(event *codexSSEEvent, result *ChatRespon
if event.Item != nil {
switch event.Item.Type {
case "message":
if event.Item.Phase != "" {
result.Phase = event.Item.Phase
}
streamState.registerMessageItem(event.ItemID, event.OutputIndex, event.Item)
streamState.flushMessage(codexEventItemKey(event.ItemID, event.Item), result, onChunk)
streamState.updateResultPhase(result)
case "function_call":
acc := toolCalls[event.Item.ID]
if acc == nil {
@@ -203,6 +212,11 @@ func (p *CodexProvider) processSSEEvent(event *codexSSEEvent, result *ChatRespon
case "response.completed", "response.incomplete", "response.failed":
if event.Response != nil {
if result.Content == "" {
streamState.ingestCompletedResponse(event.Response)
streamState.flushCompletedResponse(result, onChunk)
streamState.updateResultPhase(result)
}
if event.Response.Usage != nil {
u := event.Response.Usage
result.Usage = &Usage{
+239
View File
@@ -0,0 +1,239 @@
package providers
import (
"fmt"
"sort"
"strings"
)
// codexMessageStreamState tracks message items and their text parts during
// SSE streaming from the Codex/Responses API. It enables recovery of assistant
// text from multiple completion event types (output_text.done, output_item.done,
// response.completed) and filters out commentary-phase messages.
//
// Not thread-safe — designed for sequential SSE processing in a single goroutine.
type codexMessageStreamState struct {
messages map[string]*codexMessageState
}
type codexMessageState struct {
id string
outputIndex int
phase string
parts map[int]*codexTextPartState
}
type codexTextPartState struct {
text string
emittedBytes int
}
func newCodexMessageStreamState() *codexMessageStreamState {
return &codexMessageStreamState{messages: make(map[string]*codexMessageState)}
}
func (s *codexMessageStreamState) registerMessageItem(itemID string, outputIndex int, item *codexItem) {
if item == nil || item.Type != "message" {
return
}
msg := s.ensureMessage(itemID, item.ID, outputIndex)
if item.Phase != "" {
msg.phase = item.Phase
}
if msg.outputIndex == 0 && outputIndex != 0 {
msg.outputIndex = outputIndex
}
for idx, part := range item.Content {
if part.Type != "output_text" || part.Text == "" {
continue
}
textPart := msg.ensurePart(idx)
textPart.text = part.Text
}
}
func (s *codexMessageStreamState) recordTextDelta(itemID string, outputIndex, contentIndex int, delta string, result *ChatResponse, onChunk func(StreamChunk)) {
if delta == "" {
return
}
msg := s.ensureMessage(itemID, "", outputIndex)
part := msg.ensurePart(contentIndex)
part.text += delta
if !shouldEmitCodexPhase(msg.phase) {
return
}
msg.flushContiguous(result, onChunk)
}
func (s *codexMessageStreamState) recordFinalText(itemID string, outputIndex, contentIndex int, text string, result *ChatResponse, onChunk func(StreamChunk)) {
if text == "" {
return
}
msg := s.ensureMessage(itemID, "", outputIndex)
part := msg.ensurePart(contentIndex)
prev := part.text
part.text = text
if !shouldEmitCodexPhase(msg.phase) {
return
}
part.reconcileCompleted(prev)
msg.flushContiguous(result, onChunk)
}
func (s *codexMessageStreamState) flushMessage(itemID string, result *ChatResponse, onChunk func(StreamChunk)) {
msg, ok := s.messages[itemID]
if !ok || !shouldEmitCodexPhase(msg.phase) {
return
}
msg.flushContiguous(result, onChunk)
}
func (s *codexMessageStreamState) ingestCompletedResponse(resp *codexAPIResponse) {
if resp == nil {
return
}
for i := range resp.Output {
item := &resp.Output[i]
if item.Type != "message" {
continue
}
s.registerMessageItem(item.ID, i, item)
}
}
func (s *codexMessageStreamState) flushCompletedResponse(result *ChatResponse, onChunk func(StreamChunk)) {
ordered := s.preferredMessages()
for _, msg := range ordered {
msg.flushContiguous(result, onChunk)
}
}
func (s *codexMessageStreamState) updateResultPhase(result *ChatResponse) {
if result == nil || result.Phase == "final_answer" {
return
}
for _, msg := range s.messages {
if msg.phase == "final_answer" {
result.Phase = "final_answer"
return
}
}
}
// preferredMessages returns messages ordered by outputIndex, preferring
// final_answer phase. Falls back to non-commentary messages if no
// final_answer is found.
func (s *codexMessageStreamState) preferredMessages() []*codexMessageState {
if len(s.messages) == 0 {
return nil
}
ordered := make([]*codexMessageState, 0, len(s.messages))
for _, msg := range s.messages {
if msg.phase == "final_answer" {
ordered = append(ordered, msg)
}
}
if len(ordered) == 0 {
for _, msg := range s.messages {
if msg.phase != "commentary" {
ordered = append(ordered, msg)
}
}
}
sort.SliceStable(ordered, func(i, j int) bool {
return ordered[i].outputIndex < ordered[j].outputIndex
})
return ordered
}
func codexEventItemKey(eventItemID string, item *codexItem) string {
if eventItemID != "" {
return eventItemID
}
if item != nil {
return item.ID
}
return ""
}
func (s *codexMessageStreamState) ensureMessage(itemID, fallbackID string, outputIndex int) *codexMessageState {
key := itemID
if key == "" {
key = fallbackID
}
if key == "" {
key = fmt.Sprintf("output:%d", outputIndex)
}
if msg, ok := s.messages[key]; ok {
if msg.outputIndex == 0 && outputIndex != 0 {
msg.outputIndex = outputIndex
}
return msg
}
msg := &codexMessageState{
id: key,
outputIndex: outputIndex,
parts: make(map[int]*codexTextPartState),
}
s.messages[key] = msg
return msg
}
func (m *codexMessageState) ensurePart(contentIndex int) *codexTextPartState {
if part, ok := m.parts[contentIndex]; ok {
return part
}
part := &codexTextPartState{}
m.parts[contentIndex] = part
return part
}
func (m *codexMessageState) flushContiguous(result *ChatResponse, onChunk func(StreamChunk)) {
for nextIndex := 0; ; nextIndex++ {
part, ok := m.parts[nextIndex]
if !ok {
return
}
part.emitMissing(result, onChunk)
}
}
func (p *codexTextPartState) emitMissing(result *ChatResponse, onChunk func(StreamChunk)) {
if p.text == "" || p.emittedBytes >= len(p.text) {
return
}
suffix := p.text[p.emittedBytes:]
appendCodexContent(result, suffix, onChunk)
p.emittedBytes = len(p.text)
}
// reconcileCompleted adjusts emittedBytes after the final text replaces delta-
// accumulated text. The SSE protocol guarantees that output_text.done carries
// the same content as the concatenated deltas, so the emitted prefix is always
// a valid prefix of the final text. The guard handles rare edge cases where
// the provider omits some deltas.
func (p *codexTextPartState) reconcileCompleted(previous string) {
if p.text == "" {
return
}
if len(previous) >= p.emittedBytes && strings.HasPrefix(p.text, previous[:p.emittedBytes]) {
return
}
if p.emittedBytes > len(p.text) {
p.emittedBytes = len(p.text)
}
}
func shouldEmitCodexPhase(phase string) bool {
return phase == "" || phase == "final_answer"
}
func appendCodexContent(result *ChatResponse, text string, onChunk func(StreamChunk)) {
if text == "" {
return
}
result.Content += text
if onChunk != nil {
onChunk(StreamChunk{Content: text})
}
}
+349
View File
@@ -345,6 +345,355 @@ func TestCodexProviderChatStream(t *testing.T) {
}
}
func TestCodexProviderChatFallsBackToOutputTextDone(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprintf(w, "data: %s\n\n", mustJSON(codexSSEEvent{
Type: "response.output_text.done",
Text: "<file name=\"SOUL.md\">hello</file>",
}))
fmt.Fprintf(w, "data: %s\n\n", mustJSON(codexSSEEvent{
Type: "response.completed",
Response: &codexAPIResponse{
ID: "resp-done", Status: "completed",
Usage: &codexUsage{InputTokens: 5, OutputTokens: 4, TotalTokens: 9},
},
}))
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer server.Close()
p := NewCodexProvider("test", &staticTokenSource{token: "test"}, server.URL, "gpt-4o")
p.retryConfig.Attempts = 1
result, err := p.Chat(context.Background(), ChatRequest{
Messages: []Message{{Role: "user", Content: "Generate a file"}},
})
if err != nil {
t.Fatalf("Chat: %v", err)
}
if result.Content != "<file name=\"SOUL.md\">hello</file>" {
t.Errorf("Content = %q, want fallback output_text.done content", result.Content)
}
}
func TestCodexProviderChatFallsBackToOutputItemMessageContent(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprintf(w, "data: %s\n\n", mustJSON(codexSSEEvent{
Type: "response.output_item.done",
Item: &codexItem{
Type: "message",
Role: "assistant",
Phase: "final_answer",
Content: []codexContent{
{Type: "output_text", Text: "<file name=\"IDENTITY.md\">world</file>"},
},
},
}))
fmt.Fprintf(w, "data: %s\n\n", mustJSON(codexSSEEvent{
Type: "response.completed",
Response: &codexAPIResponse{
ID: "resp-item", Status: "completed",
Usage: &codexUsage{InputTokens: 5, OutputTokens: 4, TotalTokens: 9},
},
}))
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer server.Close()
p := NewCodexProvider("test", &staticTokenSource{token: "test"}, server.URL, "gpt-4o")
p.retryConfig.Attempts = 1
result, err := p.Chat(context.Background(), ChatRequest{
Messages: []Message{{Role: "user", Content: "Generate a file"}},
})
if err != nil {
t.Fatalf("Chat: %v", err)
}
if result.Content != "<file name=\"IDENTITY.md\">world</file>" {
t.Errorf("Content = %q, want fallback output_item.done content", result.Content)
}
if result.Phase != "final_answer" {
t.Errorf("Phase = %q, want final_answer", result.Phase)
}
}
func TestCodexProviderChatIgnoresCommentaryAndKeepsFinalAnswer(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
flusher := w.(http.Flusher)
events := []codexSSEEvent{
{
Type: "response.output_item.added",
OutputIndex: 0,
Item: &codexItem{
ID: "msg_commentary",
Type: "message",
Phase: "commentary",
},
},
{
Type: "response.output_text.delta",
ItemID: "msg_commentary",
OutputIndex: 0,
ContentIndex: 0,
Delta: "Planning out the answer...",
},
{
Type: "response.output_item.done",
OutputIndex: 0,
Item: &codexItem{
ID: "msg_commentary",
Type: "message",
Phase: "commentary",
},
},
{
Type: "response.output_item.added",
OutputIndex: 1,
Item: &codexItem{
ID: "msg_final",
Type: "message",
Phase: "final_answer",
},
},
{
Type: "response.output_text.delta",
ItemID: "msg_final",
OutputIndex: 1,
ContentIndex: 0,
Delta: "<file name=\"SOUL.md\">hello</file>",
},
{
Type: "response.output_item.done",
OutputIndex: 1,
Item: &codexItem{
ID: "msg_final",
Type: "message",
Phase: "final_answer",
},
},
{
Type: "response.completed",
Response: &codexAPIResponse{
ID: "resp-commentary", Status: "completed",
Usage: &codexUsage{InputTokens: 5, OutputTokens: 4, TotalTokens: 9},
},
},
}
for _, event := range events {
fmt.Fprintf(w, "data: %s\n\n", mustJSON(event))
flusher.Flush()
}
fmt.Fprint(w, "data: [DONE]\n\n")
flusher.Flush()
}))
defer server.Close()
p := NewCodexProvider("test", &staticTokenSource{token: "test"}, server.URL, "gpt-4o")
p.retryConfig.Attempts = 1
var chunks []string
result, err := p.ChatStream(context.Background(), ChatRequest{
Messages: []Message{{Role: "user", Content: "Generate files"}},
}, func(chunk StreamChunk) {
if chunk.Content != "" {
chunks = append(chunks, chunk.Content)
}
})
if err != nil {
t.Fatalf("ChatStream: %v", err)
}
if result.Content != "<file name=\"SOUL.md\">hello</file>" {
t.Errorf("Content = %q, want final_answer only", result.Content)
}
if len(chunks) != 1 || chunks[0] != "<file name=\"SOUL.md\">hello</file>" {
t.Errorf("chunks = %v, want only final_answer chunk", chunks)
}
if result.Phase != "final_answer" {
t.Errorf("Phase = %q, want final_answer", result.Phase)
}
}
func TestCodexProviderChatConcatenatesMultipleOutputTextDoneParts(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
flusher := w.(http.Flusher)
events := []codexSSEEvent{
{
Type: "response.output_item.added",
OutputIndex: 0,
Item: &codexItem{
ID: "msg_final",
Type: "message",
Phase: "final_answer",
},
},
{
Type: "response.output_text.done",
ItemID: "msg_final",
OutputIndex: 0,
ContentIndex: 1,
Text: "hello</file>",
},
{
Type: "response.output_text.done",
ItemID: "msg_final",
OutputIndex: 0,
ContentIndex: 0,
Text: "<file name=\"SOUL.md\">",
},
{
Type: "response.completed",
Response: &codexAPIResponse{
ID: "resp-multipart", Status: "completed",
Usage: &codexUsage{InputTokens: 5, OutputTokens: 4, TotalTokens: 9},
},
},
}
for _, event := range events {
fmt.Fprintf(w, "data: %s\n\n", mustJSON(event))
flusher.Flush()
}
fmt.Fprint(w, "data: [DONE]\n\n")
flusher.Flush()
}))
defer server.Close()
p := NewCodexProvider("test", &staticTokenSource{token: "test"}, server.URL, "gpt-4o")
p.retryConfig.Attempts = 1
result, err := p.Chat(context.Background(), ChatRequest{
Messages: []Message{{Role: "user", Content: "Generate a file"}},
})
if err != nil {
t.Fatalf("Chat: %v", err)
}
if result.Content != "<file name=\"SOUL.md\">hello</file>" {
t.Errorf("Content = %q, want concatenated output_text.done parts", result.Content)
}
}
func TestCodexProviderChatFallsBackToContentPartDone(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
flusher := w.(http.Flusher)
events := []codexSSEEvent{
{
Type: "response.output_item.added",
OutputIndex: 0,
Item: &codexItem{
ID: "msg_final",
Type: "message",
Phase: "final_answer",
},
},
{
Type: "response.content_part.done",
ItemID: "msg_final",
OutputIndex: 0,
ContentIndex: 0,
Part: &codexContentPart{
Type: "output_text",
Text: "<frontmatter>summary</frontmatter>",
},
},
{
Type: "response.completed",
Response: &codexAPIResponse{
ID: "resp-content-part", Status: "completed",
Usage: &codexUsage{InputTokens: 5, OutputTokens: 4, TotalTokens: 9},
},
},
}
for _, event := range events {
fmt.Fprintf(w, "data: %s\n\n", mustJSON(event))
flusher.Flush()
}
fmt.Fprint(w, "data: [DONE]\n\n")
flusher.Flush()
}))
defer server.Close()
p := NewCodexProvider("test", &staticTokenSource{token: "test"}, server.URL, "gpt-4o")
p.retryConfig.Attempts = 1
result, err := p.Chat(context.Background(), ChatRequest{
Messages: []Message{{Role: "user", Content: "Generate a summary"}},
})
if err != nil {
t.Fatalf("Chat: %v", err)
}
if result.Content != "<frontmatter>summary</frontmatter>" {
t.Errorf("Content = %q, want content_part.done fallback content", result.Content)
}
}
func TestCodexProviderChatFallsBackToCompletedResponseOutput(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprintf(w, "data: %s\n\n", mustJSON(codexSSEEvent{
Type: "response.completed",
Response: &codexAPIResponse{
ID: "resp-completed",
Status: "completed",
Output: []codexItem{
{
ID: "msg_commentary",
Type: "message",
Role: "assistant",
Phase: "commentary",
Content: []codexContent{
{Type: "output_text", Text: "Draft commentary"},
},
},
{
ID: "msg_final",
Type: "message",
Role: "assistant",
Phase: "final_answer",
Content: []codexContent{
{Type: "output_text", Text: "<frontmatter>summary</frontmatter>"},
},
},
},
Usage: &codexUsage{InputTokens: 6, OutputTokens: 3, TotalTokens: 9},
},
}))
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer server.Close()
p := NewCodexProvider("test", &staticTokenSource{token: "test"}, server.URL, "gpt-4o")
p.retryConfig.Attempts = 1
result, err := p.Chat(context.Background(), ChatRequest{
Messages: []Message{{Role: "user", Content: "Generate a summary"}},
})
if err != nil {
t.Fatalf("Chat: %v", err)
}
if result.Content != "<frontmatter>summary</frontmatter>" {
t.Errorf("Content = %q, want fallback response.completed output content", result.Content)
}
if result.Phase != "final_answer" {
t.Errorf("Phase = %q, want final_answer", result.Phase)
}
}
func TestCodexProviderChatStreamToolCalls(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
+33 -24
View File
@@ -3,24 +3,24 @@ package providers
// Wire types for the OpenAI Responses API (Codex flow).
type codexAPIResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Model string `json:"model"`
Output []codexItem `json:"output"`
Usage *codexUsage `json:"usage,omitempty"`
Status string `json:"status"`
ID string `json:"id"`
Object string `json:"object"`
Model string `json:"model"`
Output []codexItem `json:"output"`
Usage *codexUsage `json:"usage,omitempty"`
Status string `json:"status"`
}
type codexItem struct {
ID string `json:"id"`
Type string `json:"type"` // "message", "function_call", "reasoning"
Role string `json:"role,omitempty"`
Phase string `json:"phase,omitempty"` // gpt-5.3-codex: "commentary" or "final_answer"
Content []codexContent `json:"content,omitempty"`
CallID string `json:"call_id,omitempty"`
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
Summary []codexSummary `json:"summary,omitempty"`
ID string `json:"id"`
Type string `json:"type"` // "message", "function_call", "reasoning"
Role string `json:"role,omitempty"`
Phase string `json:"phase,omitempty"` // gpt-5.3-codex: "commentary" or "final_answer"
Content []codexContent `json:"content,omitempty"`
CallID string `json:"call_id,omitempty"`
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
Summary []codexSummary `json:"summary,omitempty"`
}
type codexContent struct {
@@ -34,10 +34,10 @@ type codexSummary struct {
}
type codexUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
OutputTokensDetails *codexTokensDetails `json:"output_tokens_details,omitempty"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
OutputTokensDetails *codexTokensDetails `json:"output_tokens_details,omitempty"`
}
type codexTokensDetails struct {
@@ -47,11 +47,15 @@ type codexTokensDetails struct {
// SSE streaming types
type codexSSEEvent struct {
Type string `json:"type"`
Delta string `json:"delta,omitempty"`
ItemID string `json:"item_id,omitempty"`
Item *codexItem `json:"item,omitempty"`
Response *codexAPIResponse `json:"response,omitempty"`
Type string `json:"type"`
Delta string `json:"delta,omitempty"`
Text string `json:"text,omitempty"`
ItemID string `json:"item_id,omitempty"`
OutputIndex int `json:"output_index,omitempty"`
ContentIndex int `json:"content_index,omitempty"`
Item *codexItem `json:"item,omitempty"`
Part *codexContentPart `json:"part,omitempty"`
Response *codexAPIResponse `json:"response,omitempty"`
}
type codexToolCallAcc struct {
@@ -59,3 +63,8 @@ type codexToolCallAcc struct {
name string
rawArgs string
}
type codexContentPart struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
}