diff --git a/internal/runtime/executor/codex_websockets_connect_retry_test.go b/internal/runtime/executor/codex_websockets_connect_retry_test.go new file mode 100644 index 00000000000..469c2fc3ab6 --- /dev/null +++ b/internal/runtime/executor/codex_websockets_connect_retry_test.go @@ -0,0 +1,81 @@ +package executor + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +func TestCodexWebsocketConnectRetriesHandshakeEOF(t *testing.T) { + for _, tc := range []struct { + name string + failures int32 + status int + wantDials int32 + wantErr bool + }{ + {name: "third_retry_succeeds", failures: 3, wantDials: 4}, + {name: "exhausted", failures: 10, wantDials: 4, wantErr: true}, + {name: "status_not_retried", status: http.StatusBadGateway, wantDials: 1, wantErr: true}, + } { + t.Run(tc.name, func(t *testing.T) { + var dials atomic.Int32 + upgrader := websocket.Upgrader{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempt := dials.Add(1) + if tc.status != 0 { + http.Error(w, "rejected", tc.status) + return + } + if attempt <= tc.failures { + raw, _, err := w.(http.Hijacker).Hijack() + if err != nil { + t.Error(err) + return + } + _ = raw.Close() + return + } + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Error(err) + return + } + _ = conn.Close() + })) + defer server.Close() + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + exec := NewCodexWebsocketsExecutor(&config.Config{}) + conn, closer, resp, err := exec.dialCodexWebsocket(ctx, &cliproxyauth.Auth{ProxyURL: "direct"}, "ws"+strings.TrimPrefix(server.URL, "http"), nil) + if resp != nil { + defer func() { + if errClose := resp.Body.Close(); errClose != nil { + t.Errorf("close response body: %v", errClose) + } + }() + } + if closer != nil { + defer func() { + if errClose := closer.Close(); errClose != nil { + t.Errorf("close websocket: %v", errClose) + } + }() + } + if (err != nil) != tc.wantErr || (err == nil && conn == nil) { + t.Errorf("dial error = %v, wantErr %v", err, tc.wantErr) + } + if got := dials.Load(); got != tc.wantDials { + t.Errorf("dial attempts = %d, want %d", got, tc.wantDials) + } + }) + } +} diff --git a/internal/runtime/executor/codex_websockets_connection.go b/internal/runtime/executor/codex_websockets_connection.go index 5af93548d34..64ce227d6b9 100644 --- a/internal/runtime/executor/codex_websockets_connection.go +++ b/internal/runtime/executor/codex_websockets_connection.go @@ -5,6 +5,7 @@ import ( "context" "errors" "fmt" + "io" "net" "net/http" "net/url" @@ -36,9 +37,29 @@ func (e *CodexWebsocketsExecutor) dialCodexWebsocket(ctx context.Context, auth * if ctx == nil { ctx = context.Background() } - conn, resp, err := dialer.DialContext(ctx, wsURL, headers) - if err != nil { + // An EOF before any HTTP response is retried on a fresh connection; no response.create + // has been sent yet. All attempts share the one handshake time budget. + dialCtx, cancelDial := context.WithTimeout(ctx, codexResponsesWebsocketHandshakeTO) + defer cancelDial() + var conn *websocket.Conn + var resp *http.Response + var err error + for retry := 0; ; retry++ { + conn, resp, err = dialer.DialContext(dialCtx, wsURL, headers) + if err == nil { + break + } cliproxyexecutor.MarkUpstreamAttempt(ctx) + if resp != nil || retry == 3 || !(errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF)) { + break + } + timer := time.NewTimer(time.Duration(100<