diff --git a/internal/api/server_routes.go b/internal/api/server_routes.go index 1913bef6..815f6e3b 100644 --- a/internal/api/server_routes.go +++ b/internal/api/server_routes.go @@ -457,42 +457,6 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { c.JSON(clienterror.HTTPStatusFromErrorOr(err, http.StatusBadGateway), gin.H{"error": err.Error()}) return } - if selection != nil && resp.StatusCode == http.StatusUnauthorized { - s.handlers.AuthManager.ReportHomeUnauthorized(ctx, selected, "codex", selectionModel) - helps.RecordAPIResponseMetadata(ctx, s.cfg, resp.StatusCode, resp.Header.Clone()) - _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<20)) - if errClose := resp.Body.Close(); errClose != nil { - log.Errorf("codex alpha search: close unauthorized response body error: %v", errClose) - } - refreshed, didRefresh, errRefresh := s.handlers.AuthManager.RefreshHomeSelectionAfterUnauthorized(ctx, selection, selected) - if errRefresh != nil { - selection.End("refresh_failed") - c.JSON(clienterror.HTTPStatusFromErrorOr(errRefresh, http.StatusServiceUnavailable), gin.H{"error": errRefresh.Error()}) - return - } - if !didRefresh || refreshed == nil { - selection.End("refresh_unavailable") - c.JSON(http.StatusUnauthorized, gin.H{"error": "Codex credential unauthorized"}) - return - } - selected = refreshed - logging.SetGinCPATraceID(c, selected.EnsureIndex()) - resp, err = performRequest(selected) - if err != nil { - if errors.Is(err, errMissingBaseURL) { - selection.End("missing_base_url") - c.JSON(http.StatusServiceUnavailable, gin.H{"error": err.Error()}) - return - } - selection.End("retry_failed") - helps.RecordAPIResponseError(ctx, s.cfg, err) - c.JSON(clienterror.HTTPStatusFromErrorOr(err, http.StatusBadGateway), gin.H{"error": err.Error()}) - return - } - if resp.StatusCode == http.StatusUnauthorized { - s.handlers.AuthManager.ReportHomeUnauthorized(ctx, selected, "codex", selectionModel) - } - } closeResponseBody := func() error { errClose := resp.Body.Close() if errClose != nil { @@ -502,6 +466,9 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { } if selection != nil { if errBind := selection.Bind(closeResponseBody); errBind != nil { + if resp.StatusCode == http.StatusUnauthorized { + s.handlers.AuthManager.ReportHomeUnauthorized(ctx, selected, "codex", selectionModel) + } selection.End("response_bind_failed") c.JSON(http.StatusServiceUnavailable, gin.H{"error": errBind.Error()}) return @@ -513,11 +480,19 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { helps.RecordAPIResponseMetadata(ctx, s.cfg, resp.StatusCode, resp.Header.Clone()) upstreamBody, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20)) if err != nil { + helps.AppendAPIResponseChunk(ctx, s.cfg, upstreamBody) + if selection != nil && resp.StatusCode == http.StatusUnauthorized { + s.handlers.AuthManager.ReportHomeUnauthorized(ctx, selected, "codex", selectionModel, upstreamBody) + } helps.RecordAPIResponseError(ctx, s.cfg, err) c.JSON(clienterror.HTTPStatusFromErrorOr(err, http.StatusBadGateway), gin.H{"error": "Failed to read Codex search response"}) return } helps.AppendAPIResponseChunk(ctx, s.cfg, upstreamBody) + if selection != nil && resp.StatusCode == http.StatusUnauthorized { + s.handlers.AuthManager.ReportHomeUnauthorized(ctx, selected, "codex", selectionModel, upstreamBody) + log.WithField("status", resp.StatusCode).Warnf("codex alpha search upstream request failed: %s", logging.SafeDiagnosticForLog(string(upstreamBody))) + } if contentType := resp.Header.Get("Content-Type"); contentType != "" { c.Header("Content-Type", contentType) } diff --git a/internal/api/server_test.go b/internal/api/server_test.go index 9f972b2a..8eb33721 100644 --- a/internal/api/server_test.go +++ b/internal/api/server_test.go @@ -27,6 +27,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" ) @@ -41,6 +42,7 @@ type codexSearchCaptureExecutor struct { statuses []int refreshCalls int httpCalls int + beforeReturn func() } func (e *codexSearchCaptureExecutor) Identifier() string { return "codex" } @@ -138,6 +140,9 @@ func (e *codexSearchCaptureExecutor) HttpRequest(_ context.Context, selected *au if e.httpCalls <= len(e.statuses) && e.statuses[e.httpCalls-1] > 0 { statusCode = e.statuses[e.httpCalls-1] } + if e.beforeReturn != nil { + e.beforeReturn() + } return &http.Response{ StatusCode: statusCode, Header: http.Header{"Content-Type": []string{"application/json"}}, @@ -150,6 +155,46 @@ type codexSearchHomeDispatcher struct { policy atomic.Value } +type homeUnauthorizedUsageCapture struct { + authID string + records chan coreusage.Record +} + +func (p *homeUnauthorizedUsageCapture) HandleUsage(_ context.Context, record coreusage.Record) { + if p == nil || record.ExecutorType != "home-result" || record.AuthID != p.authID { + return + } + select { + case p.records <- record: + default: + } +} + +func (p *homeUnauthorizedUsageCapture) wait(t *testing.T) coreusage.Record { + t.Helper() + select { + case record := <-p.records: + return record + case <-time.After(time.Second): + t.Fatal("timed out waiting for Home unauthorized usage record") + return coreusage.Record{} + } +} + +type noopHomeUnauthorizedUsagePlugin struct{} + +func (noopHomeUnauthorizedUsagePlugin) HandleUsage(context.Context, coreusage.Record) {} + +func registerHomeUnauthorizedUsageCapture(t *testing.T, name, authID string) *homeUnauthorizedUsageCapture { + t.Helper() + capture := &homeUnauthorizedUsageCapture{authID: authID, records: make(chan coreusage.Record, 1)} + coreusage.RegisterNamedPlugin(name, capture) + t.Cleanup(func() { + coreusage.RegisterNamedPlugin(name, noopHomeUnauthorizedUsagePlugin{}) + }) + return capture +} + func (*codexSearchHomeDispatcher) HeartbeatOK() bool { return true } func (d *codexSearchHomeDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { @@ -199,6 +244,24 @@ func (b *trackedSearchResponseBody) Close() error { return nil } +type errorSearchResponseBody struct { + payload []byte + read atomic.Bool + closed atomic.Bool +} + +func (b *errorSearchResponseBody) Read(p []byte) (int, error) { + if !b.read.CompareAndSwap(false, true) { + return 0, io.EOF + } + return copy(p, b.payload), io.ErrUnexpectedEOF +} + +func (b *errorSearchResponseBody) Close() error { + b.closed.Store(true) + return nil +} + type drainAwareSearchResponseBody struct { started chan struct{} closed chan struct{} @@ -324,31 +387,137 @@ func TestAuditHomeCodexSearchBodyCloseBeforeRelease(t *testing.T) { } } -func TestHomeCodexAlphaSearchRefreshesUnauthorizedSelectionOnce(t *testing.T) { +func TestHomeCodexAlphaSearchForwardsUnauthorizedResponseWithoutRefresh(t *testing.T) { + const upstreamError = `{"error":{"message":"access token expired"}}` server := newTestServer(t) + server.cfg.RequestLog = true dispatcher := &codexSearchHomeDispatcher{} server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}}) server.handlers.AuthManager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) - executor := &codexSearchCaptureExecutor{statuses: []int{http.StatusUnauthorized, http.StatusOK}} + executor := &codexSearchCaptureExecutor{ + statuses: []int{http.StatusUnauthorized}, + responseBody: io.NopCloser(strings.NewReader(upstreamError)), + } server.handlers.AuthManager.RegisterExecutor(executor) req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"id":"home-search-refresh","model":"gpt-5-codex","query":"test"}`)) req.Header.Set("Authorization", "Bearer test-key") rr := httptest.NewRecorder() - server.engine.ServeHTTP(rr, req) + c, _ := gin.CreateTestContext(rr) + c.Request = req + server.codexAlphaSearch(c) - if rr.Code != http.StatusOK { - t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusOK, rr.Body.String()) + if rr.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusUnauthorized, rr.Body.String()) + } + if got := rr.Body.String(); got != upstreamError { + t.Fatalf("body = %q, want original upstream error %q", got, upstreamError) } - if executor.refreshCalls != 1 || executor.httpCalls != 2 { - t.Fatalf("refresh/http calls = %d/%d, want 1/2", executor.refreshCalls, executor.httpCalls) + if executor.refreshCalls != 0 || executor.httpCalls != 1 { + t.Fatalf("refresh/http calls = %d/%d, want 0/1", executor.refreshCalls, executor.httpCalls) } - if got := executor.request.Header.Get("Authorization"); got != "Bearer refreshed-home-search-token" { - t.Fatalf("retry Authorization = %q, want refreshed token", got) + if got := executor.request.Header.Get("Authorization"); got != "Bearer home-search-token" { + t.Fatalf("Authorization = %q, want original Home token", got) } if got := dispatcher.calls.Load(); got != 1 { t.Fatalf("Home RPOP calls = %d, want 1", got) } + rawAPIResponse, okResponse := c.Get("API_RESPONSE") + if !okResponse { + t.Fatal("API_RESPONSE was not captured") + } + apiResponse, _ := rawAPIResponse.([]byte) + if !strings.Contains(string(apiResponse), "Status: 401") || !strings.Contains(string(apiResponse), upstreamError) { + t.Fatalf("API_RESPONSE = %q, want original upstream 401", apiResponse) + } +} + +func TestHomeCodexAlphaSearchReportsUnauthorizedBeforeEarlyReturn(t *testing.T) { + const upstreamError = `{"error":{"message":"access token expired"}}` + tests := []struct { + name string + responseBody func() io.ReadCloser + beforeReturn func(*executionregistry.Registry) + wantStatus int + wantFailBody string + }{ + { + name: "response bind failure", + responseBody: func() io.ReadCloser { + return &trackedSearchResponseBody{Reader: strings.NewReader(upstreamError)} + }, + beforeReturn: func(registry *executionregistry.Registry) { + _ = registry.Close() + }, + wantStatus: http.StatusServiceUnavailable, + wantFailBody: "upstream unauthorized", + }, + { + name: "response read failure", + responseBody: func() io.ReadCloser { + return &errorSearchResponseBody{payload: []byte(upstreamError)} + }, + wantStatus: http.StatusBadGateway, + wantFailBody: upstreamError, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := newTestServer(t) + registry := executionregistry.New() + server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}}) + server.handlers.AuthManager.PublishHomeDispatch(&codexSearchHomeDispatcher{}, registry, 1) + executor := &codexSearchCaptureExecutor{ + statuses: []int{http.StatusUnauthorized}, + responseBody: test.responseBody(), + } + if test.beforeReturn != nil { + executor.beforeReturn = func() { test.beforeReturn(registry) } + } + server.handlers.AuthManager.RegisterExecutor(executor) + usageCapture := registerHomeUnauthorizedUsageCapture(t, t.Name(), "home-codex-search") + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"model":"gpt-5-codex","query":"test"}`)) + request.Header.Set("Authorization", "Bearer test-key") + server.engine.ServeHTTP(recorder, request) + + if recorder.Code != test.wantStatus { + t.Fatalf("status = %d, want %d; body=%s", recorder.Code, test.wantStatus, recorder.Body.String()) + } + record := usageCapture.wait(t) + if record.Fail.StatusCode != http.StatusUnauthorized || record.Fail.Body != test.wantFailBody { + t.Fatalf("Home unauthorized failure = status %d body %q, want status 401 body %q", record.Fail.StatusCode, record.Fail.Body, test.wantFailBody) + } + }) + } +} + +func TestHomeCodexAlphaSearchRequestLogPreservesBodyReturnedWithReadError(t *testing.T) { + const upstreamError = `{"error":{"message":"access token expired"}}` + server := newTestServer(t) + server.cfg.RequestLog = true + server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}}) + server.handlers.AuthManager.PublishHomeDispatch(&codexSearchHomeDispatcher{}, executionregistry.New(), 1) + server.handlers.AuthManager.RegisterExecutor(&codexSearchCaptureExecutor{ + statuses: []int{http.StatusUnauthorized}, + responseBody: &errorSearchResponseBody{payload: []byte(upstreamError)}, + }) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"model":"gpt-5-codex","query":"test"}`)) + server.codexAlphaSearch(c) + + if recorder.Code != http.StatusBadGateway { + t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String()) + } + rawAPIResponse, okResponse := c.Get("API_RESPONSE") + apiResponse, _ := rawAPIResponse.([]byte) + if !okResponse || !strings.Contains(string(apiResponse), upstreamError) || !strings.Contains(string(apiResponse), io.ErrUnexpectedEOF.Error()) { + t.Fatalf("API_RESPONSE = %q, want upstream body and read error", apiResponse) + } } func TestHomeCodexAlphaSearchEndsSelectionAcrossDirectHTTPPaths(t *testing.T) { diff --git a/internal/client/codex/live/capabilities.go b/internal/client/codex/live/capabilities.go index 6a4e44ed..9c38d1be 100644 --- a/internal/client/codex/live/capabilities.go +++ b/internal/client/codex/live/capabilities.go @@ -2,7 +2,6 @@ package live import ( "context" - "io" "net/http" "net/url" "strings" @@ -149,46 +148,27 @@ func (h *Handler) HandleHangup(c *gin.Context) { writeRealtimeError(c, clienterror.HTTPStatusFromErrorOr(errRequest, http.StatusBadGateway), errRequest.Error(), "api_error", "realtime_upstream_unavailable") return } - if activeSelection != nil && response.StatusCode == http.StatusUnauthorized { - h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", session.model) - _, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 1<<20)) - if errClose := response.Body.Close(); errClose != nil { - log.Errorf("codex realtime hangup: close unauthorized response body error: %v", errClose) - } - refreshed, didRefresh, errRefresh := h.authManager.RefreshHomeSelectionAfterUnauthorized(ctx, activeSelection, selected) - if errRefresh != nil { - writeSelectionError(c, errRefresh) - return - } - if !didRefresh || refreshed == nil { - writeRealtimeError(c, http.StatusUnauthorized, "Codex credential unauthorized", "authentication_error", "realtime_upstream_unauthorized") - return - } - selected = refreshed - logging.SetGinCPATraceID(c, selected.EnsureIndex()) - response, errRequest = performRequest(selected) - if errRequest != nil { - helps.RecordAPIResponseError(ctx, runtimeConfig, errRequest) - writeRealtimeError(c, clienterror.HTTPStatusFromErrorOr(errRequest, http.StatusBadGateway), errRequest.Error(), "api_error", "realtime_upstream_unavailable") - return - } - if response.StatusCode == http.StatusUnauthorized { - h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", session.model) - } - } defer func() { if errClose := response.Body.Close(); errClose != nil { log.Errorf("codex realtime hangup: close response body error: %v", errClose) } }() + helps.RecordAPIResponseMetadata(ctx, runtimeConfig, response.StatusCode, callResponseHeaders(response.Header)) responseBody, errResponse := readLimitedBody(response.Body) if errResponse != nil { + helps.AppendAPIResponseChunk(ctx, runtimeConfig, responseBody) + if activeSelection != nil && response.StatusCode == http.StatusUnauthorized { + h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", session.model, responseBody) + } helps.RecordAPIResponseError(ctx, runtimeConfig, errResponse) writeRealtimeError(c, http.StatusBadGateway, "Failed to read Realtime hangup response", "api_error", "realtime_upstream_unavailable") return } - helps.RecordAPIResponseMetadata(ctx, runtimeConfig, response.StatusCode, callResponseHeaders(response.Header)) helps.AppendAPIResponseChunk(ctx, runtimeConfig, responseBody) + if activeSelection != nil && response.StatusCode == http.StatusUnauthorized { + h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", session.model, responseBody) + log.WithField("status", response.StatusCode).Warnf("codex realtime hangup upstream request failed: %s", logging.SafeDiagnosticForLog(string(responseBody))) + } if response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices { if selectionRelease != nil { selectionRelease() diff --git a/internal/client/codex/live/capabilities_test.go b/internal/client/codex/live/capabilities_test.go index 680ed0ef..6ba18953 100644 --- a/internal/client/codex/live/capabilities_test.go +++ b/internal/client/codex/live/capabilities_test.go @@ -8,7 +8,9 @@ import ( "testing" "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" ) func TestHandleHangupForwardsPinnedOAuthCall(t *testing.T) { @@ -53,6 +55,84 @@ func TestHandleHangupForwardsPinnedOAuthCall(t *testing.T) { } } +func TestHandleHangupForwardsUnauthorizedHomeResponseWithoutRefresh(t *testing.T) { + const upstreamError = `{"error":{"message":"access token expired"}}` + gin.SetMode(gin.TestMode) + manager := auth.NewManager(nil, nil, nil) + manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(&homeDispatcher{}, executionregistry.New(), 1) + executor := &captureExecutor{ + statusCode: http.StatusUnauthorized, + responseBody: io.NopCloser(strings.NewReader(upstreamError)), + } + manager.RegisterExecutor(executor) + handler := NewHandler(manager, nil) + handler.sessions.put("call-123", liveSession{authID: "home-codex-live", model: defaultLiveModel}) + + router := gin.New() + router.POST("/v1/realtime/calls/:call_id/hangup", handler.HandleHangup) + request := httptest.NewRequest(http.MethodPost, "/v1/realtime/calls/call-123/hangup", nil) + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusUnauthorized || recorder.Body.String() != upstreamError { + t.Fatalf("response = status %d body %q, want original upstream 401", recorder.Code, recorder.Body.String()) + } + if executor.refreshCalls.Load() != 0 || executor.httpCalls.Load() != 1 { + t.Fatalf("refresh/http calls = %d/%d, want 0/1", executor.refreshCalls.Load(), executor.httpCalls.Load()) + } + if got := executor.request.Header.Get("Authorization"); got != "Bearer home-live-token" { + t.Fatalf("Authorization = %q, want original Home token", got) + } +} + +func TestHandleHangupReportsUnauthorizedWhenResponseReadFails(t *testing.T) { + const upstreamError = `{"error":{"message":"access token expired"}}` + gin.SetMode(gin.TestMode) + runtimeConfig := &config.Config{ + Home: config.HomeConfig{Enabled: true}, + SDKConfig: config.SDKConfig{RequestLog: true}, + } + manager := auth.NewManager(nil, nil, nil) + manager.SetConfig(runtimeConfig) + manager.PublishHomeDispatch(&homeDispatcher{}, executionregistry.New(), 1) + executor := &captureExecutor{ + statusCode: http.StatusUnauthorized, + responseBody: &errorResponseBody{payload: []byte(upstreamError)}, + } + manager.RegisterExecutor(executor) + usageCapture := registerHomeUnauthorizedUsageCapture(t, t.Name(), "home-codex-live") + handler := NewHandler(manager, runtimeConfig) + handler.sessions.put("call-123", liveSession{authID: "home-codex-live", model: defaultLiveModel}) + + router := gin.New() + var apiResponse []byte + router.Use(func(c *gin.Context) { + c.Next() + if raw, exists := c.Get("API_RESPONSE"); exists { + apiResponse, _ = raw.([]byte) + } + }) + router.POST("/v1/realtime/calls/:call_id/hangup", handler.HandleHangup) + request := httptest.NewRequest(http.MethodPost, "/v1/realtime/calls/call-123/hangup", nil) + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusBadGateway { + t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String()) + } + record := usageCapture.wait(t) + if record.Fail.StatusCode != http.StatusUnauthorized || record.Fail.Body != upstreamError { + t.Fatalf("Home unauthorized failure = status %d body %q, want status 401 body %q", record.Fail.StatusCode, record.Fail.Body, upstreamError) + } + if !strings.Contains(string(apiResponse), upstreamError) || !strings.Contains(string(apiResponse), io.ErrUnexpectedEOF.Error()) { + t.Fatalf("API_RESPONSE = %q, want upstream body and read error", apiResponse) + } + if executor.refreshCalls.Load() != 0 || executor.httpCalls.Load() != 1 { + t.Fatalf("refresh/http calls = %d/%d, want 0/1", executor.refreshCalls.Load(), executor.httpCalls.Load()) + } +} + func TestHandleHangupRejectsDifferentAPIPrincipal(t *testing.T) { gin.SetMode(gin.TestMode) handler := NewHandler(auth.NewManager(nil, nil, nil), nil) diff --git a/internal/client/codex/live/live.go b/internal/client/codex/live/live.go index 4a862e92..29a33836 100644 --- a/internal/client/codex/live/live.go +++ b/internal/client/codex/live/live.go @@ -321,38 +321,6 @@ func (h *Handler) Handle(c *gin.Context) { writeLiveError(c, clienterror.HTTPStatusFromErrorOr(errRequest, http.StatusBadGateway), errRequest.Error()) return } - if selection != nil && resp.StatusCode == http.StatusUnauthorized { - h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", model) - helps.RecordAPIResponseMetadata(ctx, runtimeConfig, resp.StatusCode, callResponseHeaders(resp.Header)) - _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<20)) - if errClose := resp.Body.Close(); errClose != nil { - log.Errorf("codex live: close unauthorized response body error: %v", errClose) - } - refreshed, didRefresh, errRefresh := h.authManager.RefreshHomeSelectionAfterUnauthorized(ctx, selection, selected) - if errRefresh != nil { - selection.End("refresh_failed") - writeSelectionError(c, errRefresh) - return - } - if !didRefresh || refreshed == nil { - selection.End("refresh_unavailable") - writeLiveError(c, http.StatusUnauthorized, "Codex credential unauthorized") - return - } - selected = refreshed - logging.SetGinCPATraceID(c, selected.EnsureIndex()) - resp, errRequest = performRequest(selected) - if errRequest != nil { - selection.End("retry_failed") - helps.RecordAPIResponseError(ctx, runtimeConfig, errRequest) - writeLiveError(c, clienterror.HTTPStatusFromErrorOr(errRequest, http.StatusBadGateway), errRequest.Error()) - return - } - if resp.StatusCode == http.StatusUnauthorized { - h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", model) - } - } - var closeResponseOnce sync.Once var closeResponseErr error closeResponseBody := func() error { @@ -367,6 +335,9 @@ func (h *Handler) Handle(c *gin.Context) { defer func() { _ = closeResponseBody() }() if selection != nil { if errBind := selection.Bind(closeResponseBody); errBind != nil { + if resp.StatusCode == http.StatusUnauthorized { + h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", model) + } selection.End("response_bind_failed") writeLiveError(c, http.StatusServiceUnavailable, errBind.Error()) return @@ -377,6 +348,10 @@ func (h *Handler) Handle(c *gin.Context) { helps.RecordAPIResponseMetadata(ctx, runtimeConfig, resp.StatusCode, responseHeaders) responseBody, errResponse := readLimitedBody(resp.Body) if errResponse != nil { + helps.AppendAPIResponseChunk(ctx, runtimeConfig, responseBody) + if selection != nil && resp.StatusCode == http.StatusUnauthorized { + h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", model, responseBody) + } helps.RecordAPIResponseError(ctx, runtimeConfig, errResponse) message := "Failed to read Codex live response" status := clienterror.HTTPStatusFromErrorOr(errResponse, http.StatusBadGateway) @@ -388,6 +363,10 @@ func (h *Handler) Handle(c *gin.Context) { return } helps.AppendAPIResponseChunk(ctx, runtimeConfig, responseBody) + if selection != nil && resp.StatusCode == http.StatusUnauthorized { + h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", model, responseBody) + log.WithField("status", resp.StatusCode).Warnf("codex live upstream request failed: %s", logging.SafeDiagnosticForLog(string(responseBody))) + } responseBodyToWrite := responseBody success := resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices callID := "" @@ -522,10 +501,13 @@ func readLimitedBody(body io.Reader) ([]byte, error) { } payload, errRead := io.ReadAll(io.LimitReader(body, maxBodySize+1)) if errRead != nil { - return nil, errRead + if len(payload) > maxBodySize { + payload = payload[:maxBodySize] + } + return payload, errRead } if len(payload) > maxBodySize { - return nil, errBodyTooLarge + return payload[:maxBodySize], errBodyTooLarge } return payload, nil } diff --git a/internal/client/codex/live/live_test.go b/internal/client/codex/live/live_test.go index 3dcbff76..f69a2955 100644 --- a/internal/client/codex/live/live_test.go +++ b/internal/client/codex/live/live_test.go @@ -18,6 +18,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" ) type apiKeyFirstSelector struct{} @@ -43,6 +44,7 @@ type captureExecutor struct { statuses []int httpCalls atomic.Int32 refreshCalls atomic.Int32 + beforeReturn func() } func (*captureExecutor) Identifier() string { return "codex" } @@ -95,6 +97,9 @@ func (e *captureExecutor) HttpRequest(_ context.Context, credential *auth.Auth, if statusCode == http.StatusUnauthorized && httpCall < len(e.statuses) { responseBody = io.NopCloser(strings.NewReader("unauthorized")) } + if e.beforeReturn != nil { + e.beforeReturn() + } return &http.Response{ StatusCode: statusCode, Header: http.Header{ @@ -110,26 +115,71 @@ func (e *captureExecutor) HttpRequest(_ context.Context, credential *auth.Auth, } type homeDispatcher struct { - model string + model string + authID string +} + +type homeUnauthorizedUsageCapture struct { + authID string + records chan coreusage.Record +} + +func (p *homeUnauthorizedUsageCapture) HandleUsage(_ context.Context, record coreusage.Record) { + if p == nil || record.ExecutorType != "home-result" || record.AuthID != p.authID { + return + } + select { + case p.records <- record: + default: + } +} + +func (p *homeUnauthorizedUsageCapture) wait(t *testing.T) coreusage.Record { + t.Helper() + select { + case record := <-p.records: + return record + case <-time.After(time.Second): + t.Fatal("timed out waiting for Home unauthorized usage record") + return coreusage.Record{} + } +} + +type noopHomeUnauthorizedUsagePlugin struct{} + +func (noopHomeUnauthorizedUsagePlugin) HandleUsage(context.Context, coreusage.Record) {} + +func registerHomeUnauthorizedUsageCapture(t *testing.T, name, authID string) *homeUnauthorizedUsageCapture { + t.Helper() + capture := &homeUnauthorizedUsageCapture{authID: authID, records: make(chan coreusage.Record, 1)} + coreusage.RegisterNamedPlugin(name, capture) + t.Cleanup(func() { + coreusage.RegisterNamedPlugin(name, noopHomeUnauthorizedUsagePlugin{}) + }) + return capture } func (*homeDispatcher) HeartbeatOK() bool { return true } func (d *homeDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { d.model = model + authID := d.authID + if authID == "" { + authID = "home-codex-live" + } return json.Marshal(map[string]any{ "model": model, "provider": "codex", - "auth_index": "home-codex-live", + "auth_index": authID, "auth": map[string]any{ - "id": "home-codex-live", + "id": authID, "provider": "codex", "status": "active", "metadata": map[string]any{"access_token": "home-live-token"}, }, "concurrency": map[string]any{ "accounted": true, - "credential_id": "home-codex-live", + "credential_id": authID, "model": model, }, }) @@ -164,6 +214,32 @@ func (b *trackedResponseBody) Close() error { return nil } +type errorResponseBody struct { + payload []byte + read atomic.Bool + closed atomic.Bool +} + +func (b *errorResponseBody) Read(p []byte) (int, error) { + if !b.read.CompareAndSwap(false, true) { + return 0, io.EOF + } + return copy(p, b.payload), io.ErrUnexpectedEOF +} + +func (b *errorResponseBody) Close() error { + b.closed.Store(true) + return nil +} + +func TestReadLimitedBodyPreservesPayloadOnReadError(t *testing.T) { + const upstreamError = `{"error":{"message":"access token expired"}}` + payload, errRead := readLimitedBody(&errorResponseBody{payload: []byte(upstreamError)}) + if !errors.Is(errRead, io.ErrUnexpectedEOF) || string(payload) != upstreamError { + t.Fatalf("readLimitedBody() = %q, %v; want partial payload and unexpected EOF", payload, errRead) + } +} + type fakeMediaRelay struct { clientOffer string route mediaSessionRoute @@ -606,15 +682,16 @@ func TestHandlerClosesMediaWhenResponseWriteFails(t *testing.T) { } } -func TestHandlerRefreshesUnauthorizedHomeSelectionOnce(t *testing.T) { +func TestHandlerForwardsUnauthorizedHomeResponseWithoutRefresh(t *testing.T) { + const upstreamError = `{"error":{"message":"access token expired"}}` gin.SetMode(gin.TestMode) manager := auth.NewManager(nil, nil, nil) manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) registry := executionregistry.New() manager.PublishHomeDispatch(&homeDispatcher{}, registry, 1) executor := &captureExecutor{ - statuses: []int{http.StatusUnauthorized, http.StatusCreated}, - responseBody: &trackedResponseBody{Reader: strings.NewReader("v=0\r\n")}, + statuses: []int{http.StatusUnauthorized}, + responseBody: &trackedResponseBody{Reader: strings.NewReader(upstreamError)}, } manager.RegisterExecutor(executor) handler := NewHandler(manager, nil) @@ -626,20 +703,104 @@ func TestHandlerRefreshesUnauthorizedHomeSelectionOnce(t *testing.T) { recorder := httptest.NewRecorder() router.ServeHTTP(recorder, req) - if recorder.Code != http.StatusCreated { - t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusCreated, recorder.Body.String()) + if recorder.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusUnauthorized, recorder.Body.String()) } - if executor.refreshCalls.Load() != 1 || executor.httpCalls.Load() != 2 { - t.Fatalf("refresh/http calls = %d/%d, want 1/2", executor.refreshCalls.Load(), executor.httpCalls.Load()) + if got := recorder.Body.String(); got != upstreamError { + t.Fatalf("body = %q, want original upstream error %q", got, upstreamError) } - if got := executor.request.Header.Get("Authorization"); got != "Bearer refreshed-home-live-token" { - t.Fatalf("retry Authorization = %q, want refreshed token", got) + if executor.refreshCalls.Load() != 0 || executor.httpCalls.Load() != 1 { + t.Fatalf("refresh/http calls = %d/%d, want 0/1", executor.refreshCalls.Load(), executor.httpCalls.Load()) + } + if got := executor.request.Header.Get("Authorization"); got != "Bearer home-live-token" { + t.Fatalf("Authorization = %q, want original Home token", got) } if errDrain := registry.Drain(context.Background()); errDrain != nil { t.Fatalf("Drain() error = %v", errDrain) } } +func TestHandlerReportsUnauthorizedBeforeEarlyReturn(t *testing.T) { + const upstreamError = `{"error":{"message":"access token expired"}}` + tests := []struct { + name string + responseBody func() io.ReadCloser + beforeReturn func(*executionregistry.Registry) + wantStatus int + wantFailBody string + }{ + { + name: "response bind failure", + responseBody: func() io.ReadCloser { + return &trackedResponseBody{Reader: strings.NewReader(upstreamError)} + }, + beforeReturn: func(registry *executionregistry.Registry) { + _ = registry.Close() + }, + wantStatus: http.StatusServiceUnavailable, + wantFailBody: "upstream unauthorized", + }, + { + name: "response read failure", + responseBody: func() io.ReadCloser { + return &errorResponseBody{payload: []byte(upstreamError)} + }, + wantStatus: http.StatusBadGateway, + wantFailBody: upstreamError, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + authID := "home-codex-live-" + strings.ReplaceAll(test.name, " ", "-") + runtimeConfig := &config.Config{ + Home: config.HomeConfig{Enabled: true}, + SDKConfig: config.SDKConfig{RequestLog: true}, + } + manager := auth.NewManager(nil, nil, nil) + manager.SetConfig(runtimeConfig) + registry := executionregistry.New() + manager.PublishHomeDispatch(&homeDispatcher{authID: authID}, registry, 1) + executor := &captureExecutor{ + statuses: []int{http.StatusUnauthorized}, + responseBody: test.responseBody(), + } + if test.beforeReturn != nil { + executor.beforeReturn = func() { test.beforeReturn(registry) } + } + manager.RegisterExecutor(executor) + usageCapture := registerHomeUnauthorizedUsageCapture(t, t.Name(), authID) + handler := NewHandler(manager, runtimeConfig) + router := gin.New() + var apiResponse []byte + router.Use(func(c *gin.Context) { + c.Next() + if raw, exists := c.Get("API_RESPONSE"); exists { + apiResponse, _ = raw.([]byte) + } + }) + router.POST("/v1/live", handler.Handle) + + request := httptest.NewRequest(http.MethodPost, "/v1/live", strings.NewReader(`{"model":"gpt-live-1-codex","sdp":"v=0"}`)) + request.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, request) + + if recorder.Code != test.wantStatus { + t.Fatalf("status = %d, want %d; body=%s", recorder.Code, test.wantStatus, recorder.Body.String()) + } + record := usageCapture.wait(t) + if record.Fail.StatusCode != http.StatusUnauthorized || record.Fail.Body != test.wantFailBody { + t.Fatalf("Home unauthorized failure = status %d body %q, want status 401 body %q", record.Fail.StatusCode, record.Fail.Body, test.wantFailBody) + } + if test.name == "response read failure" && (!strings.Contains(string(apiResponse), upstreamError) || !strings.Contains(string(apiResponse), io.ErrUnexpectedEOF.Error())) { + t.Fatalf("API_RESPONSE = %q, want upstream body and read error", apiResponse) + } + }) + } +} + func TestHandlerUsesLiveModelForHomeDispatch(t *testing.T) { gin.SetMode(gin.TestMode) @@ -818,77 +979,122 @@ func TestHandleSidebandPinsAuthAndRelaysBidirectionally(t *testing.T) { } } -func TestHandleSidebandRefreshesUnauthorizedHomeHandshakeOnce(t *testing.T) { +func TestHandleSidebandForwardsUnauthorizedHomeHandshakeWithoutRefresh(t *testing.T) { gin.SetMode(gin.TestMode) - var upstreamCalls atomic.Int32 - upstreamHeaders := make(chan http.Header, 2) - upstreamServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - upstreamCalls.Add(1) - upstreamHeaders <- request.Header.Clone() - if request.Header.Get("Authorization") != "Bearer refreshed-home-live-token" { - writer.WriteHeader(http.StatusUnauthorized) - return - } - upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} - conn, errUpgrade := upgrader.Upgrade(writer, request, nil) - if errUpgrade != nil { - return - } - defer func() { _ = conn.Close() }() - messageType, payload, errRead := conn.ReadMessage() - if errRead == nil { - _ = conn.WriteMessage(messageType, append([]byte("echo:"), payload...)) - } - })) - defer upstreamServer.Close() - - manager := auth.NewManager(nil, nil, nil) - manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) - registry := executionregistry.New() - manager.PublishHomeDispatch(&homeDispatcher{}, registry, 1) - executor := &captureExecutor{} - manager.RegisterExecutor(executor) - selection, errSelect := manager.SelectHomeAuthByKind(context.Background(), "codex", defaultLiveModel, auth.AuthKindOAuth, coreexecutor.Options{}) - if errSelect != nil { - t.Fatalf("SelectHomeAuthByKind() error = %v", errSelect) - } - selection.Retain() - defer selection.End("test_complete") + for _, tc := range []struct { + name string + upstreamBody string + }{ + {name: "response body", upstreamBody: `{"error":{"message":"access token expired"}}`}, + {name: "empty response body"}, + } { + t.Run(tc.name, func(t *testing.T) { + var upstreamCalls atomic.Int32 + upstreamHeaders := make(chan http.Header, 1) + upstreamServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + upstreamCalls.Add(1) + upstreamHeaders <- request.Header.Clone() + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusUnauthorized) + _, _ = writer.Write([]byte(tc.upstreamBody)) + })) + defer upstreamServer.Close() - handler := NewHandler(manager, nil) - handler.sidebandAPIBaseURL = "ws" + strings.TrimPrefix(upstreamServer.URL, "http") + "/v1" - handler.sessions.put("call-home-refresh", liveSession{authID: "home-codex-live", model: defaultLiveModel, homeSelection: selection}) - router := gin.New() - router.GET("/v1/live/:call_id", handler.HandleSideband) - downstreamServer := httptest.NewServer(router) - defer downstreamServer.Close() + manager := auth.NewManager(nil, nil, nil) + manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + manager.PublishHomeDispatch(&homeDispatcher{}, registry, 1) + executor := &captureExecutor{} + manager.RegisterExecutor(executor) + selection, errSelect := manager.SelectHomeAuthByKind(context.Background(), "codex", defaultLiveModel, auth.AuthKindOAuth, coreexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectHomeAuthByKind() error = %v", errSelect) + } + selection.Retain() + defer selection.End("test_complete") - wsURL := "ws" + strings.TrimPrefix(downstreamServer.URL, "http") + "/v1/live/call-home-refresh" - client, response, errDial := websocket.DefaultDialer.Dial(wsURL, nil) - if errDial != nil { - if response != nil && response.Body != nil { - _ = response.Body.Close() - } - t.Fatalf("dial downstream sideband: %v", errDial) - } - if response != nil && response.Body != nil { - _ = response.Body.Close() + handler := NewHandler(manager, nil) + handler.sidebandAPIBaseURL = "ws" + strings.TrimPrefix(upstreamServer.URL, "http") + "/v1" + handler.sessions.put("call-home-refresh", liveSession{authID: "home-codex-live", model: defaultLiveModel, homeSelection: selection}) + router := gin.New() + router.GET("/v1/live/:call_id", handler.HandleSideband) + downstreamServer := httptest.NewServer(router) + defer downstreamServer.Close() + + wsURL := "ws" + strings.TrimPrefix(downstreamServer.URL, "http") + "/v1/live/call-home-refresh" + client, response, errDial := websocket.DefaultDialer.Dial(wsURL, nil) + if client != nil { + _ = client.Close() + } + if errDial == nil || response == nil { + t.Fatalf("dial downstream sideband = response %#v err %v, want rejected handshake", response, errDial) + } + defer func() { _ = response.Body.Close() }() + responseBody, errRead := io.ReadAll(response.Body) + if errRead != nil { + t.Fatalf("read downstream rejection: %v", errRead) + } + if response.StatusCode != http.StatusUnauthorized || string(responseBody) != tc.upstreamBody { + t.Fatalf("downstream rejection = status %d body %q, want original upstream 401 body %q", response.StatusCode, responseBody, tc.upstreamBody) + } + if executor.refreshCalls.Load() != 0 || upstreamCalls.Load() != 1 { + t.Fatalf("refresh/upstream calls = %d/%d, want 0/1", executor.refreshCalls.Load(), upstreamCalls.Load()) + } + if got := (<-upstreamHeaders).Get("Authorization"); got != "Bearer home-live-token" { + t.Fatalf("upstream Authorization = %q, want original Home token", got) + } + }) } - defer func() { _ = client.Close() }() - if errWrite := client.WriteMessage(websocket.TextMessage, []byte("ping")); errWrite != nil { - t.Fatalf("write sideband message: %v", errWrite) +} + +func TestHandleSidebandDialErrorPreservesBodyReturnedWithReadError(t *testing.T) { + const upstreamError = `{"error":{"message":"access token expired"}}` + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/live/call-123", nil) + ctx := context.WithValue(c.Request.Context(), "gin", c) + response := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &errorResponseBody{payload: []byte(upstreamError)}, } - _, payload, errRead := client.ReadMessage() - if errRead != nil || string(payload) != "echo:ping" { - t.Fatalf("read sideband message = %q, %v", string(payload), errRead) + + responseBody := handleSidebandDialError(c, ctx, &config.Config{SDKConfig: config.SDKConfig{RequestLog: true}}, response, errors.New("websocket: bad handshake")) + if recorder.Code != http.StatusUnauthorized || recorder.Body.String() != upstreamError || string(responseBody) != upstreamError { + t.Fatalf("forwarded response = status %d body %q returned %q, want original upstream 401", recorder.Code, recorder.Body.String(), responseBody) } - if executor.refreshCalls.Load() != 1 || upstreamCalls.Load() != 2 { - t.Fatalf("refresh/upstream calls = %d/%d, want 1/2", executor.refreshCalls.Load(), upstreamCalls.Load()) + rawTimeline, okTimeline := c.Get("API_WEBSOCKET_TIMELINE") + timeline, _ := rawTimeline.([]byte) + if !okTimeline || !strings.Contains(string(timeline), upstreamError) { + t.Fatalf("API_WEBSOCKET_TIMELINE = %q, want original upstream error", timeline) } - first := <-upstreamHeaders - second := <-upstreamHeaders - if first.Get("Authorization") != "Bearer home-live-token" || second.Get("Authorization") != "Bearer refreshed-home-live-token" { - t.Fatalf("upstream Authorization sequence = %q, %q", first.Get("Authorization"), second.Get("Authorization")) +} + +func TestHandleSidebandDialErrorDoesNotForwardNonUnauthorizedBody(t *testing.T) { + const upstreamError = `proxy-01.internal authentication failed` + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/live/call-123", nil) + ctx := context.WithValue(c.Request.Context(), "gin", c) + response := &http.Response{ + StatusCode: http.StatusBadGateway, + Header: http.Header{"Content-Type": []string{"text/html"}}, + Body: io.NopCloser(strings.NewReader(upstreamError)), + } + + responseBody := handleSidebandDialError(c, ctx, &config.Config{SDKConfig: config.SDKConfig{RequestLog: true}}, response, errors.New("websocket: bad handshake")) + if recorder.Code != http.StatusBadGateway || len(responseBody) != 0 { + t.Fatalf("response = status %d returned %q, want generic 502", recorder.Code, responseBody) + } + if strings.Contains(recorder.Body.String(), upstreamError) || !strings.Contains(recorder.Body.String(), "Codex live sideband upstream unavailable") { + t.Fatalf("downstream body = %q, want generic error without proxy detail", recorder.Body.String()) + } + rawTimeline, okTimeline := c.Get("API_WEBSOCKET_TIMELINE") + timeline, _ := rawTimeline.([]byte) + if !okTimeline || !strings.Contains(string(timeline), upstreamError) { + t.Fatalf("API_WEBSOCKET_TIMELINE = %q, want original upstream error", timeline) } } diff --git a/internal/client/codex/live/sideband.go b/internal/client/codex/live/sideband.go index 28b00535..2144f9a9 100644 --- a/internal/client/codex/live/sideband.go +++ b/internal/client/codex/live/sideband.go @@ -425,32 +425,20 @@ func (h *Handler) HandleSideband(c *gin.Context) { } upstream, handshakeResponse, errDial := dialUpstream(selected) - if errDial != nil && selection != nil && handshakeResponse != nil && handshakeResponse.StatusCode == http.StatusUnauthorized { - h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", session.model) - helps.RecordAPIWebsocketHandshake(ctx, runtimeConfig, handshakeResponse.StatusCode, callResponseHeaders(handshakeResponse.Header)) - if handshakeResponse.Body != nil { - if errClose := handshakeResponse.Body.Close(); errClose != nil { - log.Errorf("codex live sideband: close unauthorized handshake body error: %v", errClose) + if errDial != nil { + handshakeStatus := clienterror.HTTPStatusFromErrorOr(errDial, http.StatusBadGateway) + if handshakeResponse != nil && handshakeResponse.StatusCode > 0 { + handshakeStatus = handshakeResponse.StatusCode + } + responseBody := handleSidebandDialError(c, ctx, runtimeConfig, handshakeResponse, errDial) + if selection != nil && handshakeStatus == http.StatusUnauthorized { + diagnosticBody := responseBody + if len(diagnosticBody) == 0 { + diagnosticBody = []byte(errDial.Error()) } + h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", session.model, diagnosticBody) + log.WithField("status", handshakeStatus).Warnf("codex live sideband upstream handshake failed: %s", logging.SafeDiagnosticForLog(string(diagnosticBody))) } - refreshed, didRefresh, errRefresh := h.authManager.RefreshHomeSelectionAfterUnauthorized(ctx, selection, selected) - if errRefresh != nil { - writeSelectionError(c, errRefresh) - return - } - if !didRefresh || refreshed == nil { - writeLiveError(c, http.StatusUnauthorized, "Codex credential unauthorized") - return - } - selected = refreshed - logging.SetGinCPATraceID(c, selected.EnsureIndex()) - upstream, handshakeResponse, errDial = dialUpstream(selected) - if errDial != nil && handshakeResponse != nil && handshakeResponse.StatusCode == http.StatusUnauthorized { - h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", session.model) - } - } - if errDial != nil { - handleSidebandDialError(c, ctx, runtimeConfig, handshakeResponse, errDial) return } if handshakeResponse != nil { @@ -567,8 +555,9 @@ func callIDFromLocation(location string) string { return callID } -func handleSidebandDialError(c *gin.Context, ctx context.Context, cfg *config.Config, response *http.Response, errDial error) { +func handleSidebandDialError(c *gin.Context, ctx context.Context, cfg *config.Config, response *http.Response, errDial error) []byte { status := clienterror.HTTPStatusFromErrorOr(errDial, http.StatusBadGateway) + var responseBody []byte if response != nil { if response.StatusCode > 0 { status = response.StatusCode @@ -576,13 +565,32 @@ func handleSidebandDialError(c *gin.Context, ctx context.Context, cfg *config.Co copyRealtimeHandshakeHeaders(c.Writer.Header(), response.Header) helps.RecordAPIWebsocketHandshake(ctx, cfg, response.StatusCode, callResponseHeaders(response.Header)) if response.Body != nil { + var errRead error + responseBody, errRead = readLimitedBody(response.Body) + if errRead != nil { + log.Errorf("codex live sideband: read rejected handshake body error: %v", errRead) + } + helps.AppendAPIWebsocketResponse(ctx, cfg, responseBody) if errClose := response.Body.Close(); errClose != nil { log.Errorf("codex live sideband: close rejected handshake body error: %v", errClose) } } } helps.RecordAPIWebsocketError(ctx, cfg, "dial", errDial) + if response != nil && response.StatusCode == http.StatusUnauthorized { + if contentType := response.Header.Get("Content-Type"); contentType != "" { + c.Header("Content-Type", contentType) + } + c.Status(response.StatusCode) + if len(responseBody) > 0 { + if _, errWrite := c.Writer.Write(responseBody); errWrite != nil { + log.WithError(errWrite).Warn("codex live sideband: write rejected handshake body failed") + } + } + return responseBody + } writeLiveError(c, status, "Codex live sideband upstream unavailable") + return nil } func websocketCloseFunc(name string, conn *websocket.Conn) func() error { diff --git a/internal/client/codex/live/websocket.go b/internal/client/codex/live/websocket.go index e0a147dd..e57134de 100644 --- a/internal/client/codex/live/websocket.go +++ b/internal/client/codex/live/websocket.go @@ -111,28 +111,45 @@ func (h *Handler) HandleDirectWebsocket(c *gin.Context) { } upstream, handshakeResponse, errDial := dialUpstream(selected) - if errDial != nil && selection != nil && handshakeResponse != nil && handshakeResponse.StatusCode == http.StatusUnauthorized { - h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", selectionModel) - closeHandshakeBody(handshakeResponse, "direct websocket unauthorized") - refreshed, didRefresh, errRefresh := h.authManager.RefreshHomeSelectionAfterUnauthorized(ctx, selection, selected) - if errRefresh != nil { - writeSelectionError(c, errRefresh) - return - } - if didRefresh && refreshed != nil { - selected = refreshed - logging.SetGinCPATraceID(c, selected.EnsureIndex()) - upstream, handshakeResponse, errDial = dialUpstream(selected) - } - } if errDial != nil { status := clienterror.HTTPStatusFromErrorOr(errDial, http.StatusBadGateway) + helpConfig := h.currentConfig() + var responseBody []byte if handshakeResponse != nil && handshakeResponse.StatusCode > 0 { status = handshakeResponse.StatusCode copyRealtimeHandshakeHeaders(c.Writer.Header(), handshakeResponse.Header) + helps.RecordAPIWebsocketHandshake(ctx, helpConfig, handshakeResponse.StatusCode, callResponseHeaders(handshakeResponse.Header)) + if handshakeResponse.Body != nil { + var errRead error + responseBody, errRead = readLimitedBody(handshakeResponse.Body) + if errRead != nil { + log.Errorf("codex realtime: read rejected handshake body error: %v", errRead) + } + helps.AppendAPIWebsocketResponse(ctx, helpConfig, responseBody) + } } closeHandshakeBody(handshakeResponse, "direct websocket rejected") - helpConfig := h.currentConfig() + if selection != nil && status == http.StatusUnauthorized { + diagnosticBody := responseBody + if len(diagnosticBody) == 0 { + diagnosticBody = []byte(errDial.Error()) + } + h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", selectionModel, diagnosticBody) + log.WithField("status", status).Warnf("codex realtime websocket upstream handshake failed: %s", logging.SafeDiagnosticForLog(string(diagnosticBody))) + } + helps.RecordAPIWebsocketError(ctx, helpConfig, "dial", errDial) + if handshakeResponse != nil && handshakeResponse.StatusCode == http.StatusUnauthorized { + if contentType := handshakeResponse.Header.Get("Content-Type"); contentType != "" { + c.Header("Content-Type", contentType) + } + c.Status(handshakeResponse.StatusCode) + if len(responseBody) > 0 { + if _, errWrite := c.Writer.Write(responseBody); errWrite != nil { + log.WithError(errWrite).Warn("codex realtime: write rejected handshake body failed") + } + } + return + } helpDetails := "Codex Realtime WebSocket upstream unavailable" helpType := "api_error" if status == http.StatusNotFound || status == http.StatusNotImplemented { @@ -147,7 +164,6 @@ func (h *Handler) HandleDirectWebsocket(c *gin.Context) { helpType = "authentication_error" helpCode = "realtime_upstream_unauthorized" } - helps.RecordAPIWebsocketError(ctx, helpConfig, "dial", errDial) writeRealtimeError(c, status, helpDetails, helpType, helpCode) return } diff --git a/internal/client/codex/live/websocket_test.go b/internal/client/codex/live/websocket_test.go index 42eb78a7..b0de1b5c 100644 --- a/internal/client/codex/live/websocket_test.go +++ b/internal/client/codex/live/websocket_test.go @@ -2,16 +2,20 @@ package live import ( "encoding/json" + "io" "net/http" "net/http/httptest" "net/url" + "strconv" "strings" "testing" "time" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" ) func TestHandleDirectWebsocketRejectsClientSecretModelMismatch(t *testing.T) { @@ -33,6 +37,90 @@ func TestHandleDirectWebsocketRejectsClientSecretModelMismatch(t *testing.T) { } } +func TestHandleDirectWebsocketForwardsUnauthorizedHomeHandshakeWithoutRefresh(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, tc := range []struct { + name string + upstreamBody string + truncatedResponse bool + }{ + {name: "response body", upstreamBody: `{"error":{"message":"access token expired"}}`}, + {name: "response body with read error", upstreamBody: `{"error":{"message":"access token expired"}}`, truncatedResponse: true}, + {name: "empty response body"}, + } { + t.Run(tc.name, func(t *testing.T) { + upstreamAuthorization := make(chan string, 1) + upstreamServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + upstreamAuthorization <- request.Header.Get("Authorization") + writer.Header().Set("Content-Type", "application/json") + if tc.truncatedResponse { + writer.Header().Set("Content-Length", strconv.Itoa(len(tc.upstreamBody)+1)) + } + writer.WriteHeader(http.StatusUnauthorized) + _, _ = writer.Write([]byte(tc.upstreamBody)) + })) + defer upstreamServer.Close() + + runtimeConfig := &config.Config{ + Home: config.HomeConfig{Enabled: true}, + SDKConfig: config.SDKConfig{RequestLog: true}, + } + manager := auth.NewManager(nil, nil, nil) + manager.SetConfig(runtimeConfig) + manager.PublishHomeDispatch(&homeDispatcher{}, executionregistry.New(), 1) + executor := &captureExecutor{} + manager.RegisterExecutor(executor) + handler := NewHandler(manager, runtimeConfig) + handler.sidebandAPIBaseURL = "ws" + strings.TrimPrefix(upstreamServer.URL, "http") + "/v1" + router := gin.New() + timelineCapture := make(chan []byte, 1) + router.Use(func(c *gin.Context) { + c.Next() + if raw, exists := c.Get("API_WEBSOCKET_TIMELINE"); exists { + timeline, _ := raw.([]byte) + timelineCapture <- append([]byte(nil), timeline...) + } + }) + router.GET("/v1/realtime", handler.HandleRealtimeWebsocket) + downstreamServer := httptest.NewServer(router) + defer downstreamServer.Close() + + wsURL := "ws" + strings.TrimPrefix(downstreamServer.URL, "http") + "/v1/realtime?model=gpt-realtime" + connection, response, errDial := websocket.DefaultDialer.Dial(wsURL, nil) + if connection != nil { + _ = connection.Close() + } + if errDial == nil || response == nil { + t.Fatalf("dial downstream websocket = response %#v err %v, want rejected handshake", response, errDial) + } + defer func() { _ = response.Body.Close() }() + responseBody, errRead := io.ReadAll(response.Body) + if errRead != nil { + t.Fatalf("read downstream rejection: %v", errRead) + } + if response.StatusCode != http.StatusUnauthorized || string(responseBody) != tc.upstreamBody { + t.Fatalf("downstream rejection = status %d body %q, want original upstream 401 body %q", response.StatusCode, responseBody, tc.upstreamBody) + } + if executor.refreshCalls.Load() != 0 { + t.Fatalf("refresh calls = %d, want 0", executor.refreshCalls.Load()) + } + if got := <-upstreamAuthorization; got != "Bearer home-live-token" { + t.Fatalf("upstream Authorization = %q, want original Home token", got) + } + if tc.truncatedResponse { + select { + case timeline := <-timelineCapture: + if !strings.Contains(string(timeline), tc.upstreamBody) { + t.Fatalf("API_WEBSOCKET_TIMELINE = %q, want original upstream error", timeline) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for websocket request-log timeline") + } + } + }) + } +} + func TestHandleDirectWebsocketAppliesClientSecretSession(t *testing.T) { gin.SetMode(gin.TestMode) upstreamUpdate := make(chan []byte, 1) diff --git a/internal/logging/diagnostic.go b/internal/logging/diagnostic.go new file mode 100644 index 00000000..a55697e5 --- /dev/null +++ b/internal/logging/diagnostic.go @@ -0,0 +1,80 @@ +package logging + +import ( + "regexp" + "strings" + "unicode/utf8" +) + +const ( + diagnosticLogRuneLimit = 300 + diagnosticLogScanRuneLimit = 600 +) + +var ( + accessTokenExpiredLogPattern = regexp.MustCompile(`(?i)access token expired`) + sensitiveLogAssignmentPattern = regexp.MustCompile(`(?i)(["']?(?:access[\s_-]*token|refresh[\s_-]*token|id[\s_-]*token|api[\s_-]*key|client[\s_-]*secret|private[\s_-]*key|proxy[\s_-]*authorization|authorization|password|credential|token|secret)["']?\s*[:=]\s*)(?:(?:bearer|basic)\s+[^\s,;]+|"(?:\\.|[^"])*"|'(?:\\.|[^'])*'|[^\s,;&}\]]+)`) + authorizationLogPattern = regexp.MustCompile(`(?i)\b(bearer|basic)\s+[^\s,;]+`) + urlUserinfoLogPattern = regexp.MustCompile(`(?i)(https?://)[^/\s@]+@`) +) + +// SafeDiagnosticForLog returns a bounded, single-line diagnostic suitable for +// ordinary application logs. It preserves the access-token-expired signal while +// redacting credential values and URL userinfo. +func SafeDiagnosticForLog(message string) string { + excerpt, sourceTruncated := diagnosticRunePrefix(message, diagnosticLogScanRuneLimit) + if sourceTruncated && !accessTokenExpiredLogPattern.MatchString(excerpt) { + if marker := accessTokenExpiredLogPattern.FindString(message); marker != "" { + excerpt += " ... " + marker + } + } + + excerpt = strings.Join(strings.Fields(excerpt), " ") + if excerpt == "" { + return "" + } + excerpt = urlUserinfoLogPattern.ReplaceAllString(excerpt, `${1}[REDACTED]@`) + excerpt = sensitiveLogAssignmentPattern.ReplaceAllString(excerpt, `${1}"[REDACTED]"`) + excerpt = authorizationLogPattern.ReplaceAllString(excerpt, `${1} [REDACTED]`) + + return truncateDiagnosticLogExcerpt(excerpt, sourceTruncated) +} + +func diagnosticRunePrefix(value string, limit int) (string, bool) { + if limit <= 0 { + return "", value != "" + } + count := 0 + for index := range value { + if count == limit { + return value[:index], true + } + count++ + } + return value, false +} + +func truncateDiagnosticLogExcerpt(message string, sourceTruncated bool) string { + runes := []rune(message) + if len(runes) <= diagnosticLogRuneLimit { + if sourceTruncated { + return message + "..." + } + return message + } + + output := string(runes[:diagnosticLogRuneLimit]) + if match := accessTokenExpiredLogPattern.FindStringIndex(message); match != nil { + markerStart := utf8.RuneCountInString(message[:match[0]]) + markerEnd := markerStart + utf8.RuneCountInString(message[match[0]:match[1]]) + if markerEnd > diagnosticLogRuneLimit { + separator := " ... " + prefixLimit := diagnosticLogRuneLimit - len([]rune(separator)) - (markerEnd - markerStart) + if prefixLimit < 0 { + prefixLimit = 0 + } + output = string(runes[:prefixLimit]) + separator + string(runes[markerStart:markerEnd]) + } + } + return output + "..." +} diff --git a/internal/logging/diagnostic_test.go b/internal/logging/diagnostic_test.go new file mode 100644 index 00000000..dc200f95 --- /dev/null +++ b/internal/logging/diagnostic_test.go @@ -0,0 +1,59 @@ +package logging + +import ( + "strings" + "testing" +) + +func TestSafeDiagnosticForLogPreservesAccessTokenExpiredAndRedactsCredentials(t *testing.T) { + diagnostic := "access token expired\n" + + `access_token=access-secret refresh token: refresh-secret Authorization=Bearer bearer-secret ` + + `Post "https://user:password@oauth.example/token?access_token=query-secret"` + + got := SafeDiagnosticForLog(diagnostic) + if !strings.Contains(got, "access token expired") { + t.Fatalf("safe diagnostic lost access-token-expired signal: %q", got) + } + for _, secret := range []string{"access-secret", "refresh-secret", "bearer-secret", "query-secret", "user:password"} { + if strings.Contains(got, secret) { + t.Fatalf("safe diagnostic leaked %q: %q", secret, got) + } + } + if strings.ContainsAny(got, "\r\n") { + t.Fatalf("safe diagnostic retained a line break: %q", got) + } + if !strings.Contains(got, "[REDACTED]") { + t.Fatalf("safe diagnostic did not mark redacted values: %q", got) + } +} + +func TestSafeDiagnosticForLogKeepsPlainAccessTokenExpiredMessage(t *testing.T) { + const diagnostic = "access token expired" + if got := SafeDiagnosticForLog(diagnostic); got != diagnostic { + t.Fatalf("SafeDiagnosticForLog() = %q, want %q", got, diagnostic) + } +} + +func TestSafeDiagnosticForLogBoundsLargeMessageAndRetainsTrailingSignal(t *testing.T) { + diagnostic := strings.Repeat("upstream context ", 1000) + "access token expired\nforged log line" + got := SafeDiagnosticForLog(diagnostic) + if len([]rune(got)) > diagnosticLogRuneLimit+3 { + t.Fatalf("safe diagnostic length = %d, want at most %d", len([]rune(got)), diagnosticLogRuneLimit+3) + } + if !strings.Contains(got, "access token expired") { + t.Fatalf("safe diagnostic lost trailing access-token-expired signal: %q", got) + } + if strings.ContainsAny(got, "\r\n") { + t.Fatalf("safe diagnostic retained a line break: %q", got) + } + if !strings.HasSuffix(got, "...") { + t.Fatalf("safe diagnostic did not indicate truncation: %q", got) + } +} + +func TestSafeDiagnosticForLogBoundsLargeGenericMessage(t *testing.T) { + got := SafeDiagnosticForLog(strings.Repeat("x", 900)) + if len([]rune(got)) != diagnosticLogRuneLimit+3 || !strings.HasSuffix(got, "...") { + t.Fatalf("safe generic diagnostic length = %d, want %d with ellipsis", len([]rune(got)), diagnosticLogRuneLimit+3) + } +} diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index a020709e..590d4f1f 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -39,7 +39,7 @@ const ( antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" antigravityClientSecret = "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" antigravityAuthType = "antigravity" - refreshSkew = 3000 * time.Second + refreshSkew = 50 * time.Minute antigravityCreditsHintRefreshInterval = 10 * time.Minute antigravityCreditsHintRefreshTimeout = 5 * time.Second antigravityShortQuotaCooldownThreshold = 5 * time.Minute diff --git a/internal/runtime/executor/helps/home_refresh.go b/internal/runtime/executor/helps/home_refresh.go index 2e3318cf..1a4abfff 100644 --- a/internal/runtime/executor/helps/home_refresh.go +++ b/internal/runtime/executor/helps/home_refresh.go @@ -14,11 +14,15 @@ import ( ) type homeStatusErr struct { - code int - msg string + code int + msg string + upstream bool } func (e homeStatusErr) Error() string { + if e.upstream { + return e.msg + } if e.msg != "" { return e.msg } @@ -27,6 +31,15 @@ func (e homeStatusErr) Error() string { func (e homeStatusErr) StatusCode() int { return e.code } +func (e homeStatusErr) DirectResponse() bool { return e.upstream } + +func (e homeStatusErr) ResponseBody() []byte { + if !e.upstream { + return nil + } + return []byte(e.msg) +} + type homeErrorEnvelope struct { Error *homeErrorDetail `json:"error"` } @@ -37,9 +50,15 @@ type homeRefreshAuthEnvelope struct { } type homeErrorDetail struct { - Type string `json:"type"` - Message string `json:"message"` - Code string `json:"code,omitempty"` + Type string `json:"type"` + Message string `json:"message"` + Code string `json:"code,omitempty"` + Upstream *homeUpstreamResponse `json:"upstream,omitempty"` +} + +type homeUpstreamResponse struct { + Status int `json:"status"` + Body []byte `json:"body"` } type homeRefreshClient interface { @@ -88,6 +107,13 @@ func RefreshAuthViaHome(ctx context.Context, cfg *config.Config, auth *cliproxya var env homeErrorEnvelope if errUnmarshal := json.Unmarshal(raw, &env); errUnmarshal == nil && env.Error != nil { + if env.Error.Upstream != nil { + return nil, true, homeStatusErr{ + code: env.Error.Upstream.Status, + msg: string(env.Error.Upstream.Body), + upstream: true, + } + } code := strings.TrimSpace(env.Error.Type) if code == "" { code = strings.TrimSpace(env.Error.Code) diff --git a/internal/runtime/executor/helps/home_refresh_test.go b/internal/runtime/executor/helps/home_refresh_test.go index be33016d..00909938 100644 --- a/internal/runtime/executor/helps/home_refresh_test.go +++ b/internal/runtime/executor/helps/home_refresh_test.go @@ -1,6 +1,7 @@ package helps import ( + "bytes" "context" "encoding/json" "errors" @@ -60,7 +61,7 @@ func TestRefreshAuthViaHomePreservesContextErrors(t *testing.T) { } } -func TestRefreshAuthViaHomeMapsTransportFailureToRedacted503(t *testing.T) { +func TestRefreshAuthViaHomeMapsTransportFailureToGeneric503(t *testing.T) { client := &fakeHomeRefreshClient{err: errors.New("dial failed with provider-secret")} oldCurrentHomeRefreshClient := currentHomeRefreshClient currentHomeRefreshClient = func() homeRefreshClient { return client } @@ -71,14 +72,21 @@ func TestRefreshAuthViaHomeMapsTransportFailureToRedacted503(t *testing.T) { _, handled, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth) statusErr, okStatus := errRefresh.(interface{ StatusCode() int }) if !handled || !okStatus || statusErr.StatusCode() != http.StatusServiceUnavailable { - t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want redacted 503", handled, errRefresh) + t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want generic 503", handled, errRefresh) } if strings.Contains(errRefresh.Error(), "provider-secret") { - t.Fatalf("refresh error leaked transport detail: %v", errRefresh) + t.Fatalf("refresh error included transport detail: %v", errRefresh) + } + direct, okDirect := errRefresh.(interface { + DirectResponse() bool + ResponseBody() []byte + }) + if !okDirect || direct.DirectResponse() { + t.Fatalf("transport error direct response = %v/%v, want false", okDirect, direct) } } -func TestRefreshAuthViaHomeRedactsLegacyErrorEnvelope(t *testing.T) { +func TestRefreshAuthViaHomeUsesGenericMessageForLegacyErrorEnvelope(t *testing.T) { client := &fakeHomeRefreshClient{raw: []byte(`{"error":{"type":"error","message":"provider response: refresh_token=provider-secret"}}`)} oldCurrentHomeRefreshClient := currentHomeRefreshClient currentHomeRefreshClient = func() homeRefreshClient { return client } @@ -89,10 +97,125 @@ func TestRefreshAuthViaHomeRedactsLegacyErrorEnvelope(t *testing.T) { _, handled, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth) statusErr, okStatus := errRefresh.(interface{ StatusCode() int }) if !handled || !okStatus || statusErr.StatusCode() != http.StatusServiceUnavailable { - t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want redacted 503", handled, errRefresh) + t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want generic 503", handled, errRefresh) } if strings.Contains(errRefresh.Error(), "provider-secret") { - t.Fatalf("refresh error leaked legacy Home detail: %v", errRefresh) + t.Fatalf("refresh error included legacy Home detail: %v", errRefresh) + } +} + +func TestRefreshAuthViaHomePreservesUpstreamStatusAndBodyExactly(t *testing.T) { + tests := []struct { + name string + status int + body []byte + }{ + {name: "json", status: http.StatusBadRequest, body: []byte(`{"error":"invalid_request"}`)}, + {name: "access token expired", status: http.StatusUnauthorized, body: []byte(`{"error":{"message":"access token expired"}}`)}, + {name: "text", status: http.StatusBadGateway, body: []byte("provider unavailable")}, + {name: "multiline", status: http.StatusTooManyRequests, body: []byte("first line\r\nsecond line\n")}, + {name: "empty", status: http.StatusUnauthorized, body: []byte{}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + raw, errMarshal := json.Marshal(homeErrorEnvelope{Error: &homeErrorDetail{ + Type: "refresh_temporarily_unavailable", + Message: "credential refresh temporarily unavailable", + Upstream: &homeUpstreamResponse{ + Status: tt.status, + Body: tt.body, + }, + }}) + if errMarshal != nil { + t.Fatalf("marshal Home error envelope: %v", errMarshal) + } + client := &fakeHomeRefreshClient{raw: raw} + oldCurrentHomeRefreshClient := currentHomeRefreshClient + currentHomeRefreshClient = func() homeRefreshClient { return client } + t.Cleanup(func() { currentHomeRefreshClient = oldCurrentHomeRefreshClient }) + + cfg := &config.Config{Home: config.HomeConfig{Enabled: true}} + auth := &cliproxyauth.Auth{ID: "home-auth", Index: "home-auth", Provider: "codex"} + _, handled, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth) + statusErr, okStatus := errRefresh.(interface{ StatusCode() int }) + if !handled || !okStatus || statusErr.StatusCode() != tt.status { + t.Fatalf("RefreshAuthViaHome() = handled %v err %v status %v, want true/%d", handled, errRefresh, okStatus, tt.status) + } + if got := []byte(errRefresh.Error()); !bytes.Equal(got, tt.body) { + t.Fatalf("refresh body = %q, want exact body %q", got, tt.body) + } + direct, okDirect := errRefresh.(interface { + DirectResponse() bool + ResponseBody() []byte + }) + if !okDirect || !direct.DirectResponse() { + t.Fatalf("upstream error direct response = %v/%v, want true", okDirect, direct) + } + if got := direct.ResponseBody(); !bytes.Equal(got, tt.body) { + t.Fatalf("direct response body = %q, want exact body %q", got, tt.body) + } + }) + } +} + +func TestRefreshAuthViaHomeUsesGenericMessageForLegacyProviderDiagnostics(t *testing.T) { + tests := []struct { + name string + provider string + raw []byte + wantStatus int + wantError string + }{ + { + name: "transient transport diagnostic", + provider: "antigravity", + raw: []byte(`{"error":{"type":"refresh_temporarily_unavailable","message":"antigravity refresh: Post https://oauth.example/token?access_token=provider-secret: connection refused"}}`), + wantStatus: http.StatusServiceUnavailable, + wantError: "credential refresh temporarily unavailable", + }, + { + name: "terminal legacy diagnostic", + provider: "codex", + raw: []byte(`{"error":{"type":"authentication_error","message":"codex refresh: invalid_grant refresh_token=provider-secret"}}`), + wantStatus: http.StatusUnauthorized, + wantError: "credential unauthorized", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + client := &fakeHomeRefreshClient{raw: tt.raw} + oldCurrentHomeRefreshClient := currentHomeRefreshClient + currentHomeRefreshClient = func() homeRefreshClient { return client } + t.Cleanup(func() { currentHomeRefreshClient = oldCurrentHomeRefreshClient }) + + cfg := &config.Config{Home: config.HomeConfig{Enabled: true}} + auth := &cliproxyauth.Auth{ID: "home-auth", Index: "home-auth", Provider: tt.provider} + _, handled, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth) + statusErr, okStatus := errRefresh.(interface{ StatusCode() int }) + if !handled || !okStatus || statusErr.StatusCode() != tt.wantStatus { + t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want status %d", handled, errRefresh, tt.wantStatus) + } + if got := errRefresh.Error(); got != tt.wantError { + t.Fatalf("client refresh error = %q, want %q", got, tt.wantError) + } + if strings.Contains(errRefresh.Error(), "provider-secret") { + t.Fatalf("legacy refresh error leaked provider detail: %v", errRefresh) + } + }) + } +} + +func TestRefreshAuthViaHomeRejectsUnmarkedRefreshMessage(t *testing.T) { + client := &fakeHomeRefreshClient{raw: []byte(`{"error":{"type":"refresh_temporarily_unavailable","message":"database unavailable: provider-secret"}}`)} + oldCurrentHomeRefreshClient := currentHomeRefreshClient + currentHomeRefreshClient = func() homeRefreshClient { return client } + t.Cleanup(func() { currentHomeRefreshClient = oldCurrentHomeRefreshClient }) + + cfg := &config.Config{Home: config.HomeConfig{Enabled: true}} + auth := &cliproxyauth.Auth{ID: "home-auth", Index: "home-auth", Provider: "antigravity"} + _, _, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth) + if got, want := errRefresh.Error(), "credential refresh temporarily unavailable"; got != want { + t.Fatalf("refresh error = %q, want %q", got, want) } } diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index 247e7006..5596e68d 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -3,6 +3,7 @@ package helps import ( "bytes" "context" + "errors" "fmt" "io" "net/http" @@ -403,8 +404,18 @@ func failFromErrors(errs ...error) usage.Failure { if err == nil { continue } + body := strings.TrimSpace(err.Error()) + type responseBodyProvider interface { + ResponseBody() []byte + } + var responseErr responseBodyProvider + if errors.As(err, &responseErr) && responseErr != nil { + if responseBody := responseErr.ResponseBody(); len(responseBody) > 0 { + body = string(responseBody) + } + } return usage.Failure{ - Body: strings.TrimSpace(err.Error()), + Body: body, StatusCode: clienterror.HTTPStatusFromError(err), } } diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index afc42b89..8a561da5 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -3,6 +3,7 @@ package helps import ( "context" "errors" + "fmt" "io" "net/http" "net/url" @@ -812,6 +813,50 @@ func TestUsageReporterBuildAdditionalModelRecordSkipsZeroTokens(t *testing.T) { } } +type usageResponseBodyError struct { + status int + message string + body []byte +} + +func (e usageResponseBodyError) Error() string { + return e.message +} + +func (e usageResponseBodyError) StatusCode() int { + return e.status +} + +func (e usageResponseBodyError) ResponseBody() []byte { + return e.body +} + +func TestFailFromErrorsPrefersResponseBody(t *testing.T) { + for _, tc := range []struct { + name string + body []byte + }{ + {name: "original response body", body: []byte(" \n{\"error\":\"upstream rejected request\"}\r\n")}, + {name: "empty response body"}, + } { + t.Run(tc.name, func(t *testing.T) { + errExecute := fmt.Errorf("execute failed: %w", usageResponseBodyError{ + status: http.StatusUnauthorized, + message: "generic upstream error", + body: tc.body, + }) + failure := failFromErrors(errExecute) + wantBody := errExecute.Error() + if len(tc.body) > 0 { + wantBody = string(tc.body) + } + if failure.StatusCode != http.StatusUnauthorized || failure.Body != wantBody { + t.Fatalf("failure = %#v, want status %d body %q", failure, http.StatusUnauthorized, wantBody) + } + }) + } +} + func TestFailFromErrorsMapsContextStatuses(t *testing.T) { tests := []struct { name string diff --git a/sdk/api/handlers/handlers_error_response_test.go b/sdk/api/handlers/handlers_error_response_test.go index c5253901..8a5b1b4e 100644 --- a/sdk/api/handlers/handlers_error_response_test.go +++ b/sdk/api/handlers/handlers_error_response_test.go @@ -1,6 +1,7 @@ package handlers import ( + "bytes" "context" "errors" "net/http" @@ -18,6 +19,26 @@ import ( sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" ) +type directResponseTestError struct { + status int + body []byte + direct bool +} + +func (e directResponseTestError) Error() string { return string(e.body) } +func (e directResponseTestError) StatusCode() int { return e.status } +func (e directResponseTestError) DirectResponse() bool { return e.direct } +func (e directResponseTestError) ResponseBody() []byte { return e.body } + +type responseBodyOnlyTestError struct { + status int + body []byte +} + +func (e responseBodyOnlyTestError) Error() string { return string(e.body) } +func (e responseBodyOnlyTestError) StatusCode() int { return e.status } +func (e responseBodyOnlyTestError) ResponseBody() []byte { return e.body } + func TestWriteErrorResponse_AddonHeadersDisabledByDefault(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() @@ -86,6 +107,58 @@ func TestWriteErrorResponseDirectResponse(t *testing.T) { } } +func TestExecutionErrorMessagePreservesMarkedResponseBodyExactly(t *testing.T) { + tests := []struct { + name string + status int + body []byte + contentType string + }{ + {name: "json", status: http.StatusBadRequest, body: []byte(`{"error":"invalid_request"}`), contentType: "application/json"}, + {name: "text", status: http.StatusBadGateway, body: []byte("provider unavailable"), contentType: "text/plain; charset=utf-8"}, + {name: "multiline", status: http.StatusTooManyRequests, body: []byte("first line\r\nsecond line\n"), contentType: "text/plain; charset=utf-8"}, + {name: "empty", status: http.StatusUnauthorized, body: []byte{}, contentType: "application/json"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + + errUpstream := directResponseTestError{status: tt.status, body: tt.body, direct: true} + msg := executionErrorMessage(errUpstream) + if msg == nil || !msg.DirectResponse { + t.Fatalf("executionErrorMessage() direct response = %#v, want true", msg) + } + NewBaseAPIHandlers(nil, nil).WriteErrorResponse(c, msg) + if recorder.Code != tt.status { + t.Fatalf("status = %d, want %d", recorder.Code, tt.status) + } + if got := recorder.Body.Bytes(); !bytes.Equal(got, tt.body) { + t.Fatalf("body = %q, want exact body %q", got, tt.body) + } + if got := recorder.Header().Get("Content-Type"); got != tt.contentType { + t.Fatalf("Content-Type = %q, want %q", got, tt.contentType) + } + }) + } +} + +func TestExecutionErrorMessageDoesNotTrustResponseBodyWithoutMarker(t *testing.T) { + errUnmarked := responseBodyOnlyTestError{ + status: http.StatusBadGateway, + body: []byte("unmarked upstream body"), + } + msg := executionErrorMessage(errUnmarked) + if msg == nil { + t.Fatal("executionErrorMessage() returned nil") + } + if msg.DirectResponse { + t.Fatal("executionErrorMessage() trusted an unmarked response body") + } +} + func TestInternalConcurrencyBusyWritesRetryAfterWithoutPassthrough(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() diff --git a/sdk/api/handlers/handlers_execution.go b/sdk/api/handlers/handlers_execution.go index ac521ec9..79632acc 100644 --- a/sdk/api/handlers/handlers_execution.go +++ b/sdk/api/handlers/handlers_execution.go @@ -1,6 +1,7 @@ package handlers import ( + "encoding/json" "errors" "fmt" "net/http" @@ -329,6 +330,29 @@ func executionErrorMessage(err error) *interfaces.ErrorMessage { if code := clienterror.HTTPStatusFromError(err); code > 0 { status = code } + type directResponseError interface { + DirectResponse() bool + ResponseBody() []byte + } + var direct directResponseError + if errors.As(err, &direct) && direct != nil && direct.DirectResponse() { + body := direct.ResponseBody() + var headers http.Header + if len(body) > 0 { + contentType := http.DetectContentType(body) + if json.Valid(body) { + contentType = "application/json" + } + headers = http.Header{"Content-Type": []string{contentType}} + } + return &interfaces.ErrorMessage{ + StatusCode: status, + Error: err, + DirectResponse: true, + Body: body, + Headers: headers, + } + } var addon http.Header if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { if hdr := he.Headers(); hdr != nil { diff --git a/sdk/auth/antigravity.go b/sdk/auth/antigravity.go index ee41cbdb..8fe1c148 100644 --- a/sdk/auth/antigravity.go +++ b/sdk/auth/antigravity.go @@ -26,9 +26,9 @@ func NewAntigravityAuthenticator() Authenticator { return &AntigravityAuthentica // Provider returns the provider key for antigravity. func (AntigravityAuthenticator) Provider() string { return "antigravity" } -// RefreshLead instructs the manager to refresh five minutes before expiry. +// RefreshLead instructs the manager to refresh 50 minutes before expiry. func (AntigravityAuthenticator) RefreshLead() *time.Duration { - return new(5 * time.Minute) + return new(50 * time.Minute) } // Login launches a local OAuth flow to obtain antigravity tokens and persists them. diff --git a/sdk/auth/refresh_registry_test.go b/sdk/auth/refresh_registry_test.go new file mode 100644 index 00000000..ec4d76b1 --- /dev/null +++ b/sdk/auth/refresh_registry_test.go @@ -0,0 +1,32 @@ +package auth + +import ( + "testing" + "time" +) + +func TestProviderRefreshLeads(t *testing.T) { + tests := []struct { + name string + authenticator Authenticator + want time.Duration + }{ + {name: "codex", authenticator: NewCodexAuthenticator(), want: 5 * 24 * time.Hour}, + {name: "claude", authenticator: NewClaudeAuthenticator(), want: 4 * time.Hour}, + {name: "antigravity", authenticator: NewAntigravityAuthenticator(), want: 50 * time.Minute}, + {name: "kimi", authenticator: NewKimiAuthenticator(), want: 5 * time.Minute}, + {name: "xai", authenticator: NewXAIAuthenticator(), want: 5 * time.Minute}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := test.authenticator.Provider(); got != test.name { + t.Fatalf("Provider() = %q, want %q", got, test.name) + } + lead := test.authenticator.RefreshLead() + if lead == nil || *lead != test.want { + t.Fatalf("RefreshLead() = %v, want %v", lead, test.want) + } + }) + } +} diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index c7dcebe2..a968438e 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -685,7 +685,6 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string lastHomeAuthID := "" homeSameAuthRetryPending := false attempted := make(map[string]struct{}) - unauthorizedRefreshTried := make(map[string]struct{}) var lastErr error var roundTiming homeRetryRoundTiming for { @@ -784,14 +783,6 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } } } - if _, refreshedAlready := unauthorizedRefreshTried[auth.ID]; refreshedAlready { - homeExcludedAuthIDs[auth.ID] = struct{}{} - homeSameAuthRetryPending = false - if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "repeated_refresh_auth"); errEnd != nil { - return nil, errEnd - } - continue - } } entry := logEntryWithRequestID(ctx) @@ -848,7 +839,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string if errPrepare != nil { if selection != nil { excludeAuth := shouldExcludeHomeAuthAfterStreamError(execCtx, auth, errPrepare) - if _, refreshedAlready := unauthorizedRefreshTried[auth.ID]; refreshedAlready || homeSameAuthRetries[auth.ID] > 0 { + if homeSameAuthRetries[auth.ID] > 0 { excludeAuth = true } if excludeAuth { @@ -893,11 +884,11 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string models = models[:1] pooled = false } - streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, execOpts, routeModel, streamExecutionModel, models, pooled, aliasResult, routing, !homeMode || selection != nil, selection != nil, unauthorizedRefreshTried) + streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, execOpts, routeModel, streamExecutionModel, models, pooled, aliasResult, routing, !homeMode || selection != nil, selection != nil) if errStream != nil { if selection != nil { excludeAuth := shouldExcludeHomeAuthAfterStreamError(execCtx, auth, errStream) - if _, refreshedAlready := unauthorizedRefreshTried[auth.ID]; refreshedAlready || homeSameAuthRetries[auth.ID] > 0 { + if homeSameAuthRetries[auth.ID] > 0 { excludeAuth = true } if excludeAuth { @@ -951,17 +942,16 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } } -func shouldExcludeHomeAuthAfterStreamError(ctx context.Context, auth *Auth, err error) bool { +func shouldExcludeHomeAuthAfterStreamError(ctx context.Context, _ *Auth, err error) bool { if err == nil || isConnectionLifecycleError(err) { return false } // A 426 during a downstream websocket attempt is a transport fallback - // signal. OAuth authorization failures may also recover after a refresh. - // Both paths may retry the same credential once. + // signal and may retry the same credential once. if cliproxyexecutor.DownstreamWebsocket(ctx) && statusCodeFromError(err) == http.StatusUpgradeRequired { return false } - return !isUnauthorizedError(err) || auth == nil || auth.AuthKind() != AuthKindOAuth + return true } func cloneRequestMetadata(src map[string]any) map[string]any { @@ -1175,7 +1165,12 @@ func (m *Manager) prepareHomeRequestAuth(ctx context.Context, executor ProviderE if selection == nil { return nil, nil } - return m.prepareHomeAuthSnapshot(ctx, executor, selection.CloneAuth()) + auth := selection.CloneAuth() + prepared, errPrepare := m.prepareHomeAuthSnapshot(ctx, executor, auth) + if errPrepare != nil { + warnLogHomeCredentialFailure(ctx, "request_auth_preparation", selection.Provider, auth, errPrepare) + } + return prepared, errPrepare } func (m *Manager) prepareHomeAuthSnapshot(ctx context.Context, executor ProviderExecutor, auth *Auth) (*Auth, error) { @@ -1577,17 +1572,23 @@ func formatAuthIdentity(auth *Auth, provider string) string { } } -func summarizeErrorForLog(err error) string { +func warnLogHomeCredentialFailure(ctx context.Context, operation, provider string, auth *Auth, err error) { if err == nil { - return "" + return + } + provider = strings.TrimSpace(provider) + if provider == "" && auth != nil { + provider = strings.TrimSpace(auth.Provider) } - msg := strings.TrimSpace(err.Error()) - const maxRunes = 300 - runes := []rune(msg) - if len(runes) > maxRunes { - return string(runes[:maxRunes]) + "..." + fields := log.Fields{ + "auth": formatAuthIdentity(auth, provider), + "operation": operation, + "provider": provider, } - return msg + if statusCode := statusCodeFromError(err); statusCode != 0 { + fields["status"] = statusCode + } + logEntryWithRequestID(ctx).WithFields(fields).Warnf("Home credential operation failed: err=%s", logging.SafeDiagnosticForLog(err.Error())) } func warnLogUpstreamFailure(ctx context.Context, entry *log.Entry, provider, model string, auth *Auth, duration time.Duration, err error) { @@ -1611,8 +1612,13 @@ func warnLogUpstreamFailure(ctx context.Context, entry *log.Entry, provider, mod } } authIdent := formatAuthIdentity(auth, provider) - errSummary := summarizeErrorForLog(err) - entry.Warnf("upstream execution failed: provider=%s model=%s auth=%s duration=%s err=%s", provider, model, authIdent, duration.Round(time.Millisecond), errSummary) + errSummary := logging.SafeDiagnosticForLog(err.Error()) + duration = duration.Round(time.Millisecond) + if statusCode := statusCodeFromError(err); statusCode != 0 { + entry.Warnf("%3d | %13v | upstream execution failed: provider=%s model=%s auth=%s err=%s", statusCode, duration, provider, model, authIdent, errSummary) + return + } + entry.Warnf("upstream execution failed: provider=%s model=%s auth=%s duration=%s err=%s", provider, model, authIdent, duration, errSummary) } // InjectCredentials delegates per-provider HTTP request preparation when supported. diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index ce956a17..c33fee02 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -1403,7 +1403,7 @@ func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cl if len(models) == 0 { continue } - result, errStream := m.executeStreamWithModelPool(creditsCtx, c.executor, c.auth, c.provider, req, creditsOpts, routeModel, "", models, pooled, aliasResult, routing, true, false, nil) + result, errStream := m.executeStreamWithModelPool(creditsCtx, c.executor, c.auth, c.provider, req, creditsOpts, routeModel, "", models, pooled, aliasResult, routing, true, false) if errStream != nil { continue } diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index 9be8f52f..16e5f3b3 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -155,7 +155,6 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c roundTiming.Observe(lastErr) continue } - didRefreshOnUnauthorized := false for _, upstreamModel := range models { execCtx = newUpstreamAttemptContext(execCtx) resultModel := m.stateModelForExecution(preparedAuth, routeModel, upstreamModel, pooled) @@ -214,44 +213,13 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c startHomeExec := time.Now() response, errExecute = execute() durationHomeExec := time.Since(startHomeExec) - refreshAuth := preparedAuth if countTokens { - if observedAuth, fingerprint := getEffectiveAuth(); isUnauthorizedError(errExecute) { - m.reportHomeUnauthorized(execCtx, preparedAuth, selection.Provider, resultModel, fingerprint) - if observedAuth != nil { - refreshAuth = observedAuth - } + if _, fingerprint := getEffectiveAuth(); isUnauthorizedError(errExecute) { + m.reportHomeUnauthorized(execCtx, preparedAuth, selection.Provider, resultModel, fingerprint, extractErrorBody(errExecute)) } } if errExecute != nil { - refreshCtx := newUpstreamAttemptContext(execCtx) - if refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(refreshCtx, selection.Executor, refreshAuth, errExecute, didRefreshOnUnauthorized, true); errRefresh != nil { - errExecute = errRefresh - warnLogUpstreamFailure(execCtx, entry, selection.Provider, upstreamModel, preparedAuth, durationHomeExec, errExecute) - } else if okRefresh { - preparedAuth = refreshed - m.replaceHomeSelectionAuth(selection, preparedAuth) - didRefreshOnUnauthorized = true - publishSelectedAuthMetadata(opts.Metadata, preparedAuth) - setEffectiveAuth(preparedAuth) - execCtx = newUpstreamAttemptContext(execCtx) - executorCtx = execCtx - if countTokens { - executorCtx = withAccessTokenFingerprintObserver(execCtx, setEffectiveAuth) - } - startHomeRetry := time.Now() - response, errExecute = execute() - durationHomeRetry := time.Since(startHomeRetry) - if errExecute != nil { - warnLogUpstreamFailure(execCtx, entry, selection.Provider, upstreamModel, preparedAuth, durationHomeRetry, errExecute) - if countTokens && isUnauthorizedError(errExecute) { - _, fingerprint := getEffectiveAuth() - m.reportHomeUnauthorized(execCtx, preparedAuth, selection.Provider, resultModel, fingerprint) - } - } - } else { - warnLogUpstreamFailure(execCtx, entry, selection.Provider, upstreamModel, preparedAuth, durationHomeExec, errExecute) - } + warnLogUpstreamFailure(execCtx, entry, selection.Provider, upstreamModel, preparedAuth, durationHomeExec, errExecute) } result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil, Options: execOpts} if errExecute == nil { diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go index 06a2a27a..ea0ba572 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -377,54 +377,9 @@ func clearUnauthorizedModelStates(auth *Auth, now time.Time) []string { return resumed } -// tryRefreshExecutionAuthAfterUnauthorized refreshes OAuth credentials once for -// either a local auth or an ephemeral Home dispatch auth. -func (m *Manager) tryRefreshExecutionAuthAfterUnauthorized(ctx context.Context, executor ProviderExecutor, auth *Auth, execErr error, alreadyTried bool, homeDispatch bool) (*Auth, bool, error) { - if !homeDispatch { - refreshed, ok := m.tryRefreshAfterUnauthorized(ctx, auth, execErr, alreadyTried) - return refreshed, ok, nil - } - if m == nil || executor == nil || auth == nil || alreadyTried || execErr == nil { - return auth, false, nil - } - if !isUnauthorizedError(execErr) || auth.AuthKind() != AuthKindOAuth { - return auth, false, nil - } - - log.Debugf("unauthorized Home response for %s (%s), refreshing credentials before redispatch", auth.Provider, auth.ID) - target := auth.Clone() - updated, errRefresh := executor.Refresh(ctx, target) - if errRefresh != nil { - log.Debugf("Home credential refresh before redispatch failed for %s (%s)", auth.Provider, auth.ID) - return auth, false, errRefresh - } - if updated == nil { - updated = target - } - if updated.ID == "" { - updated.ID = auth.ID - } - if updated.Index == "" { - updated.Index = auth.Index - } - if updated.Provider == "" { - updated.Provider = auth.Provider - } - if updated.Runtime == nil { - updated.Runtime = auth.Runtime - } - preserveHomeRoutingAttributes(updated, auth) - prepared, errPrepare := m.prepareHomeAuthSnapshot(ctx, executor, updated) - if errPrepare != nil { - return auth, false, errPrepare - } - preserveHomeRoutingAttributes(prepared, auth) - return prepared, true, nil -} - -// RefreshHomeSelectionAfterUnauthorized refreshes the credential snapshot that -// received a 401, or reuses a newer token already installed on the selection. -func (m *Manager) RefreshHomeSelectionAfterUnauthorized(ctx context.Context, selection *HomeDispatchSelection, failedAuth *Auth) (*Auth, bool, error) { +// RefreshHomeSelectionAfterUnauthorized only reuses a newer snapshot already +// installed by Home. It never refreshes or mutates Home-owned credentials. +func (m *Manager) RefreshHomeSelectionAfterUnauthorized(_ context.Context, selection *HomeDispatchSelection, failedAuth *Auth) (*Auth, bool, error) { if m == nil || selection == nil { return nil, false, nil } @@ -436,25 +391,10 @@ func (m *Manager) RefreshHomeSelectionAfterUnauthorized(ctx context.Context, sel currentToken := authAccessToken(current) failedToken := authAccessToken(failedAuth) if currentToken != "" && failedToken != "" && currentToken != failedToken { - prepared, errPrepare := m.prepareHomeAuthSnapshot(ctx, selection.Executor, current) - if errPrepare != nil { - return current, false, errPrepare - } - preserveHomeRoutingAttributes(prepared, current) - m.replaceHomeSelectionAuth(selection, prepared) - return selection.CloneAuth(), true, nil + return current, true, nil } } - refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(ctx, selection.Executor, failedAuth, &Error{HTTPStatus: http.StatusUnauthorized, Message: "upstream unauthorized"}, false, true) - if errRefresh != nil || !okRefresh { - return current, false, errRefresh - } - m.replaceHomeSelectionAuth(selection, refreshed) - updated := selection.CloneAuth() - if updated == nil { - return nil, false, &Error{Code: "auth_not_found", Message: "refreshed Home auth is unavailable", HTTPStatus: http.StatusServiceUnavailable} - } - return updated, true, nil + return current, false, nil } // tryRefreshAfterUnauthorized refreshes local OAuth credentials once after a diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index cca3c3da..b2a4f72e 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -197,24 +197,13 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out} } -func (m *Manager) replaceHomeExecutionLifecycleAuth(lifecycle cliproxyexecutor.ExecutionLifecycle, auth *Auth) { - selection, ok := lifecycle.(*HomeDispatchSelection) - if !ok || selection == nil { - return - } - m.replaceHomeSelectionAuth(selection, auth) -} - -func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel, executionModel string, execModels []string, pooled bool, aliasResult OAuthModelAliasResult, routing *apiKeyModelRoutingSnapshot, allowRetry bool, ephemeralResult bool, unauthorizedRefreshTried map[string]struct{}) (*cliproxyexecutor.StreamResult, error) { +func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel, executionModel string, execModels []string, pooled bool, aliasResult OAuthModelAliasResult, routing *apiKeyModelRoutingSnapshot, allowRetry bool, ephemeralResult bool) (*cliproxyexecutor.StreamResult, error) { if executor == nil { return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} } ctx = contextWithRequestedModelAlias(ctx, opts, routeModel) var lastErr error didRefreshOnUnauthorized := false - if auth != nil && unauthorizedRefreshTried != nil { - _, didRefreshOnUnauthorized = unauthorizedRefreshTried[auth.ID] - } for idx, execModel := range execModels { ctx = newUpstreamAttemptContext(ctx) resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled) @@ -243,23 +232,11 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if errCtx := ctx.Err(); errCtx != nil { return nil, errCtx } - if allowRetry { + if allowRetry && !ephemeralResult { alreadyTried := didRefreshOnUnauthorized - willAttemptHomeRefresh := ephemeralResult && !alreadyTried && auth != nil && auth.AuthKind() == AuthKindOAuth && isUnauthorizedError(errStream) - refreshCtx := newUpstreamAttemptContext(ctx) - refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(refreshCtx, executor, auth, errStream, alreadyTried, ephemeralResult) - if willAttemptHomeRefresh { - didRefreshOnUnauthorized = true - if unauthorizedRefreshTried != nil { - unauthorizedRefreshTried[auth.ID] = struct{}{} - } - } - if errRefresh != nil { - errStream = errRefresh - warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, durationStream, errStream) - } else if okRefresh { + refreshed, okRefresh := m.tryRefreshAfterUnauthorized(newUpstreamAttemptContext(ctx), auth, errStream, alreadyTried) + if okRefresh { auth = refreshed - m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true ctx = newUpstreamAttemptContext(ctx) @@ -321,26 +298,12 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi discardStreamChunks(streamResult.Chunks) return nil, errCtx } - if allowRetry { + if allowRetry && !ephemeralResult { alreadyTried := didRefreshOnUnauthorized - willAttemptHomeRefresh := ephemeralResult && !alreadyTried && auth != nil && auth.AuthKind() == AuthKindOAuth && isUnauthorizedError(bootstrapErr) - refreshCtx := newUpstreamAttemptContext(ctx) - refreshed, okRefresh, errRefresh := m.tryRefreshExecutionAuthAfterUnauthorized(refreshCtx, executor, auth, bootstrapErr, alreadyTried, ephemeralResult) - if willAttemptHomeRefresh { - didRefreshOnUnauthorized = true - if unauthorizedRefreshTried != nil { - unauthorizedRefreshTried[auth.ID] = struct{}{} - } - } - if errRefresh != nil { - discardStreamChunks(streamResult.Chunks) - bootstrapErr = errRefresh - warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, time.Since(startStream), bootstrapErr) - streamResult = &cliproxyexecutor.StreamResult{} - } else if okRefresh { + refreshed, okRefresh := m.tryRefreshAfterUnauthorized(newUpstreamAttemptContext(ctx), auth, bootstrapErr, alreadyTried) + if okRefresh { discardStreamChunks(streamResult.Chunks) auth = refreshed - m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true ctx = newUpstreamAttemptContext(ctx) diff --git a/sdk/cliproxy/auth/conductor_warn_logging_test.go b/sdk/cliproxy/auth/conductor_warn_logging_test.go index 46d85de2..50ab5cc1 100644 --- a/sdk/cliproxy/auth/conductor_warn_logging_test.go +++ b/sdk/cliproxy/auth/conductor_warn_logging_test.go @@ -33,6 +33,76 @@ func setupTestLoggerHook(t *testing.T) *logtest.Hook { return hook } +type statusErrorLogTestError struct { + message string + statusCode int +} + +func (e statusErrorLogTestError) Error() string { return e.message } + +func (e statusErrorLogTestError) StatusCode() int { return e.statusCode } + +func TestWarnLogUpstreamFailureIncludesStructuredStatusCodeAndSafeDiagnostic(t *testing.T) { + hook := setupTestLoggerHook(t) + diagnostic := ` antigravity refresh: upstream request failed with status 400 error="invalid_request" error_description="Malformed request, retry & inspect" ` + + warnLogUpstreamFailure( + context.Background(), + nil, + "antigravity", + "gemini-3.7-flash-high", + &Auth{ID: "auth-1", Provider: "antigravity"}, + 83*time.Millisecond, + statusErrorLogTestError{message: diagnostic, statusCode: http.StatusServiceUnavailable}, + ) + + for _, entry := range hook.AllEntries() { + if entry.Level == log.WarnLevel && strings.Contains(entry.Message, "upstream execution failed") { + if !strings.HasPrefix(entry.Message, "503 | 83ms | upstream execution failed: provider=antigravity") || strings.Contains(entry.Message, "status=503") || !strings.HasSuffix(entry.Message, "err="+strings.TrimSpace(diagnostic)) { + t.Fatalf("unexpected Warn log content: %s", entry.Message) + } + return + } + } + t.Fatalf("expected upstream failure Warn log, got logs: %#v", hook.AllEntries()) +} + +func TestHomeCredentialBoundaryWarnLogUsesSafeDiagnostic(t *testing.T) { + diagnostic := "antigravity refresh: access token expired access_token=provider-secret\nforged log line" + upstreamErr := statusErrorLogTestError{ + message: diagnostic, + statusCode: http.StatusServiceUnavailable, + } + hook := setupTestLoggerHook(t) + auth := &Auth{ID: "auth-1", Provider: "antigravity"} + executor := &requestPrepareExecutor{prepareErr: upstreamErr} + selection := &HomeDispatchSelection{Auth: auth, Executor: executor, Provider: "antigravity"} + _, errPrepare := NewManager(nil, nil, nil).prepareHomeRequestAuth(context.Background(), executor, selection) + if errPrepare == nil || errPrepare.Error() != diagnostic { + t.Fatalf("operation error = %v, want original error %q", errPrepare, diagnostic) + } + + matches := 0 + for _, entry := range hook.AllEntries() { + if entry.Level != log.WarnLevel || !strings.HasPrefix(entry.Message, "Home credential operation failed: err=") { + continue + } + matches++ + if !strings.Contains(entry.Message, "access token expired") { + t.Fatalf("Warn log lost access-token-expired diagnostic: %q", entry.Message) + } + if strings.Contains(entry.Message, "provider-secret") || strings.ContainsAny(entry.Message, "\r\n") { + t.Fatalf("Warn log included unsafe provider detail: %q", entry.Message) + } + if entry.Data["operation"] != "request_auth_preparation" || entry.Data["provider"] != "antigravity" || entry.Data["status"] != http.StatusServiceUnavailable { + t.Fatalf("Warn log fields = %#v, want operation=request_auth_preparation provider=antigravity status=503", entry.Data) + } + } + if matches != 1 { + t.Fatalf("matching Warn logs = %d, want 1; logs=%#v", matches, hook.AllEntries()) + } +} + func TestWarnLogOnAuthUnavailable_SingleProvider(t *testing.T) { previousCooldown := quotaCooldownDisabled.Load() quotaCooldownDisabled.Store(false) diff --git a/sdk/cliproxy/auth/home_result.go b/sdk/cliproxy/auth/home_result.go index 3a9a636f..fad114ac 100644 --- a/sdk/cliproxy/auth/home_result.go +++ b/sdk/cliproxy/auth/home_result.go @@ -13,11 +13,15 @@ const homeResultExecutorType = "home-result" // ReportHomeUnauthorized publishes a result-only zero-token usage record for an // upstream 401 attempt that did not pass through an executor UsageReporter. -func (m *Manager) ReportHomeUnauthorized(ctx context.Context, auth *Auth, provider, model string) { - m.reportHomeUnauthorized(ctx, auth, provider, model, AccessTokenSHA256(auth)) +func (m *Manager) ReportHomeUnauthorized(ctx context.Context, auth *Auth, provider, model string, upstreamBody ...[]byte) { + failureBody := "" + if len(upstreamBody) > 0 { + failureBody = string(upstreamBody[0]) + } + m.reportHomeUnauthorized(ctx, auth, provider, model, AccessTokenSHA256(auth), failureBody) } -func (m *Manager) reportHomeUnauthorized(ctx context.Context, auth *Auth, provider, model, accessTokenSHA256 string) { +func (m *Manager) reportHomeUnauthorized(ctx context.Context, auth *Auth, provider, model, accessTokenSHA256, failureBody string) { if m == nil || auth == nil { return } @@ -29,6 +33,9 @@ func (m *Manager) reportHomeUnauthorized(ctx context.Context, auth *Auth, provid if authIndex == "" || accessTokenSHA256 == "" { return } + if failureBody == "" { + failureBody = "upstream unauthorized" + } provider = strings.TrimSpace(provider) if provider == "" { provider = strings.TrimSpace(auth.Provider) @@ -55,7 +62,7 @@ func (m *Manager) reportHomeUnauthorized(ctx context.Context, auth *Auth, provid Failed: true, Fail: coreusage.Failure{ StatusCode: http.StatusUnauthorized, - Body: "upstream unauthorized", + Body: failureBody, }, }) } diff --git a/sdk/cliproxy/auth/home_retry_contract_test.go b/sdk/cliproxy/auth/home_retry_contract_test.go index 5f62c652..819f9cdc 100644 --- a/sdk/cliproxy/auth/home_retry_contract_test.go +++ b/sdk/cliproxy/auth/home_retry_contract_test.go @@ -971,7 +971,7 @@ func TestHomeRetryRoundUsesRemoteCooldownWhenAttemptedErrorHasNoTiming(t *testin } } -func TestHomeStreamOAuthUnauthorizedRotatesAfterRefreshRetry(t *testing.T) { +func TestHomeStreamOAuthUnauthorizedRotatesWithoutRefreshRetry(t *testing.T) { dispatcher := &retryContractHomeDispatcher{ authIDs: []string{"home-retry-a", "home-retry-b"}, metadata: map[string]any{ @@ -995,8 +995,8 @@ func TestHomeStreamOAuthUnauthorizedRotatesAfterRefreshRetry(t *testing.T) { } for range result.Chunks { } - if got := executor.Calls(); len(got) != 3 || got[0] != "home-retry-a" || got[1] != "home-retry-a" || got[2] != "home-retry-b" { - t.Fatalf("executor calls = %v, want [home-retry-a home-retry-a home-retry-b]", got) + if got := executor.Calls(); len(got) != 2 || got[0] != "home-retry-a" || got[1] != "home-retry-b" { + t.Fatalf("executor calls = %v, want [home-retry-a home-retry-b]", got) } excluded := dispatcher.Excluded() if len(excluded) != 2 || len(excluded[0]) != 0 || len(excluded[1]) != 1 || excluded[1][0] != "home-retry-a" { diff --git a/sdk/cliproxy/auth/home_unauthorized_refresh_test.go b/sdk/cliproxy/auth/home_unauthorized_refresh_test.go index 80d8f7f9..00f8778c 100644 --- a/sdk/cliproxy/auth/home_unauthorized_refresh_test.go +++ b/sdk/cliproxy/auth/home_unauthorized_refresh_test.go @@ -59,7 +59,7 @@ func (e *homeUnauthorizedRefreshExecutor) Execute(_ context.Context, auth *Auth, } } if authAccessToken(auth) == "stale-access-token" { - return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"} + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "access token expired"} } return cliproxyexecutor.Response{Payload: []byte("ok")}, nil } @@ -70,17 +70,17 @@ func (e *homeUnauthorizedRefreshExecutor) ExecuteStream(_ context.Context, auth switch e.streamMode { case "bootstrap": chunks := make(chan cliproxyexecutor.StreamChunk, 1) - chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}} + chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusUnauthorized, Message: "access token expired"}} close(chunks) return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil case "started": chunks := make(chan cliproxyexecutor.StreamChunk, 2) chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("started")} - chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}} + chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusUnauthorized, Message: "access token expired"}} close(chunks) return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil default: - return nil, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"} + return nil, &Error{HTTPStatus: http.StatusUnauthorized, Message: "access token expired"} } } chunks := make(chan cliproxyexecutor.StreamChunk, 1) @@ -108,7 +108,7 @@ func (e *homeUnauthorizedRefreshExecutor) Refresh(_ context.Context, auth *Auth) func (e *homeUnauthorizedRefreshExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { e.countCalls.Add(1) if authAccessToken(auth) == "stale-access-token" { - return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"} + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "access token expired"} } return cliproxyexecutor.Response{Payload: []byte("ok")}, nil } @@ -125,7 +125,7 @@ func newHomeUnauthorizedRefreshManager(dispatcher *homeUnauthorizedRefreshDispat return manager } -func TestHomeUnauthorizedRefreshesSameSelectionBeforeRedispatch(t *testing.T) { +func TestHomeUnauthorizedReturnsOriginalErrorWithoutRefresh(t *testing.T) { for _, test := range []struct { name string run func(*Manager) error @@ -150,26 +150,24 @@ func TestHomeUnauthorizedRefreshesSameSelectionBeforeRedispatch(t *testing.T) { executor := &homeUnauthorizedRefreshExecutor{} manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) - if errRun := test.run(manager); errRun != nil { - t.Fatalf("execution error = %v", errRun) + errRun := test.run(manager) + if errRun == nil || errRun.Error() != "access token expired" || statusCodeFromError(errRun) != http.StatusUnauthorized { + t.Fatalf("execution error = %v, want original 401", errRun) } - if got := dispatcher.calls.Load(); got != 1 { - t.Fatalf("Home dispatch calls = %d, want 1", got) + if got := executor.refreshCalls.Load(); got != 0 { + t.Fatalf("refresh calls = %d, want 0", got) } - if got := executor.refreshCalls.Load(); got != 1 { - t.Fatalf("refresh calls = %d, want 1", got) + if test.name == "execute" && executor.executeCalls.Load() != 1 { + t.Fatalf("execute calls = %d, want 1", executor.executeCalls.Load()) } - if test.name == "execute" && executor.executeCalls.Load() != 2 { - t.Fatalf("execute calls = %d, want 2", executor.executeCalls.Load()) - } - if test.name == "count_tokens" && executor.countCalls.Load() != 2 { - t.Fatalf("count calls = %d, want 2", executor.countCalls.Load()) + if test.name == "count_tokens" && executor.countCalls.Load() != 1 { + t.Fatalf("count calls = %d, want 1", executor.countCalls.Load()) } }) } } -func TestHomeUnauthorizedRefreshUpdatesRetainedSelection(t *testing.T) { +func TestHomeUnauthorizedDoesNotRefreshRetainedSelection(t *testing.T) { dispatcher := &homeUnauthorizedRefreshDispatcher{} executor := &homeUnauthorizedRefreshExecutor{retainSelection: true} manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) @@ -179,19 +177,15 @@ func TestHomeUnauthorizedRefreshUpdatesRetainedSelection(t *testing.T) { cliproxyexecutor.PinnedAuthMetadataKey: "home-refresh-auth", }} - for range 2 { - if _, errExecute := manager.Execute(ctx, []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil { - t.Fatalf("Execute() error = %v", errExecute) - } - } - if got := dispatcher.calls.Load(); got != 1 { - t.Fatalf("Home dispatch calls = %d, want one retained selection", got) + _, errExecute := manager.Execute(ctx, []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, opts) + if errExecute == nil || errExecute.Error() != "access token expired" { + t.Fatalf("Execute() error = %v, want original upstream error", errExecute) } - if got := executor.refreshCalls.Load(); got != 1 { - t.Fatalf("refresh calls = %d, want refreshed token reused by retained selection", got) + if got := executor.refreshCalls.Load(); got != 0 { + t.Fatalf("refresh calls = %d, want 0", got) } - if got := executor.executeCalls.Load(); got != 3 { - t.Fatalf("execute calls = %d, want stale attempt, retry, and retained reuse", got) + if got := executor.executeCalls.Load(); got != 1 { + t.Fatalf("execute calls = %d, want 1", got) } } @@ -214,7 +208,7 @@ func TestRefreshHomeSelectionReusesConcurrentNewerToken(t *testing.T) { } } -func TestHomeUnauthorizedRefreshIsAttemptedAtMostOnce(t *testing.T) { +func TestHomeUnauthorizedDoesNotRefreshOrReplay(t *testing.T) { dispatcher := &homeUnauthorizedRefreshDispatcher{} executor := &homeUnauthorizedRefreshExecutor{keepStale: true} manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) @@ -223,23 +217,23 @@ func TestHomeUnauthorizedRefreshIsAttemptedAtMostOnce(t *testing.T) { if statusCodeFromError(errExecute) != http.StatusUnauthorized { t.Fatalf("Execute() error = %v, want original 401", errExecute) } - if got := executor.refreshCalls.Load(); got != 1 { - t.Fatalf("refresh calls = %d, want exactly 1", got) + if got := executor.refreshCalls.Load(); got != 0 { + t.Fatalf("refresh calls = %d, want 0", got) } - if got := executor.executeCalls.Load(); got != 2 { - t.Fatalf("execute calls = %d, want initial attempt and one retry", got) + if got := executor.executeCalls.Load(); got != 1 { + t.Fatalf("execute calls = %d, want 1", got) } } -func TestHomeNoCandidateAfterRefreshFailurePreservesRefreshError(t *testing.T) { - refreshErr := &Error{Code: "refresh_temporarily_unavailable", HTTPStatus: http.StatusServiceUnavailable, Message: "refresh unavailable"} +func TestHomeNoCandidatePreservesOriginalUpstreamError(t *testing.T) { + upstreamErr := &Error{HTTPStatus: http.StatusUnauthorized, Message: "access token expired"} noCandidate := &Error{Code: "auth_not_found", HTTPStatus: http.StatusServiceUnavailable, Message: "no auth available"} - if !shouldReturnLastErrorOnPickFailure(true, refreshErr, noCandidate) { - t.Fatal("Home no-candidate error would overwrite the original refresh error") + if !shouldReturnLastErrorOnPickFailure(true, upstreamErr, noCandidate) { + t.Fatal("Home no-candidate error would overwrite the original upstream error") } } -func TestHomeUnauthorizedTransientRefreshFailureIsReturned(t *testing.T) { +func TestHomeUnauthorizedIgnoresExecutorRefreshFailure(t *testing.T) { dispatcher := &homeUnauthorizedRefreshDispatcher{} executor := &homeUnauthorizedRefreshExecutor{ refreshErr: &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "Home refresh temporarily unavailable"}, @@ -247,18 +241,18 @@ func TestHomeUnauthorizedTransientRefreshFailureIsReturned(t *testing.T) { manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) _, errExecute := manager.Execute(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{}) - if statusCodeFromError(errExecute) != http.StatusServiceUnavailable { - t.Fatalf("Execute() error = %v, want transient 503", errExecute) + if statusCodeFromError(errExecute) != http.StatusUnauthorized || errExecute.Error() != "access token expired" { + t.Fatalf("Execute() error = %v, want original 401", errExecute) } if got := executor.executeCalls.Load(); got != 1 { t.Fatalf("execute calls = %d, want 1", got) } - if got := executor.refreshCalls.Load(); got != 1 { - t.Fatalf("refresh calls = %d, want 1", got) + if got := executor.refreshCalls.Load(); got != 0 { + t.Fatalf("refresh calls = %d, want 0", got) } } -func TestHomeUnauthorizedStreamRefreshesAtMostOnceAcrossRedispatch(t *testing.T) { +func TestHomeUnauthorizedStreamDoesNotRefreshOrReplay(t *testing.T) { dispatcher := &homeUnauthorizedRefreshDispatcher{} executor := &homeUnauthorizedRefreshExecutor{keepStale: true} manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) @@ -267,11 +261,11 @@ func TestHomeUnauthorizedStreamRefreshesAtMostOnceAcrossRedispatch(t *testing.T) if statusCodeFromError(errStream) != http.StatusUnauthorized { t.Fatalf("ExecuteStream() error = %v, want original 401", errStream) } - if got := executor.refreshCalls.Load(); got != 1 { - t.Fatalf("refresh calls = %d, want exactly 1", got) + if got := executor.refreshCalls.Load(); got != 0 { + t.Fatalf("refresh calls = %d, want 0", got) } - if got := executor.streamCalls.Load(); got != 2 { - t.Fatalf("stream calls = %d, want initial attempt and one retry", got) + if got := executor.streamCalls.Load(); got != 1 { + t.Fatalf("stream calls = %d, want 1", got) } } @@ -305,7 +299,7 @@ func TestHomeUnauthorizedStartedStreamDoesNotReplay(t *testing.T) { } } -func TestHomeUnauthorizedStreamRefreshesBeforeRedispatch(t *testing.T) { +func TestHomeUnauthorizedStreamReturnsOriginalErrorWithoutRefresh(t *testing.T) { for _, mode := range []string{"synchronous", "bootstrap"} { t.Run(mode, func(t *testing.T) { dispatcher := &homeUnauthorizedRefreshDispatcher{} @@ -313,27 +307,24 @@ func TestHomeUnauthorizedStreamRefreshesBeforeRedispatch(t *testing.T) { manager := newHomeUnauthorizedRefreshManager(dispatcher, executor) result, errStream := manager.ExecuteStream(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}) - if errStream != nil { - t.Fatalf("ExecuteStream() error = %v", errStream) - } - var payload string - for chunk := range result.Chunks { - if chunk.Err != nil { - t.Fatalf("stream chunk error = %v", chunk.Err) + if mode == "synchronous" { + if errStream == nil || errStream.Error() != "access token expired" || statusCodeFromError(errStream) != http.StatusUnauthorized { + t.Fatalf("ExecuteStream() error = %v, want original 401", errStream) + } + } else { + if errStream != nil { + t.Fatalf("ExecuteStream() error = %v", errStream) + } + chunk, ok := <-result.Chunks + if !ok || chunk.Err == nil || chunk.Err.Error() != "access token expired" || statusCodeFromError(chunk.Err) != http.StatusUnauthorized { + t.Fatalf("stream chunk = %#v, open=%v; want original 401", chunk, ok) } - payload += string(chunk.Payload) - } - if payload != "ok" { - t.Fatalf("stream payload = %q, want ok", payload) - } - if got := dispatcher.calls.Load(); got != 1 { - t.Fatalf("Home dispatch calls = %d, want 1", got) } - if got := executor.refreshCalls.Load(); got != 1 { - t.Fatalf("refresh calls = %d, want 1", got) + if got := executor.refreshCalls.Load(); got != 0 { + t.Fatalf("refresh calls = %d, want 0", got) } - if got := executor.streamCalls.Load(); got != 2 { - t.Fatalf("stream calls = %d, want 2", got) + if got := executor.streamCalls.Load(); got != 1 { + t.Fatalf("stream calls = %d, want 1", got) } }) } -- 2.51.2 From e4a8f9891344d5762aeb5a0c5d4788a951146c0f Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Fri, 28 Aug 2026 09:59:24 +0800 Subject: [PATCH 2/4] fix(logging): enhance error diagnostics and logging for home refresh operations --- internal/logging/diagnostic.go | 94 ++++++++++++++++++- internal/logging/diagnostic_test.go | 56 ++++++++++- .../runtime/executor/helps/home_refresh.go | 48 +++++++--- .../executor/helps/home_refresh_test.go | 56 ++++++++++- .../runtime/executor/helps/logging_helpers.go | 31 ++++-- .../executor/helps/logging_helpers_test.go | 72 ++++++++++++++ sdk/cliproxy/auth/conductor_execution.go | 21 ++++- .../auth/conductor_warn_logging_test.go | 54 ++++++++++- 8 files changed, 401 insertions(+), 31 deletions(-) diff --git a/internal/logging/diagnostic.go b/internal/logging/diagnostic.go index a55697e5..80928084 100644 --- a/internal/logging/diagnostic.go +++ b/internal/logging/diagnostic.go @@ -1,6 +1,11 @@ package logging import ( + "context" + "errors" + "fmt" + "io" + "net" "regexp" "strings" "unicode/utf8" @@ -15,7 +20,8 @@ var ( accessTokenExpiredLogPattern = regexp.MustCompile(`(?i)access token expired`) sensitiveLogAssignmentPattern = regexp.MustCompile(`(?i)(["']?(?:access[\s_-]*token|refresh[\s_-]*token|id[\s_-]*token|api[\s_-]*key|client[\s_-]*secret|private[\s_-]*key|proxy[\s_-]*authorization|authorization|password|credential|token|secret)["']?\s*[:=]\s*)(?:(?:bearer|basic)\s+[^\s,;]+|"(?:\\.|[^"])*"|'(?:\\.|[^'])*'|[^\s,;&}\]]+)`) authorizationLogPattern = regexp.MustCompile(`(?i)\b(bearer|basic)\s+[^\s,;]+`) - urlUserinfoLogPattern = regexp.MustCompile(`(?i)(https?://)[^/\s@]+@`) + urlUserinfoLogPattern = regexp.MustCompile(`(?i)([a-z][a-z0-9+.-]*://)[^/\s@]+@`) + diagnosticStatusPattern = regexp.MustCompile(`(?i)\bstatus(?:\s+code)?\s*[:=]?\s*([1-5][0-9]{2})\b`) ) // SafeDiagnosticForLog returns a bounded, single-line diagnostic suitable for @@ -40,6 +46,92 @@ func SafeDiagnosticForLog(message string) string { return truncateDiagnosticLogExcerpt(excerpt, sourceTruncated) } +// SafeErrorDiagnostic extracts only allowlisted failure signals from an +// arbitrary error. It never carries the original free-form error text. +func SafeErrorDiagnostic(err error) string { + if err == nil { + return "" + } + parts := make([]string, 0, 4) + appendPart := func(part string) { + if part == "" { + return + } + for _, existing := range parts { + if existing == part { + return + } + } + parts = append(parts, part) + } + + switch { + case errors.Is(err, io.ErrUnexpectedEOF): + appendPart("unexpected_EOF") + case errors.Is(err, io.EOF): + appendPart("EOF") + } + if errors.Is(err, context.DeadlineExceeded) { + appendPart("timeout") + } + if errors.Is(err, context.Canceled) { + appendPart("canceled") + } + var netErr net.Error + if errors.As(err, &netErr) && netErr != nil && netErr.Timeout() { + appendPart("timeout") + } + + rawOriginal := err.Error() + raw := strings.ToLower(rawOriginal) + if strings.EqualFold(strings.TrimSpace(rawOriginal), "EOF") { + appendPart("EOF") + } + if strings.Contains(raw, "socks") && (strings.Contains(raw, "authentication failed") || strings.Contains(raw, "authentication required")) { + appendPart("proxy_authentication_failed") + } + for _, signal := range []struct { + needle string + label string + }{ + {needle: "socks", label: "proxy=socks"}, + {needle: "proxyconnect", label: "proxy_connect_failed"}, + {needle: "proxy connect", label: "proxy_connect_failed"}, + {needle: "dial ", label: "dial_failed"}, + {needle: "dial failed", label: "dial_failed"}, + {needle: "connection refused", label: "connection_refused"}, + {needle: "connection reset", label: "connection_reset"}, + {needle: "connection aborted", label: "connection_aborted"}, + {needle: "stream reset", label: "stream_reset"}, + {needle: "network is unreachable", label: "network_unreachable"}, + {needle: "no route to host", label: "network_unreachable"}, + {needle: "no such host", label: "dns_not_found"}, + {needle: "server misbehaving", label: "dns_failure"}, + {needle: "tls handshake timeout", label: "tls_handshake_timeout"}, + {needle: "i/o timeout", label: "timeout"}, + {needle: "deadline exceeded", label: "timeout"}, + {needle: "unexpected eof", label: "unexpected_EOF"}, + {needle: "certificate", label: "tls_certificate_error"}, + {needle: "invalid character", label: "invalid_response_json"}, + {needle: "cannot unmarshal", label: "invalid_response_json"}, + {needle: "invalid_grant", label: "oauth_error=invalid_grant"}, + {needle: "refresh_token_expired", label: "oauth_error=refresh_token_expired"}, + {needle: "refresh_token_revoked", label: "oauth_error=refresh_token_revoked"}, + {needle: "refresh_token_reused", label: "oauth_error=refresh_token_reused"}, + } { + if strings.Contains(raw, signal.needle) { + appendPart(signal.label) + } + } + if match := diagnosticStatusPattern.FindStringSubmatch(rawOriginal); len(match) == 2 { + appendPart("status=" + match[1]) + } + if len(parts) == 0 { + appendPart(fmt.Sprintf("error_type=%T", err)) + } + return strings.Join(parts, " ") +} + func diagnosticRunePrefix(value string, limit int) (string, bool) { if limit <= 0 { return "", value != "" diff --git a/internal/logging/diagnostic_test.go b/internal/logging/diagnostic_test.go index dc200f95..9021f493 100644 --- a/internal/logging/diagnostic_test.go +++ b/internal/logging/diagnostic_test.go @@ -1,6 +1,9 @@ package logging import ( + "errors" + "io" + "net/url" "strings" "testing" ) @@ -8,13 +11,13 @@ import ( func TestSafeDiagnosticForLogPreservesAccessTokenExpiredAndRedactsCredentials(t *testing.T) { diagnostic := "access token expired\n" + `access_token=access-secret refresh token: refresh-secret Authorization=Bearer bearer-secret ` + - `Post "https://user:password@oauth.example/token?access_token=query-secret"` + `Post "https://user:password@oauth.example/token?access_token=query-secret" via socks5://proxy-user:proxy-password@127.0.0.1:1080` got := SafeDiagnosticForLog(diagnostic) if !strings.Contains(got, "access token expired") { t.Fatalf("safe diagnostic lost access-token-expired signal: %q", got) } - for _, secret := range []string{"access-secret", "refresh-secret", "bearer-secret", "query-secret", "user:password"} { + for _, secret := range []string{"access-secret", "refresh-secret", "bearer-secret", "query-secret", "user:password", "proxy-user", "proxy-password"} { if strings.Contains(got, secret) { t.Fatalf("safe diagnostic leaked %q: %q", secret, got) } @@ -57,3 +60,52 @@ func TestSafeDiagnosticForLogBoundsLargeGenericMessage(t *testing.T) { t.Fatalf("safe generic diagnostic length = %d, want %d with ellipsis", len([]rune(got)), diagnosticLogRuneLimit+3) } } + +func TestSafeErrorDiagnosticExtractsOnlyAllowlistedSignals(t *testing.T) { + tests := []struct { + name string + err error + wantParts []string + }{ + {name: "EOF", err: io.EOF, wantParts: []string{"EOF"}}, + {name: "SOCKS refused", err: errors.New("socks connect with unlabeled-secret: connection refused"), wantParts: []string{"proxy=socks", "connection_refused"}}, + {name: "OAuth response", err: errors.New(`upstream status 400 error="invalid_request" request_id="req-123" unlabeled-secret`), wantParts: []string{"status=400"}}, + {name: "unknown", err: errors.New("unlabeled-secret")}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := SafeErrorDiagnostic(tt.err) + for _, want := range tt.wantParts { + if !strings.Contains(got, want) { + t.Fatalf("SafeErrorDiagnostic() = %q, want %q", got, want) + } + } + if strings.Contains(got, "unlabeled-secret") { + t.Fatalf("SafeErrorDiagnostic() leaked arbitrary detail: %q", got) + } + }) + } +} + +func TestSafeErrorDiagnosticDoesNotExtractURLQueryValues(t *testing.T) { + err := &url.Error{ + Op: "Post", + URL: "https://oauth.example/token?code=oauth-secret&error=error-secret&request_id=request-secret", + Err: io.EOF, + } + + got := SafeErrorDiagnostic(err) + if !strings.Contains(got, "EOF") { + t.Fatalf("SafeErrorDiagnostic() = %q, want EOF signal", got) + } + for _, secret := range []string{"oauth-secret", "error-secret", "request-secret"} { + if strings.Contains(got, secret) { + t.Fatalf("SafeErrorDiagnostic() leaked %q: %q", secret, got) + } + } + for _, dynamicField := range []string{"oauth_error=", "request_id="} { + if strings.Contains(got, dynamicField) { + t.Fatalf("SafeErrorDiagnostic() extracted dynamic field %q: %q", dynamicField, got) + } + } +} diff --git a/internal/runtime/executor/helps/home_refresh.go b/internal/runtime/executor/helps/home_refresh.go index 1a4abfff..486279fc 100644 --- a/internal/runtime/executor/helps/home_refresh.go +++ b/internal/runtime/executor/helps/home_refresh.go @@ -10,13 +10,15 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" ) type homeStatusErr struct { - code int - msg string - upstream bool + code int + msg string + diagnostic string + upstream bool } func (e homeStatusErr) Error() string { @@ -31,6 +33,16 @@ func (e homeStatusErr) Error() string { func (e homeStatusErr) StatusCode() int { return e.code } +func (e homeStatusErr) LogDiagnostic() string { + if strings.TrimSpace(e.diagnostic) != "" { + return logging.SafeDiagnosticForLog(e.diagnostic) + } + if e.upstream { + return fmt.Sprintf("Home refresh upstream response: status=%d", e.code) + } + return logging.SafeDiagnosticForLog(e.Error()) +} + func (e homeStatusErr) DirectResponse() bool { return e.upstream } func (e homeStatusErr) ResponseBody() []byte { @@ -50,10 +62,11 @@ type homeRefreshAuthEnvelope struct { } type homeErrorDetail struct { - Type string `json:"type"` - Message string `json:"message"` - Code string `json:"code,omitempty"` - Upstream *homeUpstreamResponse `json:"upstream,omitempty"` + Type string `json:"type"` + Message string `json:"message"` + Code string `json:"code,omitempty"` + Diagnostic string `json:"diagnostic,omitempty"` + Upstream *homeUpstreamResponse `json:"upstream,omitempty"` } type homeUpstreamResponse struct { @@ -102,16 +115,21 @@ func RefreshAuthViaHome(ctx context.Context, cfg *config.Config, auth *cliproxya if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return nil, true, err } - return nil, true, homeStatusErr{code: http.StatusServiceUnavailable, msg: "home refresh temporarily unavailable"} + return nil, true, homeStatusErr{ + code: http.StatusServiceUnavailable, + msg: "home refresh temporarily unavailable", + diagnostic: "Home refresh transport failed: " + logging.SafeErrorDiagnostic(err), + } } var env homeErrorEnvelope if errUnmarshal := json.Unmarshal(raw, &env); errUnmarshal == nil && env.Error != nil { if env.Error.Upstream != nil { return nil, true, homeStatusErr{ - code: env.Error.Upstream.Status, - msg: string(env.Error.Upstream.Body), - upstream: true, + code: env.Error.Upstream.Status, + msg: string(env.Error.Upstream.Body), + diagnostic: env.Error.Diagnostic, + upstream: true, } } code := strings.TrimSpace(env.Error.Type) @@ -126,12 +144,16 @@ func RefreshAuthViaHome(ctx context.Context, cfg *config.Config, auth *cliproxya case http.StatusNotFound: message = "credential refresh target not found" } - return nil, true, homeStatusErr{code: statusCode, msg: message} + return nil, true, homeStatusErr{code: statusCode, msg: message, diagnostic: env.Error.Diagnostic} } updated, returnedIndex, errParse := parseHomeRefreshAuth(raw) if errParse != nil { - return nil, true, homeStatusErr{code: http.StatusBadGateway, msg: "home returned invalid auth payload"} + return nil, true, homeStatusErr{ + code: http.StatusBadGateway, + msg: "home returned invalid auth payload", + diagnostic: "Home refresh response decode failed: " + logging.SafeErrorDiagnostic(errParse), + } } if updated.Disabled || updated.Status == cliproxyauth.StatusDisabled { return nil, true, homeStatusErr{code: http.StatusUnauthorized, msg: "credential unauthorized"} diff --git a/internal/runtime/executor/helps/home_refresh_test.go b/internal/runtime/executor/helps/home_refresh_test.go index 00909938..362e7251 100644 --- a/internal/runtime/executor/helps/home_refresh_test.go +++ b/internal/runtime/executor/helps/home_refresh_test.go @@ -61,6 +61,21 @@ func TestRefreshAuthViaHomePreservesContextErrors(t *testing.T) { } } +func TestHomeStatusErrLogDiagnosticSanitizesUpstreamFallback(t *testing.T) { + errRefresh := homeStatusErr{ + code: http.StatusBadGateway, + msg: "upstream EOF access_token=provider-secret", + upstream: true, + } + diagnostic := errRefresh.LogDiagnostic() + if diagnostic != "Home refresh upstream response: status=502" || strings.Contains(diagnostic, "provider-secret") { + t.Fatalf("LogDiagnostic() = %q, want safe upstream fallback", diagnostic) + } + if errRefresh.Error() != "upstream EOF access_token=provider-secret" { + t.Fatalf("Error() = %q, want exact upstream response", errRefresh.Error()) + } +} + func TestRefreshAuthViaHomeMapsTransportFailureToGeneric503(t *testing.T) { client := &fakeHomeRefreshClient{err: errors.New("dial failed with provider-secret")} oldCurrentHomeRefreshClient := currentHomeRefreshClient @@ -84,6 +99,13 @@ func TestRefreshAuthViaHomeMapsTransportFailureToGeneric503(t *testing.T) { if !okDirect || direct.DirectResponse() { t.Fatalf("transport error direct response = %v/%v, want false", okDirect, direct) } + diagnosticErr, okDiagnostic := errRefresh.(interface{ LogDiagnostic() string }) + if !okDiagnostic || !strings.Contains(diagnosticErr.LogDiagnostic(), "dial_failed") { + t.Fatalf("transport log diagnostic = %T/%v, want allowlisted transport cause", errRefresh, errRefresh) + } + if strings.Contains(diagnosticErr.LogDiagnostic(), "provider-secret") { + t.Fatalf("transport log diagnostic leaked provider detail: %q", diagnosticErr.LogDiagnostic()) + } } func TestRefreshAuthViaHomeUsesGenericMessageForLegacyErrorEnvelope(t *testing.T) { @@ -102,6 +124,31 @@ func TestRefreshAuthViaHomeUsesGenericMessageForLegacyErrorEnvelope(t *testing.T if strings.Contains(errRefresh.Error(), "provider-secret") { t.Fatalf("refresh error included legacy Home detail: %v", errRefresh) } + if diagnosticErr, ok := errRefresh.(interface{ LogDiagnostic() string }); !ok || diagnosticErr.LogDiagnostic() != errRefresh.Error() { + t.Fatalf("legacy Home message became trusted diagnostic: %T/%v", errRefresh, errRefresh) + } +} + +func TestRefreshAuthViaHomeUsesDedicatedDiagnosticOnlyForLogs(t *testing.T) { + const diagnostic = "antigravity refresh failed: stage=transport err=EOF" + client := &fakeHomeRefreshClient{raw: []byte(`{"error":{"type":"refresh_temporarily_unavailable","message":"untrusted provider-secret","diagnostic":"` + diagnostic + `"}}`)} + oldCurrentHomeRefreshClient := currentHomeRefreshClient + currentHomeRefreshClient = func() homeRefreshClient { return client } + t.Cleanup(func() { currentHomeRefreshClient = oldCurrentHomeRefreshClient }) + + cfg := &config.Config{Home: config.HomeConfig{Enabled: true}} + auth := &cliproxyauth.Auth{ID: "home-auth", Index: "home-auth", Provider: "antigravity"} + _, handled, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth) + if !handled || errRefresh == nil || errRefresh.Error() != "credential refresh temporarily unavailable" { + t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want generic client error", handled, errRefresh) + } + diagnosticErr, ok := errRefresh.(interface{ LogDiagnostic() string }) + if !ok || diagnosticErr.LogDiagnostic() != diagnostic { + t.Fatalf("log diagnostic = %T/%v, want %q", errRefresh, errRefresh, diagnostic) + } + if strings.Contains(errRefresh.Error(), diagnostic) || strings.Contains(errRefresh.Error(), "provider-secret") { + t.Fatalf("client error exposed internal detail: %v", errRefresh) + } } func TestRefreshAuthViaHomePreservesUpstreamStatusAndBodyExactly(t *testing.T) { @@ -119,8 +166,9 @@ func TestRefreshAuthViaHomePreservesUpstreamStatusAndBodyExactly(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { raw, errMarshal := json.Marshal(homeErrorEnvelope{Error: &homeErrorDetail{ - Type: "refresh_temporarily_unavailable", - Message: "credential refresh temporarily unavailable", + Type: "refresh_temporarily_unavailable", + Message: "credential refresh temporarily unavailable", + Diagnostic: "antigravity refresh failed: stage=upstream_response status=400", Upstream: &homeUpstreamResponse{ Status: tt.status, Body: tt.body, @@ -154,6 +202,10 @@ func TestRefreshAuthViaHomePreservesUpstreamStatusAndBodyExactly(t *testing.T) { if got := direct.ResponseBody(); !bytes.Equal(got, tt.body) { t.Fatalf("direct response body = %q, want exact body %q", got, tt.body) } + diagnosticErr, okDiagnostic := errRefresh.(interface{ LogDiagnostic() string }) + if !okDiagnostic || !strings.Contains(diagnosticErr.LogDiagnostic(), "stage=upstream_response") { + t.Fatalf("upstream log diagnostic = %T/%v", errRefresh, errRefresh) + } }) } } diff --git a/internal/runtime/executor/helps/logging_helpers.go b/internal/runtime/executor/helps/logging_helpers.go index da94bf30..7f9b6789 100644 --- a/internal/runtime/executor/helps/logging_helpers.go +++ b/internal/runtime/executor/helps/logging_helpers.go @@ -54,6 +54,7 @@ type upstreamAttempt struct { bodyHasContent bool prevWasSSEEvent bool errorWritten bool + trailingNewlines int } func requestLogCaptureEnabled(cfg *config.Config) bool { @@ -474,6 +475,17 @@ func ensureResponseIntro(ginCtx *gin.Context, attempt *upstreamAttempt) { if attempt == nil || attempt.response == nil || attempt.responseIntroWritten { return } + attempts := getAttempts(ginCtx) + for i := len(attempts) - 1; i >= 0; i-- { + previousAttempt := attempts[i] + if previousAttempt == nil || previousAttempt == attempt || !previousAttempt.responseIntroWritten { + continue + } + if missingNewlines := 2 - previousAttempt.trailingNewlines; missingNewlines > 0 { + writeAttemptResponse(ginCtx, attempt, []byte(strings.Repeat("\n", missingNewlines))) + } + break + } writeAttemptResponse(ginCtx, attempt, []byte(fmt.Sprintf("=== API RESPONSE %d ===\n", attempt.index))) writeAttemptResponse(ginCtx, attempt, []byte(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))) writeAttemptResponse(ginCtx, attempt, []byte("\n")) @@ -484,6 +496,14 @@ func writeAttemptResponse(ginCtx *gin.Context, attempt *upstreamAttempt, payload if attempt == nil || len(payload) == 0 { return } + trailingNewlines := 0 + for i := len(payload) - 1; i >= 0 && payload[i] == '\n'; i-- { + trailingNewlines++ + } + if trailingNewlines == len(payload) { + trailingNewlines += attempt.trailingNewlines + } + attempt.trailingNewlines = trailingNewlines if attempt.responseSource == nil { attempt.responseSource = apiResponseSourceOrNil(ginCtx) } @@ -527,7 +547,7 @@ func updateAggregatedResponse(ginCtx *gin.Context, attempts []*upstreamAttempt) return } var builder strings.Builder - for idx, attempt := range attempts { + for _, attempt := range attempts { if attempt == nil || attempt.response == nil { continue } @@ -536,12 +556,9 @@ func updateAggregatedResponse(ginCtx *gin.Context, attempts []*upstreamAttempt) continue } builder.WriteString(responseText) - if !strings.HasSuffix(responseText, "\n") { - builder.WriteString("\n") - } - if idx < len(attempts)-1 { - builder.WriteString("\n") - } + } + if responseText := builder.String(); responseText != "" && !strings.HasSuffix(responseText, "\n") { + builder.WriteString("\n") } ginCtx.Set(apiResponseKey, []byte(builder.String())) } diff --git a/internal/runtime/executor/helps/logging_helpers_test.go b/internal/runtime/executor/helps/logging_helpers_test.go index d80e87a4..b6070af3 100644 --- a/internal/runtime/executor/helps/logging_helpers_test.go +++ b/internal/runtime/executor/helps/logging_helpers_test.go @@ -2,6 +2,7 @@ package helps import ( "context" + "errors" "net/http" "net/http/httptest" "strings" @@ -53,3 +54,74 @@ func TestRecordAPIResponseMetadataStoresHeadersWhenRequestLogDisabled(t *testing t.Fatalf("response header = %q, want %q", got.Get("X-Upstream-Request-Id"), "upstream-req-1") } } + +func TestAPIResponseAttemptsAreSeparated(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + fileBacked bool + firstResponseBody []byte + }{ + {name: "memory backed error"}, + {name: "file backed error", fileBacked: true}, + {name: "memory backed partial body", firstResponseBody: []byte("partial")}, + {name: "file backed partial body", fileBacked: true, firstResponseBody: []byte("partial")}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(recorder) + var responseSource *logging.FileBodySource + if tt.fileBacked { + var errSource error + responseSource, errSource = logging.NewFileBodySourceInDir(t.TempDir(), "api-response") + if errSource != nil { + t.Fatalf("NewFileBodySourceInDir: %v", errSource) + } + t.Cleanup(func() { + if errCleanup := responseSource.Cleanup(); errCleanup != nil { + t.Errorf("Cleanup: %v", errCleanup) + } + }) + ginCtx.Set(logging.APIResponseSourceContextKey, responseSource) + } + + ctx := context.WithValue(context.Background(), "gin", ginCtx) + cfg := &config.Config{SDKConfig: config.SDKConfig{RequestLog: true}} + RecordAPIRequest(ctx, cfg, UpstreamRequestLog{URL: "https://api.example.com/first", Method: http.MethodPost}) + if len(tt.firstResponseBody) > 0 { + AppendAPIResponseChunk(ctx, cfg, tt.firstResponseBody) + } else { + RecordAPIResponseError(ctx, cfg, errors.New("EOF")) + } + RecordAPIRequest(ctx, cfg, UpstreamRequestLog{URL: "https://api.example.com/second", Method: http.MethodPost}) + RecordAPIResponseError(ctx, cfg, errors.New("retry failed")) + + var response []byte + if responseSource != nil { + var errBytes error + response, errBytes = responseSource.Bytes() + if errBytes != nil { + t.Fatalf("responseSource.Bytes: %v", errBytes) + } + } else { + value, exists := ginCtx.Get(apiResponseKey) + if !exists { + t.Fatal("API_RESPONSE was not captured") + } + response, _ = value.([]byte) + } + + previousEnd := "Error: EOF" + if len(tt.firstResponseBody) > 0 { + previousEnd = string(tt.firstResponseBody) + } + wantBoundary := previousEnd + "\n\n=== API RESPONSE 2 ===" + if !strings.Contains(string(response), wantBoundary) { + t.Fatalf("API response attempts are not separated by one blank line:\n%q\nwant boundary %q", response, wantBoundary) + } + }) + } +} diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index a968438e..4d821980 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -1588,7 +1588,24 @@ func warnLogHomeCredentialFailure(ctx context.Context, operation, provider strin if statusCode := statusCodeFromError(err); statusCode != 0 { fields["status"] = statusCode } - logEntryWithRequestID(ctx).WithFields(fields).Warnf("Home credential operation failed: err=%s", logging.SafeDiagnosticForLog(err.Error())) + logEntryWithRequestID(ctx).WithFields(fields).Warnf("Home credential operation failed: err=%s", safeErrorDiagnosticForLog(err)) +} + +func safeErrorDiagnosticForLog(err error) string { + if err == nil { + return "" + } + diagnostic := err.Error() + type logDiagnosticError interface { + LogDiagnostic() string + } + var diagnosticErr logDiagnosticError + if errors.As(err, &diagnosticErr) && diagnosticErr != nil { + if markedDiagnostic := strings.TrimSpace(diagnosticErr.LogDiagnostic()); markedDiagnostic != "" { + diagnostic = markedDiagnostic + } + } + return logging.SafeDiagnosticForLog(diagnostic) } func warnLogUpstreamFailure(ctx context.Context, entry *log.Entry, provider, model string, auth *Auth, duration time.Duration, err error) { @@ -1612,7 +1629,7 @@ func warnLogUpstreamFailure(ctx context.Context, entry *log.Entry, provider, mod } } authIdent := formatAuthIdentity(auth, provider) - errSummary := logging.SafeDiagnosticForLog(err.Error()) + errSummary := safeErrorDiagnosticForLog(err) duration = duration.Round(time.Millisecond) if statusCode := statusCodeFromError(err); statusCode != 0 { entry.Warnf("%3d | %13v | upstream execution failed: provider=%s model=%s auth=%s err=%s", statusCode, duration, provider, model, authIdent, errSummary) diff --git a/sdk/cliproxy/auth/conductor_warn_logging_test.go b/sdk/cliproxy/auth/conductor_warn_logging_test.go index 50ab5cc1..a769be4e 100644 --- a/sdk/cliproxy/auth/conductor_warn_logging_test.go +++ b/sdk/cliproxy/auth/conductor_warn_logging_test.go @@ -42,6 +42,18 @@ func (e statusErrorLogTestError) Error() string { return e.message } func (e statusErrorLogTestError) StatusCode() int { return e.statusCode } +type markedDiagnosticStatusError struct { + message string + diagnostic string + statusCode int +} + +func (e markedDiagnosticStatusError) Error() string { return e.message } + +func (e markedDiagnosticStatusError) StatusCode() int { return e.statusCode } + +func (e markedDiagnosticStatusError) LogDiagnostic() string { return e.diagnostic } + func TestWarnLogUpstreamFailureIncludesStructuredStatusCodeAndSafeDiagnostic(t *testing.T) { hook := setupTestLoggerHook(t) diagnostic := ` antigravity refresh: upstream request failed with status 400 error="invalid_request" error_description="Malformed request, retry & inspect" ` @@ -67,10 +79,44 @@ func TestWarnLogUpstreamFailureIncludesStructuredStatusCodeAndSafeDiagnostic(t * t.Fatalf("expected upstream failure Warn log, got logs: %#v", hook.AllEntries()) } +func TestWarnLogUpstreamFailureUsesMarkedHomeDiagnostic(t *testing.T) { + hook := setupTestLoggerHook(t) + errRefresh := markedDiagnosticStatusError{ + message: "credential refresh temporarily unavailable", + diagnostic: "antigravity refresh failed: stage=transport err=EOF access_token=provider-secret", + statusCode: http.StatusServiceUnavailable, + } + + warnLogUpstreamFailure( + context.Background(), + nil, + "antigravity", + "gemini-3.7-flash-high", + &Auth{ID: "auth-1", Provider: "antigravity"}, + 129*time.Millisecond, + errRefresh, + ) + + for _, entry := range hook.AllEntries() { + if entry.Level != log.WarnLevel || !strings.Contains(entry.Message, "upstream execution failed") { + continue + } + if !strings.Contains(entry.Message, "stage=transport") || !strings.Contains(entry.Message, "err=EOF") { + t.Fatalf("Warn log lost marked Home diagnostic: %q", entry.Message) + } + if strings.Contains(entry.Message, "provider-secret") || strings.Contains(entry.Message, errRefresh.message) { + t.Fatalf("Warn log exposed secret or used generic client message: %q", entry.Message) + } + return + } + t.Fatalf("expected upstream failure Warn log, got logs: %#v", hook.AllEntries()) +} + func TestHomeCredentialBoundaryWarnLogUsesSafeDiagnostic(t *testing.T) { diagnostic := "antigravity refresh: access token expired access_token=provider-secret\nforged log line" - upstreamErr := statusErrorLogTestError{ - message: diagnostic, + upstreamErr := markedDiagnosticStatusError{ + message: "credential refresh temporarily unavailable", + diagnostic: diagnostic, statusCode: http.StatusServiceUnavailable, } hook := setupTestLoggerHook(t) @@ -78,8 +124,8 @@ func TestHomeCredentialBoundaryWarnLogUsesSafeDiagnostic(t *testing.T) { executor := &requestPrepareExecutor{prepareErr: upstreamErr} selection := &HomeDispatchSelection{Auth: auth, Executor: executor, Provider: "antigravity"} _, errPrepare := NewManager(nil, nil, nil).prepareHomeRequestAuth(context.Background(), executor, selection) - if errPrepare == nil || errPrepare.Error() != diagnostic { - t.Fatalf("operation error = %v, want original error %q", errPrepare, diagnostic) + if errPrepare == nil || errPrepare.Error() != upstreamErr.message { + t.Fatalf("operation error = %v, want generic error %q", errPrepare, upstreamErr.message) } matches := 0 -- 2.51.2 From bc918ab27664f0752a0891a61457cc456fdb283b Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Fri, 28 Aug 2026 11:07:52 +0800 Subject: [PATCH 3/4] fix(runtime): log safe home refresh error types --- internal/runtime/executor/helps/home_refresh.go | 15 ++++++++++----- .../runtime/executor/helps/home_refresh_test.go | 17 +++++++++++++++-- 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/internal/runtime/executor/helps/home_refresh.go b/internal/runtime/executor/helps/home_refresh.go index 486279fc..5efb9229 100644 --- a/internal/runtime/executor/helps/home_refresh.go +++ b/internal/runtime/executor/helps/home_refresh.go @@ -18,6 +18,7 @@ type homeStatusErr struct { code int msg string diagnostic string + errorType string upstream bool } @@ -40,6 +41,9 @@ func (e homeStatusErr) LogDiagnostic() string { if e.upstream { return fmt.Sprintf("Home refresh upstream response: status=%d", e.code) } + if errorType := strings.ToLower(strings.TrimSpace(e.errorType)); errorType != "" { + return logging.SafeDiagnosticForLog("Home refresh failed: type=" + errorType) + } return logging.SafeDiagnosticForLog(e.Error()) } @@ -124,18 +128,19 @@ func RefreshAuthViaHome(ctx context.Context, cfg *config.Config, auth *cliproxya var env homeErrorEnvelope if errUnmarshal := json.Unmarshal(raw, &env); errUnmarshal == nil && env.Error != nil { + code := strings.TrimSpace(env.Error.Type) + if code == "" { + code = strings.TrimSpace(env.Error.Code) + } if env.Error.Upstream != nil { return nil, true, homeStatusErr{ code: env.Error.Upstream.Status, msg: string(env.Error.Upstream.Body), diagnostic: env.Error.Diagnostic, + errorType: code, upstream: true, } } - code := strings.TrimSpace(env.Error.Type) - if code == "" { - code = strings.TrimSpace(env.Error.Code) - } statusCode := statusFromHomeErrorCode(code) message := "credential refresh temporarily unavailable" switch statusCode { @@ -144,7 +149,7 @@ func RefreshAuthViaHome(ctx context.Context, cfg *config.Config, auth *cliproxya case http.StatusNotFound: message = "credential refresh target not found" } - return nil, true, homeStatusErr{code: statusCode, msg: message, diagnostic: env.Error.Diagnostic} + return nil, true, homeStatusErr{code: statusCode, msg: message, diagnostic: env.Error.Diagnostic, errorType: code} } updated, returnedIndex, errParse := parseHomeRefreshAuth(raw) diff --git a/internal/runtime/executor/helps/home_refresh_test.go b/internal/runtime/executor/helps/home_refresh_test.go index 362e7251..cb0f67ca 100644 --- a/internal/runtime/executor/helps/home_refresh_test.go +++ b/internal/runtime/executor/helps/home_refresh_test.go @@ -124,8 +124,8 @@ func TestRefreshAuthViaHomeUsesGenericMessageForLegacyErrorEnvelope(t *testing.T if strings.Contains(errRefresh.Error(), "provider-secret") { t.Fatalf("refresh error included legacy Home detail: %v", errRefresh) } - if diagnosticErr, ok := errRefresh.(interface{ LogDiagnostic() string }); !ok || diagnosticErr.LogDiagnostic() != errRefresh.Error() { - t.Fatalf("legacy Home message became trusted diagnostic: %T/%v", errRefresh, errRefresh) + if diagnosticErr, ok := errRefresh.(interface{ LogDiagnostic() string }); !ok || diagnosticErr.LogDiagnostic() != "Home refresh failed: type=error" { + t.Fatalf("legacy Home error type log diagnostic = %T/%v", errRefresh, errRefresh) } } @@ -217,6 +217,7 @@ func TestRefreshAuthViaHomeUsesGenericMessageForLegacyProviderDiagnostics(t *tes raw []byte wantStatus int wantError string + wantLog string }{ { name: "transient transport diagnostic", @@ -224,6 +225,7 @@ func TestRefreshAuthViaHomeUsesGenericMessageForLegacyProviderDiagnostics(t *tes raw: []byte(`{"error":{"type":"refresh_temporarily_unavailable","message":"antigravity refresh: Post https://oauth.example/token?access_token=provider-secret: connection refused"}}`), wantStatus: http.StatusServiceUnavailable, wantError: "credential refresh temporarily unavailable", + wantLog: "Home refresh failed: type=refresh_temporarily_unavailable", }, { name: "terminal legacy diagnostic", @@ -231,6 +233,7 @@ func TestRefreshAuthViaHomeUsesGenericMessageForLegacyProviderDiagnostics(t *tes raw: []byte(`{"error":{"type":"authentication_error","message":"codex refresh: invalid_grant refresh_token=provider-secret"}}`), wantStatus: http.StatusUnauthorized, wantError: "credential unauthorized", + wantLog: "Home refresh failed: type=authentication_error", }, } for _, tt := range tests { @@ -253,6 +256,16 @@ func TestRefreshAuthViaHomeUsesGenericMessageForLegacyProviderDiagnostics(t *tes if strings.Contains(errRefresh.Error(), "provider-secret") { t.Fatalf("legacy refresh error leaked provider detail: %v", errRefresh) } + diagnosticErr, okDiagnostic := errRefresh.(interface{ LogDiagnostic() string }) + if !okDiagnostic { + t.Fatalf("legacy refresh log diagnostic type = %T, want LogDiagnostic", errRefresh) + } + if got := diagnosticErr.LogDiagnostic(); got != tt.wantLog { + t.Fatalf("legacy refresh log diagnostic = %q, want %q", got, tt.wantLog) + } + if strings.Contains(diagnosticErr.LogDiagnostic(), "provider-secret") { + t.Fatalf("legacy refresh log diagnostic leaked provider detail: %q", diagnosticErr.LogDiagnostic()) + } }) } } -- 2.51.2 From 6a489fa84d93de0fc025ad3a0adf4cf9d6d57df1 Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:50:27 +0800 Subject: [PATCH 4/4] fix(auth): prefer errors from upstream attempts Track when executor calls cross an upstream transport boundary and use that signal to keep model/provider errors from being replaced by later local preparation, selection, or internal failures. Mark HTTP, websocket, relay, and usage-tracked transports as upstream attempts, while avoiding marks for local validation, logging, missing sessions, and successful websocket handshakes before request send. Parse relative auth expiry metadata and adjust Antigravity refresh timing. --- internal/pluginhost/http_bridge.go | 2 + internal/pluginhost/http_bridge_test.go | 57 ++++ .../executor/aistudio_executor_test.go | 68 +++- .../runtime/executor/antigravity_executor.go | 2 +- .../executor/antigravity_executor_auth.go | 48 +-- .../antigravity_executor_keepalive_test.go | 6 +- .../executor/antigravity_executor_tokens.go | 1 + .../executor/antigravity_refresh_test.go | 54 ++++ .../claude_executor_ratelimit_test.go | 6 +- .../executor/claude_executor_request.go | 2 + .../executor/codex_websockets_connection.go | 3 + .../executor/codex_websockets_execute.go | 18 +- .../codex_websockets_executor_test.go | 135 +++++++- .../executor/codex_websockets_stream.go | 18 +- internal/runtime/executor/gemini_executor.go | 1 + .../gemini_executor_signature_test.go | 6 +- .../runtime/executor/gemini_executor_test.go | 6 +- .../executor/gemini_vertex_executor.go | 2 + .../executor/helps/logging_helpers_test.go | 30 ++ .../runtime/executor/helps/usage_helpers.go | 2 + .../executor/helps/usage_helpers_test.go | 9 +- .../executor/xai_websockets_executor.go | 5 + .../executor/xai_websockets_executor_test.go | 71 +++- internal/wsrelay/session.go | 2 + sdk/auth/antigravity.go | 4 +- sdk/auth/refresh_registry_test.go | 2 +- sdk/cliproxy/auth/auto_refresh_loop_test.go | 28 ++ sdk/cliproxy/auth/conductor_execution.go | 174 ++++++++-- ...conductor_execution_error_priority_test.go | 302 ++++++++++++++++++ sdk/cliproxy/auth/conductor_home.go | 5 + sdk/cliproxy/auth/conductor_home_execution.go | 25 +- .../conductor_request_scoped_errors_test.go | 31 ++ sdk/cliproxy/auth/conductor_stream.go | 53 ++- sdk/cliproxy/auth/home_retry_contract_test.go | 195 ++++++++++- sdk/cliproxy/auth/openai_compat_pool_test.go | 40 ++- .../auth/request_auth_prepare_test.go | 88 ++++- sdk/cliproxy/auth/types.go | 35 +- sdk/cliproxy/executor/context.go | 39 ++- 38 files changed, 1444 insertions(+), 131 deletions(-) create mode 100644 internal/pluginhost/http_bridge_test.go create mode 100644 sdk/cliproxy/auth/conductor_execution_error_priority_test.go diff --git a/internal/pluginhost/http_bridge.go b/internal/pluginhost/http_bridge.go index edd279b1..a012f2b9 100644 --- a/internal/pluginhost/http_bridge.go +++ b/internal/pluginhost/http_bridge.go @@ -10,6 +10,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" log "github.com/sirupsen/logrus" ) @@ -127,6 +128,7 @@ func (c *hostHTTPClient) doHTTP(ctx context.Context, req pluginapi.HTTPRequest) if client == nil { client = &http.Client{} } + cliproxyexecutor.MarkUpstreamAttempt(ctx) resp, errDo := client.Do(httpReq) if errDo != nil { helps.RecordAPIResponseError(ctx, cfg, errDo) diff --git a/internal/pluginhost/http_bridge_test.go b/internal/pluginhost/http_bridge_test.go new file mode 100644 index 00000000..abac90e8 --- /dev/null +++ b/internal/pluginhost/http_bridge_test.go @@ -0,0 +1,57 @@ +package pluginhost + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +func TestHostHTTPClientMarksUpstreamAttempt(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, "ok") + })) + t.Cleanup(server.Close) + + client := New().newHTTPClient(nil) + for _, test := range []struct { + name string + do func(context.Context) error + }{ + { + name: "buffered", + do: func(ctx context.Context) error { + _, errDo := client.Do(ctx, pluginapi.HTTPRequest{URL: server.URL}) + return errDo + }, + }, + { + name: "stream", + do: func(ctx context.Context) error { + response, errDo := client.DoStream(ctx, pluginapi.HTTPRequest{URL: server.URL}) + if errDo != nil { + return errDo + } + for range response.Chunks { + } + return nil + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + if errDo := test.do(ctx); errDo != nil { + t.Fatalf("host HTTP request error = %v", errDo) + } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("host HTTP request did not mark an upstream attempt") + } + }) + } +} diff --git a/internal/runtime/executor/aistudio_executor_test.go b/internal/runtime/executor/aistudio_executor_test.go index c0543f37..e14aff80 100644 --- a/internal/runtime/executor/aistudio_executor_test.go +++ b/internal/runtime/executor/aistudio_executor_test.go @@ -51,6 +51,68 @@ func TestAIStudioTranslateRequestPrependsLeadingUserForIssue4959ResponsesHistory assertIssue4959LeadingUserContents(t, gjson.GetBytes(body.payload, "contents").Array()) } +func TestAIStudioExecutorWithoutRelaySessionDoesNotMarkUpstreamAttempt(t *testing.T) { + const authID = "aistudio-not-connected" + relay := wsrelay.NewManager(wsrelay.Options{}) + exec := NewAIStudioExecutor(&config.Config{}, "aistudio", relay) + auth := &cliproxyauth.Auth{ID: authID, Provider: "aistudio"} + req := cliproxyexecutor.Request{ + Model: "gemini-3.1-pro-preview", + Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`), + } + opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini} + + tests := []struct { + name string + run func(context.Context) error + }{ + { + name: "HTTP request", + run: func(ctx context.Context) error { + httpReq, errRequest := http.NewRequestWithContext(ctx, http.MethodPost, "https://example.com/generate", strings.NewReader(`{"contents":[]}`)) + if errRequest != nil { + return errRequest + } + _, errRequest = exec.HttpRequest(ctx, auth, httpReq) + return errRequest + }, + }, + { + name: "execute", + run: func(ctx context.Context) error { + _, errExecute := exec.Execute(ctx, auth, req, opts) + return errExecute + }, + }, + { + name: "stream", + run: func(ctx context.Context) error { + _, errStream := exec.ExecuteStream(ctx, auth, req, opts) + return errStream + }, + }, + { + name: "count tokens", + run: func(ctx context.Context) error { + _, errCount := exec.CountTokens(ctx, auth, req, opts) + return errCount + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + errRun := test.run(ctx) + if errRun == nil || !strings.Contains(errRun.Error(), "not connected") { + t.Fatalf("request error = %v, want provider not connected", errRun) + } + if cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("missing relay session was marked as an upstream attempt") + } + }) + } +} + func TestAIStudioExecutorExecuteStartsTTFTBeforeRelayWait(t *testing.T) { const authID = "aistudio-ttft-auth" delay := 40 * time.Millisecond @@ -123,13 +185,17 @@ func TestAIStudioExecutorExecuteStartsTTFTBeforeRelayWait(t *testing.T) { plugin := &captureAIStudioUsagePlugin{records: make(chan usage.Record, 16)} usage.RegisterPlugin(plugin) exec := NewAIStudioExecutor(&config.Config{}, "aistudio", relay) - _, errExecute := exec.Execute(context.Background(), &cliproxyauth.Auth{ID: authID, Provider: "aistudio"}, cliproxyexecutor.Request{ + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + _, errExecute := exec.Execute(ctx, &cliproxyauth.Auth{ID: authID, Provider: "aistudio"}, cliproxyexecutor.Request{ Model: "gemini-3.1-pro-preview", Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`), }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini}) if errExecute != nil { t.Fatalf("Execute() error = %v", errExecute) } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("relay request did not mark an upstream attempt") + } if errClient := <-clientDone; errClient != nil { t.Fatal(errClient) } diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index 590d4f1f..3386fc5f 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -39,7 +39,7 @@ const ( antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" antigravityClientSecret = "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" antigravityAuthType = "antigravity" - refreshSkew = 50 * time.Minute + antigravityRequestTokenSafetyWindow = 5 * time.Minute antigravityCreditsHintRefreshInterval = 10 * time.Minute antigravityCreditsHintRefreshTimeout = 5 * time.Second antigravityShortQuotaCooldownThreshold = 5 * time.Minute diff --git a/internal/runtime/executor/antigravity_executor_auth.go b/internal/runtime/executor/antigravity_executor_auth.go index 108eb914..e0b52fcc 100644 --- a/internal/runtime/executor/antigravity_executor_auth.go +++ b/internal/runtime/executor/antigravity_executor_auth.go @@ -7,7 +7,6 @@ import ( "io" "net/http" "net/url" - "strconv" "strings" "time" @@ -72,8 +71,8 @@ func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *clipr return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"} } accessToken := metaStringValue(auth.Metadata, "access_token") - expiry := tokenExpiry(auth.Metadata) - if accessToken != "" && expiry.After(time.Now().Add(refreshSkew)) { + expiry, _ := auth.ExpirationTime() + if accessToken != "" && expiry.After(time.Now().Add(antigravityRequestTokenSafetyWindow)) { e.maybeRefreshAntigravityCreditsHint(ctx, auth, accessToken) return accessToken, nil, nil } @@ -261,26 +260,6 @@ func missingAntigravityProjectIDError(cause error) statusErr { return statusErr{code: http.StatusBadRequest, msg: msg} } -func tokenExpiry(metadata map[string]any) time.Time { - if metadata == nil { - return time.Time{} - } - if expStr, ok := metadata["expired"].(string); ok { - expStr = strings.TrimSpace(expStr) - if expStr != "" { - if parsed, errParse := time.Parse(time.RFC3339, expStr); errParse == nil { - return parsed - } - } - } - expiresIn, hasExpires := int64Value(metadata["expires_in"]) - tsMs, hasTimestamp := int64Value(metadata["timestamp"]) - if hasExpires && hasTimestamp { - return time.Unix(0, tsMs*int64(time.Millisecond)).Add(time.Duration(expiresIn) * time.Second) - } - return time.Time{} -} - func metaStringValue(metadata map[string]any, key string) string { if metadata == nil { return "" @@ -295,26 +274,3 @@ func metaStringValue(metadata map[string]any, key string) string { } return "" } - -func int64Value(value any) (int64, bool) { - switch typed := value.(type) { - case int: - return int64(typed), true - case int64: - return typed, true - case float64: - return int64(typed), true - case json.Number: - if i, errParse := typed.Int64(); errParse == nil { - return i, true - } - case string: - if strings.TrimSpace(typed) == "" { - return 0, false - } - if i, errParse := strconv.ParseInt(strings.TrimSpace(typed), 10, 64); errParse == nil { - return i, true - } - } - return 0, false -} diff --git a/internal/runtime/executor/antigravity_executor_keepalive_test.go b/internal/runtime/executor/antigravity_executor_keepalive_test.go index 8451a8c6..aecf9cfc 100644 --- a/internal/runtime/executor/antigravity_executor_keepalive_test.go +++ b/internal/runtime/executor/antigravity_executor_keepalive_test.go @@ -150,7 +150,8 @@ func TestAntigravityCountTokensReusesUpstreamConnection(t *testing.T) { const requests = 4 payload := []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`) for i := 0; i < requests; i++ { - if _, errCount := exec.CountTokens(context.Background(), auth, cliproxyexecutor.Request{ + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + if _, errCount := exec.CountTokens(ctx, auth, cliproxyexecutor.Request{ Model: "gemini-3.6-flash-high", Payload: payload, }, cliproxyexecutor.Options{ @@ -160,6 +161,9 @@ func TestAntigravityCountTokensReusesUpstreamConnection(t *testing.T) { }); errCount != nil { t.Fatalf("request %d: CountTokens() error = %v", i, errCount) } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatalf("request %d: CountTokens() did not mark the upstream attempt", i) + } } mu.Lock() diff --git a/internal/runtime/executor/antigravity_executor_tokens.go b/internal/runtime/executor/antigravity_executor_tokens.go index 65412441..78e859c0 100644 --- a/internal/runtime/executor/antigravity_executor_tokens.go +++ b/internal/runtime/executor/antigravity_executor_tokens.go @@ -113,6 +113,7 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut AuthValue: authValue, }) + cliproxyexecutor.MarkUpstreamAttempt(ctx) httpResp, errDo := httpClient.Do(httpReq) if errDo != nil { helps.RecordAPIResponseError(ctx, e.cfg, errDo) diff --git a/internal/runtime/executor/antigravity_refresh_test.go b/internal/runtime/executor/antigravity_refresh_test.go index 647b6996..3b045db7 100644 --- a/internal/runtime/executor/antigravity_refresh_test.go +++ b/internal/runtime/executor/antigravity_refresh_test.go @@ -42,6 +42,60 @@ func useAntigravityRefreshTestTransport(t *testing.T, targetHost string) { }) } +func TestAntigravityEnsureAccessTokenUsesFiveMinuteSafetyWindow(t *testing.T) { + t.Parallel() + + executor := &AntigravityExecutor{} + now := time.Now() + + t.Run("uses token outside safety window", func(t *testing.T) { + auth := &cliproxyauth.Auth{Metadata: map[string]any{ + "access_token": "still-valid-access", + "expired": now.Add(antigravityRequestTokenSafetyWindow + time.Minute).Format(time.RFC3339), + }} + + token, updated, errToken := executor.ensureAccessToken(context.Background(), auth) + if errToken != nil { + t.Fatalf("ensureAccessToken() error = %v", errToken) + } + if token != "still-valid-access" || updated != nil { + t.Fatalf("ensureAccessToken() = %q, %#v, want existing token and nil update", token, updated) + } + }) + + t.Run("uses relative expiry with issued_at seconds", func(t *testing.T) { + issuedAt := now.Add(-10 * time.Minute).Truncate(time.Second) + auth := &cliproxyauth.Auth{Metadata: map[string]any{ + "access_token": "relative-expiry-access", + "expires_in": 3600, + "issued_at": issuedAt.Unix(), + }} + + token, updated, errToken := executor.ensureAccessToken(context.Background(), auth) + if errToken != nil { + t.Fatalf("ensureAccessToken() error = %v", errToken) + } + if token != "relative-expiry-access" || updated != nil { + t.Fatalf("ensureAccessToken() = %q, %#v, want existing token and nil update", token, updated) + } + }) + + t.Run("refreshes token inside safety window", func(t *testing.T) { + auth := &cliproxyauth.Auth{Metadata: map[string]any{ + "access_token": "expiring-access", + "expired": now.Add(antigravityRequestTokenSafetyWindow - time.Minute).Format(time.RFC3339), + }} + + token, updated, errToken := executor.ensureAccessToken(context.Background(), auth) + if errToken == nil || !strings.Contains(errToken.Error(), "missing refresh token") { + t.Fatalf("ensureAccessToken() error = %v, want refresh attempt", errToken) + } + if token != "" || updated != nil { + t.Fatalf("ensureAccessToken() = %q, %#v, want empty result after failed refresh", token, updated) + } + }) +} + func TestAntigravityRefresh_DeduplicatesConcurrentRefresh(t *testing.T) { resetAntigravityRefreshGroupForTest() t.Cleanup(resetAntigravityRefreshGroupForTest) diff --git a/internal/runtime/executor/claude_executor_ratelimit_test.go b/internal/runtime/executor/claude_executor_ratelimit_test.go index 0c1a852f..be458162 100644 --- a/internal/runtime/executor/claude_executor_ratelimit_test.go +++ b/internal/runtime/executor/claude_executor_ratelimit_test.go @@ -216,13 +216,17 @@ func TestClaudeExecutor_RateLimit_CountTokensHonorsRateLimitReset(t *testing.T) } payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) - _, err := executor.countTokensUpstream(context.Background(), auth, cliproxyexecutor.Request{ + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + _, err := executor.countTokensUpstream(ctx, auth, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err == nil { t.Fatal("expected error from CountTokens, got nil") } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("CountTokens() did not mark the HTTP request as an upstream attempt") + } var rap retryAfterProvider if !errors.As(err, &rap) || rap == nil || rap.RetryAfter() == nil { diff --git a/internal/runtime/executor/claude_executor_request.go b/internal/runtime/executor/claude_executor_request.go index 1bc2d6da..caf1c9a4 100644 --- a/internal/runtime/executor/claude_executor_request.go +++ b/internal/runtime/executor/claude_executor_request.go @@ -25,6 +25,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" "github.com/router-for-me/CLIProxyAPI/v7/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "github.com/tidwall/sjson" @@ -979,6 +980,7 @@ func applyClaudeHeadersWithNativeProfile( // exactly how the streaming and non-streaming beta sets diverged before. func doClaudeUpstreamRequest(client *http.Client, req *http.Request) (*http.Response, error) { applyClaudeWireHeaderCasing(req) + cliproxyexecutor.MarkUpstreamAttempt(req.Context()) return client.Do(req) } diff --git a/internal/runtime/executor/codex_websockets_connection.go b/internal/runtime/executor/codex_websockets_connection.go index b1115907..3ad0604c 100644 --- a/internal/runtime/executor/codex_websockets_connection.go +++ b/internal/runtime/executor/codex_websockets_connection.go @@ -35,6 +35,9 @@ func (e *CodexWebsocketsExecutor) dialCodexWebsocket(ctx context.Context, auth * ctx = context.Background() } conn, resp, err := dialer.DialContext(ctx, wsURL, headers) + if err != nil { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + } closer := newWebsocketConnectionCloser(conn) if conn != nil { // Avoid gorilla/websocket flate tail validation issues on some upstreams/Go versions. diff --git a/internal/runtime/executor/codex_websockets_execute.go b/internal/runtime/executor/codex_websockets_execute.go index a7daf3d0..38747638 100644 --- a/internal/runtime/executor/codex_websockets_execute.go +++ b/internal/runtime/executor/codex_websockets_execute.go @@ -130,13 +130,15 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut var closer *websocketConnectionCloser var respHS *http.Response var errDial error + dialCtx := ctx if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { conn, closer = existingWebsocketSessionConn(sess, authID, wsURL) if conn == nil { return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() } } else { - conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + dialCtx = cliproxyexecutor.WithUpstreamAttemptTracker(ctx) + conn, closer, respHS, errDial = e.ensureUpstreamConn(dialCtx, auth, sess, authID, wsURL, wsHeaders) } if errDial != nil { bodyErr := websocketHandshakeBody(respHS) @@ -144,10 +146,16 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) } if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { - if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) { - return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + if opts.ExecutionLifecycle == nil && !cliproxyexecutor.DownstreamWebsocket(ctx) { + return e.CodexExecutor.Execute(ctx, auth, req, opts) } - return e.CodexExecutor.Execute(ctx, auth, req, opts) + if cliproxyexecutor.UpstreamAttempted(dialCtx) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + } + return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + } + if cliproxyexecutor.UpstreamAttempted(dialCtx) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) } if respHS != nil && respHS.StatusCode > 0 { return resp, newCodexStatusErr(respHS.StatusCode, bodyErr) @@ -185,6 +193,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut } restoreMultiAgentV2 := !multiAgentV2Conflict && (optimizeMultiAgentV2 || sess.isMultiAgentV2Optimized(conn)) + cliproxyexecutor.MarkUpstreamAttempt(ctx) if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { errSend = mapCodexWebsocketWriteError(sess, conn, errSend) if sess != nil { @@ -232,6 +241,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut }) recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) reporter.StartResponseTTFT() + cliproxyexecutor.MarkUpstreamAttempt(ctx) if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry == nil { wsReqBody = wsReqBodyRetry } else { diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go index 68f2dafe..43f75011 100644 --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -467,20 +467,24 @@ func TestCodexWebsocketsExecuteStreamHandshakeErrorReturnsWithoutLockingSession( } for i := 0; i < 2; i++ { + attemptCtx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) done := make(chan error, 1) - go func() { - _, errExecute := exec.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + go func(ctx context.Context) { + _, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ Model: "gpt-5.4", Payload: []byte(`{"model":"gpt-5.4","input":[{"type":"message","id":"msg-1"}]}`), }, opts) done <- errExecute - }() + }(attemptCtx) select { case errExecute := <-done: statusErr, ok := errExecute.(interface{ StatusCode() int }) if !ok || statusErr.StatusCode() != http.StatusUnauthorized { t.Fatalf("attempt %d error = %T %v, want status 401", i+1, errExecute, errExecute) } + if !cliproxyexecutor.UpstreamAttempted(attemptCtx) { + t.Fatalf("attempt %d websocket handshake was not marked as upstream", i+1) + } case <-time.After(5 * time.Second): t.Fatalf("attempt %d timed out; execution session remained locked", i+1) } @@ -544,6 +548,100 @@ func TestCodexAutoExecutorRequiredUpstreamWebsocketRejectsHTTPFallback(t *testin } } +func TestCodexWebsocketUpgradeFallbackLocalErrorDoesNotMarkUpstreamAttempt(t *testing.T) { + tests := []struct { + name string + execute func(*CodexWebsocketsExecutor, context.Context, *cliproxyauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) error + }{ + { + name: "non-stream", + execute: func(exec *CodexWebsocketsExecutor, ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) error { + _, errExecute := exec.Execute(ctx, auth, req, opts) + return errExecute + }, + }, + { + name: "stream", + execute: func(exec *CodexWebsocketsExecutor, ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) error { + _, errExecute := exec.ExecuteStream(ctx, auth, req, opts) + return errExecute + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var upgradeAttempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") { + t.Errorf("unexpected HTTP fallback request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusInternalServerError) + return + } + upgradeAttempts.Add(1) + w.WriteHeader(http.StatusUpgradeRequired) + _, _ = w.Write([]byte(`{"error":{"message":"websocket unavailable"}}`)) + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + auth := &cliproxyauth.Auth{ + ID: "codex-fallback-local-error", + Provider: "codex", + Attributes: map[string]string{ + "api_key": "sk-test", + "base_url": server.URL, + }, + } + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + opts := codexOpenAIImageTestOptions(codexImagesGenerationsPath, tc.name == "stream") + errExecute := tc.execute(exec, ctx, auth, cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte("not-json"), + }, opts) + if errExecute == nil || !strings.Contains(errExecute.Error(), "invalid OpenAI image generation request JSON") { + t.Fatalf("Execute() error = %v, want local image request validation error", errExecute) + } + if got := upgradeAttempts.Load(); got != 1 { + t.Fatalf("websocket upgrade attempts = %d, want 1", got) + } + if cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("transparent 426 fallback marked a local HTTP preparation error as an upstream attempt") + } + }) + } +} + +func TestCodexWebsocketMissingRequiredSessionDoesNotMarkUpstreamAttempt(t *testing.T) { + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ + ID: "codex-required-session", + Provider: "codex", + Attributes: map[string]string{ + "api_key": "sk-test", + }, + } + ctx := cliproxyexecutor.WithUpstreamAttemptTracker( + cliproxyexecutor.WithRequiredUpstreamWebsocket(context.Background()), + ) + _, errExecute := exec.Execute(ctx, auth, cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","role":"user","content":"hello"}]}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "missing-codex-session", + }, + }) + if !cliproxyexecutor.IsUpstreamWebsocketReplayRequired(errExecute) { + t.Fatalf("Execute() error = %T %v, want replay-required", errExecute, errExecute) + } + if cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("missing retained websocket connection was marked as an upstream attempt") + } +} + func TestCodexWebsocketsExecuteStreamPassesThroughUpstreamWebsocketPayloadForDownstreamWebsocket(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} capturedPayload := make(chan []byte, 1) @@ -584,12 +682,17 @@ func TestCodexWebsocketsExecuteStreamPassesThroughUpstreamWebsocketPayloadForDow SourceFormat: sdktranslator.FromString("openai-response"), ResponseFormat: sdktranslator.FromString("openai-response"), } - ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + ctx := cliproxyexecutor.WithUpstreamAttemptTracker( + cliproxyexecutor.WithDownstreamWebsocket(context.Background()), + ) result, err := exec.ExecuteStream(ctx, auth, req, opts) if err != nil { t.Fatalf("ExecuteStream() error = %v", err) } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("websocket write did not mark an upstream attempt") + } select { case chunk, ok := <-result.Chunks: @@ -2065,9 +2168,13 @@ func TestCodexWebsocketNonstreamLifecycleBindFailureDetachesConnection(t *testin cliproxyexecutor.ExecutionSessionMetadataKey: "nonstream-bind-failed", }, } - if _, errExecute := exec.Execute(context.Background(), auth, req, opts); errExecute == nil { + attemptCtx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + if _, errExecute := exec.Execute(attemptCtx, auth, req, opts); errExecute == nil { t.Fatal("Execute() error = nil, want lifecycle bind failure") } + if cliproxyexecutor.UpstreamAttempted(attemptCtx) { + t.Fatal("successful handshake marked a lifecycle bind failure as an upstream request attempt") + } select { case <-closed: case <-time.After(time.Second): @@ -2123,9 +2230,13 @@ func TestCodexWebsocketLifecycleBindFailureReleasesSessionRequestLock(t *testing cliproxyexecutor.ExecutionSessionMetadataKey: "bind-failed", }, } - if _, errExecute := exec.ExecuteStream(context.Background(), auth, req, opts); errExecute == nil { + attemptCtx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + if _, errExecute := exec.ExecuteStream(attemptCtx, auth, req, opts); errExecute == nil { t.Fatal("ExecuteStream() error = nil, want lifecycle bind failure") } + if cliproxyexecutor.UpstreamAttempted(attemptCtx) { + t.Fatal("successful handshake marked a lifecycle bind failure as an upstream request attempt") + } select { case <-closed: case <-time.After(time.Second): @@ -2337,10 +2448,14 @@ func TestCodexWebsocketsExecuteHandshakeUsageLimitReachedSetsRetryAfter(t *testi ResponseFormat: sdktranslator.FromString("openai-response"), } - _, errExecute := exec.Execute(context.Background(), auth, req, opts) + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + _, errExecute := exec.Execute(ctx, auth, req, opts) if errExecute == nil { t.Fatal("Execute() error = nil, want handshake rejection") } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("429 websocket handshake was not marked as an upstream attempt") + } statusErr, ok := errExecute.(interface{ StatusCode() int }) if !ok || statusErr.StatusCode() != http.StatusTooManyRequests { t.Fatalf("status = %#v, want 429", errExecute) @@ -2384,10 +2499,14 @@ func TestCodexWebsocketsExecuteStreamHandshakeUsageLimitReachedSetsRetryAfter(t ResponseFormat: sdktranslator.FromString("openai-response"), } - _, errExecuteStream := exec.ExecuteStream(context.Background(), auth, req, opts) + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + _, errExecuteStream := exec.ExecuteStream(ctx, auth, req, opts) if errExecuteStream == nil { t.Fatal("ExecuteStream() error = nil, want handshake rejection") } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("429 streaming websocket handshake was not marked as an upstream attempt") + } statusErr, ok := errExecuteStream.(interface{ StatusCode() int }) if !ok || statusErr.StatusCode() != http.StatusTooManyRequests { t.Fatalf("status = %#v, want 429", errExecuteStream) diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index a158118a..217d6e6a 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -125,6 +125,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr var closer *websocketConnectionCloser var respHS *http.Response var errDial error + dialCtx := ctx if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { conn, closer = existingWebsocketSessionConn(sess, authID, wsURL) if conn == nil { @@ -134,7 +135,8 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() } } else { - conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + dialCtx = cliproxyexecutor.WithUpstreamAttemptTracker(ctx) + conn, closer, respHS, errDial = e.ensureUpstreamConn(dialCtx, auth, sess, authID, wsURL, wsHeaders) } var upstreamHeaders http.Header if respHS != nil { @@ -149,10 +151,16 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr if sess != nil { sess.reqMu.Unlock() } - if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) { - return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + if opts.ExecutionLifecycle == nil && !cliproxyexecutor.DownstreamWebsocket(ctx) { + return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts) } - return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts) + if cliproxyexecutor.UpstreamAttempted(dialCtx) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + } + return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + } + if cliproxyexecutor.UpstreamAttempted(dialCtx) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) } if respHS != nil && respHS.StatusCode > 0 { if sess != nil { @@ -186,6 +194,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } restoreMultiAgentV2 := !multiAgentV2Conflict && (optimizeMultiAgentV2 || sess.isMultiAgentV2Optimized(conn)) + cliproxyexecutor.MarkUpstreamAttempt(ctx) if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { errSend = mapCodexWebsocketWriteError(sess, conn, errSend) helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) @@ -240,6 +249,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr }) recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) reporter.StartResponseTTFT() + cliproxyexecutor.MarkUpstreamAttempt(ctx) if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry != nil { errSendRetry = mapCodexWebsocketWriteError(sess, conn, errSendRetry) helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) diff --git a/internal/runtime/executor/gemini_executor.go b/internal/runtime/executor/gemini_executor.go index 5c577f95..186edf5f 100644 --- a/internal/runtime/executor/gemini_executor.go +++ b/internal/runtime/executor/gemini_executor.go @@ -684,6 +684,7 @@ func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut }) httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) + cliproxyexecutor.MarkUpstreamAttempt(ctx) resp, err := httpClient.Do(httpReq) if err != nil { helps.RecordAPIResponseError(ctx, e.cfg, err) diff --git a/internal/runtime/executor/gemini_executor_signature_test.go b/internal/runtime/executor/gemini_executor_signature_test.go index 25e9137e..1990e1fd 100644 --- a/internal/runtime/executor/gemini_executor_signature_test.go +++ b/internal/runtime/executor/gemini_executor_signature_test.go @@ -473,10 +473,14 @@ func TestGeminiVertexExecutorCountTokens_GeminiPayload_SanitizesClaudeCAISSignat req, opts := geminiRequestWithThinkingSignature(testClaudeCAISSample) - _, err := executor.CountTokens(context.Background(), auth, req, opts) + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + _, err := executor.CountTokens(ctx, auth, req, opts) if err != nil { t.Fatalf("CountTokens() error = %v", err) } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("CountTokens() did not mark the HTTP request as an upstream attempt") + } if bytes.Contains(upstreamBody, []byte(testClaudeCAISSample)) { t.Fatalf("vertex countTokens upstream request leaked raw Claude CAIS signature: %s", upstreamBody) diff --git a/internal/runtime/executor/gemini_executor_test.go b/internal/runtime/executor/gemini_executor_test.go index 9f1acc0c..bab5c076 100644 --- a/internal/runtime/executor/gemini_executor_test.go +++ b/internal/runtime/executor/gemini_executor_test.go @@ -179,9 +179,13 @@ func TestGeminiExecutorCountTokensPrependsLeadingUser(t *testing.T) { Payload: []byte(`{"contents":[{"role":"model","parts":[{"text":"prior output"}]}]}`), } - if _, errCount := executor.CountTokens(context.Background(), auth, request, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini}); errCount != nil { + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + if _, errCount := executor.CountTokens(ctx, auth, request, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini}); errCount != nil { t.Fatalf("CountTokens() error = %v", errCount) } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("CountTokens() did not mark the HTTP request as an upstream attempt") + } contents := gjson.GetBytes(upstreamBody, "contents").Array() if len(contents) != 2 || contents[0].Get("role").String() != "user" || contents[1].Get("role").String() != "model" { t.Fatalf("countTokens roles malformed: %s", upstreamBody) diff --git a/internal/runtime/executor/gemini_vertex_executor.go b/internal/runtime/executor/gemini_vertex_executor.go index 2c138778..97ff7951 100644 --- a/internal/runtime/executor/gemini_vertex_executor.go +++ b/internal/runtime/executor/gemini_vertex_executor.go @@ -926,6 +926,7 @@ func (e *GeminiVertexExecutor) countTokensWithServiceAccount(ctx context.Context }) httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) + cliproxyexecutor.MarkUpstreamAttempt(ctx) httpResp, errDo := httpClient.Do(httpReq) if errDo != nil { helps.RecordAPIResponseError(ctx, e.cfg, errDo) @@ -1019,6 +1020,7 @@ func (e *GeminiVertexExecutor) countTokensWithAPIKey(ctx context.Context, auth * }) httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) + cliproxyexecutor.MarkUpstreamAttempt(ctx) httpResp, errDo := httpClient.Do(httpReq) if errDo != nil { helps.RecordAPIResponseError(ctx, e.cfg, errDo) diff --git a/internal/runtime/executor/helps/logging_helpers_test.go b/internal/runtime/executor/helps/logging_helpers_test.go index b6070af3..b2d2e27c 100644 --- a/internal/runtime/executor/helps/logging_helpers_test.go +++ b/internal/runtime/executor/helps/logging_helpers_test.go @@ -11,8 +11,38 @@ import ( "github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) +func TestRequestLoggingDoesNotMarkUpstreamAttempt(t *testing.T) { + tests := []struct { + name string + record func(context.Context) + }{ + { + name: "HTTP", + record: func(ctx context.Context) { + RecordAPIRequest(ctx, &config.Config{}, UpstreamRequestLog{URL: "https://api.example.com", Method: http.MethodPost}) + }, + }, + { + name: "websocket", + record: func(ctx context.Context) { + RecordAPIWebsocketRequest(ctx, &config.Config{}, UpstreamRequestLog{URL: "wss://api.example.com", Method: "WEBSOCKET"}) + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + test.record(ctx) + if cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("request logging marked an upstream attempt before transport") + } + }) + } +} + func TestRecordAPIRequestClonesDeferredBodyWhenRequestLogDisabled(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index 5596e68d..cfb2dd2d 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -17,6 +17,7 @@ import ( internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/tidwall/gjson" "github.com/tidwall/sjson" @@ -473,6 +474,7 @@ type usageTTFTRoundTripper struct { } func (t usageTTFTRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + cliproxyexecutor.MarkUpstreamAttempt(req.Context()) t.reporter.StartResponseTTFT() resp, errRoundTrip := t.base.RoundTrip(req) if errRoundTrip != nil { diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index 8a561da5..496deb33 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -12,6 +12,7 @@ import ( "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" ) @@ -557,7 +558,8 @@ func TestUsageReporterBuildRecordIncludesLatency(t *testing.T) { func TestUsageReporterTrackHTTPClientStartsTTFTBeforeRoundTrip(t *testing.T) { delay := 40 * time.Millisecond - reporter := NewUsageReporter(context.Background(), "openai", "gpt-5.4", nil) + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil) client := reporter.TrackHTTPClient(&http.Client{ Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { time.Sleep(delay) @@ -571,7 +573,7 @@ func TestUsageReporterTrackHTTPClientStartsTTFTBeforeRoundTrip(t *testing.T) { }), }) - req, errNewRequest := http.NewRequestWithContext(context.Background(), http.MethodPost, "https://example.invalid/v1/chat/completions", strings.NewReader("{}")) + req, errNewRequest := http.NewRequestWithContext(ctx, http.MethodPost, "https://example.invalid/v1/chat/completions", strings.NewReader("{}")) if errNewRequest != nil { t.Fatalf("NewRequestWithContext() error = %v", errNewRequest) } @@ -588,6 +590,9 @@ func TestUsageReporterTrackHTTPClientStartsTTFTBeforeRoundTrip(t *testing.T) { if got := reporter.ttftDuration(); got < delay { t.Fatalf("ttft = %v, want >= %v", got, delay) } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("HTTP RoundTrip did not mark an upstream attempt") + } } func TestUsageReporterTrackHTTPClientRoundTripOnly_DoesNotTriggerOnBodyRead(t *testing.T) { diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index 16b31e1e..5a4b8819 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -583,6 +583,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox readCh = sess.activate(conn) } + cliproxyexecutor.MarkUpstreamAttempt(ctx) if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { errSend = mapXAIWebsocketWriteError(sess, conn, errSend) helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) @@ -639,6 +640,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox logXAIWebsocketRequest(executionSessionID, authID, wsURL, wsReqBodyRetry) recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) reporter.StartResponseTTFT() + cliproxyexecutor.MarkUpstreamAttempt(ctx) if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry != nil { errSendRetry = mapXAIWebsocketWriteError(sess, connRetry, errSendRetry) helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) @@ -1050,6 +1052,9 @@ func (e *XAIWebsocketsExecutor) dialXAIWebsocket(ctx context.Context, auth *clip ctx = context.Background() } conn, resp, err := dialer.DialContext(ctx, wsURL, headers) + if err != nil { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + } closer := newWebsocketConnectionCloser(conn) if conn != nil { // Avoid gorilla/websocket flate tail validation issues on some upstreams/Go versions. diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index c05ba0a3..26715ad8 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -78,6 +78,64 @@ func TestXAIWebsocketsRequiredUpstreamRejectsCompactionHTTPFallback(t *testing.T } } +func TestXAIWebsocketMissingRequiredSessionDoesNotMarkUpstreamAttempt(t *testing.T) { + exec := NewXAIWebsocketsExecutor(&config.Config{}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ + ID: "xai-required-session", + Provider: "xai", + Attributes: map[string]string{ + "api_key": "xai-key", + }, + } + ctx := cliproxyexecutor.WithUpstreamAttemptTracker( + cliproxyexecutor.WithRequiredUpstreamWebsocket(context.Background()), + ) + _, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "grok-4", + Payload: []byte(`{"model":"grok-4","previous_response_id":"resp-1","input":[{"type":"message","role":"user","content":"hello"}]}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "missing-xai-session", + }, + }) + if !cliproxyexecutor.IsUpstreamWebsocketReplayRequired(errExecute) { + t.Fatalf("ExecuteStream() error = %T %v, want replay-required", errExecute, errExecute) + } + if cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("missing retained websocket connection was marked as an upstream attempt") + } +} + +func TestXAIWebsocketSuccessfulHandshakeDoesNotMarkRequestAttempt(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) { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + t.Errorf("upgrade websocket: %v", errUpgrade) + return + } + defer func() { _ = conn.Close() }() + _, _, _ = conn.ReadMessage() + })) + defer server.Close() + + exec := NewXAIWebsocketsExecutor(&config.Config{}) + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + conn, closer, resp, errDial := exec.dialXAIWebsocket(ctx, &cliproxyauth.Auth{}, strings.Replace(server.URL, "http", "ws", 1), http.Header{}) + if errDial != nil || conn == nil { + t.Fatalf("dialXAIWebsocket() = (%p, %v), want successful connection", conn, errDial) + } + if resp != nil && resp.Body != nil { + defer func() { _ = resp.Body.Close() }() + } + defer func() { _ = closer.Close() }() + if cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("successful handshake was marked before an upstream request was sent") + } +} + func TestMapXAIWebsocketWriteErrorStopsRetryForMessageTooBig(t *testing.T) { networkWriteErr := errors.New("write: broken pipe") tests := []struct { @@ -261,12 +319,17 @@ func TestXAIWebsocketsExecuteStreamSendsResponseCreateWithPreviousResponseID(t * cliproxyexecutor.ExecutionSessionMetadataKey: "execution-session-1", }, } - ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + ctx := cliproxyexecutor.WithUpstreamAttemptTracker( + cliproxyexecutor.WithDownstreamWebsocket(context.Background()), + ) result, err := exec.ExecuteStream(ctx, auth, req, opts) if err != nil { t.Fatalf("ExecuteStream() error = %v", err) } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("websocket write did not mark an upstream attempt") + } select { case payload := <-capturedPayload: @@ -1557,10 +1620,14 @@ func TestXAIWebsocketsExecuteStreamHandshakeFreeUsageExhaustedSetsRetryAfter(t * ResponseFormat: sdktranslator.FormatOpenAIResponse, } - _, err := exec.ExecuteStream(context.Background(), auth, req, opts) + ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) + _, err := exec.ExecuteStream(ctx, auth, req, opts) if err == nil { t.Fatal("ExecuteStream() error = nil, want handshake rejection") } + if !cliproxyexecutor.UpstreamAttempted(ctx) { + t.Fatal("429 websocket handshake was not marked as an upstream attempt") + } status, ok := err.(interface{ StatusCode() int }) if !ok || status.StatusCode() != http.StatusTooManyRequests { t.Fatalf("status = %#v, want 429", err) diff --git a/internal/wsrelay/session.go b/internal/wsrelay/session.go index a728cbc3..1898b7c6 100644 --- a/internal/wsrelay/session.go +++ b/internal/wsrelay/session.go @@ -8,6 +8,7 @@ import ( "time" "github.com/gorilla/websocket" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) const ( @@ -132,6 +133,7 @@ func (s *session) send(ctx context.Context, msg Message) error { if err := s.conn.SetWriteDeadline(time.Now().Add(writeTimeout)); err != nil { return fmt.Errorf("set write deadline: %w", err) } + cliproxyexecutor.MarkUpstreamAttempt(ctx) if err := s.conn.WriteJSON(msg); err != nil { return fmt.Errorf("write json: %w", err) } diff --git a/sdk/auth/antigravity.go b/sdk/auth/antigravity.go index 8fe1c148..40a9243a 100644 --- a/sdk/auth/antigravity.go +++ b/sdk/auth/antigravity.go @@ -26,9 +26,9 @@ func NewAntigravityAuthenticator() Authenticator { return &AntigravityAuthentica // Provider returns the provider key for antigravity. func (AntigravityAuthenticator) Provider() string { return "antigravity" } -// RefreshLead instructs the manager to refresh 50 minutes before expiry. +// RefreshLead instructs the manager to refresh 30 minutes before expiry. func (AntigravityAuthenticator) RefreshLead() *time.Duration { - return new(50 * time.Minute) + return new(30 * time.Minute) } // Login launches a local OAuth flow to obtain antigravity tokens and persists them. diff --git a/sdk/auth/refresh_registry_test.go b/sdk/auth/refresh_registry_test.go index ec4d76b1..0734cbdd 100644 --- a/sdk/auth/refresh_registry_test.go +++ b/sdk/auth/refresh_registry_test.go @@ -13,7 +13,7 @@ func TestProviderRefreshLeads(t *testing.T) { }{ {name: "codex", authenticator: NewCodexAuthenticator(), want: 5 * 24 * time.Hour}, {name: "claude", authenticator: NewClaudeAuthenticator(), want: 4 * time.Hour}, - {name: "antigravity", authenticator: NewAntigravityAuthenticator(), want: 50 * time.Minute}, + {name: "antigravity", authenticator: NewAntigravityAuthenticator(), want: 30 * time.Minute}, {name: "kimi", authenticator: NewKimiAuthenticator(), want: 5 * time.Minute}, {name: "xai", authenticator: NewXAIAuthenticator(), want: 5 * time.Minute}, } diff --git a/sdk/cliproxy/auth/auto_refresh_loop_test.go b/sdk/cliproxy/auth/auto_refresh_loop_test.go index e4edb2df..9fb7d8fd 100644 --- a/sdk/cliproxy/auth/auto_refresh_loop_test.go +++ b/sdk/cliproxy/auth/auto_refresh_loop_test.go @@ -139,6 +139,34 @@ func TestNextRefreshCheckAt_ProviderLead_Expiry(t *testing.T) { } } +func TestNextRefreshCheckAt_RelativeExpiry(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 4, 12, 0, 0, 0, 0, time.UTC) + issuedAt := now.Add(-15 * time.Minute) + lead := 30 * time.Minute + setRefreshLeadFactory(t, "relative-expiry", func() *time.Duration { + d := lead + return &d + }) + + auth := &Auth{ + ID: "relative-expiry-auth", + Provider: "relative-expiry", + Metadata: map[string]any{ + "access_token": "test-access", + "expires_in": 3600, + "timestamp": int(issuedAt.UnixMilli()), + }, + } + + got, ok := nextRefreshCheckAt(now, auth, 15*time.Minute) + want := issuedAt.Add(time.Hour - lead) + if !ok || !got.Equal(want) { + t.Fatalf("nextRefreshCheckAt() = (%s, %t), want (%s, true)", got, ok, want) + } +} + func TestNextRefreshCheckAt_RefreshEvaluatorFallback(t *testing.T) { now := time.Date(2026, 4, 12, 0, 0, 0, 0, time.UTC) interval := 15 * time.Minute diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 4d821980..9949ac54 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -24,7 +24,8 @@ import ( ) func newUpstreamAttemptContext(ctx context.Context) context.Context { - return logging.WithFreshResponseHeadersHolder(ctx) + ctx = logging.WithFreshResponseHeadersHolder(ctx) + return cliproxyexecutor.WithUpstreamAttemptTracker(ctx) } func claudeOAuthRequestCancellation(ctx context.Context, auth *Auth, err error) error { @@ -40,6 +41,81 @@ func claudeOAuthRequestCancellation(ctx context.Context, auth *Auth, err error) return nil } +type upstreamExecutionAttemptError struct { + cause error +} + +func (e *upstreamExecutionAttemptError) Error() string { + if e == nil || e.cause == nil { + return "" + } + return e.cause.Error() +} + +func (e *upstreamExecutionAttemptError) Unwrap() error { + if e == nil { + return nil + } + return e.cause +} + +func markUpstreamExecutionAttempt(err error) error { + if err == nil { + return nil + } + if hasUpstreamExecutionAttempt(err) { + return err + } + return &upstreamExecutionAttemptError{cause: err} +} + +func markUpstreamExecutionAttemptFromContext(ctx context.Context, err error) error { + if err == nil || !cliproxyexecutor.UpstreamAttempted(ctx) { + return err + } + return markUpstreamExecutionAttempt(err) +} + +func hasUpstreamExecutionAttempt(err error) bool { + var marked *upstreamExecutionAttemptError + return errors.As(err, &marked) && marked != nil +} + +func unwrapUpstreamExecutionAttempt(err error) error { + marked, ok := err.(*upstreamExecutionAttemptError) + if !ok || marked == nil || marked.cause == nil { + return err + } + return marked.cause +} + +func unwrapExecutionBoundaryError(err error) error { + err = unwrapRequestStopError(err) + return unwrapUpstreamExecutionAttempt(err) +} + +func preferredExecutionAttemptError(fallback, upstream error) error { + if errors.Is(fallback, context.Canceled) || errors.Is(fallback, context.DeadlineExceeded) { + return fallback + } + if upstream == nil { + return fallback + } + var exhausted *homeRetryRoundExhaustedError + if errors.As(fallback, &exhausted) && exhausted != nil { + upstream = unwrapUpstreamExecutionAttempt(upstream) + var previousRound *homeRetryRoundExhaustedError + if errors.As(upstream, &previousRound) && previousRound != nil && previousRound.cause != nil { + upstream = previousRound.cause + } + // Keep the current Home round marker because it owns authoritative retry timing. + preferred := *exhausted + preferred.cause = upstream + return markUpstreamExecutionAttempt(&preferred) + } + return markUpstreamExecutionAttempt(upstream) +} + // Execute performs a non-streaming execution using the configured selector and executor. // It supports multiple providers for the same model and round-robins the starting provider per model. func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { @@ -50,12 +126,13 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye } if m.HomeEnabled() { resp, errHome := m.executeHome(ctx, normalized, req, opts, false) - return resp, unwrapRequestStopError(errHome) + return resp, unwrapExecutionBoundaryError(errHome) } defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error + var preferredUpstreamErr error retryModel := authSelectionModelFromOptions(opts, req.Model) for attempt := 0; ; attempt++ { resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, attempt, defaultRequestRetry) @@ -63,7 +140,10 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye return resp, nil } if isRequestTerminatedError(errExec) || isRequestStopError(errExec) { - return cliproxyexecutor.Response{}, unwrapRequestStopError(errExec) + return cliproxyexecutor.Response{}, unwrapExecutionBoundaryError(errExec) + } + if hasUpstreamExecutionAttempt(errExec) { + preferredUpstreamErr = errExec } lastErr = errExec wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry) @@ -75,7 +155,13 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye } } if lastErr != nil { - lastErr = unwrapRequestStopError(lastErr) + if ctx != nil { + if errCtx := ctx.Err(); errCtx != nil { + return cliproxyexecutor.Response{}, errCtx + } + } + lastErr = preferredExecutionAttemptError(lastErr, preferredUpstreamErr) + lastErr = unwrapExecutionBoundaryError(lastErr) if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { if resp, ok, errCredits := m.tryAntigravityCreditsExecute(ctx, req, opts); errCredits != nil { return cliproxyexecutor.Response{}, errCredits @@ -97,12 +183,13 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip } if m.HomeEnabled() { resp, errHome := m.executeHome(ctx, normalized, req, opts, true) - return resp, unwrapRequestStopError(errHome) + return resp, unwrapExecutionBoundaryError(errHome) } defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error + var preferredUpstreamErr error retryModel := authSelectionModelFromOptions(opts, req.Model) for attempt := 0; ; attempt++ { resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, attempt, defaultRequestRetry) @@ -110,7 +197,10 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip return resp, nil } if isRequestTerminatedError(errExec) || isRequestStopError(errExec) { - return cliproxyexecutor.Response{}, unwrapRequestStopError(errExec) + return cliproxyexecutor.Response{}, unwrapExecutionBoundaryError(errExec) + } + if hasUpstreamExecutionAttempt(errExec) { + preferredUpstreamErr = errExec } lastErr = errExec wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry) @@ -122,7 +212,13 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip } } if lastErr != nil { - return cliproxyexecutor.Response{}, unwrapRequestStopError(lastErr) + if ctx != nil { + if errCtx := ctx.Err(); errCtx != nil { + return cliproxyexecutor.Response{}, errCtx + } + } + lastErr = preferredExecutionAttemptError(lastErr, preferredUpstreamErr) + return cliproxyexecutor.Response{}, unwrapExecutionBoundaryError(lastErr) } return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} } @@ -144,6 +240,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error + var preferredUpstreamErr error homeRetryLimit := -1 retryModel := authSelectionModelFromOptions(opts, req.Model) attempt := 0 @@ -154,10 +251,13 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli if errStream == nil { return result, nil } + if hasUpstreamExecutionAttempt(errStream) { + preferredUpstreamErr = errStream + } if m.HomeEnabled() && retryRoundPending { if wait, okWait := pendingHomeRetryRoundDelay(errStream, maxWait, &homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == ""); okWait && m.homeRetryAllowed(attempt-1, homeRetryLimit) { if retryRoundWaited { - return nil, errStream + return nil, unwrapExecutionBoundaryError(errStream) } if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { return nil, errWait @@ -169,7 +269,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli retryRoundPending = false retryRoundWaited = false if isRequestTerminatedError(errStream) || isRequestStopError(errStream) { - return nil, unwrapRequestStopError(errStream) + return nil, unwrapExecutionBoundaryError(errStream) } lastErr = errStream wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errStream, attempt, normalized, retryModel, maxWait, homeRetryLimit, defaultRequestRetry) @@ -184,7 +284,15 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli retryRoundWaited = false } if lastErr != nil { - lastErr = unwrapRequestStopError(lastErr) + if ctx != nil { + if errCtx := ctx.Err(); errCtx != nil { + return nil, errCtx + } + } + if preferredUpstreamErr != nil && (!m.HomeEnabled() || isHomeRetryRoundExhausted(lastErr)) { + lastErr = preferredExecutionAttemptError(lastErr, preferredUpstreamErr) + } + lastErr = unwrapExecutionBoundaryError(lastErr) if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { if result, ok, errCredits := m.tryAntigravityCreditsExecuteStream(ctx, req, opts); errCredits != nil { return nil, errCredits @@ -321,10 +429,11 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } attempted := make(map[string]struct{}) var lastErr error + var upstreamErr error for { if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { if lastErr != nil { - return cliproxyexecutor.Response{}, lastErr + return cliproxyexecutor.Response{}, preferredExecutionAttemptError(lastErr, upstreamErr) } return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} } @@ -337,7 +446,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) if errPick != nil { if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { - return cliproxyexecutor.Response{}, lastErr + return cliproxyexecutor.Response{}, preferredExecutionAttemptError(lastErr, upstreamErr) } return cliproxyexecutor.Response{}, errPick } @@ -392,8 +501,12 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } startExec := time.Now() resp, errExec := executor.Execute(execCtx, auth, execReq, execOpts) + errExec = markUpstreamExecutionAttemptFromContext(execCtx, errExec) durationExec := time.Since(startExec) if errExec != nil { + if hasUpstreamExecutionAttempt(errExec) { + upstreamErr = errExec + } if errCtx := execCtx.Err(); errCtx != nil { return cliproxyexecutor.Response{}, errCtx } @@ -404,8 +517,12 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req execCtx = newUpstreamAttemptContext(execCtx) startRetry := time.Now() resp, errExec = executor.Execute(execCtx, auth, execReq, execOpts) + errExec = markUpstreamExecutionAttemptFromContext(execCtx, errExec) durationRetry := time.Since(startRetry) if errExec != nil { + if hasUpstreamExecutionAttempt(errExec) { + upstreamErr = errExec + } warnLogUpstreamFailure(execCtx, entry, provider, upstreamModel, auth, durationRetry, errExec) if errCtx := execCtx.Err(); errCtx != nil { return cliproxyexecutor.Response{}, errCtx @@ -499,10 +616,11 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } attempted := make(map[string]struct{}) var lastErr error + var upstreamErr error for { if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { if lastErr != nil { - return cliproxyexecutor.Response{}, lastErr + return cliproxyexecutor.Response{}, preferredExecutionAttemptError(lastErr, upstreamErr) } return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} } @@ -515,7 +633,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) if errPick != nil { if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { - return cliproxyexecutor.Response{}, lastErr + return cliproxyexecutor.Response{}, preferredExecutionAttemptError(lastErr, upstreamErr) } return cliproxyexecutor.Response{}, errPick } @@ -570,8 +688,12 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } startExec := time.Now() resp, errExec := executor.CountTokens(execCtx, auth, execReq, execOpts) + errExec = markUpstreamExecutionAttemptFromContext(execCtx, errExec) durationExec := time.Since(startExec) if errExec != nil { + if hasUpstreamExecutionAttempt(errExec) { + upstreamErr = errExec + } if errCtx := execCtx.Err(); errCtx != nil { return cliproxyexecutor.Response{}, errCtx } @@ -582,8 +704,12 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, execCtx = newUpstreamAttemptContext(execCtx) startRetry := time.Now() resp, errExec = executor.CountTokens(execCtx, auth, execReq, execOpts) + errExec = markUpstreamExecutionAttemptFromContext(execCtx, errExec) durationRetry := time.Since(startRetry) if errExec != nil { + if hasUpstreamExecutionAttempt(errExec) { + upstreamErr = errExec + } warnLogUpstreamFailure(execCtx, entry, provider, upstreamModel, auth, durationRetry, errExec) if errCtx := execCtx.Err(); errCtx != nil { return cliproxyexecutor.Response{}, errCtx @@ -686,15 +812,17 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string homeSameAuthRetryPending := false attempted := make(map[string]struct{}) var lastErr error + var upstreamErr error var roundTiming homeRetryRoundTiming for { allowSameAuthRetry := homeMode && homeSameAuthRetryPending && lastHomeAuthID != "" && homeSameAuthRetries[lastHomeAuthID] == 0 if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials && !allowSameAuthRetry { if lastErr != nil { + preferredErr := preferredExecutionAttemptError(lastErr, upstreamErr) if homeMode { - return nil, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), true) + return nil, markHomeRetryRoundExhausted(preferredErr, roundTiming.RetryAfter(), true) } - return nil, lastErr + return nil, preferredErr } return nil, &Error{Code: "auth_not_found", Message: "no auth available"} } @@ -721,16 +849,17 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string auth, executor, provider, errPick = m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) } if errPick != nil { + preferredErr := preferredExecutionAttemptError(lastErr, upstreamErr) var homeCooldown *homeDispatchRetryAfterError if homeMode && lastErr != nil && errors.As(errPick, &homeCooldown) && homeCooldown != nil { observeHomeCooldownRetryLimit(homeCooldown, homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == "") - return nil, markHomeRetryRoundExhausted(lastErr, homeCooldown.RetryAfter(), false) + return nil, markHomeRetryRoundExhausted(preferredErr, homeCooldown.RetryAfter(), false) } if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { if homeMode { - return nil, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), isHomeNextRoundImmediatelyAvailable(errPick)) + return nil, markHomeRetryRoundExhausted(preferredErr, roundTiming.RetryAfter(), isHomeNextRoundImmediatelyAvailable(errPick)) } - return nil, lastErr + return nil, preferredErr } return nil, errPick } @@ -748,7 +877,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string return nil, errEnd } if lastErr != nil { - return nil, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), true) + return nil, markHomeRetryRoundExhausted(preferredExecutionAttemptError(lastErr, upstreamErr), roundTiming.RetryAfter(), true) } return nil, &Error{Code: "auth_not_found", Message: "no auth available"} } @@ -766,7 +895,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string return nil, errEnd } if lastErr != nil { - return nil, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), false) + return nil, markHomeRetryRoundExhausted(preferredExecutionAttemptError(lastErr, upstreamErr), roundTiming.RetryAfter(), false) } return nil, repeatedHomeAuthError() } else { @@ -886,6 +1015,9 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, execOpts, routeModel, streamExecutionModel, models, pooled, aliasResult, routing, !homeMode || selection != nil, selection != nil) if errStream != nil { + if hasUpstreamExecutionAttempt(errStream) { + upstreamErr = errStream + } if selection != nil { excludeAuth := shouldExcludeHomeAuthAfterStreamError(execCtx, auth, errStream) if homeSameAuthRetries[auth.ID] > 0 { diff --git a/sdk/cliproxy/auth/conductor_execution_error_priority_test.go b/sdk/cliproxy/auth/conductor_execution_error_priority_test.go new file mode 100644 index 00000000..6add7883 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_execution_error_priority_test.go @@ -0,0 +1,302 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "sync/atomic" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + internalhome "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type retryTerminalPriorityExecutor struct { + identifier string + upstreamErr error + terminalErr error + cancel context.CancelFunc + calls atomic.Int32 +} + +type retryTerminalPriorityHomeDispatcher struct { + finalPayload []byte +} + +func (*retryTerminalPriorityHomeDispatcher) HeartbeatOK() bool { return true } + +func (d *retryTerminalPriorityHomeDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) { + return d.RPopAuthWithRetryRoundConstraints(ctx, model, sessionID, headers, count, 0, nil, "") +} + +func (d *retryTerminalPriorityHomeDispatcher) RPopAuthWithRetryRoundConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, retryRound int, excludedAuthIDs []string, _ string) ([]byte, error) { + if len(excludedAuthIDs) == 0 { + authID := "home-retry-a" + if retryRound > 0 { + authID = "home-retry-b" + } + requestRetry := 1 + return json.Marshal(homeAuthDispatchResponse{ + RequestRetry: &requestRetry, + Auth: Auth{ + ID: authID, + Provider: "home-retry-contract", + Status: StatusActive, + }, + }) + } + if retryRound == 0 { + return []byte(`{"error":{"type":"auth_unavailable","message":"a credential is immediately available next round"}}`), nil + } + if len(d.finalPayload) > 0 { + return append([]byte(nil), d.finalPayload...), nil + } + return nil, internalhome.ErrAuthNotFound +} + +func (*retryTerminalPriorityHomeDispatcher) AbortAmbiguousDispatch() {} + +func (e *retryTerminalPriorityExecutor) Identifier() string { return e.identifier } + +func (e *retryTerminalPriorityExecutor) Execute(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, e.nextError(ctx) +} + +func (e *retryTerminalPriorityExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, e.nextError(ctx) +} + +func (*retryTerminalPriorityExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (e *retryTerminalPriorityExecutor) CountTokens(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, e.nextError(ctx) +} + +func (*retryTerminalPriorityExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *retryTerminalPriorityExecutor) nextError(ctx context.Context) error { + if e.calls.Add(1) == 1 { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + return e.upstreamErr + } + if e.cancel != nil { + e.cancel() + } + return e.terminalErr +} + +func TestManagerRetryPreservesTerminalContextErrorAfterUpstreamFailure(t *testing.T) { + paths := []struct { + name string + invoke func(context.Context, *Manager, string, string) error + }{ + { + name: "execute", + invoke: func(ctx context.Context, manager *Manager, provider, model string) error { + _, errExecute := manager.Execute(ctx, []string{provider}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "count-tokens", + invoke: func(ctx context.Context, manager *Manager, provider, model string) error { + _, errCount := manager.ExecuteCount(ctx, []string{provider}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + return errCount + }, + }, + { + name: "stream", + invoke: func(ctx context.Context, manager *Manager, provider, model string) error { + _, errStream := manager.ExecuteStream(ctx, []string{provider}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + return errStream + }, + }, + } + terminalCases := []struct { + name string + err error + cancelParent bool + }{ + {name: "canceled", err: context.Canceled, cancelParent: true}, + {name: "deadline-exceeded", err: context.DeadlineExceeded}, + } + + for _, path := range paths { + for _, terminalCase := range terminalCases { + t.Run(path.name+"/"+terminalCase.name, func(t *testing.T) { + provider := "retry-terminal-priority-" + path.name + "-" + terminalCase.name + model := provider + "-model" + authID := provider + "-auth" + upstreamErr := &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "first retry round reached upstream"} + + ctx := context.Background() + var cancel context.CancelFunc + if terminalCase.cancelParent { + ctx, cancel = context.WithCancel(ctx) + t.Cleanup(cancel) + } + + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(1, 0, 0) + executor := &retryTerminalPriorityExecutor{ + identifier: provider, + upstreamErr: upstreamErr, + terminalErr: terminalCase.err, + cancel: cancel, + } + manager.RegisterExecutor(executor) + registerRetryRoundLocalAuths(t, manager, provider, model, map[string]int{authID: 1}) + + errExecute := path.invoke(ctx, manager, provider, model) + if !errors.Is(errExecute, terminalCase.err) { + t.Fatalf("execution error = %v, want %v", errExecute, terminalCase.err) + } + if errors.Is(errExecute, upstreamErr) { + t.Fatalf("execution error = %v, earlier upstream error took priority", errExecute) + } + if got := executor.calls.Load(); got != 2 { + t.Fatalf("executor calls = %d, want one upstream failure and one terminal retry", got) + } + }) + } + } +} + +func TestHomeRetryPreservesTerminalContextErrorAfterUpstreamFailure(t *testing.T) { + paths := []struct { + name string + invoke func(context.Context, *Manager) error + }{ + { + name: "execute", + invoke: func(ctx context.Context, manager *Manager) error { + _, errExecute := manager.Execute(ctx, []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "count-tokens", + invoke: func(ctx context.Context, manager *Manager) error { + _, errCount := manager.ExecuteCount(ctx, []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + return errCount + }, + }, + { + name: "stream", + invoke: func(ctx context.Context, manager *Manager) error { + _, errStream := manager.ExecuteStream(ctx, []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + return errStream + }, + }, + } + terminalCases := []struct { + name string + err error + }{ + {name: "canceled", err: context.Canceled}, + {name: "deadline-exceeded", err: context.DeadlineExceeded}, + } + + for _, path := range paths { + for _, terminalCase := range terminalCases { + t.Run(path.name+"/"+terminalCase.name, func(t *testing.T) { + upstreamErr := &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "first retry round reached upstream"} + executor := &retryTerminalPriorityExecutor{ + identifier: "home-retry-contract", + upstreamErr: upstreamErr, + terminalErr: terminalCase.err, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, 0, 0) + manager.PublishHomeDispatch(&retryTerminalPriorityHomeDispatcher{}, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + errExecute := path.invoke(context.Background(), manager) + if !errors.Is(errExecute, terminalCase.err) { + t.Fatalf("execution error = %v, want %v", errExecute, terminalCase.err) + } + if errors.Is(errExecute, upstreamErr) { + t.Fatalf("execution error = %v, earlier upstream error took priority", errExecute) + } + if got := executor.calls.Load(); got < 2 { + t.Fatalf("executor calls = %d, want an upstream failure followed by a terminal retry", got) + } + }) + } + } +} + +func TestHomePreferredUpstreamErrorPreservesCurrentRetryAfter(t *testing.T) { + paths := []struct { + name string + invoke func(*Manager) error + }{ + { + name: "execute", + invoke: func(manager *Manager) error { + _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "count-tokens", + invoke: func(manager *Manager) error { + _, errCount := manager.ExecuteCount(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + return errCount + }, + }, + { + name: "stream", + invoke: func(manager *Manager) error { + _, errStream := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + return errStream + }, + }, + } + + for _, path := range paths { + t.Run(path.name, func(t *testing.T) { + upstreamErr := &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "first retry round reached upstream"} + internalErr := errors.New("second retry round failed before upstream") + executor := &retryTerminalPriorityExecutor{ + identifier: "home-retry-contract", + upstreamErr: upstreamErr, + terminalErr: internalErr, + } + dispatcher := &retryTerminalPriorityHomeDispatcher{ + finalPayload: []byte(`{"error":{"type":"model_cooldown","message":"current Home retry window","retryable":true,"retry_after_ms":10000,"request_retry":1}}`), + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, 20*time.Second, 0) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + errExecute := path.invoke(manager) + if !errors.Is(errExecute, upstreamErr) { + t.Fatalf("execution error = %v, want earlier upstream error as cause", errExecute) + } + if errors.Is(errExecute, internalErr) { + t.Fatalf("execution error = %v, current internal failure remained the cause", errExecute) + } + if !isHomeRetryRoundExhausted(errExecute) { + t.Fatalf("execution error = %T %v, want current Home retry-round wrapper", errExecute, errExecute) + } + if retryAfter := retryAfterFromError(errExecute); retryAfter == nil || *retryAfter != 10*time.Second { + t.Fatalf("retry after = %v, want current Home delay 10s", retryAfter) + } + if got := SafeResponseHeaders(errExecute).Get("Retry-After"); got != "10" { + t.Fatalf("safe Retry-After header = %q, want 10", got) + } + }) + } +} diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index c33fee02..727304e7 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -170,11 +170,16 @@ func markHomeRetryRoundExhausted(err error, retryAfter *time.Duration, retryNow if err == nil { return nil } + upstreamAttempt := hasUpstreamExecutionAttempt(err) + err = unwrapUpstreamExecutionAttempt(err) marked := &homeRetryRoundExhaustedError{cause: err, retryNow: retryNow} if retryAfter != nil { marked.retryAfter = *retryAfter marked.hasRetryAfter = true } + if upstreamAttempt { + return markUpstreamExecutionAttempt(marked) + } return marked } diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index 16e5f3b3..30098387 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -21,11 +21,15 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr attempt := 0 retryRoundPending := false retryRoundWaited := false + var preferredUpstreamErr error for { response, errExecute := m.executeHomeOnce(ctx, providers, req, opts, countTokens, maxRetryCredentials, &homeRetryLimit, attempt) if errExecute == nil { return response, nil } + if hasUpstreamExecutionAttempt(errExecute) { + preferredUpstreamErr = errExecute + } if retryRoundPending { if wait, okWait := pendingHomeRetryRoundDelay(errExecute, maxWait, &homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == ""); okWait && m.homeRetryAllowed(attempt-1, homeRetryLimit) { if retryRoundWaited { @@ -41,11 +45,14 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr retryRoundPending = false retryRoundWaited = false if isRequestTerminatedError(errExecute) || isRequestStopError(errExecute) { - return cliproxyexecutor.Response{}, unwrapRequestStopError(errExecute) + return cliproxyexecutor.Response{}, unwrapExecutionBoundaryError(errExecute) } wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExecute, attempt, providers, retryModel, maxWait, homeRetryLimit, defaultRequestRetry) if !shouldRetry { - return cliproxyexecutor.Response{}, unwrapRequestStopError(errExecute) + if preferredUpstreamErr != nil && isHomeRetryRoundExhausted(errExecute) { + errExecute = preferredExecutionAttemptError(errExecute, preferredUpstreamErr) + } + return cliproxyexecutor.Response{}, unwrapExecutionBoundaryError(errExecute) } if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { return cliproxyexecutor.Response{}, errWait @@ -68,11 +75,12 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c tried := make(map[string]struct{}) attempted := make(map[string]struct{}) var lastErr error + var upstreamErr error var roundTiming homeRetryRoundTiming for homeAuthCount := 1; ; homeAuthCount++ { if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { if lastErr != nil { - return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), true) + return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(preferredExecutionAttemptError(lastErr, upstreamErr), roundTiming.RetryAfter(), true) } return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} } @@ -81,13 +89,14 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c pickOpts = withHomeExcludedAuthIDs(pickOpts, tried) selection, errSelection := m.pickHomeDispatchSelection(ctx, routeModel, pickOpts) if errSelection != nil { + preferredErr := preferredExecutionAttemptError(lastErr, upstreamErr) var homeCooldown *homeDispatchRetryAfterError if lastErr != nil && errors.As(errSelection, &homeCooldown) && homeCooldown != nil { observeHomeCooldownRetryLimit(homeCooldown, homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == "") - return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, homeCooldown.RetryAfter(), false) + return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(preferredErr, homeCooldown.RetryAfter(), false) } if shouldReturnLastErrorOnPickFailure(true, lastErr, errSelection) { - return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), isHomeNextRoundImmediatelyAvailable(errSelection)) + return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(preferredErr, roundTiming.RetryAfter(), isHomeNextRoundImmediatelyAvailable(errSelection)) } return cliproxyexecutor.Response{}, errSelection } @@ -102,7 +111,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c return cliproxyexecutor.Response{}, errEnd } if lastErr != nil { - return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), false) + return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(preferredExecutionAttemptError(lastErr, upstreamErr), roundTiming.RetryAfter(), false) } return cliproxyexecutor.Response{}, repeatedHomeAuthError() } @@ -212,6 +221,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c } startHomeExec := time.Now() response, errExecute = execute() + errExecute = markUpstreamExecutionAttemptFromContext(execCtx, errExecute) durationHomeExec := time.Since(startHomeExec) if countTokens { if _, fingerprint := getEffectiveAuth(); isUnauthorizedError(errExecute) { @@ -219,6 +229,9 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c } } if errExecute != nil { + if hasUpstreamExecutionAttempt(errExecute) { + upstreamErr = errExecute + } warnLogUpstreamFailure(execCtx, entry, selection.Provider, upstreamModel, preparedAuth, durationHomeExec, errExecute) } result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil, Options: execOpts} diff --git a/sdk/cliproxy/auth/conductor_request_scoped_errors_test.go b/sdk/cliproxy/auth/conductor_request_scoped_errors_test.go index 4d0137b3..49f1baf0 100644 --- a/sdk/cliproxy/auth/conductor_request_scoped_errors_test.go +++ b/sdk/cliproxy/auth/conductor_request_scoped_errors_test.go @@ -57,6 +57,25 @@ type customStatusError struct { retryAfter *time.Duration } +func TestUnwrapExecutionBoundaryErrorRemovesInternalMarkers(t *testing.T) { + t.Parallel() + + base := &Error{Code: "request_failed", Message: "request failed", HTTPStatus: http.StatusBadRequest} + for _, test := range []struct { + name string + err error + }{ + {name: "attempt inside stop", err: wrapRequestStopError(markUpstreamExecutionAttempt(base))}, + {name: "stop inside attempt", err: markUpstreamExecutionAttempt(wrapRequestStopError(base))}, + } { + t.Run(test.name, func(t *testing.T) { + if got := unwrapExecutionBoundaryError(test.err); got != base { + t.Fatalf("unwrapExecutionBoundaryError() = %T/%v, want original %T/%v", got, got, base, base) + } + }) + } +} + func (e customStatusError) StatusCode() int { return e.code } @@ -123,6 +142,7 @@ func TestRequestScopedErrors_ActionStop(t *testing.T) { identifier: "claude", executeFn: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { execCount++ + cliproxyexecutor.MarkUpstreamAttempt(ctx) return cliproxyexecutor.Response{}, customStatusError{ code: 400, msg: `{"error": {"message": "maximum_context_length exceeded"}}`, @@ -135,6 +155,9 @@ func TestRequestScopedErrors_ActionStop(t *testing.T) { if errExec == nil { t.Fatalf("expected error, got resp: %+v", resp) } + if _, okStatus := errExec.(customStatusError); !okStatus { + t.Fatalf("Execute() error = %T/%v, want original customStatusError", errExec, errExec) + } // Action: stop should return immediately on the first credential and not rotate to auth2. if execCount != 1 { t.Fatalf("execCount = %d, want 1 (should stop immediately)", execCount) @@ -515,6 +538,7 @@ func TestRequestScopedErrors_Stream_ActionStop(t *testing.T) { identifier: "claude", streamFn: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { execCount++ + cliproxyexecutor.MarkUpstreamAttempt(ctx) return nil, customStatusError{code: 400, msg: "stream_context_overflow"} }, } @@ -524,6 +548,9 @@ func TestRequestScopedErrors_Stream_ActionStop(t *testing.T) { if errStream == nil { t.Fatal("expected error, got nil") } + if _, okStatus := errStream.(customStatusError); !okStatus { + t.Fatalf("ExecuteStream() error = %T/%v, want original customStatusError", errStream, errStream) + } if execCount != 1 { t.Fatalf("execCount = %d, want 1 (should stop immediately)", execCount) } @@ -754,6 +781,7 @@ func TestRequestScopedErrors_CountTokens_StopAndCooldown(t *testing.T) { exec := &mockCustomErrorExecutor{ identifier: "claude", countFn: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + cliproxyexecutor.MarkUpstreamAttempt(ctx) return cliproxyexecutor.Response{}, customStatusError{ code: 404, msg: "count_endpoint_cooldown", @@ -766,6 +794,9 @@ func TestRequestScopedErrors_CountTokens_StopAndCooldown(t *testing.T) { if errCount == nil { t.Fatal("expected error, got nil") } + if _, okStatus := errCount.(customStatusError); !okStatus { + t.Fatalf("ExecuteCount() error = %T/%v, want original customStatusError", errCount, errCount) + } a1, _ := m.GetByID("auth-count-1") if !a1.Unavailable || a1.NextRetryAfter.IsZero() { diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index b2a4f72e..b4b3cb96 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -35,10 +35,16 @@ func newStreamBootstrapError(err error, headers http.Header) error { if err == nil { return nil } - return &streamBootstrapError{ + upstreamAttempt := hasUpstreamExecutionAttempt(err) + err = unwrapUpstreamExecutionAttempt(err) + bootstrapErr := &streamBootstrapError{ cause: err, headers: cloneHTTPHeader(headers), } + if upstreamAttempt { + return markUpstreamExecutionAttempt(bootstrapErr) + } + return bootstrapErr } func (e *streamBootstrapError) Error() string { @@ -203,6 +209,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } ctx = contextWithRequestedModelAlias(ctx, opts, routeModel) var lastErr error + var upstreamErr error didRefreshOnUnauthorized := false for idx, execModel := range execModels { ctx = newUpstreamAttemptContext(ctx) @@ -227,6 +234,10 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi entry := logEntryWithRequestID(ctx) startStream := time.Now() streamResult, errStream := executor.ExecuteStream(ctx, auth, execReq, execOpts) + errStream = markUpstreamExecutionAttemptFromContext(ctx, errStream) + if hasUpstreamExecutionAttempt(errStream) { + upstreamErr = errStream + } durationStream := time.Since(startStream) if errStream != nil { if errCtx := ctx.Err(); errCtx != nil { @@ -242,6 +253,10 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi ctx = newUpstreamAttemptContext(ctx) startRetry := time.Now() streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts) + errStream = markUpstreamExecutionAttemptFromContext(ctx, errStream) + if hasUpstreamExecutionAttempt(errStream) { + upstreamErr = errStream + } durationRetry := time.Since(startRetry) if errStream != nil { warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, durationRetry, errStream) @@ -262,6 +277,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } } streamResult, errStream = validateStreamResult(streamResult, errStream) + errStream = markUpstreamExecutionAttemptFromContext(ctx, errStream) if errStream != nil { rerr := resultErrorFromError(errStream) action, okAction := matchRequestScopedErrorAction(auth, errStream, m.runtimeConfigSnapshot()) @@ -278,7 +294,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } lastErr = errStream if result.CredentialScope { - return nil, errStream + return nil, preferredExecutionAttemptError(errStream, upstreamErr) } continue } @@ -287,12 +303,16 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } lastErr = errStream if result.CredentialScope { - return nil, errStream + return nil, preferredExecutionAttemptError(errStream, upstreamErr) } continue } buffered, closed, bootstrapErr := readStreamBootstrap(ctx, streamResult.Chunks) + bootstrapErr = markUpstreamExecutionAttemptFromContext(ctx, bootstrapErr) + if hasUpstreamExecutionAttempt(bootstrapErr) { + upstreamErr = newStreamBootstrapError(bootstrapErr, streamResult.Headers) + } if bootstrapErr != nil { if errCtx := ctx.Err(); errCtx != nil { discardStreamChunks(streamResult.Chunks) @@ -309,7 +329,9 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi ctx = newUpstreamAttemptContext(ctx) startRetry := time.Now() retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) + retryErr = markUpstreamExecutionAttemptFromContext(ctx, retryErr) retryStream, retryErr = validateStreamResult(retryStream, retryErr) + retryErr = markUpstreamExecutionAttemptFromContext(ctx, retryErr) if retryErr != nil { if errCtx := ctx.Err(); errCtx != nil { return nil, errCtx @@ -320,6 +342,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } else { streamResult = retryStream buffered, closed, bootstrapErr = readStreamBootstrap(ctx, streamResult.Chunks) + bootstrapErr = markUpstreamExecutionAttemptFromContext(ctx, bootstrapErr) if bootstrapErr != nil { warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, time.Since(startRetry), bootstrapErr) } @@ -330,6 +353,9 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } else { warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, time.Since(startStream), bootstrapErr) } + if hasUpstreamExecutionAttempt(bootstrapErr) { + upstreamErr = newStreamBootstrapError(bootstrapErr, streamResult.Headers) + } } if !ephemeralResult { if errCancel := claudeOAuthRequestCancellation(ctx, auth, bootstrapErr); errCancel != nil { @@ -354,7 +380,8 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } lastErr = bootstrapErr if result.CredentialScope { - return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) + currentErr := newStreamBootstrapError(bootstrapErr, streamResult.Headers) + return nil, preferredExecutionAttemptError(currentErr, upstreamErr) } continue } @@ -380,7 +407,8 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi discardStreamChunks(streamResult.Chunks) lastErr = bootstrapErr if result.CredentialScope { - return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) + currentErr := newStreamBootstrapError(bootstrapErr, streamResult.Headers) + return nil, preferredExecutionAttemptError(currentErr, upstreamErr) } continue } @@ -392,19 +420,24 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } m.recordExecutionResult(ctx, result, auth, ephemeralResult) discardStreamChunks(streamResult.Chunks) - return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) + currentErr := newStreamBootstrapError(bootstrapErr, streamResult.Headers) + return nil, preferredExecutionAttemptError(currentErr, upstreamErr) } if closed && len(buffered) == 0 { - emptyErr := &Error{Code: "empty_stream", Message: "upstream stream closed before first payload", Retryable: true} + emptyErr := markUpstreamExecutionAttemptFromContext(ctx, &Error{Code: "empty_stream", Message: "upstream stream closed before first payload", Retryable: true}) + currentErr := newStreamBootstrapError(emptyErr, streamResult.Headers) + if hasUpstreamExecutionAttempt(emptyErr) { + upstreamErr = currentErr + } warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, time.Since(startStream), emptyErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: emptyErr, Options: execOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: resultErrorFromError(emptyErr), Options: execOpts} m.recordExecutionResult(ctx, result, auth, ephemeralResult) if idx < len(execModels)-1 { lastErr = emptyErr continue } - return nil, newStreamBootstrapError(emptyErr, streamResult.Headers) + return nil, preferredExecutionAttemptError(currentErr, upstreamErr) } remaining := streamResult.Chunks @@ -419,5 +452,5 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if lastErr == nil { lastErr = &Error{Code: "auth_not_found", Message: "no upstream model available"} } - return nil, lastErr + return nil, preferredExecutionAttemptError(lastErr, upstreamErr) } diff --git a/sdk/cliproxy/auth/home_retry_contract_test.go b/sdk/cliproxy/auth/home_retry_contract_test.go index 819f9cdc..a3a10e61 100644 --- a/sdk/cliproxy/auth/home_retry_contract_test.go +++ b/sdk/cliproxy/auth/home_retry_contract_test.go @@ -27,6 +27,7 @@ type retryContractHomeDispatcher struct { requestRetry *int websocket bool exhaustedPayload []byte + exhaustedErr error } type legacyRepeatedStreamDispatcher struct { @@ -45,6 +46,11 @@ type retryRoundLimitDownshiftDispatcher struct { calls atomic.Int32 } +type retryRoundSelectionFailureDispatcher struct { + selectionPayload []byte + selectionErr error +} + type aggregateRetryHomeDispatcher struct { calls atomic.Int32 } @@ -128,6 +134,31 @@ func (d *retryRoundLimitDownshiftDispatcher) RPopAuthWithConstraints(_ context.C func (*retryRoundLimitDownshiftDispatcher) AbortAmbiguousDispatch() {} +func (*retryRoundSelectionFailureDispatcher) HeartbeatOK() bool { return true } + +func (d *retryRoundSelectionFailureDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) { + return d.RPopAuthWithRetryRoundConstraints(ctx, model, sessionID, headers, count, 0, nil, "") +} + +func (d *retryRoundSelectionFailureDispatcher) RPopAuthWithRetryRoundConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, retryRound int, excludedAuthIDs []string, _ string) ([]byte, error) { + if retryRound == 0 { + if len(excludedAuthIDs) > 0 { + return nil, home.ErrAuthNotFound + } + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "home-retry-a", + Provider: "home-retry-contract", + Status: StatusActive, + }}) + } + if len(d.selectionPayload) > 0 { + return append([]byte(nil), d.selectionPayload...), nil + } + return nil, d.selectionErr +} + +func (*retryRoundSelectionFailureDispatcher) AbortAmbiguousDispatch() {} + func (*aggregateRetryHomeDispatcher) HeartbeatOK() bool { return true } func (d *aggregateRetryHomeDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) { @@ -194,6 +225,9 @@ func (d *retryContractHomeDispatcher) RPopAuthWithConstraints(_ context.Context, if len(d.exhaustedPayload) > 0 { return append([]byte(nil), d.exhaustedPayload...), nil } + if d.exhaustedErr != nil { + return nil, d.exhaustedErr + } return nil, home.ErrAuthNotFound } @@ -210,32 +244,35 @@ func (d *retryContractHomeDispatcher) Excluded() [][]string { } type retryContractHomeExecutor struct { - mu sync.Mutex - calls []string - failAll bool - failure error - failures map[string]error - streamBootstrap bool - streamHeaders http.Header + mu sync.Mutex + calls []string + failAll bool + failure error + failures map[string]error + internalFailures map[string]bool + streamBootstrap bool + streamHeaders http.Header } func (*retryContractHomeExecutor) Identifier() string { return "home-retry-contract" } -func (e *retryContractHomeExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { +func (e *retryContractHomeExecutor) Execute(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { e.mu.Lock() e.calls = append(e.calls, auth.ID) e.mu.Unlock() if auth.ID == "home-retry-a" || e.failAll { + e.markFailureAttempt(ctx, auth.ID) return cliproxyexecutor.Response{}, e.failureError(auth.ID) } return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil } -func (e *retryContractHomeExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { +func (e *retryContractHomeExecutor) ExecuteStream(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { e.mu.Lock() e.calls = append(e.calls, auth.ID) e.mu.Unlock() if auth.ID == "home-retry-a" || e.failAll { + e.markFailureAttempt(ctx, auth.ID) errFailure := e.failureError(auth.ID) if e.streamBootstrap { chunks := make(chan cliproxyexecutor.StreamChunk, 1) @@ -263,11 +300,12 @@ func (e *retryContractHomeExecutor) failureError(authID string) error { func (*retryContractHomeExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } -func (e *retryContractHomeExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { +func (e *retryContractHomeExecutor) CountTokens(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { e.mu.Lock() e.calls = append(e.calls, auth.ID) e.mu.Unlock() if auth.ID == "home-retry-a" || e.failAll { + e.markFailureAttempt(ctx, auth.ID) return cliproxyexecutor.Response{}, e.failureError(auth.ID) } return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil @@ -283,6 +321,12 @@ func (e *retryContractHomeExecutor) Calls() []string { return append([]string(nil), e.calls...) } +func (e *retryContractHomeExecutor) markFailureAttempt(ctx context.Context, authID string) { + if e != nil && !e.internalFailures[authID] { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + } +} + type retryContractRateLimitError struct { retryAfter time.Duration } @@ -682,16 +726,23 @@ func TestHomePendingRetryRoundStopsAfterRepeatedCooldown(t *testing.T) { manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) manager.RegisterExecutor(executor) + var errExecute error if stream { - result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + var result *cliproxyexecutor.StreamResult + result, errExecute = manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) if errExecute == nil || result != nil { t.Fatalf("ExecuteStream() = result %#v, error %v; want terminal cooldown error", result, errExecute) } } else { - if _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}); errExecute == nil { + _, errExecute = manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + if errExecute == nil { t.Fatal("Execute() error = nil, want terminal cooldown error") } } + var cooldown *homeDispatchRetryAfterError + if !errors.As(errExecute, &cooldown) || cooldown == nil || cooldown.cause == nil || cooldown.cause.Code != "model_cooldown" { + t.Fatalf("terminal error = %#v, want current Home model cooldown", errExecute) + } if got := executor.Calls(); len(got) != 1 { t.Fatalf("executor calls = %v, want only the initial round execution", got) } @@ -728,6 +779,123 @@ func TestHomeRetryRoundUsesEarliestCredentialRetryAfter(t *testing.T) { } } +func TestHomePreferredErrorIgnoresLaterInternalExecutorFailure(t *testing.T) { + upstreamErr := &Error{Message: "model upstream failed", HTTPStatus: http.StatusBadGateway} + internalErr := errors.New("Home KV lookup failed") + dispatcher := &retryContractHomeDispatcher{authIDs: []string{"home-retry-a", "home-retry-b"}} + executor := &retryContractHomeExecutor{ + failAll: true, + failures: map[string]error{ + "home-retry-a": upstreamErr, + "home-retry-b": internalErr, + }, + internalFailures: map[string]bool{"home-retry-b": true}, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, 0, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + if !errors.Is(errExecute, upstreamErr) { + t.Fatalf("Execute() error = %v, want earlier model upstream error", errExecute) + } + if errors.Is(errExecute, internalErr) { + t.Fatalf("Execute() error = %v, later internal executor failure replaced upstream error", errExecute) + } +} + +func TestHomeSelectionFailureIsNotHiddenByEarlierUpstreamAttempt(t *testing.T) { + tests := []struct { + name string + exhaustedPayload []byte + exhaustedErr error + wantCode string + }{ + {name: "Home unavailable", exhaustedErr: errors.New("Home transport failed"), wantCode: "home_unavailable"}, + {name: "invalid auth payload", exhaustedPayload: []byte(`{}`), wantCode: "invalid_auth"}, + { + name: "malformed concurrency payload", + exhaustedPayload: []byte(`{"concurrency":{"accounted":true,"credential_id":"home-retry-a","model":"gpt"},"error":"busy","auth":{"id":"home-retry-a","provider":"home-retry-contract"}}`), + wantCode: "invalid_home_concurrency", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + upstreamErr := &Error{Message: "model upstream failed", HTTPStatus: http.StatusBadGateway} + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-a"}, + exhaustedPayload: test.exhaustedPayload, + exhaustedErr: test.exhaustedErr, + } + executor := &retryContractHomeExecutor{failAll: true, failure: upstreamErr} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + retryLimit := -1 + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, &retryLimit) + var authErr *Error + if !errors.As(errExecute, &authErr) || authErr == nil || authErr.Code != test.wantCode { + t.Fatalf("executeHomeOnce() error = %#v, want selection error %q", errExecute, test.wantCode) + } + if errors.Is(errExecute, upstreamErr) { + t.Fatalf("executeHomeOnce() error = %v, earlier upstream error hid selection failure", errExecute) + } + }) + } +} + +func TestHomeNextRoundSelectionFailureSupersedesEarlierUpstreamAttempt(t *testing.T) { + selectionFailures := []struct { + name string + payload []byte + err error + code string + }{ + {name: "Home unavailable", err: errors.New("Home transport failed"), code: "home_unavailable"}, + {name: "invalid auth payload", payload: []byte(`{}`), code: "invalid_auth"}, + } + for _, selectionFailure := range selectionFailures { + for _, stream := range []bool{false, true} { + name := selectionFailure.name + "/" + map[bool]string{false: "nonstream", true: "stream"}[stream] + t.Run(name, func(t *testing.T) { + upstreamErr := &Error{Message: "model upstream failed", HTTPStatus: http.StatusBadGateway} + dispatcher := &retryRoundSelectionFailureDispatcher{ + selectionPayload: selectionFailure.payload, + selectionErr: selectionFailure.err, + } + executor := &retryContractHomeExecutor{failAll: true, failure: upstreamErr} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 1) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + var errExecute error + if stream { + var result *cliproxyexecutor.StreamResult + result, errExecute = manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if result != nil { + t.Fatalf("ExecuteStream() result = %#v, want nil", result) + } + } else { + _, errExecute = manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + } + var authErr *Error + if !errors.As(errExecute, &authErr) || authErr == nil || authErr.Code != selectionFailure.code { + t.Fatalf("terminal error = %#v, want current selection error %q", errExecute, selectionFailure.code) + } + if errors.Is(errExecute, upstreamErr) { + t.Fatalf("terminal error = %v, earlier upstream error hid current selection failure", errExecute) + } + }) + } + } +} + func TestHomeStreamBootstrapErrorPreservesAggregatedRetryAfter(t *testing.T) { dispatcher := &retryContractHomeDispatcher{authIDs: []string{"home-retry-a", "home-retry-b"}} executor := &retryContractHomeExecutor{ @@ -759,6 +927,9 @@ func TestHomeStreamBootstrapErrorPreservesAggregatedRetryAfter(t *testing.T) { if !isHomeRetryRoundExhausted(chunk.Err) { t.Fatalf("stream bootstrap error = %v, want exhausted retry round", chunk.Err) } + if hasUpstreamExecutionAttempt(chunk.Err) { + t.Fatalf("stream bootstrap error leaked internal upstream-attempt marker: %T/%v", chunk.Err, chunk.Err) + } if got := SafeResponseHeaders(chunk.Err).Get("Retry-After"); got != "2" { t.Fatalf("safe Retry-After header = %q, want aggregated delay rounded to 2 seconds", got) } diff --git a/sdk/cliproxy/auth/openai_compat_pool_test.go b/sdk/cliproxy/auth/openai_compat_pool_test.go index bce2306a..24fdc96e 100644 --- a/sdk/cliproxy/auth/openai_compat_pool_test.go +++ b/sdk/cliproxy/auth/openai_compat_pool_test.go @@ -2,6 +2,7 @@ package auth import ( "context" + "errors" "net/http" "strings" "sync" @@ -25,6 +26,7 @@ type openAICompatPoolExecutor struct { executeErrors map[string]error countErrors map[string]error streamFirstErrors map[string]error + streamAttempts map[string]bool streamPayloads map[string][]cliproxyexecutor.StreamChunk } @@ -49,15 +51,18 @@ func (e *openAICompatPoolExecutor) Execute(ctx context.Context, auth *Auth, req } func (e *openAICompatPoolExecutor) ExecuteStream(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { - _ = ctx _ = auth _ = opts e.mu.Lock() e.streamModels = append(e.streamModels, req.Model) err := e.streamFirstErrors[req.Model] + attempted := e.streamAttempts[req.Model] payloadChunks, hasCustomChunks := e.streamPayloads[req.Model] chunks := append([]cliproxyexecutor.StreamChunk(nil), payloadChunks...) e.mu.Unlock() + if attempted { + cliproxyexecutor.MarkUpstreamAttempt(ctx) + } ch := make(chan cliproxyexecutor.StreamChunk, max(1, len(chunks))) if err != nil { ch <- cliproxyexecutor.StreamChunk{Err: err} @@ -835,3 +840,36 @@ func TestManagerExecuteStream_OpenAICompatAliasPoolStopsOnInvalidBootstrap(t *te t.Fatalf("stream calls = %v, want only first upstream model", got) } } + +func TestManagerExecuteStream_OpenAICompatAliasPoolPreservesEarlierUpstreamError(t *testing.T) { + alias := "claude-opus-4.66" + upstreamErr := &Error{HTTPStatus: http.StatusBadGateway, Message: "first model upstream failed"} + internalErr := errors.New("second model failed before upstream") + executor := &openAICompatPoolExecutor{ + id: openAICompatPoolProviderKey, + streamFirstErrors: map[string]error{ + "deepseek-v3.1": upstreamErr, + "glm-5": internalErr, + }, + streamAttempts: map[string]bool{"deepseek-v3.1": true}, + } + m := newOpenAICompatPoolTestManager(t, alias, []internalconfig.OpenAICompatibilityModel{ + {Name: "deepseek-v3.1", Alias: alias}, + {Name: "glm-5", Alias: alias}, + }, executor) + + result, errExecute := m.ExecuteStream(context.Background(), []string{openAICompatPoolProviderKey}, cliproxyexecutor.Request{Model: alias}, cliproxyexecutor.Options{}) + if result == nil || errExecute != nil { + t.Fatalf("ExecuteStream() = (%#v, %v), want an error stream", result, errExecute) + } + chunk, ok := <-result.Chunks + if !ok || !errors.Is(chunk.Err, upstreamErr) { + t.Fatalf("stream error = %v, want earlier upstream error %v", chunk.Err, upstreamErr) + } + if errors.Is(chunk.Err, internalErr) { + t.Fatalf("stream error = %v, later internal error replaced upstream error", chunk.Err) + } + if hasUpstreamExecutionAttempt(chunk.Err) { + t.Fatalf("stream error leaked internal upstream-attempt marker: %T/%v", chunk.Err, chunk.Err) + } +} diff --git a/sdk/cliproxy/auth/request_auth_prepare_test.go b/sdk/cliproxy/auth/request_auth_prepare_test.go index 9f2eee5d..cf5ef6fe 100644 --- a/sdk/cliproxy/auth/request_auth_prepare_test.go +++ b/sdk/cliproxy/auth/request_auth_prepare_test.go @@ -81,21 +81,23 @@ func (e *requestPrepareExecutor) recordPreparedAuth(auth *Auth) error { return nil } -func (e *requestPrepareExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { +func (e *requestPrepareExecutor) Execute(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { if errPrepared := e.recordPreparedAuth(auth); errPrepared != nil { return cliproxyexecutor.Response{}, errPrepared } if e.executeErr != nil { + cliproxyexecutor.MarkUpstreamAttempt(ctx) return cliproxyexecutor.Response{}, e.executeErr } return cliproxyexecutor.Response{Payload: []byte("ok")}, nil } -func (e *requestPrepareExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { +func (e *requestPrepareExecutor) ExecuteStream(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { if errPrepared := e.recordPreparedAuth(auth); errPrepared != nil { return nil, errPrepared } if e.executeErr != nil { + cliproxyexecutor.MarkUpstreamAttempt(ctx) return nil, e.executeErr } chunks := make(chan cliproxyexecutor.StreamChunk, 1) @@ -108,11 +110,12 @@ func (e *requestPrepareExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, return auth, nil } -func (e *requestPrepareExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { +func (e *requestPrepareExecutor) CountTokens(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { if errPrepared := e.recordPreparedAuth(auth); errPrepared != nil { return cliproxyexecutor.Response{}, errPrepared } if e.executeErr != nil { + cliproxyexecutor.MarkUpstreamAttempt(ctx) return cliproxyexecutor.Response{}, e.executeErr } return cliproxyexecutor.Response{Payload: []byte("ok")}, nil @@ -151,6 +154,85 @@ func (d *homeRequestPrepareDispatcher) RPopAuth(context.Context, string, string, func (*homeRequestPrepareDispatcher) AbortAmbiguousDispatch() {} +type homeErrorPriorityDispatcher struct{} + +func (*homeErrorPriorityDispatcher) HeartbeatOK() bool { return true } + +func (*homeErrorPriorityDispatcher) RPopAuth(_ context.Context, _ string, _ string, _ http.Header, count int) ([]byte, error) { + switch count { + case 1: + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "model-error-auth", + Provider: "antigravity", + Status: StatusActive, + Metadata: map[string]any{"access_token": "model-token", "project_id": "prepared-project"}, + }}) + case 2: + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "refresh-error-auth", + Provider: "antigravity", + Status: StatusActive, + Metadata: map[string]any{"access_token": "refresh-token"}, + }}) + default: + return json.Marshal(homeErrorEnvelope{Error: &homeErrorDetail{Code: homeRequestRetryExceededErrorCode, Message: "no more Home auths"}}) + } +} + +func (*homeErrorPriorityDispatcher) AbortAmbiguousDispatch() {} + +func TestHomeModelErrorOutranksLaterRefreshPreparationError(t *testing.T) { + paths := []struct { + name string + run func(*Manager) error + }{ + { + name: "Execute", + run: func(manager *Manager) error { + _, errExecute := manager.Execute(context.Background(), []string{"antigravity"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "Count", + run: func(manager *Manager) error { + _, errCount := manager.ExecuteCount(context.Background(), []string{"antigravity"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + return errCount + }, + }, + { + name: "Stream", + run: func(manager *Manager) error { + _, errStream := manager.ExecuteStream(context.Background(), []string{"antigravity"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{Stream: true}) + return errStream + }, + }, + } + + for _, path := range paths { + t.Run(path.name, func(t *testing.T) { + modelErr := &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "MODEL_CAPACITY_EXHAUSTED"} + refreshErr := &Error{Code: "refresh_temporarily_unavailable", HTTPStatus: http.StatusServiceUnavailable, Message: "credential refresh temporarily unavailable"} + executor := &requestPrepareExecutor{prepareErr: refreshErr, executeErr: modelErr} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(&homeErrorPriorityDispatcher{}, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + errRun := path.run(manager) + if errRun == nil || errRun.Error() != modelErr.Message || statusCodeFromError(errRun) != http.StatusServiceUnavailable { + t.Fatalf("%s error = %v, want upstream model error %q", path.name, errRun, modelErr.Message) + } + if strings.Contains(errRun.Error(), refreshErr.Message) { + t.Fatalf("%s error was overwritten by refresh failure: %v", path.name, errRun) + } + if executor.executeCalls.Load() != 1 || executor.prepareCalls.Load() != 1 { + t.Fatalf("%s calls = execute %d prepare %d, want 1/1", path.name, executor.executeCalls.Load(), executor.prepareCalls.Load()) + } + }) + } +} + func TestHomePrepareUsesEphemeralDispatchAuthAcrossExecutionPaths(t *testing.T) { for _, path := range []struct { name string diff --git a/sdk/cliproxy/auth/types.go b/sdk/cliproxy/auth/types.go index e8d14c3a..96dcb6c8 100644 --- a/sdk/cliproxy/auth/types.go +++ b/sdk/cliproxy/auth/types.go @@ -601,8 +601,8 @@ func (a *Auth) AccountInfo() (string, string) { } // ExpirationTime attempts to extract the credential expiration timestamp from metadata. -// It inspects common keys such as "expired", "expire", "expires_at", and also -// nested "token" objects to remain compatible with legacy auth file formats. +// It inspects common absolute expiry keys, expires_in plus timestamp, and nested +// token objects to remain compatible with legacy auth file formats. func (a *Auth) ExpirationTime() (time.Time, bool) { if a == nil { return time.Time{}, false @@ -641,6 +641,11 @@ func expirationFromMap(meta map[string]any) (time.Time, bool) { } } } + if expiresIn, okExpiresIn := parseRelativeExpirySeconds(meta); okExpiresIn { + if timestamp, okTimestamp := parseRelativeExpiryTimestamp(meta); okTimestamp { + return timestamp.Add(time.Duration(expiresIn) * time.Second), true + } + } for _, nestedKey := range []string{"token", "Token"} { if nested, ok := meta[nestedKey]; ok { switch val := nested.(type) { @@ -662,6 +667,28 @@ func expirationFromMap(meta map[string]any) (time.Time, bool) { return time.Time{}, false } +func parseRelativeExpirySeconds(meta map[string]any) (int, bool) { + for _, key := range []string{"expires_in", "expiresIn"} { + if value, ok := meta[key]; ok { + if seconds, okSeconds := parseIntAny(value); okSeconds && seconds > 0 { + return seconds, true + } + } + } + return 0, false +} + +func parseRelativeExpiryTimestamp(meta map[string]any) (time.Time, bool) { + for _, key := range []string{"timestamp", "issued_at", "issuedAt"} { + if value, ok := meta[key]; ok { + if timestamp, okTimestamp := parseTimeValue(value); okTimestamp && !timestamp.IsZero() { + return timestamp, true + } + } + } + return time.Time{}, false +} + func ProviderRefreshLead(provider string, runtime any) *time.Duration { provider = strings.ToLower(strings.TrimSpace(provider)) if runtime != nil { @@ -707,6 +734,10 @@ func parseTimeValue(v any) (time.Time, bool) { } case float64: return normaliseUnix(int64(value)), true + case int: + return normaliseUnix(int64(value)), true + case int32: + return normaliseUnix(int64(value)), true case int64: return normaliseUnix(value), true case json.Number: diff --git a/sdk/cliproxy/executor/context.go b/sdk/cliproxy/executor/context.go index c18d3f68..30cbebf5 100644 --- a/sdk/cliproxy/executor/context.go +++ b/sdk/cliproxy/executor/context.go @@ -1,9 +1,17 @@ package executor -import "context" +import ( + "context" + "sync/atomic" +) type downstreamWebsocketContextKey struct{} type requireUpstreamWebsocketContextKey struct{} +type upstreamAttemptTrackerContextKey struct{} + +type upstreamAttemptTracker struct { + attempted atomic.Bool +} // WithDownstreamWebsocket marks the current request as coming from a downstream websocket connection. func WithDownstreamWebsocket(ctx context.Context) context.Context { @@ -40,3 +48,32 @@ func RequiredUpstreamWebsocket(ctx context.Context) bool { enabled, ok := raw.(bool) return ok && enabled } + +// WithUpstreamAttemptTracker installs a fresh tracker for one provider execution attempt. +func WithUpstreamAttemptTracker(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, upstreamAttemptTrackerContextKey{}, &upstreamAttemptTracker{}) +} + +// MarkUpstreamAttempt records that the provider execution reached an upstream transport boundary. +func MarkUpstreamAttempt(ctx context.Context) { + if ctx == nil { + return + } + tracker, ok := ctx.Value(upstreamAttemptTrackerContextKey{}).(*upstreamAttemptTracker) + if !ok || tracker == nil { + return + } + tracker.attempted.Store(true) +} + +// UpstreamAttempted reports whether the tracked provider execution reached an upstream transport boundary. +func UpstreamAttempted(ctx context.Context) bool { + if ctx == nil { + return false + } + tracker, ok := ctx.Value(upstreamAttemptTrackerContextKey{}).(*upstreamAttemptTracker) + return ok && tracker != nil && tracker.attempted.Load() +}