From aa3e9b8ee93ede5a11a4ad7bc93a7c47b684f62a Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 6 Oct 2026 14:30:01 -0700 Subject: [PATCH] fix(tracing): scope tool ownership and bound actor grouping --- .../crates/traces/src/resolve/actors.rs | 11 ++- .../crates/traces/src/resolve/capture.rs | 8 +- .../crates/traces/src/resolve/resolution.rs | 4 +- litellm-rust/crates/traces/tests/resolve.rs | 76 +++++++++++++++++++ .../detail/conversation/conversation.test.ts | 54 +++++++++++++ .../detail/conversation/conversation.ts | 56 ++++++++------ 6 files changed, 176 insertions(+), 33 deletions(-) diff --git a/litellm-rust/crates/traces/src/resolve/actors.rs b/litellm-rust/crates/traces/src/resolve/actors.rs index ca12dfa18b6..75803eba139 100644 --- a/litellm-rust/crates/traces/src/resolve/actors.rs +++ b/litellm-rust/crates/traces/src/resolve/actors.rs @@ -78,7 +78,7 @@ impl Actors { let rows = graph.rows; let direct: Vec<_> = rows.iter().map(explicit).collect(); let mut boundaries: HashMap> = HashMap::new(); - let mut tools: HashMap<&str, BTreeSet<&str>> = HashMap::new(); + let mut tools: HashMap<(&str, &str), BTreeSet<&str>> = HashMap::new(); for (index, owner) in direct.iter().enumerate() { let Some(owner) = owner.as_deref() else { continue; @@ -90,7 +90,10 @@ impl Actors { boundaries.entry(parent).or_default().insert(owner); } if !row.tool_call_id.is_empty() { - tools.entry(&row.tool_call_id).or_default().insert(owner); + tools + .entry((&row.session_id, &row.tool_call_id)) + .or_default() + .insert(owner); } } let unique = |set: Option<&BTreeSet<&str>>| { @@ -126,7 +129,9 @@ impl Actors { direct[index] .clone() .or_else(|| unique(boundaries.get(&index))) - .or_else(|| unique(tools.get(row.tool_call_id.as_str()))) + .or_else(|| { + unique(tools.get(&(row.session_id.as_str(), row.tool_call_id.as_str()))) + }) .or_else(|| { let ancestor = if child_source(row) { inherited_child[index] diff --git a/litellm-rust/crates/traces/src/resolve/capture.rs b/litellm-rust/crates/traces/src/resolve/capture.rs index 1c14c43a6e0..8f690988d18 100644 --- a/litellm-rust/crates/traces/src/resolve/capture.rs +++ b/litellm-rust/crates/traces/src/resolve/capture.rs @@ -1,10 +1,7 @@ use std::collections::{BTreeMap, HashSet}; use super::{actors, resolution::Resolution}; -use crate::{ - ObservationType, - view::{ActorCapture, TraceCapture}, -}; +use crate::view::{ActorCapture, TraceCapture}; pub(super) fn coverage(resolution: &Resolution<'_>) -> Option { let rows = resolution.graph.rows; @@ -12,6 +9,7 @@ pub(super) fn coverage(resolution: &Resolution<'_>) -> Option { return None; } let calls: HashSet<_> = resolution.model_calls.iter().copied().collect(); + let tools: HashSet<_> = resolution.unique_tools().into_iter().collect(); let mut actors: BTreeMap<_, _> = resolution .actors .entries @@ -47,7 +45,7 @@ pub(super) fn coverage(resolution: &Resolution<'_>) -> Option { continue; }; actor.llm_calls += u64::from(calls.contains(&index)); - actor.tool_calls += u64::from(resolution.kind(index) == ObservationType::Tool); + actor.tool_calls += u64::from(tools.contains(&index)); actor.reply_events += u64::from(row.name == "claude_code.assistant_response"); actor.model_outputs += u64::from(calls.contains(&index) && row.has_output); actor.content_events += u64::from(event); diff --git a/litellm-rust/crates/traces/src/resolve/resolution.rs b/litellm-rust/crates/traces/src/resolve/resolution.rs index cdfb9f4b466..d3d55a832b1 100644 --- a/litellm-rust/crates/traces/src/resolve/resolution.rs +++ b/litellm-rust/crates/traces/src/resolve/resolution.rs @@ -246,7 +246,7 @@ impl<'a> Resolution<'a> { } pub(super) fn unique_tools(&self) -> Vec { - let mut by_call: IndexMap<&str, usize> = IndexMap::new(); + let mut by_call: IndexMap<(&str, &str), usize> = IndexMap::new(); for index in (0..self.graph.rows.len()).filter(|index| self.kind(*index) == ObservationType::Tool) { @@ -256,7 +256,7 @@ impl<'a> Resolution<'a> { } else { &row.tool_call_id }; - by_call.entry(key).or_insert(index); + by_call.entry((&row.session_id, key)).or_insert(index); } by_call.into_values().collect() } diff --git a/litellm-rust/crates/traces/tests/resolve.rs b/litellm-rust/crates/traces/tests/resolve.rs index 65aaae457b2..6f1d35c305b 100644 --- a/litellm-rust/crates/traces/tests/resolve.rs +++ b/litellm-rust/crates/traces/tests/resolve.rs @@ -2359,3 +2359,79 @@ fn many_native_actors_count_only_their_own_calls() { .all(|actor| actor.llm_calls == 1 && actor.tool_calls == 0) ); } + +#[rstest] +#[case::same_session(false)] +#[case::different_sessions(true)] +fn native_tool_ownership_and_coverage_use_session_scoped_call_ids(#[case] separate: bool) { + let root = TraceSpansRow { + framework: "claude-code".into(), + session_id: "first-session".into(), + query_source: "repl_main_thread".into(), + ..row("root", "", "claude_code.interaction", "agent", "assistant") + }; + let other_root = TraceSpansRow { + session_id: if separate { + "second-session" + } else { + "first-session" + } + .into(), + span_id: "other-root".into(), + ..root.clone() + }; + let explicit = TraceSpansRow { + framework: "claude-code".into(), + session_id: root.session_id.clone(), + native_agent_id: "child".into(), + tool_call_id: "reused-call".into(), + ..row("explicit", "root", "Bash", "tool", "assistant") + }; + let inherited = TraceSpansRow { + native_agent_id: String::new(), + session_id: other_root.session_id.clone(), + span_id: "inherited".into(), + parent_span_id: "other-root".into(), + ..explicit.clone() + }; + let trace = resolve_trace( + "trace", + "ref", + &[root, other_root, explicit, inherited], + &[], + ) + .unwrap(); + let explicit_actor = &trace + .spans + .iter() + .find(|span| span.span_id == "explicit") + .unwrap() + .actor_id; + let inherited_actor = &trace + .spans + .iter() + .find(|span| span.span_id == "inherited") + .unwrap() + .actor_id; + assert_eq!(explicit_actor == inherited_actor, !separate); + let expected = if separate { 2 } else { 1 }; + assert_eq!(trace.summary.tool_calls, expected); + assert_eq!( + trace + .agents + .iter() + .map(|actor| actor.tool_calls) + .sum::(), + expected + ); + assert_eq!( + trace + .capture + .unwrap() + .actors + .iter() + .map(|actor| actor.tool_calls) + .sum::(), + expected + ); +} diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.test.ts b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.test.ts index 862e3058158..6b36095a110 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.test.ts +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.test.ts @@ -40,6 +40,60 @@ describe("trace conversation", () => { expect(pending.has("child")).toBe(true); }); + it("groups 20,000 sibling actors without rescanning unrelated conversation entries", () => { + const spans = Array.from({ length: 20_000 }, (_, index) => ({ + ...root, + span_id: `step-${index}`, + actor_id: `actor-${index}`, + parent_actor_id: "root-actor", + })); + const items = spans.map((span) => ({ + id: span.span_id, + span, + messages: [], + branchId: span.actor_id, + parentBranchId: "root-actor", + agentName: span.actor_id, + })); + const groups = groupConversation(items, spans); + expect(groups).toHaveLength(spans.length); + expect( + groups.every((group, index) => { + if (group.kind !== "branch" || group.id !== `actor-${index}`) return false; + return ( + group.children.length === 1 && group.children[0].kind === "item" && group.children[0].item === items[index] + ); + }), + ).toBe(true); + }); + + it("retains each actor once when parent identities form a cycle", () => { + const spans = ["first", "second"].map((id, index) => ({ + ...root, + span_id: id, + actor_id: id, + parent_actor_id: index ? "first" : "second", + })); + const items = spans.map((span) => ({ + id: span.span_id, + span, + messages: [], + branchId: span.actor_id, + parentBranchId: span.parent_actor_id, + })); + expect(groupConversation(items, spans)).toEqual([ + { + kind: "branch", + id: "second", + name: root.name, + children: [ + { kind: "branch", id: "first", name: root.name, children: [{ kind: "item", item: items[0] }] }, + { kind: "item", item: items[1] }, + ], + }, + ]); + }); + it("builds a conversation below 20,000 hidden framework spans", () => { const chain = Array.from( { length: 20_000 }, diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.ts b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.ts index a565612adfc..9e864014b0a 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.ts +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/conversation.ts @@ -441,30 +441,40 @@ export function groupConversation(items: readonly ConversationItem[], spans: rea return parentById.has(parent) ? [] : [parent]; }), ); - const directBranch = (item: ConversationItem, parent?: string): string | undefined => { - let id = item.branchId; - const visited = new Set(); - while (id && !visited.has(id)) { - visited.add(id); - const ancestor = parentById.get(id); - if (parent ? ancestor === parent : ancestor !== undefined && roots.has(ancestor)) return id; - id = ancestor; + const members = new Map(); + const groupsFor = (id?: string): ConversationGroup[] => { + const existing = members.get(id); + if (existing) return existing; + const groups: ConversationGroup[] = []; + members.set(id, groups); + return groups; + }; + const names = new Map(); + for (const item of items) { + if (item.branchId && item.agentName && !names.has(item.branchId)) names.set(item.branchId, item.agentName); + } + const registered = new Set(); + for (const item of items) { + const ownBranch = item.parentBranchId ? item.branchId : undefined; + groupsFor(ownBranch).push({ kind: "item", item }); + const path = new Set(); + let id = ownBranch; + while (id && !roots.has(id) && !registered.has(id)) { + path.add(id); + registered.add(id); + const parent = parentById.get(id); + const container = parent && !roots.has(parent) && !path.has(parent) ? parent : undefined; + const group: ConversationGroup = { + kind: "branch", + id, + name: names.get(id) || byId.get(id)?.name || "Subagent", + children: groupsFor(id), + }; + groupsFor(container).push(group); + id = container; } - return undefined; - }; - const build = (parent?: string, ancestors = new Set()): ConversationGroup[] => { - const seen = new Set(); - return items.flatMap((item): ConversationGroup[] => { - if (parent ? item.branchId === parent : !item.parentBranchId) return [{ kind: "item", item }]; - const id = directBranch(item, parent); - if (!id || seen.has(id) || ancestors.has(id)) return []; - seen.add(id); - const first = items.find((candidate) => candidate.branchId === id); - const name = first?.agentName || byId.get(id)?.name || "Subagent"; - return [{ kind: "branch", id, name, children: build(id, new Set([...ancestors, id])) }]; - }); - }; - return build(); + } + return groupsFor(); } export function conversationWarnings(details: ReadonlyMap, complete: boolean): string[] {