Merge pull request router-for-me#4698 from router-for-me/fix/home-401-refresh-recovery

fix(auth): retry Home OAuth requests once after upstream 401

Author: LTbinglingfeng

internal/api/server_routes.go

@@ -333,80 +333,115 @@ func (s *Server) codexAlphaSearch(c *gin.Context) {
 	}
 	logging.SetGinCPATraceID(c, selected.EnsureIndex())
 
-	headers := make(http.Header)
-	headers.Set("Content-Type", "application/json")
-	headers.Set("Accept", "application/json")
-	headers.Set("Originator", "codex_cli_rs")
+	baseHeaders := make(http.Header)
+	baseHeaders.Set("Content-Type", "application/json")
+	baseHeaders.Set("Accept", "application/json")
+	baseHeaders.Set("Originator", "codex_cli_rs")
 	for _, name := range []string{"Version", "User-Agent", "Session_id", "X-Client-Request-Id"} {
 		if value := strings.TrimSpace(c.GetHeader(name)); value != "" {
-			headers.Set(name, value)
+			baseHeaders.Set(name, value)
 		}
 	}
-	if accountID, ok := selected.Metadata["account_id"].(string); ok && strings.TrimSpace(accountID) != "" {
-		headers.Set("Chatgpt-Account-Id", accountID)
-	}
 
