feat(codex): pass Gin context to auth selection for Alpha search

- Introduced `codexSearchGinContextSelector` to capture and validate Gin context during credential selection.
- Added test case to verify context propagation and query parameter usage (`TestCodexAlphaSearchPassesGinContextToAuthSelection`).
- Updated `SelectAuth` implementation to include Gin context for enhanced compatibility with Home scheduling.
This commit is contained in:
Luis Pater
2026-07-12 23:06:38 +08:00
parent 9c3f720710
commit 55f4d6ed1d
2 changed files with 46 additions and 2 deletions
+2 -2
View File
@@ -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,
+44
View File
@@ -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{}))