diff --git a/internal/api/server.go b/internal/api/server.go index b3c9e172..920838a4 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -650,7 +650,8 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { if sessionID := strings.TrimSpace(routing.ID); sessionID != "" { selectionHeaders.Set("X-Session-ID", sessionID) } - selected, err := s.handlers.AuthManager.SelectAuth(c.Request.Context(), "codex", strings.TrimSpace(routing.Model), coreexecutor.Options{ + ctx := context.WithValue(c.Request.Context(), "gin", c) + selected, err := s.handlers.AuthManager.SelectAuth(ctx, "codex", strings.TrimSpace(routing.Model), coreexecutor.Options{ Headers: selectionHeaders, OriginalRequest: body, }) @@ -676,7 +677,6 @@ func (s *Server) codexAlphaSearch(c *gin.Context) { headers.Set("Chatgpt-Account-Id", accountID) } - ctx := context.WithValue(c.Request.Context(), "gin", c) const upstreamURL = "https://chatgpt.com/backend-api/codex/alpha/search" req, err := s.handlers.AuthManager.NewHttpRequest( ctx, selected, http.MethodPost, upstreamURL, body, headers, diff --git a/internal/api/server_test.go b/internal/api/server_test.go index a645844a..36f61383 100644 --- a/internal/api/server_test.go +++ b/internal/api/server_test.go @@ -55,6 +55,18 @@ func (e *codexSearchCaptureExecutor) PrepareRequest(req *http.Request, a *auth.A return nil } +type codexSearchGinContextSelector struct { + ginContext *gin.Context +} + +func (s *codexSearchGinContextSelector) Pick(ctx context.Context, _ string, _ string, _ coreexecutor.Options, auths []*auth.Auth) (*auth.Auth, error) { + s.ginContext, _ = ctx.Value("gin").(*gin.Context) + if len(auths) == 0 { + return nil, nil + } + return auths[0], nil +} + func (e *codexSearchCaptureExecutor) HttpRequest(_ context.Context, selected *auth.Auth, req *http.Request) (*http.Response, error) { e.request = req.Clone(req.Context()) e.authIDs = append(e.authIDs, selected.ID) @@ -188,6 +200,38 @@ func TestCodexAlphaSearchForwardsRequest(t *testing.T) { } } +func TestCodexAlphaSearchPassesGinContextToAuthSelection(t *testing.T) { + server := newTestServer(t) + selector := &codexSearchGinContextSelector{} + server.handlers.AuthManager.SetSelector(selector) + executor := &codexSearchCaptureExecutor{} + server.handlers.AuthManager.RegisterExecutor(executor) + credential := &auth.Auth{ + ID: "codex-auth", + Provider: "codex", + Status: auth.StatusActive, + Metadata: map[string]any{"access_token": "codex-token"}, + } + if _, errRegister := server.handlers.AuthManager.Register(context.Background(), credential); errRegister != nil { + t.Fatalf("register Codex auth: %v", errRegister) + } + + req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search?key=home-query-key", strings.NewReader(`{"query":"GPT-5.6"}`)) + 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 selector.ginContext == nil { + t.Fatal("auth selection did not receive the Gin context required by Home scheduling") + } + if got := selector.ginContext.Query("key"); got != "home-query-key" { + t.Fatalf("Gin query key = %q, want %q", got, "home-query-key") + } +} + func TestCodexAlphaSearchUsesRequestIDForSessionAffinity(t *testing.T) { server := newTestServer(t) server.handlers.AuthManager.SetSelector(auth.NewSessionAffinitySelector(&auth.RoundRobinSelector{}))