From 3b9611905094fa34bee560b0b004b2b6b2345b53 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Mon, 15 Jun 2026 01:03:19 +0800 Subject: [PATCH] feat(websockets): handle terminal events and improve error propagation - Enhanced Codex Websockets Executor to capture `response.done` as a terminal event, alongside `response.completed` and `error`. - Improved error propagation for upstream websocket errors with comprehensive message handling. - Introduced utility functions for recognizing terminal events and extracting error messages. - Expanded tests to validate new websocket event logic, including terminal event handling and upstream error propagation. --- .../executor/codex_websockets_executor.go | 24 ++- .../codex_websockets_executor_test.go | 122 +++++++++++++ .../openai/openai_responses_websocket.go | 39 ++++- .../openai/openai_responses_websocket_test.go | 165 ++++++++++++++++++ 4 files changed, 347 insertions(+), 3 deletions(-) diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index 30ae848e..35d6fc94 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -650,15 +650,35 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr return } - payload = normalizeCodexWebsocketCompletion(payload) eventType := gjson.GetBytes(payload, "type").String() + isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error" + clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) + if cliproxyexecutor.DownstreamWebsocket(ctx) { + if eventType == "response.completed" || eventType == "response.done" { + if detail, ok := helps.ParseCodexUsage(payload); ok { + reporter.Publish(ctx, detail) + } + } + if !send(cliproxyexecutor.StreamChunk{Payload: clientPayload}) { + terminateReason = "context_done" + terminateErr = ctx.Err() + return + } + if isTerminalEvent { + return + } + continue + } + + payload = normalizeCodexWebsocketCompletion(payload) + eventType = gjson.GetBytes(payload, "type").String() if eventType == "response.completed" || eventType == "response.done" { if detail, ok := helps.ParseCodexUsage(payload); ok { reporter.Publish(ctx, detail) } } - clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) + clientPayload = applyCodexIdentityExposeResponsePayload(payload, identityState) line := encodeCodexWebsocketAsSSE(clientPayload) chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, clientBody, clientBody, line, ¶m) for i := range chunks { diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go index a3d3a552..b0093542 100644 --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -93,6 +93,128 @@ func TestCodexWebsocketsExecutePreservesPreviousResponseIDUpstream(t *testing.T) } } +func TestCodexWebsocketsExecuteStreamPassesThroughUpstreamWebsocketPayloadForDownstreamWebsocket(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + delta := []byte(`{"type":"response.output_text.delta","delta":"hello"}`) + completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + defer func() { _ = conn.Close() }() + + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Errorf("read upstream websocket message: %v", errRead) + return + } + if errWrite := conn.WriteMessage(websocket.TextMessage, delta); errWrite != nil { + t.Errorf("write delta websocket message: %v", errWrite) + return + } + if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { + t.Errorf("write completed websocket message: %v", errWrite) + return + } + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} + req := cliproxyexecutor.Request{ + Model: "gpt-5-codex", + Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + ResponseFormat: sdktranslator.FromString("openai-response"), + } + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + + result, err := exec.ExecuteStream(ctx, auth, req, opts) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + + select { + case chunk, ok := <-result.Chunks: + if !ok { + t.Fatal("stream closed before first chunk") + } + if chunk.Err != nil { + t.Fatalf("first chunk error = %v", chunk.Err) + } + if !bytes.Equal(bytes.TrimSpace(chunk.Payload), delta) { + t.Fatalf("first chunk = %q, want raw upstream websocket payload %q", chunk.Payload, delta) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for first stream chunk") + } +} + +func TestCodexWebsocketsExecuteStreamPropagatesUpstreamErrorForDownstreamWebsocket(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + errorPayload := []byte(`{"type":"error","status":429,"error":{"code":"websocket_connection_limit_reached","message":"too many websockets"}}`) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + defer func() { _ = conn.Close() }() + + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Errorf("read upstream websocket message: %v", errRead) + return + } + if errWrite := conn.WriteMessage(websocket.TextMessage, errorPayload); errWrite != nil { + t.Errorf("write error websocket message: %v", errWrite) + return + } + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} + req := cliproxyexecutor.Request{ + Model: "gpt-5-codex", + Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + ResponseFormat: sdktranslator.FromString("openai-response"), + } + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + + result, err := exec.ExecuteStream(ctx, auth, req, opts) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + + select { + case chunk, ok := <-result.Chunks: + if !ok { + t.Fatal("stream closed before error chunk") + } + if len(bytes.TrimSpace(chunk.Payload)) != 0 { + t.Fatalf("error chunk payload = %q, want empty", chunk.Payload) + } + if chunk.Err == nil { + t.Fatal("error chunk Err = nil, want upstream error") + } + statusErr, ok := chunk.Err.(interface{ StatusCode() int }) + if !ok { + t.Fatalf("error type %T does not expose StatusCode", chunk.Err) + } + if got := statusErr.StatusCode(); got != http.StatusTooManyRequests { + t.Fatalf("status = %d, want %d", got, http.StatusTooManyRequests) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for error stream chunk") + } +} + func TestCodexWebsocketsUpstreamDisconnectChanSignalsOnInvalidate(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/sdk/api/handlers/openai/openai_responses_websocket.go b/sdk/api/handlers/openai/openai_responses_websocket.go index 3537f5ed..8113cdbb 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket.go +++ b/sdk/api/handlers/openai/openai_responses_websocket.go @@ -33,6 +33,7 @@ const ( wsRequestTypeAppend = "response.append" wsEventTypeError = "error" wsEventTypeCompleted = "response.completed" + wsEventTypeDone = "response.done" wsDoneMarker = "[DONE]" wsTurnStateHeader = "x-codex-turn-state" wsTimelineBodyKey = "WEBSOCKET_TIMELINE_OVERRIDE" @@ -1284,7 +1285,13 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i]) recordPendingToolCallIDsFromPayload(pendingToolCallIDs, payloads[i]) eventType := gjson.GetBytes(payloads[i], "type").String() - if eventType == wsEventTypeCompleted { + var payloadErrMsg *interfaces.ErrorMessage + if eventType == wsEventTypeError { + payloadErrMsg = responsesWebsocketErrorMessageFromPayload(payloads[i]) + if h != nil { + h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), payloadErrMsg) + } + } else if isResponsesWebsocketCompletionEvent(eventType) { completed = true completedOutput = responseCompletedOutputFromPayload(payloads[i]) completedResponseID = responseCompletedIDFromPayload(payloads[i]) @@ -1307,6 +1314,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( cancel(errWrite) return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, errWrite } + if payloadErrMsg != nil { + cancel(payloadErrMsg.Error) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, nil + } } } } @@ -1530,6 +1541,32 @@ func websocketPayloadPreview(payload []byte) string { return previewText } +func isResponsesWebsocketCompletionEvent(eventType string) bool { + return eventType == wsEventTypeCompleted || eventType == wsEventTypeDone +} + +func responsesWebsocketErrorMessageFromPayload(payload []byte) *interfaces.ErrorMessage { + status := int(gjson.GetBytes(payload, "status").Int()) + if status <= 0 { + status = int(gjson.GetBytes(payload, "status_code").Int()) + } + if status <= 0 { + status = http.StatusInternalServerError + } + + errText := strings.TrimSpace(gjson.GetBytes(payload, "error.message").String()) + if errText == "" { + errText = strings.TrimSpace(gjson.GetBytes(payload, "message").String()) + } + if errText == "" { + errText = strings.TrimSpace(string(payload)) + } + if errText == "" { + errText = http.StatusText(status) + } + return &interfaces.ErrorMessage{StatusCode: status, Error: fmt.Errorf("%s", errText)} +} + func setWebsocketTimelineBody(c *gin.Context, body string) { setWebsocketBody(c, wsTimelineBodyKey, body) } diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go index cefffcc9..b67147f0 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_test.go @@ -1156,6 +1156,171 @@ func TestForwardResponsesWebsocketPreservesCompletedEvent(t *testing.T) { } } +func TestForwardResponsesWebsocketTreatsResponseDoneAsTerminalWithoutRewriting(t *testing.T) { + gin.SetMode(gin.TestMode) + + serverErrCh := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := responsesWebsocketUpgrader.Upgrade(w, r, nil) + if err != nil { + serverErrCh <- err + return + } + defer func() { + errClose := conn.Close() + if errClose != nil { + serverErrCh <- errClose + } + }() + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = r + + data := make(chan []byte, 1) + errCh := make(chan *interfaces.ErrorMessage) + data <- []byte(`{"type":"response.done","response":{"id":"resp-1","output":[{"type":"message","id":"out-1"}]}}`) + close(data) + close(errCh) + + timelineLog := newInMemoryWebsocketTimelineLog() + completedOutput, completedResponseID, pendingToolCallIDs, errMsg, err := (*OpenAIResponsesAPIHandler)(nil).forwardResponsesWebsocket( + ctx, + conn, + func(...interface{}) {}, + data, + errCh, + timelineLog, + "session-1", + ) + if err != nil { + serverErrCh <- err + return + } + if errMsg != nil { + serverErrCh <- fmt.Errorf("unexpected websocket error message: %v", errMsg.Error) + return + } + if gjson.GetBytes(completedOutput, "0.id").String() != "out-1" { + serverErrCh <- errors.New("done output not captured") + return + } + if completedResponseID != "resp-1" { + serverErrCh <- fmt.Errorf("completed response id = %q, want resp-1", completedResponseID) + return + } + if len(pendingToolCallIDs) != 0 { + serverErrCh <- fmt.Errorf("pending tool call ids = %v, want empty", pendingToolCallIDs) + return + } + serverErrCh <- nil + })) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { + errClose := conn.Close() + if errClose != nil { + t.Fatalf("close websocket: %v", errClose) + } + }() + + _, payload, errReadMessage := conn.ReadMessage() + if errReadMessage != nil { + t.Fatalf("read websocket message: %v", errReadMessage) + } + if got := gjson.GetBytes(payload, "type").String(); got != "response.done" { + t.Fatalf("payload type = %s, want response.done; payload=%s", got, payload) + } + + if errServer := <-serverErrCh; errServer != nil { + t.Fatalf("server error: %v", errServer) + } +} + +func TestForwardResponsesWebsocketTreatsErrorPayloadAsTerminal(t *testing.T) { + gin.SetMode(gin.TestMode) + + serverErrCh := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := responsesWebsocketUpgrader.Upgrade(w, r, nil) + if err != nil { + serverErrCh <- err + return + } + defer func() { + errClose := conn.Close() + if errClose != nil { + serverErrCh <- errClose + } + }() + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = r + + data := make(chan []byte, 1) + errCh := make(chan *interfaces.ErrorMessage) + data <- []byte(`{"type":"error","status":429,"error":{"message":"upstream failed"}}`) + close(data) + close(errCh) + + _, _, _, errMsg, err := (*OpenAIResponsesAPIHandler)(nil).forwardResponsesWebsocket( + ctx, + conn, + func(...interface{}) {}, + data, + errCh, + newInMemoryWebsocketTimelineLog(), + "session-1", + ) + if err != nil { + serverErrCh <- err + return + } + if errMsg == nil { + serverErrCh <- errors.New("expected websocket error message") + return + } + if errMsg.StatusCode != http.StatusTooManyRequests { + serverErrCh <- fmt.Errorf("websocket error status = %d, want %d", errMsg.StatusCode, http.StatusTooManyRequests) + return + } + if errMsg.Error == nil || !strings.Contains(errMsg.Error.Error(), "upstream failed") { + serverErrCh <- fmt.Errorf("websocket error = %v, want upstream failed", errMsg.Error) + return + } + serverErrCh <- nil + })) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { + errClose := conn.Close() + if errClose != nil { + t.Fatalf("close websocket: %v", errClose) + } + }() + + _, payload, errReadMessage := conn.ReadMessage() + if errReadMessage != nil { + t.Fatalf("read websocket message: %v", errReadMessage) + } + if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeError { + t.Fatalf("payload type = %s, want %s; payload=%s", got, wsEventTypeError, payload) + } + + if errServer := <-serverErrCh; errServer != nil { + t.Fatalf("server error: %v", errServer) + } +} + func TestRecordPendingToolCallIDsFromPayloadDropsSatisfiedCalls(t *testing.T) { pending := map[string]struct{}{} payload := []byte(`{"type":"response.completed","response":{"output":[{"type":"function_call","call_id":"call-1","id":"fc-1"},{"type":"function_call_output","call_id":"call-1","id":"out-1"},{"type":"custom_tool_call","call_id":"call-2","id":"ctc-1"},{"type":"custom_tool_call_output","call_id":"call-2","id":"custom-out-1"}]}}`) -- 2.51.2