From 5755d00b1bedca1e96dacaac1de4c90e3e0f4c38 Mon Sep 17 00:00:00 2001 From: sususu Date: Mon, 31 Aug 2026 19:59:16 +0800 Subject: [PATCH] refactor(session): prune in-memory session tree store in favor of flat KV affinity and pure info extraction - Replace complex materialized tree structures with lightweight (session_id, parent_id) extraction - Maintain O(1) flat KV cache for subagent-to-parent prompt cache inheritance - Add tree_compat.go with deprecated stubs to preserve SDK compatibility - Add comprehensive regression tests for multi-protocol session headers and payloads --- sdk/cliproxy/auth/selector.go | 21 +- sdk/cliproxy/auth/selector_lcp_test.go | 55 +- sdk/cliproxy/session/info.go | 524 ++++++++++++++ sdk/cliproxy/session/info_test.go | 434 +++++++++++ sdk/cliproxy/session/tree.go | 953 ------------------------- sdk/cliproxy/session/tree_compat.go | 204 ++++++ sdk/cliproxy/session/tree_test.go | 736 ------------------- 7 files changed, 1184 insertions(+), 1743 deletions(-) create mode 100644 sdk/cliproxy/session/info.go create mode 100644 sdk/cliproxy/session/info_test.go delete mode 100644 sdk/cliproxy/session/tree.go create mode 100644 sdk/cliproxy/session/tree_compat.go delete mode 100644 sdk/cliproxy/session/tree_test.go diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 80cc89b4..e488878a 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -661,7 +661,6 @@ type SessionAffinitySelector struct { fallback Selector cache *SessionCache matcher *cliproxysession.MerklePrefixMatcher - trees *cliproxysession.InMemorySessionTreeStore } // SessionAffinityConfig configures the session affinity selector. @@ -690,16 +689,13 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff fallback: cfg.Fallback, cache: NewSessionCache(cfg.TTL), matcher: cliproxysession.NewMerklePrefixMatcher(cfg.TTL), - trees: cliproxysession.NewInMemorySessionTreeStore(0, cfg.TTL), } } -// Trees returns the in-memory session tree store. +// Trees returns a backward-compatible in-memory session tree store. +// Deprecated: Session tree management has moved to Home. func (s *SessionAffinitySelector) Trees() *cliproxysession.InMemorySessionTreeStore { - if s == nil { - return nil - } - return s.trees + return cliproxysession.NewInMemorySessionTreeStore(0, time.Hour) } // Pick selects an auth with session affinity when possible. @@ -773,14 +769,6 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } else { s.cache.Set(cacheKey, authID) } - if s.trees != nil { - if treeInfo, ok := cliproxysession.ExtractTreeInfo(opts.Headers, opts.OriginalRequest, opts.Metadata); ok { - treeInfo.AuthID = authID - treeInfo.Provider = provider - treeInfo.Model = model - s.trees.RecordNode(treeInfo) - } - } } if cachedAuthID, ok := s.cache.GetAndRefresh(cacheKey); ok { @@ -962,9 +950,6 @@ func (s *SessionAffinitySelector) Stop() { if s.matcher != nil { s.matcher.Clear() } - if s.trees != nil { - s.trees.Clear() - } } // InvalidateAuth removes all session bindings for a specific auth. diff --git a/sdk/cliproxy/auth/selector_lcp_test.go b/sdk/cliproxy/auth/selector_lcp_test.go index 82f9316d..740adae5 100644 --- a/sdk/cliproxy/auth/selector_lcp_test.go +++ b/sdk/cliproxy/auth/selector_lcp_test.go @@ -612,7 +612,7 @@ func TestSessionAffinityExtendedHeadersAndPayloadIdentities(t *testing.T) { } } -func TestSessionAffinitySelectorSessionTreeIntegration(t *testing.T) { +func TestSessionAffinitySelectorSubagentInheritance(t *testing.T) { t.Parallel() selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ @@ -636,7 +636,7 @@ func TestSessionAffinitySelectorSessionTreeIntegration(t *testing.T) { t.Fatalf("root Pick() error = %v", errRoot) } - // 2. Subagent 1 Request (Claude Code Subagent) + // 2. Subagent 1 Request (Claude Code Subagent) inherits Root Auth sub1Opts := cliproxyexecutor.Options{ Headers: http.Header{ "X-Claude-Code-Session-Id": []string{"tree-root-sess"}, @@ -652,44 +652,27 @@ func TestSessionAffinitySelectorSessionTreeIntegration(t *testing.T) { t.Fatalf("sub1 Pick() error = %v", errSub1) } if sub1Auth.ID != rootAuth.ID { - t.Fatalf("sub1 did not inherit root auth") + t.Fatalf("sub1 did not inherit root auth: got %s, want %s", sub1Auth.ID, rootAuth.ID) } - // 3. Verify SessionTreeStore in selector - trees := selector.Trees() - if trees == nil { - t.Fatalf("expected selector.Trees() to be non-nil") - } - - rootNode, ok := trees.GetNode("claude:tree-root-sess") - if !ok || rootNode == nil { - t.Fatalf("expected rootNode to be found in trees") - } - if rootNode.TreeDepth != 0 || rootNode.TreePath != "claude:tree-root-sess" { - t.Errorf("rootNode hierarchy mismatch: depth=%d path=%q", rootNode.TreeDepth, rootNode.TreePath) - } - if rootNode.LastAuthID != rootAuth.ID || rootNode.LastProvider != "claude" { - t.Errorf("rootNode affinity mismatch: %+v", rootNode) - } - - subNode, ok := trees.GetNode("claude:tree-root-sess:agent:checker-agent") - if !ok || subNode == nil { - t.Fatalf("expected subNode to be found in trees") - } - if subNode.TreeDepth != 1 || subNode.TreePath != "claude:tree-root-sess/claude:tree-root-sess:agent:checker-agent" { - t.Errorf("subNode hierarchy mismatch: depth=%d path=%q", subNode.TreeDepth, subNode.TreePath) - } - if subNode.ParentSessionID != "claude:tree-root-sess" || subNode.RootSessionID != "claude:tree-root-sess" { - t.Errorf("subNode parent/root mismatch: %+v", subNode) + // 3. Subagent 2 Request with explicit parent agent inherits same auth + sub2Opts := cliproxyexecutor.Options{ + Headers: http.Header{ + "X-Claude-Code-Session-Id": []string{"tree-root-sess"}, + "X-Claude-Code-Agent-Id": []string{"leaf-agent"}, + "X-Claude-Code-Parent-Agent-Id": []string{"checker-agent"}, + }, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"run leaf checker"}]}`), + Metadata: map[string]any{ + cliproxyexecutor.CallerScopeMetadataKey: "scope-corp", + }, } - - // 4. Test GetTree for root - fullTree := trees.GetTree("claude:tree-root-sess") - if len(fullTree) != 2 { - t.Fatalf("GetTree returned %d nodes, want 2", len(fullTree)) + sub2Auth, errSub2 := selector.Pick(context.Background(), "claude", "claude-3-7-sonnet", sub2Opts, auths) + if errSub2 != nil { + t.Fatalf("sub2 Pick() error = %v", errSub2) } - if fullTree[0].SessionID != "claude:tree-root-sess" || fullTree[1].SessionID != "claude:tree-root-sess:agent:checker-agent" { - t.Errorf("GetTree nodes ordering mismatch: %+v", fullTree) + if sub2Auth.ID != rootAuth.ID { + t.Fatalf("sub2 did not inherit root auth: got %s, want %s", sub2Auth.ID, rootAuth.ID) } } diff --git a/sdk/cliproxy/session/info.go b/sdk/cliproxy/session/info.go new file mode 100644 index 00000000..2e80c63f --- /dev/null +++ b/sdk/cliproxy/session/info.go @@ -0,0 +1,524 @@ +// Package session derives stable conversation identities and extracts hierarchical session relationships. +package session + +import ( + "net/http" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/tidwall/gjson" +) + +// SessionInfo encapsulates incoming request details needed for session affinity and upstream reporting. +type SessionInfo struct { + SessionID string `json:"session_id"` + ParentSessionID string `json:"parent_session_id,omitempty"` + AgentName string `json:"agent_name,omitempty"` + ClientType string `json:"client_type,omitempty"` + CallerScope string `json:"caller_scope,omitempty"` + AuthID string `json:"auth_id,omitempty"` + Provider string `json:"provider,omitempty"` + Model string `json:"model,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` +} + +// SessionTreeInfo is an alias for SessionInfo for backward compatibility. +type SessionTreeInfo = SessionInfo + +// ExtractSessionInfo extracts session hierarchy and client identification from request attributes. +// Priority matches selector.go: +// 1. X-Claude-Code-Session-Id +// 2. Claude Code metadata.user_id session +// 3. Session-Id / Session_id (Codex and compatible clients) +// 4. X-Http-Session-Id (Antigravity CLI) +// 5. X-Session-ID / X-Session-Affinity / X-Slot-Session-Id +// 6. X-Conversation-Id / X-Thread-Id / X-Client-Request-Id +// 7. Gemini cachedContent +// 8. OpenAI thread_id +// 9. session_id / sessionId +// 10. prompt_cache_key (pck:), conversation.id (conv:), metadata.user_id (user:) +// 11. conversation_id / chat_id +// 12. execution_session_id metadata +func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string]any) (SessionInfo, bool) { + var info SessionInfo + if metadata != nil { + if scope, ok := metadata[cliproxyexecutor.CallerScopeMetadataKey].(string); ok { + info.CallerScope = strings.TrimSpace(scope) + } + } + + var root gjson.Result + var reqRoot gjson.Result + var hasNestedReq bool + var parentCandidate string + + if len(payload) > 0 { + root = util.ParseGJSONBytesNoCopy(payload) + reqRoot = root + req := root.Get("request") + hasNestedReq = req.Exists() && !root.Get("contents").Exists() + if hasNestedReq { + reqRoot = req + } + for _, p := range []string{ + "parent_session_id", "parentSessionId", + "parent_thread_id", "parentThreadId", + "forked_from_thread_id", "forked_from_id", + "parent_conversation_id", "parentConversationId", + "metadata.parent_session_id", "metadata.parent_thread_id", + "extra_body.parent_session_id", "extra_body.parent_thread_id", + } { + if val := normalizedSessionCandidate(root.Get(p).String()); val != "" { + parentCandidate = val + break + } + if hasNestedReq { + if val := normalizedSessionCandidate(reqRoot.Get(p).String()); val != "" { + parentCandidate = val + break + } + } + } + if parentCandidate == "" { + parentCandidate = ClaudeMetadataParentSessionID(payload) + } + } + + // 1. Anthropic / Claude Code Headers + if sid := sessionHeaderValue(headers, "X-Claude-Code-Session-Id"); sid != "" { + info.ClientType = "claude" + agentID := sessionHeaderValue(headers, "X-Claude-Code-Agent-Id") + parentAgentID := sessionHeaderValue(headers, "X-Claude-Code-Parent-Agent-Id") + if agentID != "" && agentID != "main" { + info.AgentName = agentID + info.ParentSessionID = "claude:" + sid + if parentAgentID != "" && parentAgentID != "main" && parentAgentID != agentID { + info.ParentSessionID = "claude:" + sid + ":agent:" + parentAgentID + } else if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "claude:" + parentCandidate + } + info.SessionID = "claude:" + sid + ":agent:" + agentID + } else { + info.AgentName = "main" + info.SessionID = "claude:" + sid + if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "claude:" + parentCandidate + info.AgentName = "subagent" + } + } + return finalizeSessionInfo(info) + } + + // 2. Claude Code metadata.user_id in payload (outranks generic headers) + if len(payload) > 0 { + if sid, parentSID, agentID := ClaudeMetadataIdentities(payload); sid != "" { + info.ClientType = "claude" + if agentID == "" && root.Exists() { + agentID = normalizedSessionCandidate(root.Get("metadata.agent_id").String()) + } + if agentID != "" && agentID != "main" { + info.SessionID = "claude:" + sid + ":agent:" + agentID + info.ParentSessionID = "claude:" + sid + if parentSID != "" && parentSID != sid { + info.ParentSessionID = "claude:" + parentSID + } else if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "claude:" + parentCandidate + } + info.AgentName = agentID + } else { + info.SessionID = "claude:" + sid + if parentSID != "" && parentSID != sid { + info.ParentSessionID = "claude:" + parentSID + info.AgentName = "subagent" + } else if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "claude:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + } + return finalizeSessionInfo(info) + } + } + + // 3. OpenAI / Codex CLI Headers + if sid := sessionHeaderValue(headers, "Session-Id"); sid != "" { + info.ClientType = "codex" + info.SessionID = "codex:" + sid + parentThread := sessionHeaderValue(headers, "x-codex-parent-thread-id") + if parentThread == "" { + parentThread = sessionHeaderValue(headers, "X-Codex-Parent-Thread-Id") + } + if parentThread != "" && parentThread != sid { + info.ParentSessionID = "codex:" + parentThread + info.AgentName = "subagent" + } else if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "codex:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + if sid := sessionHeaderValue(headers, "Session_id"); sid != "" { + info.ClientType = "codex" + info.SessionID = "codex:" + sid + parentThread := sessionHeaderValue(headers, "x-codex-parent-thread-id") + if parentThread == "" { + parentThread = sessionHeaderValue(headers, "X-Codex-Parent-Thread-Id") + } + if parentThread != "" && parentThread != sid { + info.ParentSessionID = "codex:" + parentThread + info.AgentName = "subagent" + } else if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "codex:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + + // 4. Antigravity CLI (agy) Headers + if sid := sessionHeaderValue(headers, "X-Http-Session-Id"); sid != "" { + info.ClientType = "agy" + info.SessionID = "agy:" + sid + parentSID := sessionHeaderValue(headers, "X-Parent-Session-ID") + if parentSID == "" { + parentSID = sessionHeaderValue(headers, "X-Parent-Session-Id") + } + if parentSID != "" && parentSID != sid { + info.ParentSessionID = "agy:" + parentSID + info.AgentName = "subagent" + } else if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "agy:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + + // 5. OpenCode / Pi Slot / Generic Headers + if sid := sessionHeaderValue(headers, "X-Session-ID"); sid != "" { + info.ClientType = "generic" + info.SessionID = "header:" + sid + parentSID := sessionHeaderValue(headers, "X-Parent-Session-ID") + if parentSID == "" { + parentSID = sessionHeaderValue(headers, "X-Parent-Session-Id") + } + if parentSID != "" && parentSID != sid { + info.ParentSessionID = "header:" + parentSID + info.AgentName = "subagent" + } else if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "header:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + if sid := sessionHeaderValue(headers, "X-Session-Affinity"); sid != "" { + info.ClientType = "opencode" + info.SessionID = "affinity:" + sid + parentAffinity := sessionHeaderValue(headers, "X-Parent-Session-Affinity") + if parentAffinity == "" { + parentAffinity = sessionHeaderValue(headers, "X-Parent-Session-ID") + } + if parentAffinity != "" && parentAffinity != sid { + info.ParentSessionID = "affinity:" + parentAffinity + info.AgentName = "subagent" + } else if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "affinity:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + if sid := sessionHeaderValue(headers, "X-Slot-Session-Id"); sid != "" { + info.ClientType = "pi" + info.SessionID = "slot:" + sid + parentSID := sessionHeaderValue(headers, "X-Parent-Session-ID") + if parentSID == "" { + parentSID = sessionHeaderValue(headers, "X-Parent-Session-Id") + } + if parentSID != "" && parentSID != sid { + info.ParentSessionID = "slot:" + parentSID + info.AgentName = "subagent" + } else if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "slot:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "slot" + } + return finalizeSessionInfo(info) + } + if sid := sessionHeaderValue(headers, "X-Conversation-Id"); sid != "" { + info.ClientType = "conv" + info.SessionID = "conv:" + sid + if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "conv:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + if sid := sessionHeaderValue(headers, "X-Conversation-ID"); sid != "" { + info.ClientType = "conv" + info.SessionID = "conv:" + sid + if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "conv:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + if sid := sessionHeaderValue(headers, "X-Thread-Id"); sid != "" { + info.ClientType = "openai-thread" + info.SessionID = "thread:" + sid + if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "thread:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + if sid := sessionHeaderValue(headers, "X-Thread-ID"); sid != "" { + info.ClientType = "openai-thread" + info.SessionID = "thread:" + sid + if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "thread:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + if sid := sessionHeaderValue(headers, "Thread-Id"); sid != "" { + info.ClientType = "openai-thread" + info.SessionID = "thread:" + sid + if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "thread:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + if sid := sessionHeaderValue(headers, "X-Client-Request-Id"); sid != "" { + info.ClientType = "generic" + info.SessionID = "clientreq:" + sid + info.AgentName = "main" + return finalizeSessionInfo(info) + } + + // 6. Payload inspection + if len(payload) > 0 && root.Exists() { + // Gemini context caching + for _, cachePath := range []string{"cachedContent", "cached_content"} { + cacheID := normalizedSessionCandidate(root.Get(cachePath).String()) + if cacheID == "" && hasNestedReq { + cacheID = normalizedSessionCandidate(reqRoot.Get(cachePath).String()) + } + if cacheID != "" { + info.ClientType = "gemini" + info.SessionID = "geminicache:" + cacheID + if parentCandidate != "" && parentCandidate != cacheID { + info.ParentSessionID = "geminicache:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + } + + // OpenAI thread in payload + for _, threadPath := range []string{"thread_id", "threadId", "metadata.thread_id"} { + tid := normalizedSessionCandidate(root.Get(threadPath).String()) + if tid == "" && hasNestedReq { + tid = normalizedSessionCandidate(reqRoot.Get(threadPath).String()) + } + if tid != "" { + info.ClientType = "openai-thread" + info.SessionID = "thread:" + tid + if parentCandidate != "" && parentCandidate != tid { + info.ParentSessionID = "thread:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + } + + // Generic session in payload + agentID := normalizedSessionCandidate(root.Get("metadata.agent_id").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()) + } + } + + for _, path := range []string{"session_id", "sessionId", "sessionID", "metadata.session_id", "extra_body.session_id"} { + sid := normalizedSessionCandidate(root.Get(path).String()) + if sid == "" && hasNestedReq { + sid = normalizedSessionCandidate(reqRoot.Get(path).String()) + } + if sid != "" { + info.ClientType = "generic" + if agentID != "" && agentID != "main" { + info.SessionID = "session:" + sid + ":agent:" + agentID + info.ParentSessionID = "session:" + sid + if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "session:" + parentCandidate + } + info.AgentName = agentID + } else { + info.SessionID = "session:" + sid + if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "session:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + } + return finalizeSessionInfo(info) + } + } + + // Prompt cache key & Conversation object + var conversationID string + conversation := root.Get("conversation") + if !conversation.Exists() && hasNestedReq { + conversation = reqRoot.Get("conversation") + } + if sid := normalizedSessionCandidate(conversation.Get("id").String()); sid != "" { + conversationID = "conv:" + sid + } else if conversation.Type == gjson.String { + if sid := normalizedSessionCandidate(conversation.String()); sid != "" { + conversationID = "conv:" + sid + } + } + pck := root.Get("prompt_cache_key") + if !pck.Exists() { + pck = root.Get("promptCacheKey") + } + if sid := normalizedSessionCandidate(pck.String()); sid != "" { + info.ClientType = "generic" + info.SessionID = "pck:" + sid + if conversationID != "" { + info.ParentSessionID = conversationID + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + if conversationID != "" { + info.ClientType = "conv" + info.SessionID = conversationID + if parentCandidate != "" && ("conv:"+parentCandidate) != conversationID { + info.ParentSessionID = "conv:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + + // Plain metadata.user_id + if userID := normalizedSessionCandidate(root.Get("metadata.user_id").String()); userID != "" { + info.ClientType = "generic" + info.SessionID = "user:" + userID + info.AgentName = "main" + return finalizeSessionInfo(info) + } + + // Legacy conversation string paths + for _, convPath := range []string{"conversation_id", "conversationId", "chat_id", "chatId", "metadata.conversation_id", "extra_body.conversation_id"} { + cid := normalizedSessionCandidate(root.Get(convPath).String()) + if cid == "" && hasNestedReq { + cid = normalizedSessionCandidate(reqRoot.Get(convPath).String()) + } + if cid != "" { + info.ClientType = "conv" + info.SessionID = "conv:" + cid + if parentCandidate != "" && ("conv:"+parentCandidate) != ("conv:"+cid) { + info.ParentSessionID = "conv:" + parentCandidate + info.AgentName = "subagent" + } else { + info.AgentName = "main" + } + return finalizeSessionInfo(info) + } + } + } + + // 7. ExecutionSessionMetadataKey + if executionID, ok := metadata[cliproxyexecutor.ExecutionSessionMetadataKey].(string); ok { + if executionID = normalizedSessionCandidate(executionID); executionID != "" { + info.ClientType = "generic" + info.SessionID = "execution:" + executionID + info.AgentName = "main" + return finalizeSessionInfo(info) + } + } + + return SessionInfo{}, false +} + +func finalizeSessionInfo(info SessionInfo) (SessionInfo, bool) { + if info.SessionID == "" { + return SessionInfo{}, false + } + if info.AgentName == "" { + info.AgentName = "main" + } + if info.ClientType == "" { + info.ClientType = "generic" + } + // Self-referential loop protection + if info.ParentSessionID == info.SessionID { + info.ParentSessionID = "" + } + return info, true +} + +// ExtractTreeInfo is an alias for ExtractSessionInfo for backward compatibility. +func ExtractTreeInfo(headers http.Header, payload []byte, metadata map[string]any) (SessionInfo, bool) { + return ExtractSessionInfo(headers, payload, metadata) +} + +func sessionHeaderValue(headers http.Header, name string) string { + if headers == nil { + return "" + } + if value := normalizedSessionCandidate(headers.Get(name)); value != "" { + return value + } + for key, values := range headers { + if !strings.EqualFold(key, name) { + continue + } + for _, value := range values { + if normalized := normalizedSessionCandidate(value); normalized != "" { + return normalized + } + } + } + return "" +} + +func normalizedSessionCandidate(raw string) string { + return NormalizeExplicitID(raw) +} diff --git a/sdk/cliproxy/session/info_test.go b/sdk/cliproxy/session/info_test.go new file mode 100644 index 00000000..a0ab4cf6 --- /dev/null +++ b/sdk/cliproxy/session/info_test.go @@ -0,0 +1,434 @@ +package session + +import ( + "net/http" + "testing" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func TestExtractSessionInfoAllClients(t *testing.T) { + t.Parallel() + + // 1. Claude Code with Subagent + claudeHeaders := http.Header{ + "X-Claude-Code-Session-Id": []string{"claude-root-123"}, + "X-Claude-Code-Agent-Id": []string{"subagent-checker"}, + } + info, ok := ExtractSessionInfo(claudeHeaders, nil, nil) + if !ok { + t.Fatalf("ExtractSessionInfo failed for Claude Code") + } + if info.ClientType != "claude" { + t.Errorf("ClientType = %q, want claude", info.ClientType) + } + if info.SessionID != "claude:claude-root-123:agent:subagent-checker" { + t.Errorf("SessionID = %q", info.SessionID) + } + if info.ParentSessionID != "claude:claude-root-123" { + t.Errorf("ParentSessionID = %q", info.ParentSessionID) + } + if info.AgentName != "subagent-checker" { + t.Errorf("AgentName = %q", info.AgentName) + } + + // 2. Codex CLI with parent thread + codexHeaders := http.Header{ + "Session-Id": []string{"codex-child-555"}, + "x-codex-parent-thread-id": []string{"codex-parent-111"}, + } + info, ok = ExtractSessionInfo(codexHeaders, nil, nil) + if !ok || info.ClientType != "codex" { + t.Fatalf("ExtractSessionInfo failed for Codex CLI") + } + if info.SessionID != "codex:codex-child-555" || info.ParentSessionID != "codex:codex-parent-111" { + t.Errorf("Codex session ids mismatch: %+v", info) + } + + // 3. Pi Slot Session + piHeaders := http.Header{ + "X-Slot-Session-Id": []string{"pi-slot-777"}, + } + info, ok = ExtractSessionInfo(piHeaders, nil, nil) + if !ok || info.ClientType != "pi" || info.SessionID != "slot:pi-slot-777" { + t.Errorf("Pi slot mismatch: %+v", info) + } + + // 4. OpenCode Session Affinity & Parent + opencodeHeaders := http.Header{ + "X-Session-Affinity": []string{"oc-child-999"}, + "X-Parent-Session-Affinity": []string{"oc-parent-333"}, + } + info, ok = ExtractSessionInfo(opencodeHeaders, nil, nil) + if !ok || info.ClientType != "opencode" || info.SessionID != "affinity:oc-child-999" || info.ParentSessionID != "affinity:oc-parent-333" { + t.Errorf("OpenCode mismatch: %+v", info) + } + + // 5. Antigravity X-Http-Session-Id + agyHeaders := http.Header{ + "X-Http-Session-Id": []string{"agy-sess-888"}, + } + info, ok = ExtractSessionInfo(agyHeaders, nil, nil) + if !ok || info.ClientType != "agy" || info.SessionID != "agy:agy-sess-888" { + t.Errorf("Antigravity mismatch: %+v", info) + } + + // 6. Payload with metadata.agent_id and parent_session_id + payload := []byte(`{ + "session_id": "payload-child-10", + "parent_session_id": "payload-parent-01", + "metadata": { + "agent_id": "analyzer" + } + }`) + metadata := map[string]any{ + cliproxyexecutor.CallerScopeMetadataKey: "test-scope", + } + info, ok = ExtractSessionInfo(nil, payload, metadata) + if !ok || info.ClientType != "generic" { + t.Fatalf("ExtractSessionInfo failed for payload") + } + if info.SessionID != "session:payload-child-10:agent:analyzer" { + t.Errorf("SessionID = %q", info.SessionID) + } + if info.ParentSessionID != "session:payload-parent-01" { + t.Errorf("ParentSessionID = %q", info.ParentSessionID) + } + if info.CallerScope != "test-scope" { + t.Errorf("CallerScope = %q", info.CallerScope) + } +} + +func TestExtractSessionInfoCanonicalizesPayloadParentForHeaderSessions(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + headers http.Header + wantParent string + }{ + { + name: "generic header", + headers: http.Header{"X-Session-ID": []string{"child"}}, + wantParent: "header:parent", + }, + { + name: "codex header", + headers: http.Header{"Session-Id": []string{"child"}}, + wantParent: "codex:parent", + }, + { + name: "claude header", + headers: http.Header{"X-Claude-Code-Session-Id": []string{"child"}}, + wantParent: "claude:parent", + }, + } + + payload := []byte(`{"parent_session_id":"parent"}`) + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + info, ok := ExtractSessionInfo(test.headers, payload, nil) + if !ok { + t.Fatal("ExtractSessionInfo() returned no session") + } + if info.ParentSessionID != test.wantParent { + t.Fatalf("ParentSessionID = %q, want %q", info.ParentSessionID, test.wantParent) + } + }) + } +} + +func TestExtractSessionInfoRejectsControlCharacters(t *testing.T) { + t.Parallel() + + payloadWithNewline := []byte(`{"session_id": "test\nsession"}`) + if _, ok := ExtractSessionInfo(nil, payloadWithNewline, nil); ok { + t.Fatal("ExtractSessionInfo accepted session_id with newline") + } + + payloadWithNull := []byte(`{"session_id": "test\u0000session"}`) + if _, ok := ExtractSessionInfo(nil, payloadWithNull, nil); ok { + t.Fatal("ExtractSessionInfo accepted session_id with null byte") + } +} + +func TestExtractSessionInfoClaudeNestedParent(t *testing.T) { + t.Parallel() + + // Nested metadata.user_id containing session_id and parent_session_id + payload := []byte(`{ + "metadata": { + "user_id": "{\"session_id\":\"child-session-123\",\"parent_session_id\":\"parent-session-456\",\"agent_id\":\"subagent-worker\"}" + } + }`) + info, ok := ExtractSessionInfo(nil, payload, nil) + if !ok { + t.Fatal("ExtractSessionInfo failed on nested Claude user_id") + } + if info.SessionID != "claude:child-session-123:agent:subagent-worker" { + t.Fatalf("expected child session with agent, got %q", info.SessionID) + } + if info.ParentSessionID != "claude:parent-session-456" { + t.Fatalf("expected parent claude:parent-session-456, got %q", info.ParentSessionID) + } + + // Explicit header + payload body parent_session_id + headerPayload := []byte(`{"parent_session_id":"parent-session-789"}`) + headers := http.Header{} + headers.Set("X-Claude-Code-Session-Id", "header-child-123") + info2, ok := ExtractSessionInfo(headers, headerPayload, nil) + if !ok { + t.Fatal("ExtractSessionInfo failed on header + body parent") + } + if info2.SessionID != "claude:header-child-123" { + t.Fatalf("expected session claude:header-child-123, got %q", info2.SessionID) + } + if info2.ParentSessionID != "claude:parent-session-789" { + t.Fatalf("expected parent claude:parent-session-789, got %q", info2.ParentSessionID) + } +} + +func TestExtractSessionInfoGeminiAndAntigravityHierarchy(t *testing.T) { + t.Parallel() + + // Gemini cachedContent with parent_session_id + geminiPayload := []byte(`{"cachedContent":"cache-child-1","parent_session_id":"cache-parent-1"}`) + info, ok := ExtractSessionInfo(nil, geminiPayload, nil) + if !ok { + t.Fatal("ExtractSessionInfo failed on Gemini cachedContent") + } + if info.SessionID != "geminicache:cache-child-1" { + t.Fatalf("expected session geminicache:cache-child-1, got %q", info.SessionID) + } + if info.ParentSessionID != "geminicache:cache-parent-1" { + t.Fatalf("expected parent geminicache:cache-parent-1, got %q", info.ParentSessionID) + } + if info.AgentName != "subagent" { + t.Fatalf("expected subagent, got %q", info.AgentName) + } + + // Antigravity headers with parent + headers := http.Header{} + headers.Set("X-Http-Session-Id", "agy-child-2") + headers.Set("X-Parent-Session-ID", "agy-parent-2") + info2, ok := ExtractSessionInfo(headers, nil, nil) + if !ok { + t.Fatal("ExtractSessionInfo failed on Antigravity headers") + } + if info2.SessionID != "agy:agy-child-2" { + t.Fatalf("expected session agy:agy-child-2, got %q", info2.SessionID) + } + if info2.ParentSessionID != "agy:agy-parent-2" { + t.Fatalf("expected parent agy:agy-parent-2, got %q", info2.ParentSessionID) + } + if info2.AgentName != "subagent" { + t.Fatalf("expected subagent, got %q", info2.AgentName) + } +} + +func TestExtractSessionInfoNestedAntigravityRequest(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "project_id": "proj-123", + "request": { + "parentSessionId": "parent-sess-456", + "sessionId": "child-sess-789" + } + }`) + info, ok := ExtractSessionInfo(nil, payload, nil) + if !ok { + t.Fatal("ExtractSessionInfo failed on nested Antigravity payload") + } + if info.SessionID != "session:child-sess-789" { + t.Fatalf("SessionID = %q, want session:child-sess-789", info.SessionID) + } + if info.ParentSessionID != "session:parent-sess-456" { + t.Fatalf("ParentSessionID = %q, want session:parent-sess-456", info.ParentSessionID) + } +} + +func TestExtractSessionInfoThreadAndConversation(t *testing.T) { + t.Parallel() + + // 1. Thread Header + threadHeaders := http.Header{ + "X-Thread-Id": []string{"thread-abc-123"}, + } + info, ok := ExtractSessionInfo(threadHeaders, nil, nil) + if !ok || info.ClientType != "openai-thread" || info.SessionID != "thread:thread-abc-123" { + t.Fatalf("thread header failed: %+v", info) + } + + // 2. Conversation Header + convHeaders := http.Header{ + "X-Conversation-Id": []string{"conv-xyz-789"}, + } + info, ok = ExtractSessionInfo(convHeaders, nil, nil) + if !ok || info.ClientType != "conv" || info.SessionID != "conv:conv-xyz-789" { + t.Fatalf("conversation header failed: %+v", info) + } + + // 3. Payload Thread with parent + threadPayload := []byte(`{ + "thread_id": "thread-child-1", + "parent_thread_id": "thread-parent-1" + }`) + info, ok = ExtractSessionInfo(nil, threadPayload, nil) + if !ok || info.ClientType != "openai-thread" || info.SessionID != "thread:thread-child-1" || info.ParentSessionID != "thread:thread-parent-1" { + t.Fatalf("thread payload failed: %+v", info) + } + + // 4. Payload Conversation with parent + convPayload := []byte(`{ + "conversation_id": "conv-child-2", + "parent_conversation_id": "conv-parent-2" + }`) + info, ok = ExtractSessionInfo(nil, convPayload, nil) + if !ok || info.ClientType != "conv" || info.SessionID != "conv:conv-child-2" || info.ParentSessionID != "conv:conv-parent-2" { + t.Fatalf("conv payload failed: %+v", info) + } +} + +func TestExtractSessionInfoSelfReferentialParentFiltered(t *testing.T) { + t.Parallel() + + headers := http.Header{ + "X-Session-ID": []string{"self-session-123"}, + } + payload := []byte(`{"parent_session_id": "self-session-123"}`) + info, ok := ExtractSessionInfo(headers, payload, nil) + if !ok { + t.Fatal("ExtractSessionInfo failed") + } + if info.SessionID != "header:self-session-123" { + t.Fatalf("SessionID = %q, want header:self-session-123", info.SessionID) + } + if info.ParentSessionID != "" { + t.Fatalf("ParentSessionID = %q, want empty (self-referential parent must be rejected)", info.ParentSessionID) + } +} + +func TestDeprecatedInMemorySessionTreeStoreCompatibility(t *testing.T) { + t.Parallel() + + var store SessionTreeStore = NewInMemorySessionTreeStore(100, 0) + node := store.RecordNode(SessionTreeInfo{ + SessionID: "child-1", + ParentSessionID: "root-1", + AgentName: "reviewer", + ClientType: "pi", + }) + if node == nil || node.SessionID != "child-1" || node.TreeDepth != 1 { + t.Fatalf("RecordNode returned invalid node: %+v", node) + } + + gotNode, ok := store.GetNode("child-1") + if !ok || gotNode == nil || gotNode.SessionID != "child-1" { + t.Fatalf("GetNode failed: %+v", gotNode) + } + + tree := store.GetTree("root-1") + if len(tree) != 1 || tree[0].SessionID != "child-1" { + t.Fatalf("GetTree failed: %+v", tree) + } + + ancestors := store.Ancestors("child-1") + if len(ancestors) != 1 || ancestors[0] != "root-1" { + t.Fatalf("Ancestors failed: %+v", ancestors) + } + + if !store.UpdateAffinity("child-1", "auth-99", "claude", "sonnet") { + t.Fatal("UpdateAffinity failed") + } + + if store.Len() != 1 { + t.Fatalf("Len() = %d, want 1", store.Len()) + } + + store.Clear() + if store.Len() != 0 { + t.Fatalf("Len() after Clear = %d, want 0", store.Len()) + } +} + +func TestExtractSessionInfoClaudePayloadOutranksGenericHeader(t *testing.T) { + t.Parallel() + + headers := http.Header{ + "X-Session-ID": []string{"generic-fallback-session"}, + } + payload := []byte(`{ + "metadata": { + "user_id": "{\"session_id\":\"claude-real-session\",\"agent_id\":\"reviewer\"}" + } + }`) + info, ok := ExtractSessionInfo(headers, payload, nil) + if !ok { + t.Fatal("ExtractSessionInfo failed") + } + if info.ClientType != "claude" || info.SessionID != "claude:claude-real-session:agent:reviewer" { + t.Fatalf("expected claude identity to outrank generic header, got %+v", info) + } +} + +func TestExtractSessionInfoPromptCacheKeyAndClientReqAndMetadata(t *testing.T) { + t.Parallel() + + // 1. prompt_cache_key + payloadPCK := []byte(`{"prompt_cache_key":"prompt-key-123","conversation":{"id":"conv-456"}}`) + info, ok := ExtractSessionInfo(nil, payloadPCK, nil) + if !ok || info.SessionID != "pck:prompt-key-123" || info.ParentSessionID != "conv:conv-456" { + t.Fatalf("pck extraction failed: %+v", info) + } + + // 2. plain metadata.user_id + payloadUser := []byte(`{"metadata":{"user_id":"user-999"}}`) + info, ok = ExtractSessionInfo(nil, payloadUser, nil) + if !ok || info.SessionID != "user:user-999" { + t.Fatalf("user_id extraction failed: %+v", info) + } + + // 3. X-Client-Request-Id + headersReqID := http.Header{"X-Client-Request-Id": []string{"client-req-001"}} + info, ok = ExtractSessionInfo(headersReqID, nil, nil) + if !ok || info.SessionID != "clientreq:client-req-001" { + t.Fatalf("clientreq extraction failed: %+v", info) + } + + // 4. ExecutionSessionMetadataKey + meta := map[string]any{ + "execution_session_id": "exec-777", + } + info, ok = ExtractSessionInfo(nil, nil, meta) + if !ok || info.SessionID != "execution:exec-777" { + t.Fatalf("execution extraction failed: %+v", info) + } +} + +func TestExtractSessionInfoNestedRequestAgent(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "request": { + "sessionId": "child-sess-1", + "metadata": { + "agent_id": "worker-sub" + } + } + }`) + info, ok := ExtractSessionInfo(nil, payload, nil) + if !ok { + t.Fatal("ExtractSessionInfo failed on nested request") + } + if info.SessionID != "session:child-sess-1:agent:worker-sub" { + t.Fatalf("SessionID = %q, want session:child-sess-1:agent:worker-sub", info.SessionID) + } + if info.ParentSessionID != "session:child-sess-1" { + t.Fatalf("ParentSessionID = %q, want session:child-sess-1", info.ParentSessionID) + } + if info.AgentName != "worker-sub" { + t.Fatalf("AgentName = %q, want worker-sub", info.AgentName) + } +} diff --git a/sdk/cliproxy/session/tree.go b/sdk/cliproxy/session/tree.go deleted file mode 100644 index 3821736d..00000000 --- a/sdk/cliproxy/session/tree.go +++ /dev/null @@ -1,953 +0,0 @@ -// Package session derives stable conversation identities and records hierarchical session trees. -package session - -import ( - "container/list" - "net/http" - "sort" - "strings" - "sync" - "time" - - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" - cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" -) - -const ( - defaultMaxTreeNodes = 65536 - defaultTreeTTL = 24 * time.Hour - maxTreeDepth = 128 -) - -// SessionTreeNode represents a node in the hierarchical session tree. -type SessionTreeNode struct { - SessionID string `json:"session_id"` - ParentSessionID string `json:"parent_session_id,omitempty"` - RootSessionID string `json:"root_session_id"` - TreePath string `json:"tree_path"` // e.g. "root_id/parent_id/self_id" - TreeDepth int `json:"tree_depth"` - AgentName string `json:"agent_name,omitempty"` - ClientType string `json:"client_type,omitempty"` - CallerScope string `json:"caller_scope,omitempty"` - LastAuthID string `json:"last_auth_id,omitempty"` - LastProvider string `json:"last_provider,omitempty"` - LastModel string `json:"last_model,omitempty"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` - Metadata map[string]any `json:"metadata,omitempty"` -} - -// Clone returns a deep copy of the node. -func (n *SessionTreeNode) Clone() *SessionTreeNode { - if n == nil { - return nil - } - res := *n - if n.Metadata != nil { - res.Metadata = cloneTreeMetadata(n.Metadata) - } - return &res -} - -const maxMetadataCloneDepth = 16 - -func cloneTreeMetadata(metadata map[string]any) map[string]any { - return cloneTreeMetadataWithDepth(metadata, 0) -} - -func cloneTreeMetadataWithDepth(metadata map[string]any, depth int) map[string]any { - if metadata == nil || depth > maxMetadataCloneDepth { - return nil - } - cloned := make(map[string]any, len(metadata)) - for key, value := range metadata { - cloned[key] = cloneTreeMetadataValueWithDepth(value, depth+1) - } - return cloned -} - -func cloneTreeMetadataValueWithDepth(value any, depth int) any { - if depth > maxMetadataCloneDepth { - return nil - } - switch typed := value.(type) { - case map[string]any: - return cloneTreeMetadataWithDepth(typed, depth+1) - case map[string]string: - cloned := make(map[string]string, len(typed)) - for k, v := range typed { - cloned[k] = v - } - return cloned - case []any: - cloned := make([]any, len(typed)) - for index, item := range typed { - cloned[index] = cloneTreeMetadataValueWithDepth(item, depth+1) - } - return cloned - case []string: - return append([]string(nil), typed...) - case []int: - return append([]int(nil), typed...) - case []byte: - return append([]byte(nil), typed...) - default: - return value - } -} - -// SessionTreeInfo encapsulates incoming request details needed to record or update a tree node. -type SessionTreeInfo struct { - SessionID string - ParentSessionID string - AgentName string - ClientType string - CallerScope string - AuthID string - Provider string - Model string - Metadata map[string]any -} - -// ExtractTreeInfo extracts session hierarchy and client identification from request attributes. -func ExtractTreeInfo(headers http.Header, payload []byte, metadata map[string]any) (SessionTreeInfo, bool) { - var info SessionTreeInfo - if metadata != nil { - if scope, ok := metadata[cliproxyexecutor.CallerScopeMetadataKey].(string); ok { - info.CallerScope = strings.TrimSpace(scope) - } - } - - // 1. Client Type & Explicit Header Detection - switch { - case sessionHeaderValue(headers, "X-Claude-Code-Session-Id") != "": - info.ClientType = "claude" - rawSID := sessionHeaderValue(headers, "X-Claude-Code-Session-Id") - agentID := sessionHeaderValue(headers, "X-Claude-Code-Agent-Id") - parentAgentID := sessionHeaderValue(headers, "X-Claude-Code-Parent-Agent-Id") - if agentID != "" && agentID != "main" { - info.AgentName = agentID - info.ParentSessionID = "claude:" + rawSID - if parentAgentID != "" && parentAgentID != "main" && parentAgentID != agentID { - info.ParentSessionID = "claude:" + rawSID + ":agent:" + parentAgentID - } - info.SessionID = "claude:" + rawSID + ":agent:" + agentID - } else { - info.AgentName = "main" - info.SessionID = "claude:" + rawSID - } - - case sessionHeaderValue(headers, "Session-Id") != "" || sessionHeaderValue(headers, "Session_id") != "": - info.ClientType = "codex" - sid := sessionHeaderValue(headers, "Session-Id") - if sid == "" { - sid = sessionHeaderValue(headers, "Session_id") - } - info.SessionID = "codex:" + sid - parentThread := sessionHeaderValue(headers, "x-codex-parent-thread-id") - if parentThread == "" { - parentThread = sessionHeaderValue(headers, "X-Codex-Parent-Thread-Id") - } - if parentThread != "" && parentThread != sid { - info.ParentSessionID = "codex:" + parentThread - info.AgentName = "subagent" - } else { - info.AgentName = "main" - } - - case sessionHeaderValue(headers, "X-Slot-Session-Id") != "": - info.ClientType = "pi" - sid := sessionHeaderValue(headers, "X-Slot-Session-Id") - info.SessionID = "slot:" + sid - parentSID := sessionHeaderValue(headers, "X-Parent-Session-ID") - if parentSID == "" { - parentSID = sessionHeaderValue(headers, "X-Parent-Session-Id") - } - if parentSID != "" && parentSID != sid { - info.ParentSessionID = "slot:" + parentSID - info.AgentName = "subagent" - } else { - info.AgentName = "slot" - } - - case sessionHeaderValue(headers, "X-Http-Session-Id") != "": - info.ClientType = "agy" - sid := sessionHeaderValue(headers, "X-Http-Session-Id") - info.SessionID = "agy:" + sid - parentSID := sessionHeaderValue(headers, "X-Parent-Session-ID") - if parentSID == "" { - parentSID = sessionHeaderValue(headers, "X-Parent-Session-Id") - } - if parentSID != "" && parentSID != sid { - info.ParentSessionID = "agy:" + parentSID - info.AgentName = "subagent" - } else { - info.AgentName = "main" - } - - case sessionHeaderValue(headers, "X-Session-ID") != "": - info.ClientType = "generic" - sid := sessionHeaderValue(headers, "X-Session-ID") - info.SessionID = "header:" + sid - parentSID := sessionHeaderValue(headers, "X-Parent-Session-ID") - if parentSID == "" { - parentSID = sessionHeaderValue(headers, "X-Parent-Session-Id") - } - if parentSID != "" && parentSID != sid { - info.ParentSessionID = "header:" + parentSID - info.AgentName = "subagent" - } else { - info.AgentName = "main" - } - - case sessionHeaderValue(headers, "X-Session-Affinity") != "": - info.ClientType = "opencode" - sid := sessionHeaderValue(headers, "X-Session-Affinity") - info.SessionID = "affinity:" + sid - parentAffinity := sessionHeaderValue(headers, "X-Parent-Session-Affinity") - if parentAffinity == "" { - parentAffinity = sessionHeaderValue(headers, "X-Parent-Session-ID") - } - if parentAffinity != "" && parentAffinity != sid { - info.ParentSessionID = "affinity:" + parentAffinity - info.AgentName = "subagent" - } else { - info.AgentName = "main" - } - - case sessionHeaderValue(headers, "X-Thread-Id") != "" || sessionHeaderValue(headers, "X-Thread-ID") != "" || sessionHeaderValue(headers, "Thread-Id") != "": - info.ClientType = "openai-thread" - tid := sessionHeaderValue(headers, "X-Thread-Id") - if tid == "" { - tid = sessionHeaderValue(headers, "X-Thread-ID") - } - if tid == "" { - tid = sessionHeaderValue(headers, "Thread-Id") - } - info.SessionID = "thread:" + tid - info.AgentName = "main" - - case sessionHeaderValue(headers, "X-Conversation-Id") != "" || sessionHeaderValue(headers, "X-Conversation-ID") != "": - info.ClientType = "conv" - cid := sessionHeaderValue(headers, "X-Conversation-Id") - if cid == "" { - cid = sessionHeaderValue(headers, "X-Conversation-ID") - } - info.SessionID = "conv:" + cid - info.AgentName = "main" - } - - // 2. Payload Inspection if headers didn't fully resolve - if len(payload) > 0 { - root := util.ParseGJSONBytesNoCopy(payload) - reqRoot := root - req := root.Get("request") - hasNestedReq := req.Exists() && !root.Get("contents").Exists() - if hasNestedReq { - reqRoot = req - } - var parentCandidate string - for _, p := range []string{ - "parent_session_id", "parentSessionId", - "parent_thread_id", "parentThreadId", - "forked_from_thread_id", "forked_from_id", - "parent_conversation_id", "parentConversationId", - "metadata.parent_session_id", "metadata.parent_thread_id", - "extra_body.parent_session_id", "extra_body.parent_thread_id", - } { - if val := normalizedSessionCandidate(root.Get(p).String()); val != "" { - parentCandidate = val - break - } - if hasNestedReq { - if val := normalizedSessionCandidate(reqRoot.Get(p).String()); val != "" { - parentCandidate = val - break - } - } - } - if parentCandidate == "" { - parentCandidate = ClaudeMetadataParentSessionID(payload) - } - - if info.SessionID == "" { - // Claude metadata.user_id - if sid, parentSID, agentID := ClaudeMetadataIdentities(payload); sid != "" { - info.ClientType = "claude" - if agentID == "" { - agentID = normalizedSessionCandidate(root.Get("metadata.agent_id").String()) - } - if agentID != "" && agentID != "main" { - info.SessionID = "claude:" + sid + ":agent:" + agentID - info.ParentSessionID = "claude:" + sid - if parentSID != "" && parentSID != sid { - info.ParentSessionID = "claude:" + parentSID - } - info.AgentName = agentID - } else { - info.SessionID = "claude:" + sid - if parentSID != "" && parentSID != sid { - info.ParentSessionID = "claude:" + parentSID - info.AgentName = "subagent" - } else { - info.AgentName = "main" - } - } - } - - // Gemini context caching - if info.SessionID == "" { - for _, cachePath := range []string{"cachedContent", "cached_content"} { - cacheID := normalizedSessionCandidate(root.Get(cachePath).String()) - if cacheID == "" && hasNestedReq { - cacheID = normalizedSessionCandidate(reqRoot.Get(cachePath).String()) - } - if cacheID != "" { - info.ClientType = "gemini" - info.SessionID = "geminicache:" + cacheID - if parentCandidate != "" && parentCandidate != cacheID { - info.ParentSessionID = "geminicache:" + parentCandidate - info.AgentName = "subagent" - } else { - info.AgentName = "main" - } - break - } - } - } - - // OpenAI thread in payload - if info.SessionID == "" { - for _, threadPath := range []string{"thread_id", "threadId", "metadata.thread_id"} { - tid := normalizedSessionCandidate(root.Get(threadPath).String()) - if tid == "" && hasNestedReq { - tid = normalizedSessionCandidate(reqRoot.Get(threadPath).String()) - } - if tid != "" { - info.ClientType = "openai-thread" - info.SessionID = "thread:" + tid - if parentCandidate != "" && parentCandidate != tid { - info.ParentSessionID = "thread:" + parentCandidate - info.AgentName = "subagent" - } else { - info.AgentName = "main" - } - break - } - } - } - - // Generic session in payload - if info.SessionID == "" { - agentID := normalizedSessionCandidate(root.Get("metadata.agent_id").String()) - if agentID == "" { - agentID = normalizedSessionCandidate(root.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 { - sid = normalizedSessionCandidate(reqRoot.Get(path).String()) - } - if sid != "" { - info.ClientType = "generic" - if agentID != "" && agentID != "main" { - info.SessionID = "session:" + sid + ":agent:" + agentID - info.ParentSessionID = "session:" + sid - if parentCandidate != "" && parentCandidate != sid { - info.ParentSessionID = "session:" + parentCandidate - } - info.AgentName = agentID - } else { - info.SessionID = "session:" + sid - if parentCandidate != "" && parentCandidate != sid { - info.ParentSessionID = "session:" + parentCandidate - info.AgentName = "subagent" - } else { - info.AgentName = "main" - } - } - break - } - } - } - - // Conversation paths - if info.SessionID == "" { - for _, convPath := range []string{"conversation_id", "conversationId", "chat_id", "chatId", "metadata.conversation_id", "extra_body.conversation_id"} { - cid := normalizedSessionCandidate(root.Get(convPath).String()) - if cid == "" && hasNestedReq { - cid = normalizedSessionCandidate(reqRoot.Get(convPath).String()) - } - if cid != "" { - info.ClientType = "conv" - info.SessionID = "conv:" + cid - if parentCandidate != "" && parentCandidate != cid { - info.ParentSessionID = "conv:" + parentCandidate - info.AgentName = "subagent" - } else { - info.AgentName = "main" - } - break - } - } - } - } else if info.ParentSessionID == "" && parentCandidate != "" { - info.ParentSessionID = canonicalParentSessionID(info.ClientType, parentCandidate) - if info.AgentName == "main" || info.AgentName == "" { - info.AgentName = "subagent" - } - } - } - - if info.SessionID == "" { - return SessionTreeInfo{}, false - } - if info.AgentName == "" { - info.AgentName = "main" - } - if info.ClientType == "" { - info.ClientType = "generic" - } - return info, true -} - -func sessionHeaderValue(headers http.Header, name string) string { - if headers == nil { - return "" - } - for key, values := range headers { - if !strings.EqualFold(key, name) { - continue - } - for _, value := range values { - if normalized := normalizedSessionCandidate(value); normalized != "" { - return normalized - } - } - } - return "" -} - -func normalizedSessionCandidate(raw string) string { - return NormalizeExplicitID(raw) -} - -func canonicalParentSessionID(clientType, raw string) string { - raw = normalizedSessionCandidate(raw) - if raw == "" { - return "" - } - - var scheme string - switch clientType { - case "agy": - scheme = "agy" - case "claude": - scheme = "claude" - case "codex": - scheme = "codex" - case "conv": - scheme = "conv" - case "gemini": - scheme = "geminicache" - case "generic": - scheme = "header" - case "openai-thread": - scheme = "thread" - case "opencode": - scheme = "affinity" - case "pi": - scheme = "slot" - } - if scheme != "" { - if strings.HasPrefix(raw, scheme+":") { - return raw - } - return scheme + ":" + raw - } - return raw -} - -// SessionTreeStore defines the interface for recording and querying session trees. -type SessionTreeStore interface { - RecordNode(info SessionTreeInfo) *SessionTreeNode - GetNode(sessionID string) (*SessionTreeNode, bool) - GetTree(rootSessionID string) []*SessionTreeNode - GetSubtree(sessionID string) []*SessionTreeNode - Ancestors(sessionID string) []string - UpdateAffinity(sessionID, authID, provider, model string) bool - Len() int - Clear() -} - -// InMemorySessionTreeStore is a bounded, concurrent-safe in-memory session tree store. -type InMemorySessionTreeStore struct { - mu sync.RWMutex - maxNodes int - ttl time.Duration - nodes map[string]*SessionTreeNode - rootIndex map[string]map[string]struct{} // root_session_id -> set of session_ids - parentIndex map[string]map[string]struct{} // parent_session_id -> set of direct child session_ids - lru *list.List - lruElements map[string]*list.Element -} - -type lruItem struct { - sessionID string - expiresAt time.Time -} - -// NewInMemorySessionTreeStore creates a new in-memory session tree store with bounded limits. -func NewInMemorySessionTreeStore(maxNodes int, ttl time.Duration) *InMemorySessionTreeStore { - if maxNodes <= 0 { - maxNodes = defaultMaxTreeNodes - } - if ttl <= 0 { - ttl = defaultTreeTTL - } - return &InMemorySessionTreeStore{ - maxNodes: maxNodes, - ttl: ttl, - nodes: make(map[string]*SessionTreeNode), - rootIndex: make(map[string]map[string]struct{}), - parentIndex: make(map[string]map[string]struct{}), - lru: list.New(), - lruElements: make(map[string]*list.Element), - } -} - -// RecordNode atomically creates or updates a session tree node with materialized paths. -func (s *InMemorySessionTreeStore) RecordNode(info SessionTreeInfo) *SessionTreeNode { - sessionID := strings.TrimSpace(info.SessionID) - if sessionID == "" { - return nil - } - - s.mu.Lock() - defer s.mu.Unlock() - - now := time.Now() - s.cleanupExpiredLocked(now) - expiresAt := now.Add(s.ttl) - - // Update existing node - if existing, ok := s.nodes[sessionID]; ok { - existing.UpdatedAt = now - if info.AuthID != "" { - existing.LastAuthID = info.AuthID - } - if info.Provider != "" { - existing.LastProvider = info.Provider - } - if info.Model != "" { - existing.LastModel = info.Model - } - if info.AgentName != "" && existing.AgentName == "main" { - existing.AgentName = info.AgentName - } - if info.ClientType != "" && existing.ClientType == "generic" { - existing.ClientType = info.ClientType - } - if parentID := strings.TrimSpace(info.ParentSessionID); parentID != "" && parentID != existing.ParentSessionID && parentID != existing.SessionID && !s.isDescendantLocked(parentID, existing.SessionID) { - oldParent := existing.ParentSessionID - existing.ParentSessionID = parentID - s.updateParentIndexLocked(existing.SessionID, oldParent, existing.ParentSessionID) - oldRoot, changed := s.computeNodeLineageLocked(existing) - if changed { - s.updateRootIndexLocked(existing.SessionID, oldRoot, existing.RootSessionID) - } - s.cascadeDescendantsLineageLocked(existing.SessionID) - } - if elem, ok := s.lruElements[sessionID]; ok { - elem.Value = lruItem{sessionID: sessionID, expiresAt: expiresAt} - s.lru.MoveToBack(elem) - } - return existing.Clone() - } - - // Create new node - node := &SessionTreeNode{ - SessionID: sessionID, - ParentSessionID: strings.TrimSpace(info.ParentSessionID), - AgentName: info.AgentName, - ClientType: info.ClientType, - CallerScope: info.CallerScope, - LastAuthID: info.AuthID, - LastProvider: info.Provider, - LastModel: info.Model, - CreatedAt: now, - UpdatedAt: now, - Metadata: cloneTreeMetadata(info.Metadata), - } - if node.AgentName == "" { - node.AgentName = "main" - } - if node.ClientType == "" { - node.ClientType = "generic" - } - - // Compute RootSessionID, TreePath, and TreeDepth - s.computeNodeLineageLocked(node) - - s.nodes[sessionID] = node - if node.ParentSessionID != "" { - s.updateParentIndexLocked(sessionID, "", node.ParentSessionID) - } - - // Index by root - s.updateRootIndexLocked(sessionID, "", node.RootSessionID) - - // Cascade lineage in case children were recorded before this node arrived - s.cascadeDescendantsLineageLocked(sessionID) - - // LRU tracking - elem := s.lru.PushBack(lruItem{sessionID: sessionID, expiresAt: expiresAt}) - s.lruElements[sessionID] = elem - - // Evict if capacity exceeded - for len(s.nodes) > s.maxNodes { - s.evictOldestLocked() - } - - return node.Clone() -} - -// GetNode returns a clone of the specified node if it exists. -func (s *InMemorySessionTreeStore) GetNode(sessionID string) (*SessionTreeNode, bool) { - s.mu.Lock() - defer s.mu.Unlock() - s.cleanupExpiredLocked(time.Now()) - node, ok := s.nodes[sessionID] - if !ok { - return nil, false - } - return node.Clone(), true -} - -// GetTree returns all nodes belonging to the same root session, ordered by depth and created time. -func (s *InMemorySessionTreeStore) GetTree(rootSessionID string) []*SessionTreeNode { - s.mu.Lock() - defer s.mu.Unlock() - s.cleanupExpiredLocked(time.Now()) - - set := s.rootIndex[rootSessionID] - if len(set) == 0 { - return nil - } - - res := make([]*SessionTreeNode, 0, len(set)) - for sid := range set { - if node, ok := s.nodes[sid]; ok { - res = append(res, node.Clone()) - } - } - - sort.Slice(res, func(i, j int) bool { - if res[i].TreeDepth != res[j].TreeDepth { - return res[i].TreeDepth < res[j].TreeDepth - } - return res[i].CreatedAt.Before(res[j].CreatedAt) - }) - return res -} - -// GetSubtree returns the specified node and all of its descendants by parent linkage. -// Parent-link traversal is used instead of TreePath prefix matching because session IDs -// may contain "/". The cost is O(rootSet x depth), which is acceptable for this -// management/debug-oriented accessor. -func (s *InMemorySessionTreeStore) GetSubtree(sessionID string) []*SessionTreeNode { - s.mu.Lock() - defer s.mu.Unlock() - s.cleanupExpiredLocked(time.Now()) - - target, ok := s.nodes[sessionID] - if !ok { - return nil - } - - res := []*SessionTreeNode{target.Clone()} - - set := s.rootIndex[target.RootSessionID] - for sid := range set { - if sid == sessionID { - continue - } - if node, ok := s.nodes[sid]; ok && s.isDescendantLocked(node.SessionID, sessionID) { - res = append(res, node.Clone()) - } - } - - sort.Slice(res, func(i, j int) bool { - if res[i].TreeDepth != res[j].TreeDepth { - return res[i].TreeDepth < res[j].TreeDepth - } - return res[i].CreatedAt.Before(res[j].CreatedAt) - }) - return res -} - -// Ancestors returns the sequence of ancestor session IDs from root to direct parent without -// database lookup. The chain may include IDs of ancestors that already expired or were -// evicted, so callers must not assume every returned ID resolves via GetNode. -func (s *InMemorySessionTreeStore) Ancestors(sessionID string) []string { - s.mu.Lock() - defer s.mu.Unlock() - s.cleanupExpiredLocked(time.Now()) - - node, ok := s.nodes[sessionID] - if !ok || node.ParentSessionID == "" { - return nil - } - - ancestors := make([]string, 0, node.TreeDepth) - visited := make(map[string]struct{}, node.TreeDepth+1) - visited[sessionID] = struct{}{} - parentID := node.ParentSessionID - for steps := 0; parentID != "" && steps <= len(s.nodes); steps++ { - if _, seen := visited[parentID]; seen { - break - } - visited[parentID] = struct{}{} - ancestors = append(ancestors, parentID) - parent, exists := s.nodes[parentID] - if !exists { - break - } - parentID = parent.ParentSessionID - } - for left, right := 0, len(ancestors)-1; left < right; left, right = left+1, right-1 { - ancestors[left], ancestors[right] = ancestors[right], ancestors[left] - } - return ancestors -} - -// UpdateAffinity updates the latest successful upstream binding for a session node. -func (s *InMemorySessionTreeStore) UpdateAffinity(sessionID, authID, provider, model string) bool { - s.mu.Lock() - defer s.mu.Unlock() - now := time.Now() - s.cleanupExpiredLocked(now) - - node, ok := s.nodes[sessionID] - if !ok { - return false - } - node.UpdatedAt = now - if elem, exists := s.lruElements[sessionID]; exists { - elem.Value = lruItem{sessionID: sessionID, expiresAt: now.Add(s.ttl)} - s.lru.MoveToBack(elem) - } - if authID != "" { - node.LastAuthID = authID - } - if provider != "" { - node.LastProvider = provider - } - if model != "" { - node.LastModel = model - } - return true -} - -func (s *InMemorySessionTreeStore) evictOldestLocked() { - oldest := s.lru.Front() - if oldest == nil { - return - } - item, ok := oldest.Value.(lruItem) - if !ok { - s.lru.Remove(oldest) - return - } - s.removeNodeLocked(item.sessionID) -} - -func (s *InMemorySessionTreeStore) cleanupExpiredLocked(now time.Time) { - for element := s.lru.Front(); element != nil; { - item, ok := element.Value.(lruItem) - if !ok { - next := element.Next() - s.lru.Remove(element) - element = next - continue - } - // Every touch sets expiresAt = touch + the same TTL and moves the element to the - // LRU tail, so LRU order is exactly expiresAt order and expired nodes form a - // front prefix. Stop at the first live node to keep this amortized O(expired). - if now.Before(item.expiresAt) { - return - } - next := element.Next() - s.removeNodeLocked(item.sessionID) - element = next - } -} - -func (s *InMemorySessionTreeStore) removeNodeLocked(sessionID string) { - if element := s.lruElements[sessionID]; element != nil { - s.lru.Remove(element) - delete(s.lruElements, sessionID) - } - if node, ok := s.nodes[sessionID]; ok { - delete(s.nodes, sessionID) - s.updateParentIndexLocked(sessionID, node.ParentSessionID, "") - if set := s.rootIndex[node.RootSessionID]; set != nil { - delete(set, sessionID) - if len(set) == 0 { - delete(s.rootIndex, node.RootSessionID) - } - } - } -} - -func (s *InMemorySessionTreeStore) computeNodeLineageLocked(node *SessionTreeNode) (oldRoot string, changed bool) { - oldRoot = node.RootSessionID - var newRoot, newPath string - var newDepth int - - if node.ParentSessionID != "" && node.ParentSessionID != node.SessionID && !s.isDescendantLocked(node.ParentSessionID, node.SessionID) { - if parent, ok := s.nodes[node.ParentSessionID]; ok { - if parent.TreeDepth < maxTreeDepth { - newRoot = parent.RootSessionID - newPath = parent.TreePath + "/" + node.SessionID - newDepth = parent.TreeDepth + 1 - } else { - s.updateParentIndexLocked(node.SessionID, node.ParentSessionID, "") - node.ParentSessionID = "" - newRoot = node.SessionID - newPath = node.SessionID - newDepth = 0 - } - } else { - newRoot = node.ParentSessionID - newPath = node.ParentSessionID + "/" + node.SessionID - newDepth = 1 - } - } else { - node.ParentSessionID = "" - newRoot = node.SessionID - newPath = node.SessionID - newDepth = 0 - } - - if node.RootSessionID != newRoot || node.TreePath != newPath || node.TreeDepth != newDepth { - node.RootSessionID = newRoot - node.TreePath = newPath - node.TreeDepth = newDepth - changed = true - } - return oldRoot, changed -} - -func (s *InMemorySessionTreeStore) updateRootIndexLocked(sessionID, oldRoot, newRoot string) { - if oldRoot != "" && oldRoot != newRoot { - if set := s.rootIndex[oldRoot]; set != nil { - delete(set, sessionID) - if len(set) == 0 { - delete(s.rootIndex, oldRoot) - } - } - } - if newRoot != "" { - set := s.rootIndex[newRoot] - if set == nil { - set = make(map[string]struct{}) - s.rootIndex[newRoot] = set - } - set[sessionID] = struct{}{} - } -} - -func (s *InMemorySessionTreeStore) updateParentIndexLocked(sessionID, oldParent, newParent string) { - if oldParent != "" && oldParent != newParent { - if set := s.parentIndex[oldParent]; set != nil { - delete(set, sessionID) - if len(set) == 0 { - delete(s.parentIndex, oldParent) - } - } - } - if newParent != "" { - set := s.parentIndex[newParent] - if set == nil { - set = make(map[string]struct{}) - s.parentIndex[newParent] = set - } - set[sessionID] = struct{}{} - } -} - -func (s *InMemorySessionTreeStore) cascadeDescendantsLineageLocked(parentID string) { - visited := make(map[string]struct{}) - s.cascadeDescendantsHelperLocked(parentID, visited) -} - -func (s *InMemorySessionTreeStore) cascadeDescendantsHelperLocked(parentID string, visited map[string]struct{}) { - if _, seen := visited[parentID]; seen { - return - } - visited[parentID] = struct{}{} - - childrenSet := s.parentIndex[parentID] - if len(childrenSet) == 0 { - return - } - childIDs := make([]string, 0, len(childrenSet)) - for childID := range childrenSet { - childIDs = append(childIDs, childID) - } - - for _, childID := range childIDs { - child := s.nodes[childID] - if child != nil && child.ParentSessionID == parentID { - oldRoot, changed := s.computeNodeLineageLocked(child) - if changed { - s.updateRootIndexLocked(child.SessionID, oldRoot, child.RootSessionID) - } - s.cascadeDescendantsHelperLocked(child.SessionID, visited) - } - } -} - -func (s *InMemorySessionTreeStore) isDescendantLocked(sessionID, ancestorID string) bool { - if sessionID == "" || ancestorID == "" || sessionID == ancestorID { - return false - } - visited := make(map[string]struct{}) - visited[sessionID] = struct{}{} - currentID := sessionID - for steps := 0; currentID != "" && steps <= len(s.nodes); steps++ { - node, ok := s.nodes[currentID] - if !ok { - return false - } - if node.ParentSessionID == ancestorID { - return true - } - if _, seen := visited[node.ParentSessionID]; seen { - return false - } - visited[node.ParentSessionID] = struct{}{} - currentID = node.ParentSessionID - } - return false -} - -// Len returns the current count of tracked session tree nodes. -func (s *InMemorySessionTreeStore) Len() int { - s.mu.Lock() - defer s.mu.Unlock() - s.cleanupExpiredLocked(time.Now()) - return len(s.nodes) -} - -// Clear flushes all tracked nodes. -func (s *InMemorySessionTreeStore) Clear() { - s.mu.Lock() - defer s.mu.Unlock() - s.nodes = make(map[string]*SessionTreeNode) - s.rootIndex = make(map[string]map[string]struct{}) - s.parentIndex = make(map[string]map[string]struct{}) - s.lru = list.New() - s.lruElements = make(map[string]*list.Element) -} diff --git a/sdk/cliproxy/session/tree_compat.go b/sdk/cliproxy/session/tree_compat.go new file mode 100644 index 00000000..b5f45960 --- /dev/null +++ b/sdk/cliproxy/session/tree_compat.go @@ -0,0 +1,204 @@ +package session + +import ( + "sync" + "time" +) + +// SessionTreeNode represents a node in the session hierarchy. +// Deprecated: Session tree management and multi-node lineage have moved to Home. +// This struct is maintained as a migration stub for SDK backward compatibility. +type SessionTreeNode struct { + SessionID string `json:"session_id"` + ParentSessionID string `json:"parent_session_id,omitempty"` + RootSessionID string `json:"root_session_id"` + TreePath string `json:"tree_path"` + TreeDepth int `json:"tree_depth"` + AgentName string `json:"agent_name,omitempty"` + ClientType string `json:"client_type,omitempty"` + CallerScope string `json:"caller_scope,omitempty"` + LastAuthID string `json:"last_auth_id,omitempty"` + LastProvider string `json:"last_provider,omitempty"` + LastModel string `json:"last_model,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + Metadata map[string]any `json:"metadata,omitempty"` +} + +// Clone returns a deep copy of the node. +// Deprecated: Session tree management has moved to Home. +func (n *SessionTreeNode) Clone() *SessionTreeNode { + if n == nil { + return nil + } + res := *n + if n.Metadata != nil { + res.Metadata = make(map[string]any, len(n.Metadata)) + for k, v := range n.Metadata { + res.Metadata[k] = v + } + } + return &res +} + +// SessionTreeStore defines the interface for recording and querying session trees. +// Deprecated: Session tree management has moved to Home. +type SessionTreeStore interface { + RecordNode(info SessionTreeInfo) *SessionTreeNode + GetNode(sessionID string) (*SessionTreeNode, bool) + GetTree(rootSessionID string) []*SessionTreeNode + GetSubtree(sessionID string) []*SessionTreeNode + Ancestors(sessionID string) []string + UpdateAffinity(sessionID, authID, provider, model string) bool + Len() int + Clear() +} + +// InMemorySessionTreeStore is a backward-compatible in-memory session tree stub. +// Deprecated: Standalone CPA nodes use flat KV session cache (SessionCache). +// Full hierarchical session tree management and global aggregation are handled by Home. +type InMemorySessionTreeStore struct { + mu sync.RWMutex + nodes map[string]*SessionTreeNode +} + +// NewInMemorySessionTreeStore creates a new in-memory session tree store stub. +// Deprecated: Standalone CPA nodes use flat KV session cache. +func NewInMemorySessionTreeStore(maxNodes int, ttl time.Duration) *InMemorySessionTreeStore { + return &InMemorySessionTreeStore{ + nodes: make(map[string]*SessionTreeNode), + } +} + +// RecordNode records a session tree node stub. +// Deprecated: Session tree management has moved to Home. +func (s *InMemorySessionTreeStore) RecordNode(info SessionTreeInfo) *SessionTreeNode { + if s == nil || info.SessionID == "" { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + + now := time.Now() + node := &SessionTreeNode{ + SessionID: info.SessionID, + ParentSessionID: info.ParentSessionID, + RootSessionID: info.SessionID, + TreePath: info.SessionID, + TreeDepth: 0, + AgentName: info.AgentName, + ClientType: info.ClientType, + CallerScope: info.CallerScope, + LastAuthID: info.AuthID, + LastProvider: info.Provider, + LastModel: info.Model, + CreatedAt: now, + UpdatedAt: now, + Metadata: info.Metadata, + } + if info.ParentSessionID != "" { + node.RootSessionID = info.ParentSessionID + node.TreePath = info.ParentSessionID + "/" + info.SessionID + node.TreeDepth = 1 + } + s.nodes[info.SessionID] = node + return node.Clone() +} + +// GetNode returns a recorded node. +// Deprecated: Session tree management has moved to Home. +func (s *InMemorySessionTreeStore) GetNode(sessionID string) (*SessionTreeNode, bool) { + if s == nil { + return nil, false + } + s.mu.RLock() + defer s.mu.RUnlock() + node, ok := s.nodes[sessionID] + if !ok { + return nil, false + } + return node.Clone(), true +} + +// GetTree returns nodes under rootSessionID. +// Deprecated: Session tree management has moved to Home. +func (s *InMemorySessionTreeStore) GetTree(rootSessionID string) []*SessionTreeNode { + if s == nil { + return nil + } + s.mu.RLock() + defer s.mu.RUnlock() + var res []*SessionTreeNode + for _, n := range s.nodes { + if n.RootSessionID == rootSessionID || n.SessionID == rootSessionID { + res = append(res, n.Clone()) + } + } + return res +} + +// GetSubtree returns subtree rooted at sessionID. +// Deprecated: Session tree management has moved to Home. +func (s *InMemorySessionTreeStore) GetSubtree(sessionID string) []*SessionTreeNode { + return s.GetTree(sessionID) +} + +// Ancestors returns ancestor session IDs. +// Deprecated: Session tree management has moved to Home. +func (s *InMemorySessionTreeStore) Ancestors(sessionID string) []string { + if s == nil { + return nil + } + s.mu.RLock() + defer s.mu.RUnlock() + if node, ok := s.nodes[sessionID]; ok && node.ParentSessionID != "" { + return []string{node.ParentSessionID} + } + return nil +} + +// UpdateAffinity updates the latest successful upstream binding. +// Deprecated: Session tree management has moved to Home. +func (s *InMemorySessionTreeStore) UpdateAffinity(sessionID, authID, provider, model string) bool { + if s == nil { + return false + } + s.mu.Lock() + defer s.mu.Unlock() + if node, ok := s.nodes[sessionID]; ok { + node.UpdatedAt = time.Now() + if authID != "" { + node.LastAuthID = authID + } + if provider != "" { + node.LastProvider = provider + } + if model != "" { + node.LastModel = model + } + return true + } + return false +} + +// Len returns the number of nodes. +// Deprecated: Session tree management has moved to Home. +func (s *InMemorySessionTreeStore) Len() int { + if s == nil { + return 0 + } + s.mu.RLock() + defer s.mu.RUnlock() + return len(s.nodes) +} + +// Clear flushes all nodes. +// Deprecated: Session tree management has moved to Home. +func (s *InMemorySessionTreeStore) Clear() { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.nodes = make(map[string]*SessionTreeNode) +} diff --git a/sdk/cliproxy/session/tree_test.go b/sdk/cliproxy/session/tree_test.go deleted file mode 100644 index c92e5b29..00000000 --- a/sdk/cliproxy/session/tree_test.go +++ /dev/null @@ -1,736 +0,0 @@ -package session - -import ( - "fmt" - "net/http" - "sync" - "testing" - "time" - - cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" -) - -func TestSessionTreeStoreRootAndChildHierarchy(t *testing.T) { - t.Parallel() - - store := NewInMemorySessionTreeStore(1000, time.Hour) - - // 1. Root Node - rootInfo := SessionTreeInfo{ - SessionID: "session-root", - AgentName: "main", - ClientType: "pi", - AuthID: "auth-1", - Provider: "claude", - Model: "claude-3-7-sonnet", - } - rootNode := store.RecordNode(rootInfo) - if rootNode == nil { - t.Fatalf("expected rootNode to be non-nil") - } - if rootNode.RootSessionID != "session-root" { - t.Errorf("rootNode.RootSessionID = %q, want %q", rootNode.RootSessionID, "session-root") - } - if rootNode.TreePath != "session-root" { - t.Errorf("rootNode.TreePath = %q, want %q", rootNode.TreePath, "session-root") - } - if rootNode.TreeDepth != 0 { - t.Errorf("rootNode.TreeDepth = %d, want 0", rootNode.TreeDepth) - } - - // 2. Child 1 (Subagent under root) - child1Info := SessionTreeInfo{ - SessionID: "session-child-1", - ParentSessionID: "session-root", - AgentName: "code-reviewer", - ClientType: "pi", - } - child1Node := store.RecordNode(child1Info) - if child1Node == nil { - t.Fatalf("expected child1Node to be non-nil") - } - if child1Node.RootSessionID != "session-root" { - t.Errorf("child1Node.RootSessionID = %q, want %q", child1Node.RootSessionID, "session-root") - } - if child1Node.TreePath != "session-root/session-child-1" { - t.Errorf("child1Node.TreePath = %q, want %q", child1Node.TreePath, "session-root/session-child-1") - } - if child1Node.TreeDepth != 1 { - t.Errorf("child1Node.TreeDepth = %d, want 1", child1Node.TreeDepth) - } - - // 3. Child 2 (Grandchild, subagent under Child 1) - child2Info := SessionTreeInfo{ - SessionID: "session-child-2", - ParentSessionID: "session-child-1", - AgentName: "test-runner", - ClientType: "pi", - } - child2Node := store.RecordNode(child2Info) - if child2Node == nil { - t.Fatalf("expected child2Node to be non-nil") - } - if child2Node.RootSessionID != "session-root" { - t.Errorf("child2Node.RootSessionID = %q, want %q", child2Node.RootSessionID, "session-root") - } - if child2Node.TreePath != "session-root/session-child-1/session-child-2" { - t.Errorf("child2Node.TreePath = %q, want %q", child2Node.TreePath, "session-root/session-child-1/session-child-2") - } - if child2Node.TreeDepth != 2 { - t.Errorf("child2Node.TreeDepth = %d, want 2", child2Node.TreeDepth) - } - - // 4. Child 3 (Sibling under Root) - child3Info := SessionTreeInfo{ - SessionID: "session-child-3", - ParentSessionID: "session-root", - AgentName: "explorer", - ClientType: "pi", - } - child3Node := store.RecordNode(child3Info) - if child3Node.TreePath != "session-root/session-child-3" { - t.Errorf("child3Node.TreePath = %q, want %q", child3Node.TreePath, "session-root/session-child-3") - } - if child3Node.TreeDepth != 1 { - t.Errorf("child3Node.TreeDepth = %d, want 1", child3Node.TreeDepth) - } -} - -func TestSessionTreeStoreGetTreeAndSubtree(t *testing.T) { - t.Parallel() - - store := NewInMemorySessionTreeStore(1000, time.Hour) - - store.RecordNode(SessionTreeInfo{SessionID: "root"}) - store.RecordNode(SessionTreeInfo{SessionID: "branch-a", ParentSessionID: "root", AgentName: "agent-a"}) - store.RecordNode(SessionTreeInfo{SessionID: "branch-a-sub", ParentSessionID: "branch-a", AgentName: "agent-a-sub"}) - store.RecordNode(SessionTreeInfo{SessionID: "branch-b", ParentSessionID: "root", AgentName: "agent-b"}) - - // 1. GetTree for root - fullTree := store.GetTree("root") - if len(fullTree) != 4 { - t.Fatalf("GetTree(root) returned %d nodes, want 4", len(fullTree)) - } - if fullTree[0].SessionID != "root" || fullTree[0].TreeDepth != 0 { - t.Errorf("first node should be root (depth 0), got %+v", fullTree[0]) - } - if fullTree[3].SessionID != "branch-a-sub" || fullTree[3].TreeDepth != 2 { - t.Errorf("last node should be branch-a-sub (depth 2), got %+v", fullTree[3]) - } - - // 2. GetSubtree for branch-a - subTree := store.GetSubtree("branch-a") - if len(subTree) != 2 { - t.Fatalf("GetSubtree(branch-a) returned %d nodes, want 2", len(subTree)) - } - if subTree[0].SessionID != "branch-a" || subTree[1].SessionID != "branch-a-sub" { - t.Errorf("subTree nodes mismatch: %+v", subTree) - } - - // 3. Ancestors check - ancestors := store.Ancestors("branch-a-sub") - if len(ancestors) != 2 || ancestors[0] != "root" || ancestors[1] != "branch-a" { - t.Fatalf("Ancestors(branch-a-sub) = %v, want [root, branch-a]", ancestors) - } - if got := store.Ancestors("root"); len(got) != 0 { - t.Fatalf("Ancestors(root) = %v, want empty", got) - } - - // 4. UpdateAffinity - ok := store.UpdateAffinity("branch-a-sub", "auth-updated", "openai", "gpt-4o") - if !ok { - t.Fatalf("UpdateAffinity returned false") - } - node, _ := store.GetNode("branch-a-sub") - if node.LastAuthID != "auth-updated" || node.LastProvider != "openai" { - t.Errorf("node affinity not updated: %+v", node) - } -} - -func TestSessionTreeStoreExpiresNodesLazily(t *testing.T) { - t.Parallel() - - store := NewInMemorySessionTreeStore(1000, time.Millisecond) - store.RecordNode(SessionTreeInfo{SessionID: "expiring"}) - time.Sleep(5 * time.Millisecond) - - if node, ok := store.GetNode("expiring"); ok || node != nil { - t.Fatalf("GetNode() returned expired node: %#v", node) - } - if got := store.Len(); got != 0 { - t.Fatalf("Len() = %d after lazy expiry, want 0", got) - } -} - -func TestSessionTreeStoreLazyCleanupKeepsTouchedNodes(t *testing.T) { - t.Parallel() - - store := NewInMemorySessionTreeStore(1000, 30*time.Millisecond) - store.RecordNode(SessionTreeInfo{SessionID: "touched"}) - store.RecordNode(SessionTreeInfo{SessionID: "untouched"}) - time.Sleep(10 * time.Millisecond) - if !store.UpdateAffinity("touched", "auth-1", "openai", "gpt-5") { - t.Fatal("UpdateAffinity() failed for existing node") - } - time.Sleep(22 * time.Millisecond) - - if _, ok := store.GetNode("untouched"); ok { - t.Fatal("GetNode() returned a node that should have expired") - } - if _, ok := store.GetNode("touched"); !ok { - t.Fatal("GetNode() expired a node whose TTL was refreshed by UpdateAffinity") - } - if got := store.Len(); got != 1 { - t.Fatalf("Len() = %d after mixed expiry, want 1", got) - } -} - -func TestSessionTreeStoreHierarchySupportsSlashIDs(t *testing.T) { - t.Parallel() - - store := NewInMemorySessionTreeStore(1000, time.Hour) - store.RecordNode(SessionTreeInfo{SessionID: "root/with/slash"}) - store.RecordNode(SessionTreeInfo{SessionID: "child/with/slash", ParentSessionID: "root/with/slash"}) - store.RecordNode(SessionTreeInfo{SessionID: "leaf", ParentSessionID: "child/with/slash"}) - - ancestors := store.Ancestors("leaf") - if len(ancestors) != 2 || ancestors[0] != "root/with/slash" || ancestors[1] != "child/with/slash" { - t.Fatalf("Ancestors(leaf) = %v, want exact slash-bearing IDs", ancestors) - } - subtree := store.GetSubtree("child/with/slash") - if len(subtree) != 2 || subtree[0].SessionID != "child/with/slash" || subtree[1].SessionID != "leaf" { - t.Fatalf("GetSubtree() = %#v, want child and leaf", subtree) - } -} - -func TestExtractTreeInfoAllClients(t *testing.T) { - t.Parallel() - - // 1. Claude Code with Subagent - claudeHeaders := http.Header{ - "X-Claude-Code-Session-Id": []string{"claude-root-123"}, - "X-Claude-Code-Agent-Id": []string{"subagent-checker"}, - } - info, ok := ExtractTreeInfo(claudeHeaders, nil, nil) - if !ok { - t.Fatalf("ExtractTreeInfo failed for Claude Code") - } - if info.ClientType != "claude" { - t.Errorf("ClientType = %q, want claude", info.ClientType) - } - if info.SessionID != "claude:claude-root-123:agent:subagent-checker" { - t.Errorf("SessionID = %q", info.SessionID) - } - if info.ParentSessionID != "claude:claude-root-123" { - t.Errorf("ParentSessionID = %q", info.ParentSessionID) - } - if info.AgentName != "subagent-checker" { - t.Errorf("AgentName = %q", info.AgentName) - } - - // 2. Codex CLI with parent thread - codexHeaders := http.Header{ - "Session-Id": []string{"codex-child-555"}, - "x-codex-parent-thread-id": []string{"codex-parent-111"}, - } - info, ok = ExtractTreeInfo(codexHeaders, nil, nil) - if !ok || info.ClientType != "codex" { - t.Fatalf("ExtractTreeInfo failed for Codex CLI") - } - if info.SessionID != "codex:codex-child-555" || info.ParentSessionID != "codex:codex-parent-111" { - t.Errorf("Codex session ids mismatch: %+v", info) - } - - // 3. Pi Slot Session - piHeaders := http.Header{ - "X-Slot-Session-Id": []string{"pi-slot-777"}, - } - info, ok = ExtractTreeInfo(piHeaders, nil, nil) - if !ok || info.ClientType != "pi" || info.SessionID != "slot:pi-slot-777" { - t.Errorf("Pi slot mismatch: %+v", info) - } - - // 4. OpenCode Session Affinity & Parent - opencodeHeaders := http.Header{ - "X-Session-Affinity": []string{"oc-child-999"}, - "X-Parent-Session-Affinity": []string{"oc-parent-333"}, - } - info, ok = ExtractTreeInfo(opencodeHeaders, nil, nil) - if !ok || info.ClientType != "opencode" || info.SessionID != "affinity:oc-child-999" || info.ParentSessionID != "affinity:oc-parent-333" { - t.Errorf("OpenCode mismatch: %+v", info) - } - - // 5. Antigravity X-Http-Session-Id - agyHeaders := http.Header{ - "X-Http-Session-Id": []string{"agy-sess-888"}, - } - info, ok = ExtractTreeInfo(agyHeaders, nil, nil) - if !ok || info.ClientType != "agy" || info.SessionID != "agy:agy-sess-888" { - t.Errorf("Antigravity mismatch: %+v", info) - } - - // 6. Payload with metadata.agent_id and parent_session_id - payload := []byte(`{ - "session_id": "payload-child-10", - "parent_session_id": "payload-parent-01", - "metadata": { - "agent_id": "analyzer" - } - }`) - metadata := map[string]any{ - cliproxyexecutor.CallerScopeMetadataKey: "test-scope", - } - info, ok = ExtractTreeInfo(nil, payload, metadata) - if !ok || info.ClientType != "generic" { - t.Fatalf("ExtractTreeInfo failed for payload") - } - if info.SessionID != "session:payload-child-10:agent:analyzer" { - t.Errorf("SessionID = %q", info.SessionID) - } - if info.ParentSessionID != "session:payload-parent-01" { - t.Errorf("ParentSessionID = %q", info.ParentSessionID) - } - if info.CallerScope != "test-scope" { - t.Errorf("CallerScope = %q", info.CallerScope) - } -} - -func TestExtractTreeInfoCanonicalizesPayloadParentForHeaderSessions(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - headers http.Header - wantParent string - }{ - { - name: "generic header", - headers: http.Header{"X-Session-ID": []string{"child"}}, - wantParent: "header:parent", - }, - { - name: "codex header", - headers: http.Header{"Session-Id": []string{"child"}}, - wantParent: "codex:parent", - }, - { - name: "claude header", - headers: http.Header{"X-Claude-Code-Session-Id": []string{"child"}}, - wantParent: "claude:parent", - }, - } - - payload := []byte(`{"parent_session_id":"parent"}`) - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - t.Parallel() - info, ok := ExtractTreeInfo(test.headers, payload, nil) - if !ok { - t.Fatal("ExtractTreeInfo() returned no session") - } - if info.ParentSessionID != test.wantParent { - t.Fatalf("ParentSessionID = %q, want %q", info.ParentSessionID, test.wantParent) - } - }) - } -} - -func TestSessionTreeStoreConcurrencyAndLRUEviction(t *testing.T) { - t.Parallel() - - // Max 50 nodes - store := NewInMemorySessionTreeStore(50, time.Hour) - - var wg sync.WaitGroup - for i := 0; i < 100; i++ { - wg.Add(1) - go func(idx int) { - defer wg.Done() - rootID := fmt.Sprintf("root-%d", idx%10) - childID := fmt.Sprintf("child-%d", idx) - store.RecordNode(SessionTreeInfo{SessionID: rootID}) - store.RecordNode(SessionTreeInfo{SessionID: childID, ParentSessionID: rootID}) - _ = store.GetTree(rootID) - _ = store.GetSubtree(childID) - }(i) - } - wg.Wait() - - if store.Len() > 50 { - t.Fatalf("store length %d exceeded maxNodes 50", store.Len()) - } -} - -func TestSessionTreeNodeDeepCloneEmptyMetadata(t *testing.T) { - t.Parallel() - - store := NewInMemorySessionTreeStore(10, time.Hour) - meta := map[string]any{} - store.RecordNode(SessionTreeInfo{SessionID: "session-meta", Metadata: meta}) - - node1, ok1 := store.GetNode("session-meta") - if !ok1 || node1 == nil { - t.Fatal("GetNode() returned false") - } - node1.Metadata["mutated"] = true - - node2, ok2 := store.GetNode("session-meta") - if !ok2 || node2 == nil { - t.Fatal("GetNode() returned false on second read") - } - if val, exists := node2.Metadata["mutated"]; exists || val != nil { - t.Fatalf("mutation of clone leaked into tree store: %#v", node2.Metadata) - } -} - -func TestSessionTreeStoreCycleRejection(t *testing.T) { - t.Parallel() - - store := NewInMemorySessionTreeStore(10, time.Hour) - - // Record A -> B (A parent is B) - store.RecordNode(SessionTreeInfo{SessionID: "node-a", ParentSessionID: "node-b"}) - // Record B -> A (B parent is A). This should be detected as a cycle and rejected, making B a root node. - nodeB := store.RecordNode(SessionTreeInfo{SessionID: "node-b", ParentSessionID: "node-a"}) - if nodeB.ParentSessionID != "" { - t.Fatalf("node-b parent = %q, want empty (cycle rejected)", nodeB.ParentSessionID) - } - if nodeB.RootSessionID != "node-b" { - t.Fatalf("node-b root = %q, want node-b", nodeB.RootSessionID) - } - if nodeB.TreeDepth != 0 { - t.Fatalf("node-b depth = %d, want 0", nodeB.TreeDepth) - } - - ancestorsA := store.Ancestors("node-a") - if len(ancestorsA) != 1 || ancestorsA[0] != "node-b" { - t.Fatalf("Ancestors(node-a) = %v, want [node-b]", ancestorsA) - } - ancestorsB := store.Ancestors("node-b") - if len(ancestorsB) != 0 { - t.Fatalf("Ancestors(node-b) = %v, want empty", ancestorsB) - } -} - -func TestExtractTreeInfoRejectsControlCharacters(t *testing.T) { - t.Parallel() - - payloadWithNewline := []byte(`{"session_id":"bad\nsession"}`) - if _, ok := ExtractTreeInfo(nil, payloadWithNewline, nil); ok { - t.Fatal("ExtractTreeInfo accepted session_id with newline") - } - - payloadWithNull := []byte(`{"session_id":"bad\u0000session"}`) - if _, ok := ExtractTreeInfo(nil, payloadWithNull, nil); ok { - t.Fatal("ExtractTreeInfo accepted session_id with null byte") - } -} - -func TestSessionTreeStoreLateParentAndReparenting(t *testing.T) { - t.Parallel() - - store := NewInMemorySessionTreeStore(100, time.Hour) - - // 1. Child arrives before parent - child := store.RecordNode(SessionTreeInfo{ - SessionID: "child-1", - ParentSessionID: "parent-1", - }) - if child.RootSessionID != "parent-1" || child.TreeDepth != 1 || child.TreePath != "parent-1/child-1" { - t.Fatalf("unexpected initial child state: %+v", child) - } - - // 2. Grandchild arrives before root and parent - grandchild := store.RecordNode(SessionTreeInfo{ - SessionID: "grandchild-1", - ParentSessionID: "child-1", - }) - if grandchild.RootSessionID != "parent-1" || grandchild.TreeDepth != 2 || grandchild.TreePath != "parent-1/child-1/grandchild-1" { - t.Fatalf("unexpected initial grandchild state: %+v", grandchild) - } - - // 3. Parent arrives, rooted at root-1 (which is already known or self-rooted) - parent := store.RecordNode(SessionTreeInfo{ - SessionID: "parent-1", - ParentSessionID: "root-1", - }) - if parent.RootSessionID != "root-1" || parent.TreeDepth != 1 || parent.TreePath != "root-1/parent-1" { - t.Fatalf("unexpected parent state: %+v", parent) - } - - // Verify child and grandchild were cascaded to root-1 - updatedChild, ok := store.GetNode("child-1") - if !ok || updatedChild.RootSessionID != "root-1" || updatedChild.TreeDepth != 2 || updatedChild.TreePath != "root-1/parent-1/child-1" { - t.Fatalf("child not cascaded to root-1: %+v", updatedChild) - } - - updatedGrandchild, ok := store.GetNode("grandchild-1") - if !ok || updatedGrandchild.RootSessionID != "root-1" || updatedGrandchild.TreeDepth != 3 || updatedGrandchild.TreePath != "root-1/parent-1/child-1/grandchild-1" { - t.Fatalf("grandchild not cascaded to root-1: %+v", updatedGrandchild) - } - - // Check GetTree on root-1 - tree := store.GetTree("root-1") - if len(tree) != 3 { - t.Fatalf("GetTree(root-1) returned %d nodes, want 3", len(tree)) - } - - // 4. Reparent child-1 to another parent (parent-2 under root-2) - _ = store.RecordNode(SessionTreeInfo{ - SessionID: "root-2", - }) - _ = store.RecordNode(SessionTreeInfo{ - SessionID: "parent-2", - ParentSessionID: "root-2", - }) - _ = store.RecordNode(SessionTreeInfo{ - SessionID: "child-1", - ParentSessionID: "parent-2", - }) - - reparentedChild, _ := store.GetNode("child-1") - if reparentedChild.RootSessionID != "root-2" || reparentedChild.TreeDepth != 2 || reparentedChild.TreePath != "root-2/parent-2/child-1" { - t.Fatalf("reparented child state incorrect: %+v", reparentedChild) - } - - reparentedGrandchild, _ := store.GetNode("grandchild-1") - if reparentedGrandchild.RootSessionID != "root-2" || reparentedGrandchild.TreeDepth != 3 || reparentedGrandchild.TreePath != "root-2/parent-2/child-1/grandchild-1" { - t.Fatalf("reparented grandchild state incorrect: %+v", reparentedGrandchild) - } - - tree1 := store.GetTree("root-1") - if len(tree1) != 1 { // only parent-1 left - t.Fatalf("GetTree(root-1) returned %d nodes, want 1", len(tree1)) - } - tree2 := store.GetTree("root-2") - if len(tree2) != 4 { // root-2, parent-2, child-1, grandchild-1 - t.Fatalf("GetTree(root-2) returned %d nodes, want 4", len(tree2)) - } -} - -func TestSessionTreeNodeCloneDeepCopiesTypedMapsAndSlices(t *testing.T) { - original := &SessionTreeNode{ - SessionID: "test-session", - Metadata: map[string]any{ - "stringMap": map[string]string{ - "key1": "val1", - }, - "stringSlice": []string{"a", "b"}, - "intSlice": []int{1, 2, 3}, - "byteSlice": []byte("hello"), - }, - } - - clone := original.Clone() - if clone == nil { - t.Fatal("clone returned nil") - } - - // Mutate clone - if sm, ok := clone.Metadata["stringMap"].(map[string]string); ok { - sm["key1"] = "mutated" - sm["key2"] = "val2" - } else { - t.Fatal("expected stringMap in clone metadata") - } - - if origSM, ok := original.Metadata["stringMap"].(map[string]string); ok { - if origSM["key1"] != "val1" || len(origSM) != 1 { - t.Fatalf("original stringMap was mutated: %v", origSM) - } - } -} - -func TestExtractTreeInfoClaudeNestedParent(t *testing.T) { - // Nested metadata.user_id containing session_id and parent_session_id - payload := []byte(`{ - "metadata": { - "user_id": "{\"session_id\":\"child-session-123\",\"parent_session_id\":\"parent-session-456\",\"agent_id\":\"subagent-worker\"}" - } - }`) - info, ok := ExtractTreeInfo(nil, payload, nil) - if !ok { - t.Fatal("ExtractTreeInfo failed on nested Claude user_id") - } - if info.SessionID != "claude:child-session-123:agent:subagent-worker" { - t.Fatalf("expected child session with agent, got %q", info.SessionID) - } - if info.ParentSessionID != "claude:parent-session-456" { - t.Fatalf("expected parent claude:parent-session-456, got %q", info.ParentSessionID) - } - - // Explicit header + payload body parent_session_id - headerPayload := []byte(`{"parent_session_id":"parent-session-789"}`) - headers := http.Header{} - headers.Set("X-Claude-Code-Session-Id", "header-child-123") - info2, ok := ExtractTreeInfo(headers, headerPayload, nil) - if !ok { - t.Fatal("ExtractTreeInfo failed on header + body parent") - } - if info2.SessionID != "claude:header-child-123" { - t.Fatalf("expected session claude:header-child-123, got %q", info2.SessionID) - } - if info2.ParentSessionID != "claude:parent-session-789" { - t.Fatalf("expected parent claude:parent-session-789, got %q", info2.ParentSessionID) - } -} - -func TestExtractTreeInfoGeminiAndAntigravityHierarchy(t *testing.T) { - // Gemini cachedContent with parent_session_id - geminiPayload := []byte(`{"cachedContent":"cache-child-1","parent_session_id":"cache-parent-1"}`) - info, ok := ExtractTreeInfo(nil, geminiPayload, nil) - if !ok { - t.Fatal("ExtractTreeInfo failed on Gemini cachedContent") - } - if info.SessionID != "geminicache:cache-child-1" { - t.Fatalf("expected session geminicache:cache-child-1, got %q", info.SessionID) - } - if info.ParentSessionID != "geminicache:cache-parent-1" { - t.Fatalf("expected parent geminicache:cache-parent-1, got %q", info.ParentSessionID) - } - if info.AgentName != "subagent" { - t.Fatalf("expected subagent, got %q", info.AgentName) - } - - // Antigravity headers with parent - headers := http.Header{} - headers.Set("X-Http-Session-Id", "agy-child-2") - headers.Set("X-Parent-Session-ID", "agy-parent-2") - info2, ok := ExtractTreeInfo(headers, nil, nil) - if !ok { - t.Fatal("ExtractTreeInfo failed on Antigravity headers") - } - if info2.SessionID != "agy:agy-child-2" { - t.Fatalf("expected session agy:agy-child-2, got %q", info2.SessionID) - } - if info2.ParentSessionID != "agy:agy-parent-2" { - t.Fatalf("expected parent agy:agy-parent-2, got %q", info2.ParentSessionID) - } - if info2.AgentName != "subagent" { - t.Fatalf("expected subagent, got %q", info2.AgentName) - } -} - -func TestSessionTreeStoreMaxDepthCapping(t *testing.T) { - t.Parallel() - - store := NewInMemorySessionTreeStore(500, time.Hour) - current := "node-0" - store.RecordNode(SessionTreeInfo{SessionID: current}) - - // Create a chain exceeding maxTreeDepth (128) - for i := 1; i <= 150; i++ { - next := fmt.Sprintf("node-%d", i) - node := store.RecordNode(SessionTreeInfo{ - SessionID: next, - ParentSessionID: current, - }) - if i <= 128 { - if node.TreeDepth != i { - t.Fatalf("expected depth %d, got %d", i, node.TreeDepth) - } - } else { - // Once exceeding maxTreeDepth, the chain is capped and the node becomes its own root - if node.TreeDepth > maxTreeDepth { - t.Fatalf("depth %d exceeded maxTreeDepth %d", node.TreeDepth, maxTreeDepth) - } - } - current = next - } - - // Verify that the capped node's parent index was cleaned up - cappedParent := fmt.Sprintf("node-%d", 128) - cappedNode := fmt.Sprintf("node-%d", 129) - store.mu.Lock() - if children, ok := store.parentIndex[cappedParent]; ok { - if _, exists := children[cappedNode]; exists { - store.mu.Unlock() - t.Fatalf("capped node %s still found in parentIndex[%s]", cappedNode, cappedParent) - } - } - store.mu.Unlock() -} - -func TestSessionTreeStoreMetadataDeepRecursionProtection(t *testing.T) { - t.Parallel() - - store := NewInMemorySessionTreeStore(100, time.Hour) - - // Create deeply nested map - nested := map[string]any{"level": 0} - curr := nested - for i := 1; i <= 30; i++ { - next := map[string]any{"level": i} - curr["child"] = next - curr = next - } - - // Recording node should succeed without stack overflow - node := store.RecordNode(SessionTreeInfo{ - SessionID: "sess-deep-meta", - Metadata: nested, - }) - if node.SessionID != "sess-deep-meta" { - t.Fatalf("expected sess-deep-meta, got %s", node.SessionID) - } -} - -func TestExtractTreeInfoNestedAntigravityRequest(t *testing.T) { - t.Parallel() - - payload := []byte(`{ - "project_id": "proj-123", - "request": { - "parentSessionId": "parent-sess-456", - "sessionId": "child-sess-789" - } - }`) - info, ok := ExtractTreeInfo(nil, payload, nil) - if !ok { - t.Fatal("ExtractTreeInfo failed on nested Antigravity payload") - } - if info.SessionID != "session:child-sess-789" { - t.Fatalf("SessionID = %q, want session:child-sess-789", info.SessionID) - } - if info.ParentSessionID != "session:parent-sess-456" { - t.Fatalf("ParentSessionID = %q, want session:parent-sess-456", info.ParentSessionID) - } -} - -func BenchmarkSessionTreeRecordAndCascade(b *testing.B) { - store := NewInMemorySessionTreeStore(10000, time.Hour) - // Pre-populate root and parent - store.RecordNode(SessionTreeInfo{SessionID: "root-node"}) - for i := 0; i < 50; i++ { - store.RecordNode(SessionTreeInfo{ - SessionID: fmt.Sprintf("parent-%d", i), - ParentSessionID: "root-node", - }) - } - b.ReportAllocs() - b.ResetTimer() - for i := 0; b.Loop(); i++ { - store.RecordNode(SessionTreeInfo{ - SessionID: fmt.Sprintf("child-%d", i%1000), - ParentSessionID: fmt.Sprintf("parent-%d", i%50), - }) - } -} - -func BenchmarkSessionTreeGetSubtree(b *testing.B) { - store := NewInMemorySessionTreeStore(10000, time.Hour) - store.RecordNode(SessionTreeInfo{SessionID: "root-node"}) - for i := 0; i < 20; i++ { - p := fmt.Sprintf("parent-%d", i) - store.RecordNode(SessionTreeInfo{SessionID: p, ParentSessionID: "root-node"}) - for j := 0; j < 10; j++ { - c := fmt.Sprintf("child-%d-%d", i, j) - store.RecordNode(SessionTreeInfo{SessionID: c, ParentSessionID: p}) - } - } - b.ReportAllocs() - b.ResetTimer() - for b.Loop() { - _ = store.GetSubtree("root-node") - } -} -- 2.54.0