Repository navigation
Fix Codex model catalog routing for API-key sessions - #183
Conversation
📝 WalkthroughWalkthroughCodex model-catalog requests are detected by method, path, and ChangesCodex catalog OAuth flow
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant Proxy
participant CredentialBroker
participant OAuthUpstream
Client->>Proxy: GET /models with client_version
Proxy->>CredentialBroker: Select OAuth catalog account
CredentialBroker-->>Proxy: OAuth account and token
Proxy->>OAuthUpstream: Send catalog request
OAuthUpstream-->>Proxy: Return model catalog
Proxy-->>Client: Return response
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
internal/proxy/codex_model_catalog_test.go (1)
102-104: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winSynchronize the
codexAuthcapture.The httptest handler goroutine writes
codexAuthat line 104. The test goroutine reads it at line 168. No explicit synchronization edge exists between the two. The other counters in this file useatomic.Int32for exactly this reason. Undergo test -racethis write and read pair can be reported as a data race.Use an atomic value to make the capture explicit.
♻️ Proposed fix
- var codexAuth string + var codexAuth atomic.Value codexUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { - codexAuth = request.Header.Get("Authorization") + codexAuth.Store(request.Header.Get("Authorization")) w.Header().Set("Content-Type", "application/json") _, _ = io.WriteString(w, `{"models":[]}`) }))- if codexAuth != "Bearer healthy-token" { - t.Fatalf("Codex upstream authorization = %q, want healthy OAuth token", codexAuth) + if got, _ := codexAuth.Load().(string); got != "Bearer healthy-token" { + t.Fatalf("Codex upstream authorization = %q, want healthy OAuth token", got) }Also applies to: 168-170
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/proxy/codex_model_catalog_test.go` around lines 102 - 104, Replace the plain codexAuth variable with an atomic value, store the Authorization header from the httptest handler via its atomic store operation, and load it in the test assertion around the existing read at line 168. Preserve the current captured-header assertion while using the same atomic synchronization pattern as the other counters.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/proxy/proxy.go`:
- Around line 4337-4349: The Codex catalog routing flow must neutralize
query-based model overrides on both paths. In internal/proxy/proxy.go lines
4337-4349, update codexModelCatalogRoutingRequest to remove the model query
parameter from the cloned URL as well as the routing headers; in
internal/proxy/session_lease.go lines 563-565, narrow validateRequestModel’s
catalog early return to only the body-model requirement so the existing
query-model conflict validation remains active for bound leases.
- Around line 2239-2243: Update the request handling around modelCatalogRequest,
sessionID, and the subsequent session/account selection flow to set userEmail to
"" for Codex model-catalog requests when no lease is bound. Preserve normal
caller email extraction for non-catalog requests and catalog requests within a
bound lease, preventing the shared codexModelCatalogSessionID record from being
updated with the caller’s identity.
In `@internal/proxy/session_lease.go`:
- Around line 563-565: Update the lease validation flow around the Codex
catalog-request handling so the request body requirement is bypassed without
returning before the query-model conflict check. Preserve the existing query
loop that validates the requested model against lease.Model, while retaining the
early bypass behavior only for body validation.
---
Nitpick comments:
In `@internal/proxy/codex_model_catalog_test.go`:
- Around line 102-104: Replace the plain codexAuth variable with an atomic
value, store the Authorization header from the httptest handler via its atomic
store operation, and load it in the test assertion around the existing read at
line 168. Preserve the current captured-header assertion while using the same
atomic synchronization pattern as the other counters.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: c9630fb4-e866-4cbb-b304-1b801317e3a7
📒 Files selected for processing (3)
internal/proxy/codex_model_catalog_test.gointernal/proxy/proxy.gointernal/proxy/session_lease.go
| modelCatalogRequest := requestProvider == accounts.ProviderCodex && codexModelCatalogRequest(r) | ||
| sessionID := session.ExtractID(r, s.MaxBodyBytes) | ||
| if modelCatalogRequest { | ||
| sessionID = codexModelCatalogSessionID | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect ExtractUserEmail sources and Sessions.Put write path.
set -euo pipefail
fd -t f 'extract.go' session | xargs rg -n -C5 'func ExtractUserEmail|userEmailHeaderCandidates'
rg -n -C3 'func .*Store\) Put\(' --type=goRepository: manaflow-ai/subrouter
Length of output: 674
🏁 Script executed:
#!/bin/bash
set -u
echo "== files =="
git ls-files | rg '(^|/)proxy\.go$|session' | head -100
echo "== proxy.go relevant lines =="
if [ -f internal/proxy/proxy.go ]; then
sed -n '2210,2350p' internal/proxy/proxy.go | nl -ba -v2210
echo "== around accountForSessionProviderWithOptions =="
rg -n -C8 'accountForSessionProviderWithOptions|codexModelCatalogSessionID|Sessions\.Put|userEmail' internal/proxy/proxy.go || true
fi
echo "== session package/store files =="
fd -t f . . | rg 'session' | while read -r f; do
echo "--- $f ---"
rg -n -C5 'func .*Put|Store|Sessions|Put\(' "$f" || true
done
echo "== deterministic source checks =="
python3 - <<'PY'
from pathlib import Path
s = Path('internal/proxy/proxy.go').read_text()
checks = {
"catalog session constant exists": "codexModelCatalogSessionID" in s,
"handler replaces session id when modelCatalogRequest": "sessionID = codexModelCatalogSessionID" in s,
"user email extracted before Sessions.Put account path": "userEmail := session.ExtractUserEmail(r)" in s,
"Put called with userEmail": "s.Sessions.Put(agentType, sessionID, assignment.AccountID, userEmail)" in s,
}
for k,v in checks.items():
print(f"{k}: {v}")
# Locate line ranges containing the key symbols.
lines = s.splitlines()
for needle in ["codexModelCatalogSessionID", "requestProvider == accounts.ProviderCodex && codexModelCatalogRequest(r)", "userEmail := session.ExtractUserEmail(r)", "s.Sessions.Put(agentType, sessionID, assignment.AccountID, userEmail)"]:
for i,l in enumerate(lines,1):
if needle in l:
print(f"line {i}: {l.strip()}")
start=max(1,i-8); end=min(len(lines),i+8)
for j in range(start,end+1):
print(f" {j}: {lines[j-1]}")
break
PYRepository: manaflow-ai/subrouter
Length of output: 50377
Clear caller email for catalog requests on a shared session slot.
codexModelCatalogSessionID is one global session key. Codex catalog requests still call session.ExtractUserEmail(r) and pass that value into account selection, which writes the caller’s email onto internal:codex-model-catalog when it differs from the stored UserEmail. This mixes caller identity on a shared session record and causes extra Sessions.Put calls for catalog polling. Clear userEmail to "" for model-catalog requests outside any bound lease.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/proxy/proxy.go` around lines 2239 - 2243, Update the request
handling around modelCatalogRequest, sessionID, and the subsequent
session/account selection flow to set userEmail to "" for Codex model-catalog
requests when no lease is bound. Preserve normal caller email extraction for
non-catalog requests and catalog requests within a bound lease, preventing the
shared codexModelCatalogSessionID record from being updated with the caller’s
identity.
| func codexModelCatalogRoutingRequest(r *http.Request) *http.Request { | ||
| request := r.Clone(r.Context()) | ||
| request.Header = r.Header.Clone() | ||
| for _, header := range []string{ | ||
| "X-Subrouter-Account-ID", | ||
| "X-Subrouter-Account", | ||
| "X-Subrouter-Model", | ||
| "X-Model", | ||
| } { | ||
| request.Header.Del(header) | ||
| } | ||
| return request | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
A model query parameter on a Codex catalog request is never neutralized on either routing path. session.ExtractModel reads the model query parameter in addition to headers (session/extract.go:129-139). The unbound path sanitizes only headers, and the bound-lease path skips model validation entirely, so ?model= survives in both cases and reaches scheduler pool selection and the broker Model field.
internal/proxy/proxy.go#L4337-L4349: incodexModelCatalogRoutingRequest, delete themodelquery parameter from a clonedURLin addition to the routing headers.internal/proxy/session_lease.go#L563-L565: invalidateRequestModel, restrict the catalog early return to the body-model requirement, and keep the query-model conflict loop at lines 576-581 so a bound lease still rejects a conflicting?model=.
📍 Affects 2 files
internal/proxy/proxy.go#L4337-L4349(this comment)internal/proxy/session_lease.go#L563-L565
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/proxy/proxy.go` around lines 4337 - 4349, The Codex catalog routing
flow must neutralize query-based model overrides on both paths. In
internal/proxy/proxy.go lines 4337-4349, update codexModelCatalogRoutingRequest
to remove the model query parameter from the cloned URL as well as the routing
headers; in internal/proxy/session_lease.go lines 563-565, narrow
validateRequestModel’s catalog early return to only the body-model requirement
so the existing query-model conflict validation remains active for bound leases.
| if lease.Provider == accounts.ProviderCodex && codexModelCatalogRequest(r) { | ||
| return nil | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
The model bypass also removes the query-model conflict check.
The early return skips the query loop at lines 576-581. A bound-lease catalog request such as GET /v1/models?client_version=0.146.1&model=other-model is no longer checked against lease.Model. The lease account stays pinned by allowsAccount in internal/proxy/proxy.go at line 2383, so the account cannot be changed. The value still reaches session.ExtractModel and influences scheduler pool selection and the broker Model field.
Bypass only the body requirement, and keep the query conflict check.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/proxy/session_lease.go` around lines 563 - 565, Update the lease
validation flow around the Codex catalog-request handling so the request body
requirement is bypassed without returning before the query-model conflict check.
Preserve the existing query loop that validates the requested model against
lease.Model, while retaining the early bypass behavior only for body validation.
1 participant
Summary
client_versionquery/modelsrequests on their selected API-key accountVerification
{object,data}model list because it expects{models:[...]}go test ./...reached one pre-existing macOS deployment-script timeout; the failing test also times out alone on unchangedmain, while all proxy packages passedNeed help on this PR? Tag
@codesmith-botwith what you need. Autofix is disabled.Summary by cubic
Fixes Codex model catalog routing to always use ChatGPT OAuth and preserve the caller’s session assignment, preventing fallback to the incompatible API-key
/modelsresponse.client_versionand route them to the OAuth Codex upstream, using internal sessioninternal:codex-model-catalogwithout changing response assignment./modelsrequests on the selected API-key account whenclient_versionis absent.Written for commit bb29b01. Summary will update on new commits.
Summary by CodeRabbit