From 4d6457e6ec6ac6f8fcc1c31190a64121a3e3ca86 Mon Sep 17 00:00:00 2001 From: XYenon Date: Thu, 23 Apr 2026 13:47:22 +0800 Subject: [PATCH 1/2] feat: support extracting X-Amp-Thread-Id header as session id for session affinity --- sdk/cliproxy/auth/selector.go | 20 ++++++++++----- sdk/cliproxy/auth/selector_test.go | 40 ++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 6 deletions(-) diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 51275a311..f49979ce4 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -570,9 +570,10 @@ func (s *SessionAffinitySelector) InvalidateAuth(authID string) { // Priority order: // 1. metadata.user_id (Claude Code format with _session_{uuid}) - highest priority for Claude Code clients // 2. X-Session-ID header -// 3. metadata.user_id (non-Claude Code format) -// 4. conversation_id field in request body -// 5. Stable hash from first few messages content (fallback) +// 3. X-Amp-Thread-Id header (Amp CLI thread ID) +// 4. metadata.user_id (non-Claude Code format) +// 5. conversation_id field in request body +// 6. Stable hash from first few messages content (fallback) func ExtractSessionID(headers http.Header, payload []byte, metadata map[string]any) string { primary, _ := extractSessionIDs(headers, payload, metadata) return primary @@ -608,22 +609,29 @@ func extractSessionIDs(headers http.Header, payload []byte, metadata map[string] } } + // 3. X-Amp-Thread-Id header (Amp CLI thread ID) + if headers != nil { + if tid := headers.Get("X-Amp-Thread-Id"); tid != "" { + return "amp:" + tid, "" + } + } + if len(payload) == 0 { return "", "" } - // 3. metadata.user_id (non-Claude Code format) + // 4. metadata.user_id (non-Claude Code format) userID := gjson.GetBytes(payload, "metadata.user_id").String() if userID != "" { return "user:" + userID, "" } - // 4. conversation_id field + // 5. conversation_id field if convID := gjson.GetBytes(payload, "conversation_id").String(); convID != "" { return "conv:" + convID, "" } - // 5. Hash-based fallback from message content + // 6. Hash-based fallback from message content return extractMessageHashIDs(payload) } diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index 560d3b9e9..c3041b5ba 100644 --- a/sdk/cliproxy/auth/selector_test.go +++ b/sdk/cliproxy/auth/selector_test.go @@ -776,6 +776,46 @@ func TestExtractSessionID_Headers(t *testing.T) { } } +func TestExtractSessionID_AmpThreadId(t *testing.T) { + t.Parallel() + + headers := make(http.Header) + headers.Set("X-Amp-Thread-Id", "T-7873e6bd-6354-4a9a-be2c-c7702c6e1b64") + + got := ExtractSessionID(headers, nil, nil) + want := "amp:T-7873e6bd-6354-4a9a-be2c-c7702c6e1b64" + if got != want { + t.Errorf("ExtractSessionID() with X-Amp-Thread-Id = %q, want %q", got, want) + } +} + +// TestExtractSessionID_AmpThreadIdLowerPriority verifies X-Amp-Thread-Id is lower +// priority than Claude Code metadata.user_id but higher than conversation_id. +func TestExtractSessionID_AmpThreadIdPriority(t *testing.T) { + t.Parallel() + + // X-Amp-Thread-Id should be used when no Claude Code user_id is present + headers := make(http.Header) + headers.Set("X-Amp-Thread-Id", "T-priority-test") + + payload := []byte(`{"conversation_id":"conv-12345"}`) + got := ExtractSessionID(headers, payload, nil) + want := "amp:T-priority-test" + if got != want { + t.Errorf("ExtractSessionID() = %q, want %q (Amp thread ID should take priority over conversation_id)", got, want) + } + + // Claude Code user_id should take priority over X-Amp-Thread-Id + headers2 := make(http.Header) + headers2.Set("X-Amp-Thread-Id", "T-priority-test") + payload2 := []byte(`{"metadata":{"user_id":"user_xxx_account__session_ac980658-63bd-4fb3-97ba-8da64cb1e344"}}`) + got2 := ExtractSessionID(headers2, payload2, nil) + want2 := "claude:ac980658-63bd-4fb3-97ba-8da64cb1e344" + if got2 != want2 { + t.Errorf("ExtractSessionID() = %q, want %q (Claude Code should take priority over Amp thread ID)", got2, want2) + } +} + // TestExtractSessionID_IdempotencyKey verifies that idempotency_key is intentionally // ignored for session affinity (it's auto-generated per-request, causing cache misses). func TestExtractSessionID_IdempotencyKey(t *testing.T) { From 8e49c795f5ebddc0e9ec594c654d2aa23aeb4145 Mon Sep 17 00:00:00 2001 From: XYenon Date: Thu, 23 Apr 2026 15:26:14 +0800 Subject: [PATCH 2/2] fix: forward HTTP headers to executor Options so session affinity can read X-Amp-Thread-Id --- sdk/api/handlers/handlers.go | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index 1fda8f49f..369ab5a8d 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -211,6 +211,19 @@ func requestExecutionMetadata(ctx context.Context) map[string]any { return meta } +// headersFromContext extracts the original HTTP request headers from the gin context +// embedded in the provided context. This allows session affinity selectors to read +// client headers like X-Amp-Thread-Id. +func headersFromContext(ctx context.Context) http.Header { + if ctx == nil { + return nil + } + if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + return ginCtx.Request.Header.Clone() + } + return nil +} + func pinnedAuthIDFromContext(ctx context.Context) string { if ctx == nil { return "" @@ -488,6 +501,7 @@ func (h *BaseAPIHandler) ExecuteWithAuthManager(ctx context.Context, handlerType Alt: alt, OriginalRequest: rawJSON, SourceFormat: sdktranslator.FromString(handlerType), + Headers: headersFromContext(ctx), } opts.Metadata = reqMeta resp, err := h.AuthManager.Execute(ctx, providers, req, opts) @@ -535,6 +549,7 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle Alt: alt, OriginalRequest: rawJSON, SourceFormat: sdktranslator.FromString(handlerType), + Headers: headersFromContext(ctx), } opts.Metadata = reqMeta resp, err := h.AuthManager.ExecuteCount(ctx, providers, req, opts) @@ -586,6 +601,7 @@ func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handl Alt: alt, OriginalRequest: rawJSON, SourceFormat: sdktranslator.FromString(handlerType), + Headers: headersFromContext(ctx), } opts.Metadata = reqMeta streamResult, err := h.AuthManager.ExecuteStream(ctx, providers, req, opts)