feat: support extracting X-Amp-Thread-Id header as session id for session affinity

This commit is contained in:
XYenon
2026-04-23 13:47:22 +08:00
parent a188159632
commit 4d6457e6ec
2 changed files with 54 additions and 6 deletions

View File

@@ -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)
}

View File

@@ -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) {