diff --git a/backend/internal/service/openai_image_generation_controls_test.go b/backend/internal/service/openai_image_generation_controls_test.go index de2446eb7f..300bea3763 100644 --- a/backend/internal/service/openai_image_generation_controls_test.go +++ b/backend/internal/service/openai_image_generation_controls_test.go @@ -649,6 +649,12 @@ func TestNormalizeCompletedImageGenerationStatus(t *testing.T) { want: `{"type":"response.output_item.added","item":{"type":"image_generation_call","status":"generating","result":"image-data"}}`, wantChanged: false, }, + { + name: "done preserves base64 result", + input: `{"type":"response.done","response":{"output":[{"type":"image_generation_call","status":"generating","result":"iVBORw0KGgoAAAANSUhEUg/+=="}]}}`, + want: `{"type":"response.done","response":{"output":[{"type":"image_generation_call","status":"completed","result":"iVBORw0KGgoAAAANSUhEUg/+=="}]}}`, + wantChanged: true, + }, } for _, tt := range tests { diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index f6a7dca250..67bffe6ce2 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -803,6 +803,9 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( wroteDownstream, ) } + if normalized, changed := normalizeCompletedImageGenerationStatus(upstreamMessage); changed { + upstreamMessage = normalized + } eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage) if responseID == "" && eventResponseID != "" { diff --git a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go index 8fbf168f66..c90a1ebc0c 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -38,6 +38,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT captureConn := &openAIWSCaptureConn{ events: [][]byte{ + []byte(`{"type":"response.output_item.done","item":{"id":"ig_ingress_1","type":"image_generation_call","status":"generating","result":"iVBORw0KGgoAAAANSUhEUg/+=="}}`), []byte(`{"type":"response.completed","response":{"id":"resp_ingress_turn_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), []byte(`{"type":"response.completed","response":{"id":"resp_ingress_turn_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), }, @@ -138,6 +139,10 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT } writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false}`) + firstTurnImageEvent := readMessage() + require.Equal(t, "response.output_item.done", gjson.GetBytes(firstTurnImageEvent, "type").String()) + require.Equal(t, "completed", gjson.GetBytes(firstTurnImageEvent, "item.status").String()) + require.Equal(t, "iVBORw0KGgoAAAANSUhEUg/+==", gjson.GetBytes(firstTurnImageEvent, "item.result").String()) firstTurnEvent := readMessage() require.Equal(t, "response.completed", gjson.GetBytes(firstTurnEvent, "type").String()) require.Equal(t, "resp_ingress_turn_1", gjson.GetBytes(firstTurnEvent, "response.id").String()) @@ -765,7 +770,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughModeR upstreamConn := &openAIWSCaptureConn{ events: [][]byte{ - []byte(`{"type":"response.completed","response":{"id":"resp_passthrough_turn_1","model":"gpt-5.1","usage":{"input_tokens":2,"output_tokens":3}}}`), + []byte(`{"type":"response.completed","response":{"id":"resp_passthrough_turn_1","model":"gpt-5.1","output":[{"id":"ig_passthrough_1","type":"image_generation_call","status":"generating","result":"final-image"}],"usage":{"input_tokens":2,"output_tokens":3}}}`), }, } captureDialer := &openAIWSCaptureDialer{conn: upstreamConn} @@ -858,6 +863,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughModeR require.NoError(t, readErr) require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) require.Equal(t, "resp_passthrough_turn_1", gjson.GetBytes(event, "response.id").String()) + require.Equal(t, "completed", gjson.GetBytes(event, "response.output.0.status").String()) _ = clientConn.Close(coderws.StatusNormalClosure, "done") select { @@ -1047,7 +1053,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_HTTPBridgeModeRe Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_bridge_1"}}, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n" + - "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_http_bridge_1\",\"usage\":{\"input_tokens\":2,\"output_tokens\":1,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n" + + "data: {\"type\":\"response.done\",\"response\":{\"id\":\"resp_http_bridge_1\",\"output\":[{\"id\":\"ig_bridge_1\",\"type\":\"image_generation_call\",\"status\":\"in_progress\",\"result\":\"final-image\"}],\"usage\":{\"input_tokens\":2,\"output_tokens\":1,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n" + "data: [DONE]\n\n", )), }, @@ -1145,8 +1151,9 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_HTTPBridgeModeRe _, event2, readErr2 := clientConn.Read(readCtx2) cancelRead2() require.NoError(t, readErr2) - require.Equal(t, "response.completed", gjson.GetBytes(event2, "type").String()) + require.Equal(t, "response.done", gjson.GetBytes(event2, "type").String()) require.Equal(t, "resp_http_bridge_1", gjson.GetBytes(event2, "response.id").String()) + require.Equal(t, "completed", gjson.GetBytes(event2, "response.output.0.status").String()) _ = clientConn.Close(coderws.StatusNormalClosure, "done") diff --git a/backend/internal/service/openai_ws_forwarder_success_test.go b/backend/internal/service/openai_ws_forwarder_success_test.go index 37c68a4126..0d031de189 100644 --- a/backend/internal/service/openai_ws_forwarder_success_test.go +++ b/backend/internal/service/openai_ws_forwarder_success_test.go @@ -281,6 +281,7 @@ func TestOpenAIGatewayService_Forward_WSv2_ImageGenerationCountsOutputs(t *testi "item": map[string]any{ "id": "ig_ws_1", "type": "image_generation_call", + "status": "generating", "result": "final-image", }, }); err != nil { @@ -296,6 +297,7 @@ func TestOpenAIGatewayService_Forward_WSv2_ImageGenerationCountsOutputs(t *testi map[string]any{ "id": "ig_ws_1", "type": "image_generation_call", + "status": "in_progress", "result": "final-image", }, }, @@ -375,6 +377,7 @@ func TestOpenAIGatewayService_Forward_WSv2_ImageGenerationCountsOutputs(t *testi require.Equal(t, 4, result.Usage.OutputTokens) require.True(t, result.OpenAIWSMode) require.Equal(t, "resp_ws_image_1", gjson.GetBytes(rec.Body.Bytes(), "id").String()) + require.Equal(t, "completed", gjson.GetBytes(rec.Body.Bytes(), "output.0.status").String()) } func requestToJSONString(payload map[string]any) string { diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index c406b6b337..62aee29c20 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -504,6 +504,9 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( setOpsUpstreamError(c, 0, sanitizeUpstreamErrorMessage(readErr.Error()), "") return nil, fmt.Errorf("openai ws read event: %w", readErr) } + if normalized, changed := normalizeCompletedImageGenerationStatus(message); changed { + message = normalized + } eventType, eventResponseID, responseField := parseOpenAIWSEventEnvelope(message) if eventType == "" { diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index 9140f61443..09b3dc0628 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -331,6 +331,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( } upstreamMessage := []byte(trimmedData) + if normalized, changed := normalizeCompletedImageGenerationStatus(upstreamMessage); changed { + upstreamMessage = normalized + } eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage) if responseID == "" && eventResponseID != "" { responseID = eventResponseID diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index 4505ffb665..af1bf7456a 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -212,6 +212,11 @@ func (c *openAIWSClientFrameConn) WriteFrame(ctx context.Context, msgType coderw if ctx == nil { ctx = context.Background() } + if msgType == coderws.MessageText { + if normalized, changed := normalizeCompletedImageGenerationStatus(payload); changed { + payload = normalized + } + } return c.conn.Write(ctx, msgType, payload) }