feat: support Codex/PI session headers for session affinity

Amp-Thread-ID: https://ampcode.com/threads/T-019dce25-c070-773a-ac52-11c541220b30
Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
XYenon
2026-04-27 17:08:49 +08:00
parent 34027da7f1
commit 3ac39dcc7d
4 changed files with 91 additions and 15 deletions
+54
View File
@@ -776,6 +776,46 @@ func TestExtractSessionID_Headers(t *testing.T) {
}
}
func TestExtractSessionID_CodexSessionIDHeader(t *testing.T) {
t.Parallel()
headers := make(http.Header)
headers.Set("Session_id", "codex-session-123")
got := ExtractSessionID(headers, nil, nil)
want := "codex:codex-session-123"
if got != want {
t.Errorf("ExtractSessionID() with Session_id = %q, want %q", got, want)
}
}
func TestExtractSessionID_ClientRequestIDHeader(t *testing.T) {
t.Parallel()
headers := make(http.Header)
headers.Set("X-Client-Request-Id", "pi-session-123")
got := ExtractSessionID(headers, nil, nil)
want := "clientreq:pi-session-123"
if got != want {
t.Errorf("ExtractSessionID() with X-Client-Request-Id = %q, want %q", got, want)
}
}
func TestExtractSessionID_CodexSessionIDPriorityOverClientRequestID(t *testing.T) {
t.Parallel()
headers := make(http.Header)
headers.Set("X-Client-Request-Id", "pi-session-123")
headers.Set("Session_id", "codex-session-456")
got := ExtractSessionID(headers, nil, nil)
want := "codex:codex-session-456"
if got != want {
t.Errorf("ExtractSessionID() = %q, want %q (Session_id should take priority over X-Client-Request-Id)", got, want)
}
}
func TestExtractSessionID_AmpThreadId(t *testing.T) {
t.Parallel()
@@ -789,6 +829,20 @@ func TestExtractSessionID_AmpThreadId(t *testing.T) {
}
}
func TestExtractSessionID_AmpThreadIdPriorityOverClientRequestID(t *testing.T) {
t.Parallel()
headers := make(http.Header)
headers.Set("X-Amp-Thread-Id", "T-priority-test")
headers.Set("X-Client-Request-Id", "pi-session-123")
got := ExtractSessionID(headers, nil, nil)
want := "amp:T-priority-test"
if got != want {
t.Errorf("ExtractSessionID() = %q, want %q (X-Amp-Thread-Id should take priority over X-Client-Request-Id)", 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) {