-	upstreamURL := "https://chatgpt.com/backend-api/codex/alpha/search"
-	if selected.AuthKind() == auth.AuthKindAPIKey {
-		baseURL := ""
-		if selected.Attributes != nil {
-			baseURL = strings.TrimSpace(selected.Attributes["base_url"])
+	errMissingBaseURL := errors.New("Codex Alpha Search API key base URL unavailable")
+	performRequest := func(current *auth.Auth) (*http.Response, error) {
+		headers := baseHeaders.Clone()
+		if accountID, ok := current.Metadata["account_id"].(string); ok && strings.TrimSpace(accountID) != "" {
+			headers.Set("Chatgpt-Account-Id", accountID)
 		}
-		if baseURL == "" {
-			if selection != nil {
-				selection.End("missing_base_url")
+		upstreamURL := "https://chatgpt.com/backend-api/codex/alpha/search"
+		if current.AuthKind() == auth.AuthKindAPIKey {
+			baseURL := ""
+			if current.Attributes != nil {
+				baseURL = strings.TrimSpace(current.Attributes["base_url"])
 			}
-			c.JSON(http.StatusServiceUnavailable, gin.H{"error": "Codex Alpha Search API key base URL unavailable"})
-			return
-		}
-		upstreamURL = strings.TrimRight(baseURL, "/") + "/alpha/search"
-	}
-	req, err := s.handlers.AuthManager.NewHttpRequest(
-		ctx, selected, http.MethodPost, upstreamURL, upstreamRequestBody, headers,
-	)
-	if err != nil {
-		if selection != nil {
-			selection.End("request_build_failed")
-		}
-		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
-		return
+			if baseURL == "" {
+				return nil, errMissingBaseURL
+			}
+			upstreamURL = strings.TrimRight(baseURL, "/") + "/alpha/search"
+		}
+		req, errRequest := s.handlers.AuthManager.NewHttpRequest(ctx, current, http.MethodPost, upstreamURL, upstreamRequestBody, headers)
+		if errRequest != nil {
+			return nil, errRequest
+		}
+		authType, authValue := current.AccountInfo()
+		helps.RecordAPIRequest(ctx, s.cfg, helps.UpstreamRequestLog{
+			URL:       upstreamURL,
+			Method:    http.MethodPost,
+			Headers:   req.Header.Clone(),
+			Body:      upstreamRequestBody,
+			Provider:  "codex",
+			AuthID:    current.ID,
+			AuthLabel: current.Label,
+			AuthType:  authType,
+			AuthValue: authValue,
+		})
+		return s.handlers.AuthManager.HttpRequest(ctx, current, req)
 	}
 
-	var authID, authLabel, authType, authValue string
-	if selected != nil {
-		authID = selected.ID
-		authLabel = selected.Label
-		authType, authValue = selected.AccountInfo()
-	}
-	helpHeaders := req.Header.Clone()
-	helps.RecordAPIRequest(ctx, s.cfg, helps.UpstreamRequestLog{
-		URL:       upstreamURL,
-		Method:    http.MethodPost,
-		Headers:   helpHeaders,
-		Body:      upstreamRequestBody,
-		Provider:  "codex",
-		AuthID:    authID,
-		AuthLabel: authLabel,
-		AuthType:  authType,
-		AuthValue: authValue,
-	})
-
 	if errCtx := ctx.Err(); errCtx != nil {
 		if selection != nil {
 			selection.End("attempt_canceled")
 		}
 		c.JSON(http.StatusRequestTimeout, gin.H{"error": errCtx.Error()})
 		return
 	}
-	resp, err := s.handlers.AuthManager.HttpRequest(ctx, selected, req)
+	resp, err := performRequest(selected)
 	if err != nil {
+		if errors.Is(err, errMissingBaseURL) {
+			if selection != nil {
+				selection.End("missing_base_url")
+			}
+			c.JSON(http.StatusServiceUnavailable, gin.H{"error": err.Error()})
+			return
+		}
 		if selection != nil {
 			selection.End("request_failed")
 		}
 		helps.RecordAPIResponseError(ctx, s.cfg, err)
 		c.JSON(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")
+			status := http.StatusServiceUnavailable
+			if statusError, ok := errRefresh.(interface{ StatusCode() int }); ok && statusError.StatusCode() > 0 {
+				status = statusError.StatusCode()
+			}
+			c.JSON(status, 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(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 {

internal/api/server_test.go

@@ -38,6 +38,9 @@ type codexSearchCaptureExecutor struct {
 	prepareErr   error
 	httpErr      error
 	responseBody io.ReadCloser
+	statuses     []int
+	refreshCalls int
+	httpCalls    int
 }
 
 func (e *codexSearchCaptureExecutor) Identifier() string { return "codex" }
@@ -51,7 +54,13 @@ func (e *codexSearchCaptureExecutor) ExecuteStream(context.Context, *auth.Auth,
 }
 
 func (e *codexSearchCaptureExecutor) Refresh(_ context.Context, a *auth.Auth) (*auth.Auth, error) {
-	return a, nil
+	e.refreshCalls++
+	updated := a.Clone()
+	if updated.Metadata == nil {
+		updated.Metadata = make(map[string]any)
+	}
+	updated.Metadata["access_token"] = "refreshed-home-search-token"
+	return updated, nil
 }
 
 func (e *codexSearchCaptureExecutor) CountTokens(context.Context, *auth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
@@ -115,6 +124,7 @@ func (e *codexSearchCaptureExecutor) HttpRequest(_ context.Context, selected *au
 	}
 	e.request = req.Clone(req.Context())
 	e.authIDs = append(e.authIDs, selected.ID)
+	e.httpCalls++
 	body, err := io.ReadAll(req.Body)
 	if err != nil {
 		return nil, err
@@ -124,8 +134,12 @@ func (e *codexSearchCaptureExecutor) HttpRequest(_ context.Context, selected *au
 	if responseBody == nil {
 		responseBody = io.NopCloser(strings.NewReader(`{"results":[{"url":"https://example.com"}]}`))
 	}
+	statusCode := http.StatusOK
+	if e.httpCalls <= len(e.statuses) && e.statuses[e.httpCalls-1] > 0 {
+		statusCode = e.statuses[e.httpCalls-1]
+	}
 	return &http.Response{
-		StatusCode: http.StatusOK,
+		StatusCode: statusCode,
 		Header:     http.Header{"Content-Type": []string{"application/json"}},
 		Body:       responseBody,
 	}, nil
@@ -310,6 +324,33 @@ func TestAuditHomeCodexSearchBodyCloseBeforeRelease(t *testing.T) {
 	}
 }
 
+func TestHomeCodexAlphaSearchRefreshesUnauthorizedSelectionOnce(t *testing.T) {
+	server := newTestServer(t)
+	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}}
+	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)
+
+	if rr.Code != http.StatusOK {
+		t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusOK, rr.Body.String())
+	}
+	if executor.refreshCalls != 1 || executor.httpCalls != 2 {
+		t.Fatalf("refresh/http calls = %d/%d, want 1/2", 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 := dispatcher.calls.Load(); got != 1 {
+		t.Fatalf("Home RPOP calls = %d, want 1", got)
+	}
+}
+
 func TestHomeCodexAlphaSearchEndsSelectionAcrossDirectHTTPPaths(t *testing.T) {
 	tests := []struct {
 		name       string

internal/client/codex/live/live.go

@@ -263,39 +263,38 @@ func (h *Handler) Handle(c *gin.Context) {
 		}
 	}
 
-	headers := protocolHeaders(c.Request.Header)
-	headers.Set("Content-Type", upstreamContentType)
-	setAccountHeader(headers, selected)
-	req, errRequest := h.authManager.NewHttpRequest(ctx, selected, http.MethodPost, upstreamCallURL, upstreamBody, headers)
-	if errRequest != nil {
-		if selection != nil {
-			selection.End("request_build_failed")
-		}
-		c.JSON(http.StatusBadGateway, gin.H{"error": errRequest.Error()})
-		return
+	baseHeaders := protocolHeaders(c.Request.Header)
+	baseHeaders.Set("Content-Type", upstreamContentType)
+	performRequest := func(current *auth.Auth) (*http.Response, error) {
+		headers := baseHeaders.Clone()
+		setAccountHeader(headers, current)
+		req, errRequest := h.authManager.NewHttpRequest(ctx, current, http.MethodPost, upstreamCallURL, upstreamBody, headers)
+		if errRequest != nil {
+			return nil, errRequest
+		}
+		authType, authValue := current.AccountInfo()
+		helps.RecordAPIRequest(ctx, runtimeConfig, helps.UpstreamRequestLog{
+			URL:       upstreamCallURL,
+			Method:    http.MethodPost,
+			Headers:   headersForLogging(req.Header),
+			Body:      upstreamBody,
+			Provider:  "codex",
+			AuthID:    current.ID,
+			AuthLabel: current.Label,
+			AuthType:  authType,
+			AuthValue: authValue,
+		})
+		return h.authManager.HttpRequest(ctx, current, req)
 	}
 
-	authType, authValue := selected.AccountInfo()
-	helps.RecordAPIRequest(ctx, runtimeConfig, helps.UpstreamRequestLog{
-		URL:       upstreamCallURL,
-		Method:    http.MethodPost,
-		Headers:   headersForLogging(req.Header),
-		Body:      upstreamBody,
-		Provider:  "codex",
-		AuthID:    selected.ID,
-		AuthLabel: selected.Label,
-		AuthType:  authType,
-		AuthValue: authValue,
-	})
-
 	if errContext := ctx.Err(); errContext != nil {
 		if selection != nil {
 			selection.End("attempt_canceled")
 		}
 		c.JSON(http.StatusRequestTimeout, gin.H{"error": errContext.Error()})
 		return
 	}
-	resp, errRequest := h.authManager.HttpRequest(ctx, selected, req)
+	resp, errRequest := performRequest(selected)
 	if errRequest != nil {
 		if selection != nil {
 			selection.End("request_failed")
@@ -304,6 +303,37 @@ func (h *Handler) Handle(c *gin.Context) {
 		c.JSON(http.StatusBadGateway, gin.H{"error": 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")
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "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)
+			c.JSON(http.StatusBadGateway, gin.H{"error": errRequest.Error()})
+			return
+		}
+		if resp.StatusCode == http.StatusUnauthorized {
+			h.authManager.ReportHomeUnauthorized(ctx, selected, "codex", model)
+		}
+	}
 
 	var closeResponseOnce sync.Once
 	var closeResponseErr error