feat: sticky session routing via X-Session-ID header#1998
feat: sticky session routing via X-Session-ID header#1998lichengzhe wants to merge 2 commits intorouter-for-me:mainfrom
Conversation
Route requests with the same X-Session-ID to the same auth account, leveraging the existing pinnedAuthID mechanism. Bindings expire after 1 hour and are cleared on 429/5xx errors for automatic failover. - New stickyStore (in-memory, sync.RWMutex) with TTL-based expiration - Integrated into executeMixedOnce and executeStreamMixedOnce - Explicit pinned_auth_id takes precedence over sticky binding - Periodic cleanup piggybacked on existing autoRefresh ticker - X-Session-ID header extracted in handler metadata layer Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a sticky session routing mechanism to improve consistency in multi-account setups. By allowing clients to opt-in via an Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
xkonjin
left a comment
xkonjin
left a comment
There was a problem hiding this comment.
Automated review focusing on bugs, security, and test coverage.
PR: feat: sticky session routing via X-Session-ID header (#1998)
Key observations:
- No test files changed. Consider adding or updating tests for modified code paths.
Suggestions:
- Add or update unit/integration tests for changed code paths.
- Run dependency and static security scans (pip-audit/npm audit/Semgrep) and address findings.
- Double check error handling and input validation on new code paths.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces sticky session routing using the X-Session-ID header, a valuable feature for multi-account setups, with a new stickyStore for session bindings and integrated logic in conductor.go. However, the current stickyStore implementation is vulnerable to memory exhaustion and denial-of-service (DoS) attacks due to its unbounded in-memory map. Additionally, the review identifies opportunities to improve maintainability and robustness by addressing code duplication and brittle logic in the conductor. It is recommended to implement a maximum size for the stickyStore and sanitize the X-Session-ID input header to mitigate these risks.
| func (s *stickyStore) Set(sessionID, authID string, ttl time.Duration) { | ||
| s.mu.Lock() | ||
| s.entries[sessionID] = stickyEntry{authID: authID, expiresAt: time.Now().Add(ttl)} | ||
| s.mu.Unlock() |
There was a problem hiding this comment.
The stickyStore uses an unbounded in-memory map to store session-to-auth bindings. Since the sessionID is derived from the user-controlled X-Session-ID header, an attacker can send a large number of requests with unique session IDs to exhaust the server's memory.
Furthermore, the Cleanup function (lines 52-61) iterates over the entire map while holding a write lock. If the map grows very large due to an attack, this cleanup process will block all other operations on the stickyStore for an extended period, leading to a denial of service.
There was a problem hiding this comment.
Fixed. Added maxEntries cap (10k) — new entries are silently dropped at capacity. Overwriting existing keys still works. Added TestStickyStore_MaxEntries covering both cases.
| func stickySessionIDFromHeader(ctx context.Context) string { | ||
| if ctx == nil { | ||
| return "" | ||
| } | ||
| ginCtx, ok := ctx.Value("gin").(*gin.Context) | ||
| if !ok || ginCtx == nil || ginCtx.Request == nil { | ||
| return "" | ||
| } | ||
| return strings.TrimSpace(ginCtx.GetHeader("X-Session-ID")) | ||
| } |
There was a problem hiding this comment.
The X-Session-ID header is extracted and used directly as a key in the stickyStore without any length validation or sanitization. This untrusted input contributes to the potential memory exhaustion vulnerability in the stickyStore.
There was a problem hiding this comment.
Fixed. Added 256-char length limit — headers exceeding it return empty (no binding created).
sdk/cliproxy/auth/conductor.go
Outdated
| // Sticky session: resolve session→auth binding before pick loop. | ||
| // An explicit pinned_auth_id takes precedence over sticky binding. | ||
| stickySessionID := stickySessionIDFromMetadata(opts.Metadata) | ||
| if stickySessionID != "" { | ||
| if _, alreadyPinned := opts.Metadata[cliproxyexecutor.PinnedAuthMetadataKey]; !alreadyPinned { | ||
| if boundAuth, found := m.sticky.Get(stickySessionID); found { | ||
| opts.Metadata[cliproxyexecutor.PinnedAuthMetadataKey] = boundAuth | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
This block of logic for resolving sticky sessions is duplicated in executeStreamMixedOnce (lines 1160-1169). To improve maintainability and reduce code duplication, consider extracting this logic into a private helper method on the Manager.
For example:
func (m *Manager) resolveStickySession(meta map[string]any) string {
stickySessionID := stickySessionIDFromMetadata(meta)
if stickySessionID != "" {
if _, alreadyPinned := meta[cliproxyexecutor.PinnedAuthMetadataKey]; !alreadyPinned {
if boundAuth, found := m.sticky.Get(stickySessionID); found {
meta[cliproxyexecutor.PinnedAuthMetadataKey] = boundAuth
}
}
}
return stickySessionID
}You could then call this helper in both executeMixedOnce and executeStreamMixedOnce like so: stickySessionID := m.resolveStickySession(opts.Metadata).
There was a problem hiding this comment.
Extracted resolveStickySession helper — both executeMixedOnce and executeStreamMixedOnce now call it.
sdk/cliproxy/auth/conductor.go
Outdated
| sc := 0 | ||
| if se, ok := errors.AsType[cliproxyexecutor.StatusError](errStream); ok && se != nil { | ||
| sc = se.StatusCode() | ||
| } | ||
| if sc == 429 || sc >= 500 { |
There was a problem hiding this comment.
This logic to extract the status code from an error appears to be a re-implementation of the existing statusCodeFromError helper function. Using the helper function directly would improve consistency and readability.
if sc := statusCodeFromError(errStream); sc == 429 || sc >= 500 {There was a problem hiding this comment.
Good catch — switched to statusCodeFromError(errStream).
sdk/cliproxy/auth/conductor.go
Outdated
| stickyCleanupCounter++ | ||
| if stickyCleanupCounter >= 60 { // ~every 5 min at default 5s interval | ||
| m.sticky.Cleanup() | ||
| stickyCleanupCounter = 0 | ||
| } |
There was a problem hiding this comment.
The current implementation for triggering m.sticky.Cleanup() relies on a counter and a magic number 60. As the comment indicates, this is tied to the default 5s interval and is brittle if the interval parameter changes. A more robust approach would be to base the cleanup on elapsed time, making it independent of the ticker's interval.
For example:
// Before the loop
lastStickyCleanup := time.Now()
const stickyCleanupInterval = 5 * time.Minute
// Inside the loop
if time.Since(lastStickyCleanup) >= stickyCleanupInterval {
m.sticky.Cleanup()
lastStickyCleanup = time.Now()
}There was a problem hiding this comment.
Replaced counter with time.Since(lastStickyCleanup) — now interval-independent.
- Cap stickyStore at 10k entries to prevent memory exhaustion (DoS) - Reject X-Session-ID headers longer than 256 chars - Extract resolveStickySession helper to reduce duplication - Use existing statusCodeFromError instead of manual extraction - Replace counter-based cleanup with time.Since for interval independence - Add TestStickyStore_MaxEntries covering capacity + overwrite behavior Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
luispater
left a comment
luispater
left a comment
There was a problem hiding this comment.
Summary:
This is a useful feature, but I found two blocking correctness issues in the sticky routing behavior.
Key findings:
- Blocking: sticky bindings are keyed only by session ID, not by provider/model. Reusing the same X-Session-ID across different providers or incompatible models can inject a stale pinned_auth_id and cause auth_not_found instead of falling back to normal selection.
- Blocking: the streaming path pins the session before the stream has actually succeeded. executeStreamWithModelPool can return a wrapped error stream with err == nil for bootstrap/empty-stream failures, and executeStreamMixedOnce still stores the sticky binding. Those bindings are only cleared for 429/5xx, so empty-stream and transport/bootstrap failures can poison the session.
Test plan:
- Ran: go test ./sdk/cliproxy/auth/...
Follow-ups:
- Add coverage for cross-provider/model reuse of X-Session-ID.
- Add coverage for bootstrap/empty-stream failures in the streaming path.
|
Additional concern: I recommend switching to a namespaced header such as |
3 participants
Summary
X-Session-IDheader to the same auth account, leveraging the existingpinnedAuthIDmechanismWhy
In multi-account setups, the same logical conversation can bounce between different auth accounts across requests. This causes issues with providers that track conversation state per-account. Sticky sessions give clients a simple opt-in mechanism to keep a session pinned to one account.
How
stickyStore(in-memory,sync.RWMutex) with TTL-based expiration insdk/cliproxy/auth/sticky.goexecuteMixedOnceandexecuteStreamMixedOnce— resolves binding before pick loop, binds on success, clears on 429/5xxpinned_auth_idtakes precedence over sticky bindingautoRefreshticker (~every 5 min)X-Session-IDheader extracted in handler metadata layer (sdk/api/handlers/handlers.go)Files Changed
sdk/cliproxy/auth/sticky.gosdk/cliproxy/auth/sticky_test.gosdk/cliproxy/auth/conductor.gosdk/cliproxy/executor/types.goStickySessionMetadataKeyconstantsdk/api/handlers/handlers.goX-Session-IDheader into metadataTest plan
go build ./cmd/server— compiles cleanly (Go 1.26)go test ./sdk/cliproxy/auth/ -run TestSticky— 7/7 passgo test ./sdk/cliproxy/auth/— all auth tests pass (no regression)curl -H "X-Session-ID: test-123"verifies routing affinitypr-test-build.ymlauto-triggered🤖 Generated with Claude Code