diff --git a/sdk/cliproxy/auth/home_session_alias_test.go b/sdk/cliproxy/auth/home_session_alias_test.go index 844cfe0c..b1788ac0 100644 --- a/sdk/cliproxy/auth/home_session_alias_test.go +++ b/sdk/cliproxy/auth/home_session_alias_test.go @@ -448,6 +448,50 @@ func TestPickNextViaHomePassesParentSessionIDToHierarchyDispatcher(t *testing.T) } } +func TestPickNextViaHomeNestedRequestSubagentHierarchy(t *testing.T) { + dispatcher := &hierarchyCaptureDispatcher{} + oldCurrentHomeDispatcher := currentHomeDispatcher + currentHomeDispatcher = func() homeAuthDispatcher { return dispatcher } + t.Cleanup(func() { currentHomeDispatcher = oldCurrentHomeDispatcher }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{ + Home: internalconfig.HomeConfig{Enabled: true}, + Routing: internalconfig.RoutingConfig{SessionAffinityTTL: "1h"}, + }) + manager.SetHomeExecutionRegistry(executionregistry.New()) + manager.RegisterExecutor(schedulerTestExecutor{provider: "hierarchy-provider"}) + + nestedSubOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{ + "request": { + "sessionId": "root-home-session", + "metadata": { + "agent_id": "worker-subagent" + } + } + }`), + Metadata: map[string]any{}, + } + + auth, _, _, errPick := manager.pickNextViaHome(context.Background(), "test-model", nestedSubOpts, nil) + if errPick != nil { + t.Fatalf("pickNextViaHome() error = %v", errPick) + } + if auth == nil || auth.ID != "hierarchy-auth" { + t.Fatalf("auth = %#v, want hierarchy-auth", auth) + } + + dispatcher.mu.Lock() + defer dispatcher.mu.Unlock() + if len(dispatcher.sessionIDs) != 1 || dispatcher.sessionIDs[0] != "session:root-home-session:agent:worker-subagent" { + t.Fatalf("dispatched sessionID = %v, want session:root-home-session:agent:worker-subagent", dispatcher.sessionIDs) + } + if len(dispatcher.parentIDs) != 1 || dispatcher.parentIDs[0] != "session:root-home-session" { + t.Fatalf("dispatched parentID = %v, want session:root-home-session", dispatcher.parentIDs) + } +} + func TestHomeDispatchSessionIDsExtractsParentFromHeaderPlusBody(t *testing.T) { manager := NewManager(nil, nil, nil) manager.SetConfig(&internalconfig.Config{ @@ -528,3 +572,49 @@ func TestHomeDispatchSessionIDsExtractsParentFromHeaderPlusBody(t *testing.T) { t.Fatalf("geminiParentID = %q, want geminicache:cache-parent-200", geminiParentID) } } + +func TestHomeDispatchSessionIDsNestedRequestSubagent(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{ + Home: internalconfig.HomeConfig{Enabled: true}, + Routing: internalconfig.RoutingConfig{SessionAffinityTTL: "1h"}, + }) + + // 1. Nested request with sessionId and agent_id + nestedAgentOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{ + "request": { + "sessionId": "root", + "metadata": { + "agent_id": "worker" + } + } + }`), + } + sessionID, parentID := manager.homeDispatchSessionIDs(nestedAgentOpts) + if sessionID != "session:root:agent:worker" { + t.Fatalf("sessionID = %q, want session:root:agent:worker", sessionID) + } + if parentID != "session:root" { + t.Fatalf("parentID = %q, want session:root", parentID) + } + + // 2. Nested request with sessionId and subagent_id + nestedSubagentOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{ + "request": { + "sessionId": "root", + "metadata": { + "subagent_id": "worker-sub" + } + } + }`), + } + subSessionID, subParentID := manager.homeDispatchSessionIDs(nestedSubagentOpts) + if subSessionID != "session:root:agent:worker-sub" { + t.Fatalf("subSessionID = %q, want session:root:agent:worker-sub", subSessionID) + } + if subParentID != "session:root" { + t.Fatalf("subParentID = %q, want session:root", subParentID) + } +} diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 36133f0b..bd1952eb 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -1359,12 +1359,14 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map // Extract parent candidate from payload once if payload is non-empty var root gjson.Result + var reqRoot gjson.Result + var hasNestedReq bool var parentIDCandidate string if len(payload) > 0 { root = util.ParseGJSONBytesNoCopy(payload) - reqRoot := root + reqRoot = root req := root.Get("request") - hasNestedReq := req.Exists() && !root.Get("contents").Exists() + hasNestedReq = req.Exists() && !root.Get("contents").Exists() if hasNestedReq { reqRoot = req } @@ -1400,6 +1402,12 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map if agentID == "" { agentID = normalizedSessionCandidate(root.Get("metadata.subagent_id").String()) } + if agentID == "" && hasNestedReq { + agentID = normalizedSessionCandidate(reqRoot.Get("metadata.agent_id").String()) + if agentID == "" { + agentID = normalizedSessionCandidate(reqRoot.Get("metadata.subagent_id").String()) + } + } } if agentID == "" { _, _, agentID = cliproxysession.ClaudeMetadataIdentities(payload) @@ -1430,6 +1438,12 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map if agentID == "" { agentID = normalizedSessionCandidate(root.Get("metadata.subagent_id").String()) } + if agentID == "" && hasNestedReq { + agentID = normalizedSessionCandidate(reqRoot.Get("metadata.agent_id").String()) + if agentID == "" { + agentID = normalizedSessionCandidate(reqRoot.Get("metadata.subagent_id").String()) + } + } } parentAgentID := sessionHeaderValue(headers, "X-Claude-Code-Parent-Agent-Id") if agentID != "" && agentID != "main" { @@ -1618,6 +1632,12 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map if agentID == "" { agentID = sessionHeaderValue(headers, "x-agent-id") } + if agentID == "" && hasNestedReq { + agentID = normalizedSessionCandidate(reqRoot.Get("metadata.agent_id").String()) + if agentID == "" { + agentID = normalizedSessionCandidate(reqRoot.Get("metadata.subagent_id").String()) + } + } for _, path := range []string{"session_id", "sessionId", "sessionID", "metadata.session_id", "extra_body.session_id"} { sid := normalizedSessionCandidate(root.Get(path).String()) if sid == "" && hasNestedReq { @@ -1651,12 +1671,18 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map conversationID = "conv:" + sid } } - pck := root.Get("prompt_cache_key") - if !pck.Exists() { - pck = root.Get("promptCacheKey") + pck := normalizedSessionCandidate(root.Get("prompt_cache_key").String()) + if pck == "" { + pck = normalizedSessionCandidate(root.Get("promptCacheKey").String()) } - if sid := normalizedSessionCandidate(pck.String()); sid != "" { - return "pck:" + sid, conversationID + if pck == "" && hasNestedReq { + pck = normalizedSessionCandidate(reqRoot.Get("prompt_cache_key").String()) + if pck == "" { + pck = normalizedSessionCandidate(reqRoot.Get("promptCacheKey").String()) + } + } + if pck != "" { + return "pck:" + pck, conversationID } if conversationID != "" { if parentIDCandidate != "" && ("conv:"+parentIDCandidate) != conversationID { @@ -1664,7 +1690,11 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map } return conversationID, "" } - if userID := normalizedSessionCandidate(root.Get("metadata.user_id").String()); userID != "" { + userID := normalizedSessionCandidate(root.Get("metadata.user_id").String()) + if userID == "" && hasNestedReq { + userID = normalizedSessionCandidate(reqRoot.Get("metadata.user_id").String()) + } + if userID != "" { return "user:" + userID, "" } for _, convPath := range []string{"conversation_id", "conversationId", "chat_id", "chatId", "metadata.conversation_id", "extra_body.conversation_id"} { diff --git a/sdk/cliproxy/auth/selector_antigravity_subagent_test.go b/sdk/cliproxy/auth/selector_antigravity_subagent_test.go index 1c45bcba..9c6c72d7 100644 --- a/sdk/cliproxy/auth/selector_antigravity_subagent_test.go +++ b/sdk/cliproxy/auth/selector_antigravity_subagent_test.go @@ -494,3 +494,160 @@ func TestSessionAffinityOtherGoogleProvidersSubagentIsolation(t *testing.T) { }) } } + +func TestSessionAffinityNestedAntigravitySubagentIsolation(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "auth-ag-1", Provider: "antigravity"}, + {ID: "auth-ag-2", Provider: "antigravity"}, + } + + // 1. Parent request with nested request payload binds to auth-ag-1 + parentOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{ + "request": { + "sessionId": "root-task" + } + }`), + Metadata: map[string]any{}, + } + parentAuth, errParent := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", parentOpts, auths) + if errParent != nil { + t.Fatalf("parent Pick() error = %v", errParent) + } + if parentAuth.ID != "auth-ag-1" { + t.Fatalf("parent auth = %q, want auth-ag-1", parentAuth.ID) + } + if got := parentOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]; got != "session:root-task" { + t.Fatalf("parent canonical session ID = %v, want session:root-task", got) + } + + // 2. Subagent 1 request with nested metadata.agent_id must NOT inherit parent's auth-ag-1 + subagent1Opts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{ + "request": { + "sessionId": "root-task", + "metadata": { + "agent_id": "worker-1" + } + } + }`), + Metadata: map[string]any{}, + } + subagent1Auth, errSub1 := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", subagent1Opts, auths) + if errSub1 != nil { + t.Fatalf("subagent 1 Pick() error = %v", errSub1) + } + if subagent1Auth.ID == parentAuth.ID { + t.Fatalf("subagent 1 incorrectly inherited parent auth %q; should be isolated to prevent 429 rate limit", parentAuth.ID) + } + if subagent1Auth.ID != "auth-ag-2" { + t.Fatalf("subagent 1 auth = %q, want auth-ag-2", subagent1Auth.ID) + } + if got := subagent1Opts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]; got != "session:root-task:agent:worker-1" { + t.Fatalf("subagent 1 canonical session ID = %v, want session:root-task:agent:worker-1", got) + } + + // 3. Subagent 1 turn 2 must stay sticky to auth-ag-2 + subagent1Turn2Opts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{ + "request": { + "sessionId": "root-task", + "metadata": { + "agent_id": "worker-1" + } + } + }`), + Metadata: map[string]any{}, + } + subagent1Turn2Auth, errSub1Turn2 := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", subagent1Turn2Opts, auths) + if errSub1Turn2 != nil { + t.Fatalf("subagent 1 turn 2 Pick() error = %v", errSub1Turn2) + } + if subagent1Turn2Auth.ID != "auth-ag-2" { + t.Fatalf("subagent 1 turn 2 did not stay sticky: got %q, want auth-ag-2", subagent1Turn2Auth.ID) + } + + // 4. Subagent 2 request with metadata.subagent_id must select next available credential (auth-ag-1) + subagent2Opts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{ + "request": { + "sessionId": "root-task", + "metadata": { + "subagent_id": "worker-2" + } + } + }`), + Metadata: map[string]any{}, + } + subagent2Auth, errSub2 := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", subagent2Opts, auths) + if errSub2 != nil { + t.Fatalf("subagent 2 Pick() error = %v", errSub2) + } + if subagent2Auth.ID != "auth-ag-1" { + t.Fatalf("subagent 2 auth = %q, want auth-ag-1", subagent2Auth.ID) + } + + // 5. Parent second turn must stay sticky to auth-ag-1 + parentTurn2Opts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{ + "request": { + "sessionId": "root-task" + } + }`), + Metadata: map[string]any{}, + } + parentTurn2Auth, errParent2 := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", parentTurn2Opts, auths) + if errParent2 != nil { + t.Fatalf("parent turn 2 Pick() error = %v", errParent2) + } + if parentTurn2Auth.ID != "auth-ag-1" { + t.Fatalf("parent turn 2 auth = %q, want auth-ag-1", parentTurn2Auth.ID) + } +} + +func TestSessionAffinityNestedGeminiProviderSubagentIsolation(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "gemini-auth-1", Provider: "gemini"}, + {ID: "gemini-auth-2", Provider: "gemini"}, + } + + parentOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"request":{"sessionId":"gemini-root"}}`), + Metadata: map[string]any{}, + } + parentAuth, errParent := selector.Pick(context.Background(), "gemini", "gemini-2.5-pro", parentOpts, auths) + if errParent != nil { + t.Fatalf("parent Pick() error = %v", errParent) + } + if parentAuth.ID != "gemini-auth-1" { + t.Fatalf("parent auth = %q, want gemini-auth-1", parentAuth.ID) + } + + subOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"request":{"sessionId":"gemini-root","metadata":{"agent_id":"gemini-worker"}}}`), + Metadata: map[string]any{}, + } + subAuth, errSub := selector.Pick(context.Background(), "gemini", "gemini-2.5-pro", subOpts, auths) + if errSub != nil { + t.Fatalf("subagent Pick() error = %v", errSub) + } + if subAuth.ID != "gemini-auth-2" { + t.Fatalf("subagent incorrectly picked parent gemini auth: got %q, want gemini-auth-2", subAuth.ID) + } +} diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index 546e628e..cab20798 100644 --- a/sdk/cliproxy/auth/selector_test.go +++ b/sdk/cliproxy/auth/selector_test.go @@ -1187,6 +1187,73 @@ func TestSessionAffinitySelector_FailoverWhenAuthUnavailable(t *testing.T) { } } +func TestExtractSessionID_NestedRequestSubagent(t *testing.T) { + t.Parallel() + + // 1. Nested request with sessionId and metadata.agent_id + payloadAgent := []byte(`{ + "request": { + "sessionId": "root", + "metadata": { + "agent_id": "worker" + } + } + }`) + if got := ExtractSessionID(nil, payloadAgent, nil); got != "session:root:agent:worker" { + t.Fatalf("ExtractSessionID() = %q, want session:root:agent:worker", got) + } + primary, fallback := extractExplicitSessionIDs(nil, payloadAgent, nil) + if primary != "session:root:agent:worker" || fallback != "session:root" { + t.Fatalf("extractExplicitSessionIDs() = (%q, %q), want (session:root:agent:worker, session:root)", primary, fallback) + } + + // 2. Nested request with sessionId and metadata.subagent_id + payloadSubagent := []byte(`{ + "request": { + "sessionId": "root", + "metadata": { + "subagent_id": "worker" + } + } + }`) + if got := ExtractSessionID(nil, payloadSubagent, nil); got != "session:root:agent:worker" { + t.Fatalf("ExtractSessionID() = %q, want session:root:agent:worker", got) + } + primary, fallback = extractExplicitSessionIDs(nil, payloadSubagent, nil) + if primary != "session:root:agent:worker" || fallback != "session:root" { + t.Fatalf("extractExplicitSessionIDs() = (%q, %q), want (session:root:agent:worker, session:root)", primary, fallback) + } + + // 3. Nested request with sessionId and parentSessionId + payloadParent := []byte(`{ + "request": { + "sessionId": "root", + "parentSessionId": "parent-root", + "metadata": { + "agent_id": "worker" + } + } + }`) + if got := ExtractSessionID(nil, payloadParent, nil); got != "session:root:agent:worker" { + t.Fatalf("ExtractSessionID() = %q, want session:root:agent:worker", got) + } + primary, fallback = extractExplicitSessionIDs(nil, payloadParent, nil) + if primary != "session:root:agent:worker" || fallback != "session:parent-root" { + t.Fatalf("extractExplicitSessionIDs() = (%q, %q), want (session:root:agent:worker, session:parent-root)", primary, fallback) + } + + // 4. Nested promptCacheKey when top-level prompt_cache_key is empty string + payloadPCKShadow := []byte(`{ + "prompt_cache_key": "", + "request": { + "promptCacheKey": "nested-pck-valid" + } + }`) + if got := ExtractSessionID(nil, payloadPCKShadow, nil); got != "pck:nested-pck-valid" { + t.Fatalf("ExtractSessionID() with shadowed empty top-level pck = %q, want pck:nested-pck-valid", got) + } +} + func TestExtractSessionID_ClaudeCodePriorityOverHeader(t *testing.T) { t.Parallel() diff --git a/sdk/cliproxy/session/identity.go b/sdk/cliproxy/session/identity.go index e092272e..85b99532 100644 --- a/sdk/cliproxy/session/identity.go +++ b/sdk/cliproxy/session/identity.go @@ -62,6 +62,12 @@ func ClaudeMetadataIdentities(payload []byte) (sessionID, parentSessionID, agent } root := util.ParseGJSONBytesNoCopy(payload) userID := strings.TrimSpace(root.Get("metadata.user_id").String()) + if userID == "" { + req := root.Get("request") + if req.Exists() && !root.Get("contents").Exists() { + userID = strings.TrimSpace(req.Get("metadata.user_id").String()) + } + } if userID == "" { return "", "", "" } @@ -182,6 +188,12 @@ func hasExplicitSession(headers map[string][]string, payload []byte) bool { // Parsing without copying matters here: this runs on every request and the // payload can be multiple megabytes. root := util.ParseGJSONBytesNoCopy(payload) + reqRoot := root + req := root.Get("request") + hasNestedReq := req.Exists() && !root.Get("contents").Exists() + if hasNestedReq { + reqRoot = req + } for _, path := range []string{ "session_id", "sessionId", @@ -210,15 +222,24 @@ func hasExplicitSession(headers map[string][]string, payload []byte) bool { if NormalizeExplicitID(root.Get(path).String()) != "" { return true } + if hasNestedReq && NormalizeExplicitID(reqRoot.Get(path).String()) != "" { + return true + } } if ClaudeMetadataSessionID(payload) != "" { return true } userID := strings.TrimSpace(root.Get("metadata.user_id").String()) + if userID == "" && hasNestedReq { + userID = strings.TrimSpace(reqRoot.Get("metadata.user_id").String()) + } if NormalizeExplicitID(userID) != "" { return true } conversation := root.Get("conversation") + if !conversation.Exists() && hasNestedReq { + conversation = reqRoot.Get("conversation") + } if NormalizeExplicitID(conversation.Get("id").String()) != "" { return true } diff --git a/sdk/cliproxy/session/identity_test.go b/sdk/cliproxy/session/identity_test.go index 6d441c3f..227908ab 100644 --- a/sdk/cliproxy/session/identity_test.go +++ b/sdk/cliproxy/session/identity_test.go @@ -196,6 +196,14 @@ func TestEnrichSkipsDerivationForExplicitSessions(t *testing.T) { cliproxyexecutor.DerivedSessionIDMetadataKey: "ctx:v1:stale", }, }, + { + name: "nested request sessionId", + payload: []byte(`{"request":{"sessionId":"nested-session"},"messages":[{"role":"user","content":"hello"}]}`), + }, + { + name: "nested request subagent", + payload: []byte(`{"request":{"sessionId":"nested-session","metadata":{"agent_id":"worker"}},"messages":[{"role":"user","content":"hello"}]}`), + }, } for _, test := range tests { diff --git a/sdk/cliproxy/session/info.go b/sdk/cliproxy/session/info.go index 41c8104e..79efc50f 100644 --- a/sdk/cliproxy/session/info.go +++ b/sdk/cliproxy/session/info.go @@ -94,6 +94,12 @@ func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string if agentID == "" { agentID = normalizedSessionCandidate(root.Get("metadata.subagent_id").String()) } + if agentID == "" && hasNestedReq { + agentID = normalizedSessionCandidate(reqRoot.Get("metadata.agent_id").String()) + if agentID == "" { + agentID = normalizedSessionCandidate(reqRoot.Get("metadata.subagent_id").String()) + } + } } if agentID == "" { _, _, agentID = ClaudeMetadataIdentities(payload) @@ -131,6 +137,12 @@ func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string if agentID == "" { agentID = normalizedSessionCandidate(root.Get("metadata.subagent_id").String()) } + if agentID == "" && hasNestedReq { + agentID = normalizedSessionCandidate(reqRoot.Get("metadata.agent_id").String()) + if agentID == "" { + agentID = normalizedSessionCandidate(reqRoot.Get("metadata.subagent_id").String()) + } + } } parentAgentID := sessionHeaderValue(headers, "X-Claude-Code-Parent-Agent-Id") if agentID != "" && agentID != "main" { @@ -433,13 +445,19 @@ func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string conversationID = "conv:" + sid } } - pck := root.Get("prompt_cache_key") - if !pck.Exists() { - pck = root.Get("promptCacheKey") + pck := normalizedSessionCandidate(root.Get("prompt_cache_key").String()) + if pck == "" { + pck = normalizedSessionCandidate(root.Get("promptCacheKey").String()) + } + if pck == "" && hasNestedReq { + pck = normalizedSessionCandidate(reqRoot.Get("prompt_cache_key").String()) + if pck == "" { + pck = normalizedSessionCandidate(reqRoot.Get("promptCacheKey").String()) + } } - if sid := normalizedSessionCandidate(pck.String()); sid != "" { + if pck != "" { info.ClientType = "generic" - info.SessionID = "pck:" + sid + info.SessionID = "pck:" + pck if conversationID != "" { info.ParentSessionID = conversationID info.AgentName = "subagent" @@ -461,7 +479,11 @@ func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string } // Plain metadata.user_id - if userID := normalizedSessionCandidate(root.Get("metadata.user_id").String()); userID != "" { + userID := normalizedSessionCandidate(root.Get("metadata.user_id").String()) + if userID == "" && hasNestedReq { + userID = normalizedSessionCandidate(reqRoot.Get("metadata.user_id").String()) + } + if userID != "" { info.ClientType = "generic" info.SessionID = "user:" + userID info.AgentName = "main" diff --git a/sdk/cliproxy/session/info_test.go b/sdk/cliproxy/session/info_test.go index a4980a72..00294a38 100644 --- a/sdk/cliproxy/session/info_test.go +++ b/sdk/cliproxy/session/info_test.go @@ -461,3 +461,75 @@ func TestExtractSessionInfoClaudeMetadataUserIDWithAgentHeader(t *testing.T) { t.Fatalf("AgentName = %q, want subagent-uuid-123", info.AgentName) } } + +func TestExtractSessionInfoNestedSubagentIDAndUserID(t *testing.T) { + t.Parallel() + + // 1. Nested request with subagent_id + payloadSub := []byte(`{ + "request": { + "sessionId": "main-sess-999", + "metadata": { + "subagent_id": "worker-sub-999" + } + } + }`) + info, ok := ExtractSessionInfo(nil, payloadSub, nil) + if !ok { + t.Fatal("ExtractSessionInfo failed on nested subagent_id request") + } + if info.SessionID != "session:main-sess-999:agent:worker-sub-999" { + t.Fatalf("SessionID = %q, want session:main-sess-999:agent:worker-sub-999", info.SessionID) + } + if info.ParentSessionID != "session:main-sess-999" { + t.Fatalf("ParentSessionID = %q, want session:main-sess-999", info.ParentSessionID) + } + if info.AgentName != "worker-sub-999" { + t.Fatalf("AgentName = %q, want worker-sub-999", info.AgentName) + } + + // 2. Nested request with plain user_id + payloadUser := []byte(`{ + "request": { + "metadata": { + "user_id": "nested-user-123" + } + } + }`) + infoUser, okUser := ExtractSessionInfo(nil, payloadUser, nil) + if !okUser { + t.Fatal("ExtractSessionInfo failed on nested user_id request") + } + if infoUser.SessionID != "user:nested-user-123" { + t.Fatalf("SessionID = %q, want user:nested-user-123", infoUser.SessionID) + } + + // 3. Nested request with promptCacheKey + payloadPCK := []byte(`{ + "request": { + "promptCacheKey": "nested-pck-456" + } + }`) + infoPCK, okPCK := ExtractSessionInfo(nil, payloadPCK, nil) + if !okPCK { + t.Fatal("ExtractSessionInfo failed on nested promptCacheKey request") + } + if infoPCK.SessionID != "pck:nested-pck-456" { + t.Fatalf("SessionID = %q, want pck:nested-pck-456", infoPCK.SessionID) + } + + // 4. Nested promptCacheKey when top-level prompt_cache_key is empty string + payloadPCKShadow := []byte(`{ + "prompt_cache_key": "", + "request": { + "promptCacheKey": "nested-pck-valid" + } + }`) + infoShadow, okShadow := ExtractSessionInfo(nil, payloadPCKShadow, nil) + if !okShadow { + t.Fatal("ExtractSessionInfo failed on shadowed promptCacheKey request") + } + if infoShadow.SessionID != "pck:nested-pck-valid" { + t.Fatalf("SessionID = %q, want pck:nested-pck-valid", infoShadow.SessionID) + } +}