diff --git a/litellm-rust/crates/traces-cache/src/cache.rs b/litellm-rust/crates/traces-cache/src/cache.rs index a5526a39963..980ce399574 100644 --- a/litellm-rust/crates/traces-cache/src/cache.rs +++ b/litellm-rust/crates/traces-cache/src/cache.rs @@ -58,9 +58,8 @@ impl SnapshotKey { } } -/// How long a read result stays reusable: traces still receiving spans, or read with spend -/// unavailable, are re-read after `LIVE_TTL`; traces quiet for `SETTLED_AFTER_MS` are kept for -/// `SETTLED_TTL`. +/// Native sessions can resume without a terminal record, so their reads retain `LIVE_TTL`. +/// Other traces with known spend settle after `SETTLED_AFTER_MS` of inactivity. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum Freshness { Live, @@ -69,6 +68,12 @@ pub enum Freshness { impl Freshness { pub fn of(rows: &[TraceSpansRow], spend_known: bool, snapshot_ms: u64) -> Self { + if rows + .iter() + .any(|row| matches!(row.framework.as_str(), "claude-code" | "claude-agent-sdk")) + { + return Self::Live; + } let last_end_ms = rows .iter() .map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns) / 1_000_000) @@ -297,6 +302,12 @@ mod tests { fn row(span_id: &str) -> TraceSpansRow { TraceSpansRow { + native_agent_id: String::new(), + native_parent_agent_id: String::new(), + session_id: String::new(), + query_source: String::new(), + has_output: false, + capture_warning: false, trace_id: String::new(), original_trace_id: String::new(), span_id: span_id.into(), @@ -393,4 +404,18 @@ mod tests { expected ); } + + #[rstest] + #[case::claude_code("claude-code")] + #[case::claude_agent_sdk("claude-agent-sdk")] + fn idle_native_sessions_remain_live(#[case] framework: &str) { + let native = TraceSpansRow { + framework: framework.into(), + ..row("native") + }; + assert_eq!( + Freshness::of(&[row("root"), native], true, LAST_END_MS + SETTLED_AFTER_MS), + Freshness::Live + ); + } } diff --git a/litellm-rust/crates/traces-cache/src/reader.rs b/litellm-rust/crates/traces-cache/src/reader.rs index 6aacb2febb0..96a05969d9e 100644 --- a/litellm-rust/crates/traces-cache/src/reader.rs +++ b/litellm-rust/crates/traces-cache/src/reader.rs @@ -1,4 +1,4 @@ -use std::{sync::Arc, time::Duration}; +use std::{collections::HashSet, sync::Arc, time::Duration}; use crate::{ ReadError, Snapshot, SnapshotCache, SnapshotKey, StoreError, TraceStore, @@ -345,10 +345,39 @@ fn page( let spans = &snapshot.trace().spans; let create_page = |count: usize| { let end = position.offset.saturating_add(count).min(spans.len()); + let page_spans = &spans[position.offset..end]; + let actors: HashSet<&str> = page_spans + .iter() + .filter_map(|span| span.actor_id.as_deref()) + .collect(); Trace { + capture: snapshot.trace().capture.as_ref().map(|capture| { + litellm_traces::TraceCapture { + actors: capture + .actors + .iter() + .filter(|actor| actors.contains(actor.actor_id.as_str())) + .cloned() + .collect(), + content_events: capture.content_events, + unassigned_events: capture.unassigned_events, + warning_events: capture.warning_events, + } + }), summary: snapshot.trace().summary.clone(), - agents: snapshot.trace().agents.clone(), - spans: spans[position.offset..end].to_vec(), + agents: snapshot + .trace() + .agents + .iter() + .filter(|agent| { + agent + .actor_id + .as_deref() + .is_none_or(|id| actors.contains(id)) + }) + .cloned() + .collect(), + spans: page_spans.to_vec(), next_cursor: (end < spans.len()).then(|| { encode_cursor(&SpanPosition { trace_ref: position.trace_ref.clone(), diff --git a/litellm-rust/crates/traces-cache/tests/read.rs b/litellm-rust/crates/traces-cache/tests/read.rs index 02f11fdd1b7..50fd08bef38 100644 --- a/litellm-rust/crates/traces-cache/tests/read.rs +++ b/litellm-rust/crates/traces-cache/tests/read.rs @@ -269,6 +269,12 @@ fn access() -> ReadAccessParams { fn span(index: usize) -> TraceSpansRow { TraceSpansRow { + native_agent_id: String::new(), + native_parent_agent_id: String::new(), + session_id: String::new(), + query_source: String::new(), + has_output: false, + capture_warning: false, trace_id: "trace".into(), original_trace_id: String::new(), span_id: format!("span-{index}"), @@ -495,6 +501,76 @@ async fn list_run_budget_halves_the_limit_and_cursor_requires_a_full_page() { assert_eq!(shorter.calls(Operation::ListRuns), 1); } +#[rstest] +#[case::small_response(4096)] +#[case::larger_response(8192)] +#[tokio::test] +async fn actor_metadata_pages_with_spans_without_losing_trace_totals( + #[case] response_bytes: usize, +) { + let rows = (0..120) + .map(|index| TraceSpansRow { + framework: "claude-code".into(), + session_id: "native-session".into(), + native_agent_id: format!("actor-{}", index / 2), + parent_span_id: String::new(), + kind: if index % 2 == 0 { + ObservationType::Llm + } else { + ObservationType::Framework + }, + name: if index % 2 == 0 { + "claude_code.llm_request" + } else { + "claude_code.assistant_response" + } + .into(), + ..span(index) + }) + .collect(); + let store = FakeStore::with_spans("ref", rows); + let reader = TraceReader::new(response_bytes); + let access = access(); + let mut cursor = None; + let mut actors = HashSet::new(); + let mut spans = HashSet::new(); + loop { + let page = reader + .get_trace_page(&store, &access, "trace", "ref", cursor.as_deref(), 200) + .await + .unwrap() + .unwrap(); + assert!(serde_json::to_vec(&page).unwrap().len() <= response_bytes); + assert_eq!(page.summary.agent_count, 60); + assert_eq!(page.summary.span_count, 120); + let coverage = page.capture.unwrap(); + assert_eq!(coverage.content_events, 60); + let page_actors: HashSet<_> = page + .spans + .iter() + .filter_map(|span| span.actor_id.as_deref()) + .collect(); + assert_eq!(page.agents.len(), page_actors.len()); + assert_eq!(coverage.actors.len(), page_actors.len()); + for actor in coverage.actors { + assert!(page_actors.contains(actor.actor_id.as_str())); + assert_eq!(actor.llm_calls, 1); + assert_eq!(actor.reply_events, 1); + actors.insert(actor.actor_id); + } + for span in page.spans { + assert!(spans.insert(span.span_id)); + } + cursor = page.next_cursor; + if cursor.is_none() { + break; + } + } + assert_eq!(actors.len(), 60); + assert_eq!(spans.len(), 120); + assert_eq!(store.calls(Operation::TraceSpans), 1); +} + #[rstest] #[tokio::test] async fn oversized_run_batch_falls_back_to_each_run_and_keeps_listed_summaries() { @@ -716,6 +792,54 @@ async fn failed_reads_are_not_cached() { assert_eq!(store.calls(Operation::TraceSpans), 2); } +#[rstest] +#[case::claude_code("claude-code")] +#[case::claude_agent_sdk("claude-agent-sdk")] +#[tokio::test] +async fn resumed_native_sessions_refresh_after_live_ttl(#[case] framework: &str) { + let original = TraceSpansRow { + framework: framework.into(), + ..span(0) + }; + let store = FakeStore::with_spans("ref", vec![original.clone()]); + let reader = TraceReader::new(usize::MAX); + let access = access(); + let first = reader + .get_trace(&store, &access, "trace", "ref") + .await + .unwrap() + .unwrap(); + assert_eq!(first.spans.len(), 1); + + store.state.lock().unwrap().trace_spans.insert( + "ref".into(), + vec![ + original, + TraceSpansRow { + start_ns: now_ns(), + ..span(1) + }, + ], + ); + let cached = reader + .get_trace(&store, &access, "trace", "ref") + .await + .unwrap() + .unwrap(); + assert_eq!(cached.spans.len(), 1); + assert_eq!(store.calls(Operation::TraceSpans), 1); + + tokio::time::sleep(LIVE_TTL + Duration::from_millis(200)).await; + let resumed = reader + .get_trace(&store, &access, "trace", "ref") + .await + .unwrap() + .unwrap(); + assert_eq!(resumed.spans.len(), 2); + assert_eq!(store.calls(Operation::TraceSpans), 2); + assert_eq!(first.spans.len(), 1); +} + #[rstest] #[tokio::test] async fn listed_runs_are_read_once_until_a_live_run_expires() { diff --git a/litellm-rust/crates/traces-cache/tests/snapshots.rs b/litellm-rust/crates/traces-cache/tests/snapshots.rs index 0a206004080..a77a26d59a8 100644 --- a/litellm-rust/crates/traces-cache/tests/snapshots.rs +++ b/litellm-rust/crates/traces-cache/tests/snapshots.rs @@ -16,6 +16,12 @@ const TTL: Duration = Duration::from_secs(120); fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> TraceSpansRow { TraceSpansRow { + native_agent_id: String::new(), + native_parent_agent_id: String::new(), + session_id: String::new(), + query_source: String::new(), + has_output: false, + capture_warning: false, trace_id: String::new(), original_trace_id: String::new(), span_id: span_id.into(), diff --git a/litellm-rust/crates/traces-clickhouse/query/list_traces.sql b/litellm-rust/crates/traces-clickhouse/query/list_traces.sql index c52adf7ef49..697feec35ab 100644 --- a/litellm-rust/crates/traces-clickhouse/query/list_traces.sql +++ b/litellm-rust/crates/traces-clickhouse/query/list_traces.sql @@ -1,4 +1,14 @@ -WITH page AS ( +WITH + coalesce(nullIf(SpanAttributes['query_source_safe'], ''), SpanAttributes['query_source']) AS native_source, + ((Framework = 'claude-code' OR (Framework = 'claude-agent-sdk' AND + (startsWith(SpanName, 'claude_code.') OR SpanAttributes['agent_id'] != '' OR native_source != ''))) + AND (SpanAttributes['session.id'] != '' OR SpanAttributes['agent_id'] != '')) AS native, + if(SpanAttributes['session.id'] = '', TraceId, SpanAttributes['session.id']) AS native_session, + (native AND SpanAttributes['agent_id'] != '') AS native_child, + (native AND NOT native_child AND (SpanName = 'claude_code.interaction' OR + native_source IN + ('repl_main_thread', 'sdk', 'sdk_main_thread', 'generate_session_title', 'prompt_suggestion'))) AS native_root, + page AS ( SELECT TraceId AS trace_id, hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref, if(length(groupUniqArrayArray(UserIds)) = 1, arrayElement(groupUniqArrayArray(UserIds), 1), '') AS user_id, TeamId AS team_id, ApiKeyHash AS api_key_hash, @@ -27,20 +37,34 @@ HAVING min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64}) ORDER BY start_ms DESC, trace_ref DESC LIMIT {limit:UInt32} ) -SELECT page.* EXCEPT (trace_start, trace_end), +SELECT page.* EXCEPT (trace_start, trace_end, agent_invocations), + if(identities.native_count > 0, identities.resolved_invocations, page.agent_invocations) AS agent_invocations, identities.agent_names AS agent_names, identities.agent_count AS agent_count, identities.frameworks AS frameworks FROM page LEFT JOIN ( SELECT TeamId, ApiKeyHash, TraceId, - arraySort(groupUniqArrayIf(AgentName, AgentName != '')) AS agent_names, - arraySort(groupUniqArrayIf(toString(Framework), Framework != '')) AS frameworks, - uniqExactIf(if(AgentName = '', SpanName, AgentName), ObservationType = 'agent') AS agent_count - FROM otel_traces - WHERE Timestamp >= (SELECT min(trace_start) FROM page) - AND Timestamp <= (SELECT max(trace_end) FROM page) - AND TraceId IN (SELECT trace_id FROM page) - AND (TeamId, ApiKeyHash, TraceId) IN (SELECT team_id, api_key_hash, trace_id FROM page) + arraySort(groupUniqArrayArray(agent_names)) AS agent_names, + arraySort(groupUniqArrayArray(frameworks)) AS frameworks, + countIf(native_actor) AS native_count, + native_count + uniqExactIf(legacy_name, legacy_invocations > 0) AS agent_count, + sum(if(native_actor, greatest(1, native_interactions), legacy_invocations)) AS resolved_invocations + FROM ( + SELECT TeamId, ApiKeyHash, TraceId, + native_root OR native_child AS native_actor, + (native_session, if(native_child, concat('agent:', SpanAttributes['agent_id']), 'root')) AS actor_identity, + if(native_actor, '', if(AgentName = '', SpanName, AgentName)) AS legacy_name, + groupUniqArrayIf(AgentName, AgentName != '') AS agent_names, + groupUniqArrayIf(toString(Framework), Framework != '') AS frameworks, + uniqExactIf(SpanId, native_actor AND SpanName = 'claude_code.interaction') AS native_interactions, + uniqExactIf(SpanId, ObservationType = 'agent' AND NOT native) AS legacy_invocations + FROM otel_traces + WHERE Timestamp >= (SELECT min(trace_start) FROM page) + AND Timestamp <= (SELECT max(trace_end) FROM page) + AND TraceId IN (SELECT trace_id FROM page) + AND (TeamId, ApiKeyHash, TraceId) IN (SELECT team_id, api_key_hash, trace_id FROM page) + GROUP BY TeamId, ApiKeyHash, TraceId, native_actor, actor_identity, legacy_name + ) GROUP BY TeamId, ApiKeyHash, TraceId ) AS identities ON page.team_id = identities.TeamId AND page.api_key_hash = identities.ApiKeyHash diff --git a/litellm-rust/crates/traces-clickhouse/query/trace_list_span_batch.sql b/litellm-rust/crates/traces-clickhouse/query/trace_list_span_batch.sql index d6ab416dfd3..3a6501a4f6e 100644 --- a/litellm-rust/crates/traces-clickhouse/query/trace_list_span_batch.sql +++ b/litellm-rust/crates/traces-clickhouse/query/trace_list_span_batch.sql @@ -1,6 +1,12 @@ SELECT * FROM ( SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS original_trace_id, o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name, o.ObservationType AS type, toUInt8(o.WrapperCandidate) AS wrapper_candidate, o.AgentName AS agent, + o.SpanAttributes['agent_id'] AS native_agent_id, + o.SpanAttributes['parent_agent_id'] AS native_parent_agent_id, + o.SpanAttributes['session.id'] AS session_id, + coalesce(nullIf(o.SpanAttributes['query_source_safe'], ''), o.SpanAttributes['query_source']) AS query_source, + o.Output != '' AS has_output, + o.SpanAttributes['lens.capture.warning'] != '' AS capture_warning, o.Framework AS framework, o.StatusCode AS status, substringUTF8(o.StatusMessage, 1, 128) AS status_message, lengthUTF8(o.StatusMessage) > 128 AS error_truncated, diff --git a/litellm-rust/crates/traces-clickhouse/query/trace_page_spans.sql b/litellm-rust/crates/traces-clickhouse/query/trace_page_spans.sql index 27a8e2ac0da..23857064fce 100644 --- a/litellm-rust/crates/traces-clickhouse/query/trace_page_spans.sql +++ b/litellm-rust/crates/traces-clickhouse/query/trace_page_spans.sql @@ -1,5 +1,11 @@ SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS original_trace_id, o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name, o.ObservationType AS type, toUInt8(o.WrapperCandidate) AS wrapper_candidate, o.AgentName AS agent, + o.SpanAttributes['agent_id'] AS native_agent_id, + o.SpanAttributes['parent_agent_id'] AS native_parent_agent_id, + o.SpanAttributes['session.id'] AS session_id, + coalesce(nullIf(o.SpanAttributes['query_source_safe'], ''), o.SpanAttributes['query_source']) AS query_source, + o.Output != '' AS has_output, + o.SpanAttributes['lens.capture.warning'] != '' AS capture_warning, o.Framework AS framework, o.StatusCode AS status, substringUTF8(o.StatusMessage, 1, 128) AS status_message, lengthUTF8(o.StatusMessage) > 128 AS error_truncated, diff --git a/litellm-rust/crates/traces-clickhouse/query/trace_span_batch.sql b/litellm-rust/crates/traces-clickhouse/query/trace_span_batch.sql index a084325feaa..394344841cb 100644 --- a/litellm-rust/crates/traces-clickhouse/query/trace_span_batch.sql +++ b/litellm-rust/crates/traces-clickhouse/query/trace_span_batch.sql @@ -1,6 +1,12 @@ SELECT * FROM ( SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS original_trace_id, o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name, o.ObservationType AS type, toUInt8(o.WrapperCandidate) AS wrapper_candidate, o.AgentName AS agent, + o.SpanAttributes['agent_id'] AS native_agent_id, + o.SpanAttributes['parent_agent_id'] AS native_parent_agent_id, + o.SpanAttributes['session.id'] AS session_id, + coalesce(nullIf(o.SpanAttributes['query_source_safe'], ''), o.SpanAttributes['query_source']) AS query_source, + o.Output != '' AS has_output, + o.SpanAttributes['lens.capture.warning'] != '' AS capture_warning, o.Framework AS framework, o.StatusCode AS status, substringUTF8(o.StatusMessage, 1, 128) AS status_message, lengthUTF8(o.StatusMessage) > 128 AS error_truncated, diff --git a/litellm-rust/crates/traces-clickhouse/query/trace_spans.sql b/litellm-rust/crates/traces-clickhouse/query/trace_spans.sql index 306692e9709..6da8dd9735c 100644 --- a/litellm-rust/crates/traces-clickhouse/query/trace_spans.sql +++ b/litellm-rust/crates/traces-clickhouse/query/trace_spans.sql @@ -1,5 +1,11 @@ SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS original_trace_id, o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name, o.ObservationType AS type, toUInt8(o.WrapperCandidate) AS wrapper_candidate, o.AgentName AS agent, + o.SpanAttributes['agent_id'] AS native_agent_id, + o.SpanAttributes['parent_agent_id'] AS native_parent_agent_id, + o.SpanAttributes['session.id'] AS session_id, + coalesce(nullIf(o.SpanAttributes['query_source_safe'], ''), o.SpanAttributes['query_source']) AS query_source, + o.Output != '' AS has_output, + o.SpanAttributes['lens.capture.warning'] != '' AS capture_warning, o.Framework AS framework, o.StatusCode AS status, substringUTF8(o.StatusMessage, 1, 128) AS status_message, lengthUTF8(o.StatusMessage) > 128 AS error_truncated, diff --git a/litellm-rust/crates/traces-clickhouse/src/query/named.rs b/litellm-rust/crates/traces-clickhouse/src/query/named.rs index ff64259ec33..68550af4515 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query/named.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query/named.rs @@ -97,6 +97,26 @@ struct TraceSpansRowEncoding { pub wrapper_candidate: bool, pub agent: String, #[serde(default)] + pub native_agent_id: String, + #[serde(default)] + pub native_parent_agent_id: String, + #[serde(default)] + pub session_id: String, + #[serde(default)] + pub query_source: String, + #[serde( + default, + deserialize_with = "super::number::boolean", + serialize_with = "litellm_traces::wire::serialize_flag" + )] + pub has_output: bool, + #[serde( + default, + deserialize_with = "super::number::boolean", + serialize_with = "litellm_traces::wire::serialize_flag" + )] + pub capture_warning: bool, + #[serde(default)] pub framework: String, #[serde(serialize_with = "litellm_traces::wire::serialize_status")] pub status: litellm_traces::SpanStatus, @@ -364,7 +384,7 @@ mod tests { quoted, ); round_trip::( - json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "claude-agent-sdk", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "team_id": "team", "api_key_hash": "key", "user_id": "user"}), + json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "native_agent_id": "child", "native_parent_agent_id": "parent", "session_id": "session", "query_source": "agent.builtin.general-purpose", "has_output": 1, "capture_warning": 0, "framework": "claude-agent-sdk", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "team_id": "team", "api_key_hash": "key", "user_id": "user"}), quoted, ); round_trip::( diff --git a/litellm-rust/crates/traces-clickhouse/tests/reads.rs b/litellm-rust/crates/traces-clickhouse/tests/reads.rs index 0dabf9d8727..ab2fb34c48b 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/reads.rs @@ -1052,3 +1052,274 @@ async fn native_tool_content_agrees_between_single_and_bulk_reads( assert_eq!(bulk[1].attributes["lens.content.output_source"], "reply"); Ok(()) } + +#[rstest] +#[tokio::test] +async fn historical_native_actors_are_resolved_from_raw_attributes_with_tenant_isolation( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let client = &fixture.database.client; + let writer = Connection::writer(&fixture.database.url)?; + let start_ms = 1_790_000_000_000_i64; + let records = [ + ("root", "", "agent", "", "", "team-a"), + ("first", "root", "llm", "first", "", "team-a"), + ("second", "root", "llm", "second", "", "team-a"), + ("nested", "root", "llm", "nested", "first", "team-a"), + ("tool", "root", "tool", "first", "", "team-a"), + ("foreign", "root", "llm", "foreign", "", "team-b"), + ]; + let rows = records.iter().enumerate().map(|(index, (id, parent, kind, actor, parent_actor, team))| { + BTreeMap::from([ + ("Timestamp".into(), json!((start_ms + index as i64) * 1_000_000)), + ("TraceId".into(), json!("native-run")), + ("SpanId".into(), json!(id)), + ("ParentSpanId".into(), json!(parent)), + ("SpanName".into(), json!(if *kind == "agent" { "claude_code.interaction" } else { kind })), + ("ObservationType".into(), json!(kind)), + ("Output".into(), json!(if *id == "first" { "completed" } else { "" })), + ("Framework".into(), json!("claude-code")), + ("AgentName".into(), json!("claude-code")), + ("SpanAttributes".into(), json!({ + "lens.capture.warning": if *id == "nested" { "synthetic warning" } else { "" }, + "session.id": "historical-session", "agent_id": actor, "parent_agent_id": parent_actor, + "query_source_safe": if actor.is_empty() { "repl_main_thread" } else { "agent.builtin.general-purpose" }, + })), + ("TeamId".into(), json!(team)), + ("ApiKeyHash".into(), json!("key-a")), + ("Duration".into(), json!(1_000_000)), + ]) + }).collect(); + insert_rows(client, &writer, DATABASE, InsertTable::OtelTraces, rows).await?; + let connection = fixture + .readers + .connection(client, &QueryScope::All, "fixture-secret") + .await?; + let (reader, store) = make_reader(client, connection); + let access = ReadAccessParams { + all_teams: false, + user_id: String::new(), + team_ids: vec!["team-a".into()], + }; + let listed = store + .list_runs(&litellm_traces::query::named::ListTracesParams { + access: access.clone(), + start_ms: 0, + end_ms: 2_000_000_000_000, + cursor_ms: 0, + cursor_trace_id: String::new(), + limit: 50, + }) + .await?; + assert_eq!(listed[0].agent_count, 4); + assert_eq!(listed[0].agent_invocations, 4); + let page = reader + .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .await?; + assert_eq!(page.data.len(), 1); + assert_eq!(page.data[0].agent_count, 4); + let trace = reader + .get_trace(&store, &access, "native-run", &page.data[0].trace_ref) + .await? + .ok_or("missing trace")?; + assert_eq!(trace.summary.agent_count, 4); + assert_eq!(trace.summary.agent_invocations, 4); + assert_eq!( + trace + .agents + .iter() + .map(|actor| actor.llm_calls) + .sum::(), + 3 + ); + assert_eq!( + trace + .agents + .iter() + .map(|actor| actor.tool_calls) + .sum::(), + 1 + ); + let capture = trace.capture.as_ref().ok_or("missing capture coverage")?; + assert_eq!(capture.warning_events, 1); + assert_eq!( + capture + .actors + .iter() + .map(|actor| actor.model_outputs) + .sum::(), + 1 + ); + assert_eq!( + capture + .actors + .iter() + .map(|actor| actor.llm_calls) + .sum::(), + 3 + ); + assert!(!trace.spans.iter().any(|span| span.span_id == "foreign")); + let first = trace + .spans + .iter() + .find(|span| span.span_id == "first") + .ok_or("missing first actor")?; + let nested = trace + .spans + .iter() + .find(|span| span.span_id == "nested") + .ok_or("missing nested actor")?; + assert!(first.actor_id.is_some()); + assert_eq!(nested.parent_actor_id, first.actor_id); + Ok(()) +} + +#[rstest] +#[case::claude_code("claude-code")] +#[case::agent_sdk("claude-agent-sdk")] +#[tokio::test] +async fn native_list_counts_each_root_session_and_does_not_count_children_as_roots( + #[future(awt)] migrated_database: TestResult, + #[case] framework: &str, +) -> TestResult { + let fixture = migrated_database?; + let client = &fixture.database.client; + let writer = Connection::writer(&fixture.database.url)?; + let records = [ + ("first", "session-a", "", "agent", "claude_code.interaction"), + ( + "second", + "session-a", + "", + "agent", + "claude_code.interaction", + ), + ("pending", "session-b", "", "llm", "claude_code.llm_request"), + ( + "child", + "session-a", + "child", + "agent", + "claude_code.interaction", + ), + ( + "child-resumed", + "session-a", + "child", + "agent", + "claude_code.interaction", + ), + ( + "legacy-sdk", + "sdk-session", + "", + "agent", + "ClaudeAgentSDK.query", + ), + ]; + insert_rows(client, &writer, DATABASE, InsertTable::OtelTraces, records.into_iter().map(|(id, session, actor, kind, name)| BTreeMap::from([ + ("Timestamp".into(), json!(1_790_000_000_000_000_000_i64)), + ("TraceId".into(), json!("mixed-sessions")), + ("SpanId".into(), json!(id)), + ("SpanName".into(), json!(name)), + ("ObservationType".into(), json!(kind)), + ("Framework".into(), json!(if id == "legacy-sdk" { "claude-agent-sdk" } else { framework })), + ("AgentName".into(), json!("assistant")), + ("SpanAttributes".into(), json!({"session.id": session, "agent_id": actor, "query_source_safe": if id == "legacy-sdk" { "" } else { "sdk" }})), + ("TeamId".into(), json!("team-a")), + ("ApiKeyHash".into(), json!("key-a")), + ])).collect()).await?; + let connection = fixture + .readers + .connection(client, &QueryScope::All, "fixture-secret") + .await?; + let (reader, store) = make_reader(client, connection); + let access = ReadAccessParams { + all_teams: true, + user_id: String::new(), + team_ids: Vec::new(), + }; + let listed = store + .list_runs(&litellm_traces::query::named::ListTracesParams { + access: access.clone(), + start_ms: 0, + end_ms: 2_000_000_000_000, + cursor_ms: 0, + cursor_trace_id: String::new(), + limit: 50, + }) + .await?; + let summary = &listed[0]; + let detail = reader + .get_trace(&store, &access, "mixed-sessions", &summary.trace_ref) + .await? + .ok_or("missing trace")?; + assert_eq!((summary.agent_count, summary.agent_invocations), (4, 6)); + assert_eq!( + (detail.summary.agent_count, detail.summary.agent_invocations), + (4, 6) + ); + Ok(()) +} + +#[rstest] +#[case::root_sessions(false)] +#[case::child_actors(true)] +#[tokio::test] +async fn many_native_invocations_are_aggregated_by_identity( + #[future(awt)] migrated_database: TestResult, + #[case] children: bool, +) -> TestResult { + let fixture = migrated_database?; + let client = &fixture.database.client; + let writer = Connection::writer(&fixture.database.url)?; + let rows = (0..50_000) + .map(|index| { + BTreeMap::from([ + ("Timestamp".into(), json!(1_790_000_000_000_000_000_i64)), + ("TraceId".into(), json!("many-invocations")), + ("SpanId".into(), json!(index.to_string())), + ("SpanName".into(), json!("claude_code.interaction")), + ("ObservationType".into(), json!("agent")), + ("Framework".into(), json!("claude-code")), + ("AgentName".into(), json!("assistant")), + ( + "SpanAttributes".into(), + json!({ + "session.id": if children { "shared".into() } else { index.to_string() }, + "agent_id": if children { index.to_string() } else { String::new() }, + }), + ), + ("TeamId".into(), json!("team-a")), + ("ApiKeyHash".into(), json!("key-a")), + ]) + }) + .collect(); + insert_rows(client, &writer, DATABASE, InsertTable::OtelTraces, rows).await?; + let connection = fixture + .readers + .connection(client, &QueryScope::All, "fixture-secret") + .await?; + let (_, store) = make_reader(client, connection); + let listed = store + .list_runs(&litellm_traces::query::named::ListTracesParams { + access: ReadAccessParams { + all_teams: true, + user_id: String::new(), + team_ids: Vec::new(), + }, + start_ms: 0, + end_ms: 2_000_000_000_000, + cursor_ms: 0, + cursor_trace_id: String::new(), + limit: 50, + }) + .await?; + assert_eq!(listed.len(), 1); + assert_eq!( + (listed[0].agent_count, listed[0].agent_invocations), + (50_000, 50_000) + ); + Ok(()) +} diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs index b51701638e8..4bfb553eeb8 100644 --- a/litellm-rust/crates/traces/src/lib.rs +++ b/litellm-rust/crates/traces/src/lib.rs @@ -44,6 +44,6 @@ pub use tenant::Tenant; pub use truncate::{truncate_messages, truncate_value}; pub use ui::{ChatRole, UiContent, UiField, UiMessage, UiToolCall, to_ui_content}; pub use view::{ - AgentNode, Span, SpanDetail, SpanErrorPage, SpanStatus, SpendMatch, Trace, TracePage, - TraceSummary, + ActorCapture, AgentNode, Span, SpanDetail, SpanErrorPage, SpanStatus, SpendMatch, Trace, + TraceCapture, TracePage, TraceSummary, }; diff --git a/litellm-rust/crates/traces/src/normalize/format/claude_code.rs b/litellm-rust/crates/traces/src/normalize/format/claude_code.rs index 37fb2d01330..3a59748ab50 100644 --- a/litellm-rust/crates/traces/src/normalize/format/claude_code.rs +++ b/litellm-rust/crates/traces/src/normalize/format/claude_code.rs @@ -45,12 +45,15 @@ fn span_type(name: &str, attributes: &BTreeMap) -> SpanType { } } -/// `agent:custom:search_agent` -> `search_agent`: the subagent a request ran for. fn subagent(attributes: &BTreeMap) -> Option<&str> { - let mut parts = attr(attributes, "query_source") - .strip_prefix("agent:")? - .splitn(2, ':'); - let (_kind, name) = (parts.next()?, parts.next()?); + let source = + super::super::select_attribute(attributes, &["query_source", "query_source_safe"])?; + let (rest, separator) = if let Some(rest) = source.text.strip_prefix("agent:") { + (rest, ':') + } else { + (source.text.strip_prefix("agent.")?, '.') + }; + let (_, name) = rest.split_once(separator)?; (!name.is_empty()).then_some(name) } @@ -212,7 +215,7 @@ impl Format for ClaudeCode { let kind = span_type(context.name, attributes); let base = SpanFacts { role: Some(RoleEvidence::Declared(ObservationType::Framework)), - agent_name: Some(CLAUDE_CODE_AGENT.to_owned()), + agent_name: Some(subagent(attributes).unwrap_or(CLAUDE_CODE_AGENT).to_owned()), tool_call_id: present(attributes, &["gen_ai.tool.call.id"]), ..SpanFacts::default() }; @@ -325,6 +328,32 @@ mod tests { otlp::DecodedEvent, }; + #[rstest] + #[case::log_source("query_source", "agent:builtin:general-purpose")] + #[case::span_source("query_source_safe", "agent.builtin.general-purpose")] + fn native_actor_metadata_does_not_require_session_capture( + #[case] key: &str, + #[case] source: &str, + ) { + let attributes = BTreeMap::from([ + ("span.type".into(), "llm_request".into()), + ("session.id".into(), "session".into()), + ("agent_id".into(), "child".into()), + ("parent_agent_id".into(), "parent".into()), + (key.into(), source.into()), + ]); + let span = normalization("claude_code.llm_request", &attributes, &[]) + .unwrap() + .span; + assert_eq!(span.agent_name.as_deref(), Some("general-purpose")); + assert_eq!(span.agent_metadata.agent_id.as_deref(), Some("child")); + assert_eq!( + span.agent_metadata.parent_agent_id.as_deref(), + Some("parent") + ); + assert_eq!(span.agent_metadata.session_id.as_deref(), Some("session")); + } + fn normalization( name: &str, attributes: &BTreeMap, diff --git a/litellm-rust/crates/traces/src/normalize/metadata.rs b/litellm-rust/crates/traces/src/normalize/metadata.rs index 100008714c5..0a704072926 100644 --- a/litellm-rust/crates/traces/src/normalize/metadata.rs +++ b/litellm-rust/crates/traces/src/normalize/metadata.rs @@ -70,6 +70,12 @@ impl From for String { #[derive(Debug, Default, Deserialize, Eq, PartialEq, Serialize)] #[serde(default)] pub struct AgentMetadata { + #[serde(deserialize_with = "optional")] + pub agent_id: Option, + #[serde(deserialize_with = "optional")] + pub parent_agent_id: Option, + #[serde(deserialize_with = "optional")] + pub session_id: Option, #[serde(deserialize_with = "optional")] pub lc_agent_name: Option, #[serde(deserialize_with = "optional")] @@ -109,6 +115,9 @@ pub struct AgentMetadata { impl AgentMetadata { pub(crate) fn byte_len(&self) -> usize { let strings = [ + &self.agent_id, + &self.parent_agent_id, + &self.session_id, &self.lc_agent_name, &self.ls_agent_purpose, &self.ls_agent_runtime, @@ -140,6 +149,10 @@ impl AgentMetadata { #[derive(strum::EnumString, strum::IntoStaticStr)] #[strum(serialize_all = "snake_case")] enum MetadataField { + AgentId, + ParentAgentId, + #[strum(serialize = "session.id", to_string = "session_id")] + SessionId, LcAgentName, LsIntegration, LsAgentType, diff --git a/litellm-rust/crates/traces/src/query/named.rs b/litellm-rust/crates/traces/src/query/named.rs index 9af02ecfad1..0650c824a5d 100644 --- a/litellm-rust/crates/traces/src/query/named.rs +++ b/litellm-rust/crates/traces/src/query/named.rs @@ -83,6 +83,26 @@ pub struct TraceSpansRow { pub wrapper_candidate: bool, pub agent: String, #[serde(default)] + pub native_agent_id: String, + #[serde(default)] + pub native_parent_agent_id: String, + #[serde(default)] + pub session_id: String, + #[serde(default)] + pub query_source: String, + #[serde( + default, + deserialize_with = "crate::wire::flag", + serialize_with = "crate::wire::serialize_flag" + )] + pub has_output: bool, + #[serde( + default, + deserialize_with = "crate::wire::flag", + serialize_with = "crate::wire::serialize_flag" + )] + pub capture_warning: bool, + #[serde(default)] pub framework: String, #[serde(serialize_with = "crate::wire::serialize_status")] pub status: crate::SpanStatus, diff --git a/litellm-rust/crates/traces/src/resolve/actors.rs b/litellm-rust/crates/traces/src/resolve/actors.rs new file mode 100644 index 00000000000..ca12dfa18b6 --- /dev/null +++ b/litellm-rust/crates/traces/src/resolve/actors.rs @@ -0,0 +1,186 @@ +use std::collections::{BTreeMap, BTreeSet, HashMap}; + +use sha2::{Digest, Sha256}; + +use crate::query::named::TraceSpansRow; + +use super::graph::Graph; + +pub(super) fn native(row: &TraceSpansRow) -> bool { + (row.framework == "claude-code" + || (row.framework == "claude-agent-sdk" + && (row.name.starts_with("claude_code.") + || !row.native_agent_id.is_empty() + || !row.query_source.is_empty()))) + && (!row.session_id.is_empty() || !row.native_agent_id.is_empty()) +} + +fn identity(row: &TraceSpansRow, agent: &str) -> String { + let session = if row.session_id.is_empty() { + &row.trace_id + } else { + &row.session_id + }; + let mut digest = Sha256::new(); + digest.update((session.len() as u64).to_be_bytes()); + digest.update(session); + digest.update(agent); + format!("claude-code:{:x}", digest.finalize()) +} + +fn label(value: &str) -> String { + value.chars().take(256).collect() +} + +fn root(row: &TraceSpansRow) -> String { + identity(row, "root") +} + +fn explicit(row: &TraceSpansRow) -> Option { + if !native(row) { + return None; + } + if !row.native_agent_id.is_empty() { + return Some(identity(row, &format!("agent:{}", row.native_agent_id))); + } + (row.name == "claude_code.interaction" + || matches!( + row.query_source.as_str(), + "repl_main_thread" + | "sdk" + | "sdk_main_thread" + | "generate_session_title" + | "prompt_suggestion" + )) + .then(|| root(row)) +} + +fn child_source(row: &TraceSpansRow) -> bool { + row.query_source == "agent" + || row.query_source.starts_with("agent:") + || row.query_source.starts_with("agent.") +} + +pub(super) struct Actor { + pub(super) id: String, + pub(super) parent: Option, + pub(super) name: String, + pub(super) spans: Vec, +} + +pub(super) struct Actors { + owners: Vec>, + pub(super) entries: BTreeMap, +} + +impl Actors { + pub(super) fn new(graph: &Graph<'_>) -> Self { + 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(); + for (index, owner) in direct.iter().enumerate() { + let Some(owner) = owner.as_deref() else { + continue; + }; + let row = &rows[index]; + if let Some(parent) = graph.parent(index) + && rows[parent].name == "claude_code.tool.execution" + { + boundaries.entry(parent).or_default().insert(owner); + } + if !row.tool_call_id.is_empty() { + tools.entry(&row.tool_call_id).or_default().insert(owner); + } + } + let unique = |set: Option<&BTreeSet<&str>>| { + set.filter(|set| set.len() == 1) + .and_then(|set| set.first().copied()) + .map(str::to_owned) + }; + let candidates: Vec<_> = (0..rows.len()) + .map(|index| unique(boundaries.get(&index)).or_else(|| direct[index].clone())) + .collect(); + let inherited = + graph.nearest_ancestors(&candidates.iter().map(Option::is_some).collect::>()); + let inherited_child = graph.nearest_ancestors( + &candidates + .iter() + .enumerate() + .map(|(index, candidate)| { + candidate + .as_ref() + .is_some_and(|candidate| candidate != &root(&rows[index])) + }) + .collect::>(), + ); + let owners: Vec<_> = rows + .iter() + .enumerate() + .map(|(index, row)| { + if !native(row) + && !(row.framework == "claude-agent-sdk" && !row.session_id.is_empty()) + { + return None; + } + direct[index] + .clone() + .or_else(|| unique(boundaries.get(&index))) + .or_else(|| unique(tools.get(row.tool_call_id.as_str()))) + .or_else(|| { + let ancestor = if child_source(row) { + inherited_child[index] + } else { + inherited[index] + }?; + candidates[ancestor].clone() + }) + }) + .collect(); + let mut entries: BTreeMap = BTreeMap::new(); + for (index, owner) in owners.iter().enumerate() { + let Some(id) = owner else { continue }; + let row = &rows[index]; + let actor = entries.entry(id.clone()).or_insert_with(|| Actor { + id: id.clone(), + parent: (id != &root(row)).then(|| root(row)), + name: if id == &root(row) { + label(&row.agent) + } else { + "Subagent".to_owned() + }, + spans: Vec::new(), + }); + actor.spans.push(index); + if !row.native_parent_agent_id.is_empty() { + actor.parent = Some(identity( + row, + &format!("agent:{}", row.native_parent_agent_id), + )); + } + if child_source(row) { + let name = row + .query_source + .strip_prefix("agent:") + .and_then(|source| source.split_once(':')) + .or_else(|| { + row.query_source + .strip_prefix("agent.") + .and_then(|source| source.split_once('.')) + }); + if let Some((_, name)) = name + && !name.is_empty() + { + actor.name = label(name); + } + } + } + Self { owners, entries } + } + + pub(super) fn owner(&self, index: usize) -> Option<&Actor> { + self.owners[index] + .as_ref() + .and_then(|id| self.entries.get(id)) + } +} diff --git a/litellm-rust/crates/traces/src/resolve/capture.rs b/litellm-rust/crates/traces/src/resolve/capture.rs new file mode 100644 index 00000000000..1c14c43a6e0 --- /dev/null +++ b/litellm-rust/crates/traces/src/resolve/capture.rs @@ -0,0 +1,57 @@ +use std::collections::{BTreeMap, HashSet}; + +use super::{actors, resolution::Resolution}; +use crate::{ + ObservationType, + view::{ActorCapture, TraceCapture}, +}; + +pub(super) fn coverage(resolution: &Resolution<'_>) -> Option { + let rows = resolution.graph.rows; + if !rows.iter().any(actors::native) { + return None; + } + let calls: HashSet<_> = resolution.model_calls.iter().copied().collect(); + let mut actors: BTreeMap<_, _> = resolution + .actors + .entries + .values() + .map(|actor| { + ( + actor.id.as_str(), + ActorCapture { + actor_id: actor.id.clone(), + name: actor.name.clone(), + ..Default::default() + }, + ) + }) + .collect(); + let mut capture = TraceCapture::default(); + for (index, row) in rows.iter().enumerate() { + let event = matches!( + row.name.as_str(), + "claude_code.user_prompt" + | "claude_code.assistant_response" + | "claude_code.tool_result" + | "claude_code.api_request_body" + ); + capture.content_events += u64::from(event); + capture.warning_events += u64::from(row.capture_warning); + let Some(actor) = resolution + .actors + .owner(index) + .and_then(|owner| actors.get_mut(owner.id.as_str())) + else { + capture.unassigned_events += u64::from(event); + continue; + }; + actor.llm_calls += u64::from(calls.contains(&index)); + actor.tool_calls += u64::from(resolution.kind(index) == ObservationType::Tool); + 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); + } + capture.actors = actors.into_values().collect(); + Some(capture) +} diff --git a/litellm-rust/crates/traces/src/resolve/graph.rs b/litellm-rust/crates/traces/src/resolve/graph.rs index 733f5e87213..d6fbb1188f8 100644 --- a/litellm-rust/crates/traces/src/resolve/graph.rs +++ b/litellm-rust/crates/traces/src/resolve/graph.rs @@ -64,6 +64,43 @@ impl<'a> Graph<'a> { found } + pub(super) fn nearest_ancestors(&self, matches: &[bool]) -> Vec> { + let mut nearest: Vec<_> = matches + .iter() + .enumerate() + .map(|(index, found)| found.then_some(index)) + .collect(); + let mut resolved = matches.to_vec(); + let mut visiting = vec![false; self.rows.len()]; + for index in 0..self.rows.len() { + let mut path = Vec::new(); + let mut current = Some(index); + let found = loop { + let Some(node) = current else { break None }; + if resolved[node] { + break nearest[node]; + } + if visiting[node] { + break None; + } + visiting[node] = true; + path.push(node); + current = self.parent(node); + }; + for node in path { + nearest[node] = found; + resolved[node] = true; + } + } + (0..self.rows.len()) + .map(|index| { + self.parent(index) + .and_then(|parent| nearest[parent]) + .filter(|found| *found != index) + }) + .collect() + } + pub(super) fn descendants(&self, index: usize) -> Vec { let children = |index: usize| { self.children @@ -141,6 +178,33 @@ mod tests { assert!(!graph.is_root(0)); } + #[rstest] + #[case::root(vec![false, false, false, true], vec![Some(3), Some(3), Some(3), None])] + #[case::nearer(vec![false, false, true, true], vec![Some(2), Some(3), Some(3), None])] + #[case::unassigned(vec![false; 4], vec![None; 4])] + fn nearest_candidates_follow_unordered_links( + unordered_rows: Vec, + #[case] candidates: Vec, + #[case] expected: Vec>, + ) { + assert_eq!( + Graph::new(&unordered_rows).nearest_ancestors(&candidates), + expected + ); + } + + #[rstest] + #[case::unassigned(vec![false, false], vec![None, None])] + #[case::one_candidate(vec![true, false], vec![None, Some(0)])] + #[case::both_candidates(vec![true, true], vec![Some(1), Some(0)])] + fn nearest_candidates_do_not_revisit_cycles( + #[case] candidates: Vec, + #[case] expected: Vec>, + ) { + let rows = [row("first", "second"), row("second", "first")]; + assert_eq!(Graph::new(&rows).nearest_ancestors(&candidates), expected); + } + #[rstest] #[case::missing_parent("missing", &[], &[1])] #[case::self_link("first", &[], &[1])] diff --git a/litellm-rust/crates/traces/src/resolve/mod.rs b/litellm-rust/crates/traces/src/resolve/mod.rs index a3a68f86776..1efb2659db9 100644 --- a/litellm-rust/crates/traces/src/resolve/mod.rs +++ b/litellm-rust/crates/traces/src/resolve/mod.rs @@ -1,3 +1,5 @@ +mod actors; +mod capture; mod graph; mod resolution; mod spend; diff --git a/litellm-rust/crates/traces/src/resolve/resolution.rs b/litellm-rust/crates/traces/src/resolve/resolution.rs index e19a98a1ede..cdfb9f4b466 100644 --- a/litellm-rust/crates/traces/src/resolve/resolution.rs +++ b/litellm-rust/crates/traces/src/resolve/resolution.rs @@ -9,6 +9,7 @@ use crate::{ }; use super::{ + actors::Actors, graph::Graph, spend::{self, Ownership, Requests, SpendEvidence}, }; @@ -23,6 +24,7 @@ pub(super) fn agent_label(row: &TraceSpansRow) -> &str { pub(super) struct Resolution<'a> { pub(super) graph: Graph<'a>, + pub(super) actors: Actors, ownership: Ownership<'a>, spend: &'a [SpendRow], types: HashMap<&'a str, ObservationType>, @@ -50,6 +52,7 @@ impl<'a> Resolution<'a> { }) .collect(); let resolution = Self { + actors: Actors::new(&graph), ownership: Ownership { team_id: &rows[0].team_id, api_key_hash: &rows[0].api_key_hash, diff --git a/litellm-rust/crates/traces/src/resolve/view.rs b/litellm-rust/crates/traces/src/resolve/view.rs index 9e9edd51bf4..c43389a3642 100644 --- a/litellm-rust/crates/traces/src/resolve/view.rs +++ b/litellm-rust/crates/traces/src/resolve/view.rs @@ -40,7 +40,16 @@ fn span(resolution: &Resolution<'_>, index: usize, trace_start_ns: i64) -> Span parent_span_id: optional(&row.parent_span_id), name: row.name.clone(), kind: resolution.kind(index), - agent: row.agent.clone(), + actor_id: resolution.actors.owner(index).map(|actor| actor.id.clone()), + actor_unassigned: super::actors::native(row) && resolution.actors.owner(index).is_none(), + parent_actor_id: resolution + .actors + .owner(index) + .and_then(|actor| actor.parent.clone()), + agent: resolution + .actors + .owner(index) + .map_or_else(|| row.agent.clone(), |actor| actor.name.clone()), framework: row.framework.clone(), start_offset_ms: (i128::from(row.start_ns) - i128::from(trace_start_ns)) as f64 / NANOS_PER_MS, @@ -62,7 +71,11 @@ fn span(resolution: &Resolution<'_>, index: usize, trace_start_ns: i64) -> Span fn agents(resolution: &Resolution<'_>) -> Vec { let graph = &resolution.graph; let mut entries: IndexMap<&str, Vec> = IndexMap::new(); - for index in (0..graph.rows.len()).filter(|index| resolution.is_agent(*index)) { + for index in (0..graph.rows.len()).filter(|index| { + resolution.is_agent(*index) + && !super::actors::native(resolution.row(*index)) + && resolution.actors.owner(*index).is_none() + }) { entries .entry(agent_label(resolution.row(index))) .or_default() @@ -73,7 +86,9 @@ fn agents(resolution: &Resolution<'_>) -> Vec { let parent_agent = graph .parent(index) .map(|parent| graph.rows[parent].agent.as_str()); - if !row.agent.is_empty() + if !super::actors::native(row) + && resolution.actors.owner(index).is_none() + && !row.agent.is_empty() && !explicit.contains(row.agent.as_str()) && parent_agent != Some(row.agent.as_str()) { @@ -83,40 +98,100 @@ fn agents(resolution: &Resolution<'_>) -> Vec { let calls: Vec<(&str, Option>)> = resolution .model_calls .iter() + .filter(|call| { + !super::actors::native(resolution.row(**call)) + && resolution.actors.owner(**call).is_none() + }) .map(|call| (resolution.owner(*call), resolution.call_requests(*call))) .collect(); let tools = resolution.unique_tools(); - entries - .into_iter() - .map(|(name, spans)| { - let parent_agent = graph.ancestors(spans[0]).into_iter().find_map(|ancestor| { - let label = agent_label(resolution.row(ancestor)); - (resolution.is_agent(ancestor) && label != name).then(|| label.to_owned()) - }); - let owned_calls: Vec>> = calls + let legacy = entries.into_iter().map(|(name, spans)| { + let parent_agent = graph.ancestors(spans[0]).into_iter().find_map(|ancestor| { + let label = agent_label(resolution.row(ancestor)); + (resolution.is_agent(ancestor) && label != name).then(|| label.to_owned()) + }); + let owned_calls: Vec>> = calls + .iter() + .filter(|(owner, _)| *owner == name) + .map(|(_, requests)| requests.clone()) + .collect(); + let priced = total(&owned_calls); + AgentNode { + actor_id: None, + parent_actor_id: None, + name: name.to_owned(), + parent_agent, + invocations: spans.len() as u64, + llm_calls: owned_calls.len() as u64, + tool_calls: tools .iter() - .filter(|(owner, _)| *owner == name) - .map(|(_, requests)| requests.clone()) + .filter(|tool| { + !super::actors::native(resolution.row(**tool)) + && resolution.actors.owner(**tool).is_none() + && resolution.owner(**tool) == name + }) + .count() as u64, + duration_ms: spans + .iter() + .map(|span| graph.rows[*span].duration_ns) + .sum::() as f64 + / NANOS_PER_MS, + spend: priced.spend, + priced_calls: priced.priced_calls, + } + }); + let model_calls: HashSet<_> = resolution.model_calls.iter().copied().collect(); + let native_tools: HashSet<_> = tools.iter().copied().collect(); + legacy + .chain(resolution.actors.entries.values().map(|actor| { + let calls: Vec<_> = actor + .spans + .iter() + .filter(|index| model_calls.contains(index)) + .map(|index| resolution.call_requests(*index)) .collect(); - let priced = total(&owned_calls); + let start = actor + .spans + .iter() + .map(|index| i128::from(graph.rows[*index].start_ns)) + .min() + .unwrap_or_default(); + let end = actor + .spans + .iter() + .map(|index| { + i128::from(graph.rows[*index].start_ns) + + i128::from(graph.rows[*index].duration_ns) + }) + .max() + .unwrap_or(start); + let priced = total(&calls); AgentNode { - name: name.to_owned(), - parent_agent, - invocations: spans.len() as u64, - llm_calls: owned_calls.len() as u64, - tool_calls: tools + actor_id: Some(actor.id.clone()), + parent_actor_id: actor.parent.clone(), + name: actor.name.clone(), + parent_agent: actor + .parent + .as_ref() + .and_then(|id| resolution.actors.entries.get(id)) + .map(|parent| parent.name.clone()), + invocations: actor + .spans .iter() - .filter(|tool| resolution.owner(**tool) == name) + .filter(|index| graph.rows[**index].name == "claude_code.interaction") + .count() + .max(1) as u64, + llm_calls: calls.len() as u64, + tool_calls: actor + .spans + .iter() + .filter(|index| native_tools.contains(index)) .count() as u64, - duration_ms: spans - .iter() - .map(|span| graph.rows[*span].duration_ns) - .sum::() as f64 - / NANOS_PER_MS, + duration_ms: (end - start) as f64 / NANOS_PER_MS, spend: priced.spend, priced_calls: priced.priced_calls, } - }) + })) .collect() } @@ -231,6 +306,7 @@ pub fn resolve_trace( priced_calls: priced.priced_calls, }; Some(Trace { + capture: super::capture::coverage(&resolution), summary, agents, spans, diff --git a/litellm-rust/crates/traces/src/view.rs b/litellm-rust/crates/traces/src/view.rs index 5a214d33da7..6996ab5909f 100644 --- a/litellm-rust/crates/traces/src/view.rs +++ b/litellm-rust/crates/traces/src/view.rs @@ -25,6 +25,13 @@ pub struct Span { #[serde(rename = "type")] pub kind: crate::ObservationType, pub agent: String, + #[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))] + pub actor_id: Option, + #[serde(default)] + #[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))] + pub actor_unassigned: bool, + #[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))] + pub parent_actor_id: Option, pub framework: String, pub start_offset_ms: f64, pub duration_ms: f64, @@ -56,6 +63,10 @@ pub enum SpendMatch { #[macro_rules_attribute::apply(response_type)] #[derive(Clone, Debug, PartialEq)] pub struct AgentNode { + #[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))] + pub actor_id: Option, + #[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))] + pub parent_actor_id: Option, pub name: String, pub parent_agent: Option, pub invocations: u64, @@ -100,6 +111,9 @@ pub struct TraceSummary { #[macro_rules_attribute::apply(response_type)] #[derive(Clone, Debug, PartialEq)] pub struct Trace { + #[serde(default)] + #[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))] + pub capture: Option, pub summary: TraceSummary, pub agents: Vec, pub spans: Vec, @@ -133,3 +147,24 @@ pub struct SpanErrorPage { pub total_chars: u64, pub next_cursor: Option, } + +#[macro_rules_attribute::apply(response_type)] +#[derive(Clone, Debug, Default, PartialEq)] +pub struct ActorCapture { + pub actor_id: String, + pub name: String, + pub llm_calls: u64, + pub tool_calls: u64, + pub reply_events: u64, + pub model_outputs: u64, + pub content_events: u64, +} + +#[macro_rules_attribute::apply(response_type)] +#[derive(Clone, Debug, Default, PartialEq)] +pub struct TraceCapture { + pub actors: Vec, + pub content_events: u64, + pub unassigned_events: u64, + pub warning_events: u64, +} diff --git a/litellm-rust/crates/traces/tests/captures.rs b/litellm-rust/crates/traces/tests/captures.rs index ff20db617f6..20bc696951d 100644 --- a/litellm-rust/crates/traces/tests/captures.rs +++ b/litellm-rust/crates/traces/tests/captures.rs @@ -177,6 +177,25 @@ fn trace_span(span: DecodedSpan) -> TraceSpansRow { }) .unwrap_or_default(); TraceSpansRow { + has_output: !normalized.output.is_empty(), + capture_warning: span.attributes.contains_key("lens.capture.warning"), + native_agent_id: span.attributes.get("agent_id").cloned().unwrap_or_default(), + native_parent_agent_id: span + .attributes + .get("parent_agent_id") + .cloned() + .unwrap_or_default(), + session_id: span + .attributes + .get("session.id") + .cloned() + .unwrap_or_default(), + query_source: span + .attributes + .get("query_source_safe") + .or_else(|| span.attributes.get("query_source")) + .cloned() + .unwrap_or_default(), trace_id: span.trace_id, original_trace_id: String::new(), span_id: span.span_id, @@ -268,6 +287,12 @@ fn unrelated_transport(call: &TraceSpansRow) -> TraceSpansRow { i64::try_from(i128::from(call.start_ns) + i128::from(call.duration_ns) + 1_000_000) .expect("valid unrelated transport timestamp"); TraceSpansRow { + native_agent_id: String::new(), + native_parent_agent_id: String::new(), + session_id: String::new(), + query_source: String::new(), + has_output: false, + capture_warning: false, trace_id: call.trace_id.clone(), original_trace_id: call.original_trace_id.clone(), span_id: format!("unrelated-transport-{}", call.span_id), diff --git a/litellm-rust/crates/traces/tests/query/named.rs b/litellm-rust/crates/traces/tests/query/named.rs index 3ac9668e42a..f69e07a0bd1 100644 --- a/litellm-rust/crates/traces/tests/query/named.rs +++ b/litellm-rust/crates/traces/tests/query/named.rs @@ -53,7 +53,11 @@ fn result_contracts_preserve_public_field_names() { json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "STATUS_CODE_OK", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["framework"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}), ); round_trip::( - json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "framework", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "team_id": "team", "api_key_hash": "key", "user_id": "user"}), + json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", + "native_agent_id": "child", + "native_parent_agent_id": "parent", + "session_id": "session", + "query_source": "agent.builtin.general-purpose", "has_output": 1, "capture_warning": 0, "framework": "framework", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "team_id": "team", "api_key_hash": "key", "user_id": "user"}), ); round_trip::( json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}), diff --git a/litellm-rust/crates/traces/tests/resolve.rs b/litellm-rust/crates/traces/tests/resolve.rs index 8d65adb0c0f..65aaae457b2 100644 --- a/litellm-rust/crates/traces/tests/resolve.rs +++ b/litellm-rust/crates/traces/tests/resolve.rs @@ -18,6 +18,12 @@ fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> Trac kind: kind.parse().unwrap(), wrapper_candidate: false, agent: agent.into(), + native_agent_id: String::new(), + native_parent_agent_id: String::new(), + session_id: String::new(), + query_source: String::new(), + has_output: false, + capture_warning: false, framework: String::new(), status: SpanStatus::Ok, status_message: String::new(), @@ -202,6 +208,8 @@ fn repeated_subagent_invocations_aggregate_into_one_node() { assert_eq!( trace.agents[0], AgentNode { + actor_id: None, + parent_actor_id: None, name: "deep_research_agent".into(), parent_agent: None, invocations: 1, @@ -1847,3 +1855,507 @@ fn provider_request_id_cannot_match_a_message_id_of_the_same_value() { assert_eq!(trace.summary.spend, None); assert_eq!(trace.spans[0].spend_match, Some(SpendMatch::NoSpendLog)); } + +#[rstest] +#[case::root_recorded(true)] +#[case::root_pending(false)] +fn native_actor_identity_survives_same_names_nested_and_resumed_work(#[case] root_recorded: bool) { + let native = |mut row: TraceSpansRow, actor: &str, parent: &str, source: &str| { + row.framework = "claude-code".into(); + row.session_id = "session-one".into(); + row.native_agent_id = actor.into(); + row.native_parent_agent_id = parent.into(); + row.query_source = source.into(); + row + }; + let mut rows = vec![ + native( + at( + row( + "root-call", + "root", + "claude_code.llm_request", + "llm", + "claude-code", + ), + 0, + 10, + ), + "", + "", + "repl_main_thread", + ), + native( + at( + row( + "first-call", + "launch-execution", + "claude_code.llm_request", + "llm", + "claude-code", + ), + 10, + 10, + ), + "first", + "", + "agent.builtin.general-purpose", + ), + native( + at( + row( + "second-call", + "root", + "claude_code.llm_request", + "llm", + "claude-code", + ), + 20, + 10, + ), + "second", + "", + "agent.builtin.general-purpose", + ), + native( + at( + row( + "nested-call", + "root", + "claude_code.llm_request", + "llm", + "claude-code", + ), + 30, + 10, + ), + "nested", + "first", + "agent.builtin.general-purpose", + ), + native( + at( + row( + "resumed-call", + "new-root", + "claude_code.llm_request", + "llm", + "claude-code", + ), + 100, + 10, + ), + "first", + "", + "agent.builtin.general-purpose", + ), + native( + at(row("tool", "root", "Bash", "tool", "claude-code"), 105, 2), + "first", + "", + "", + ), + native( + at( + row( + "launch-execution", + "root", + "claude_code.tool.execution", + "framework", + "claude-code", + ), + 9, + 1, + ), + "", + "", + "", + ), + native( + at( + row( + "reply", + "launch-execution", + "claude_code.assistant_response", + "chain", + "claude-code", + ), + 21, + 0, + ), + "", + "", + "agent:builtin:general-purpose", + ), + native( + at( + row( + "unassigned", + "new-root", + "claude_code.assistant_response", + "chain", + "claude-code", + ), + 115, + 0, + ), + "", + "", + "agent:builtin:general-purpose", + ), + ]; + if root_recorded { + rows.push(native( + at( + row( + "root", + "", + "claude_code.interaction", + "agent", + "claude-code", + ), + 0, + 90, + ), + "", + "", + "", + )); + rows.push(native( + at( + row( + "new-root", + "", + "claude_code.interaction", + "agent", + "claude-code", + ), + 95, + 25, + ), + "", + "", + "", + )); + } + let trace = resolve_trace("trace", "", &rows, &[]).unwrap(); + assert_eq!(trace.summary.agent_count, 4); + assert_eq!( + trace + .agents + .iter() + .map(|actor| actor.llm_calls) + .sum::(), + trace.summary.llm_calls + ); + assert_eq!( + trace + .agents + .iter() + .map(|actor| actor.tool_calls) + .sum::(), + trace.summary.tool_calls + ); + let first_id = trace + .spans + .iter() + .find(|span| span.span_id == "first-call") + .unwrap() + .actor_id + .as_ref() + .unwrap(); + let first = trace + .agents + .iter() + .find(|actor| actor.actor_id.as_ref() == Some(first_id)) + .unwrap(); + assert_eq!(first.name, "general-purpose"); + assert_eq!( + ( + first.invocations, + first.llm_calls, + first.tool_calls, + first.duration_ms + ), + (1, 2, 1, 101.0) + ); + for id in ["resumed-call", "tool", "reply"] { + assert_eq!( + trace + .spans + .iter() + .find(|span| span.span_id == id) + .unwrap() + .actor_id + .as_ref(), + Some(first_id) + ); + } + let second = trace + .spans + .iter() + .find(|span| span.span_id == "second-call") + .unwrap(); + assert_ne!(second.actor_id.as_ref(), Some(first_id)); + let nested = trace + .spans + .iter() + .find(|span| span.span_id == "nested-call") + .unwrap(); + assert_eq!(nested.parent_actor_id.as_ref(), Some(first_id)); + assert!( + trace + .spans + .iter() + .find(|span| span.span_id == "unassigned") + .unwrap() + .actor_id + .is_none() + ); + assert_eq!( + trace.summary.agent_invocations, + if root_recorded { 5 } else { 4 } + ); +} + +#[rstest] +fn native_actor_identifiers_are_scoped_to_their_session() { + let rows = ["session-one", "session-two"].map(|session| TraceSpansRow { + session_id: session.into(), + native_agent_id: "shared-child-id".into(), + framework: "claude-code".into(), + ..row(session, "", "claude_code.llm_request", "llm", "claude-code") + }); + let trace = resolve_trace("trace", "", &rows, &[]).unwrap(); + assert_eq!(trace.summary.agent_count, 2); + assert_ne!(trace.spans[0].actor_id, trace.spans[1].actor_id); + assert_ne!( + trace.spans[0].parent_actor_id, + trace.spans[1].parent_actor_id + ); +} + +#[rstest] +#[case::legacy("", "", false)] +#[case::unconfirmed_child("session", "agent.builtin.reader", true)] +fn unresolved_native_actors_are_distinguished_from_legacy_grouping( + #[case] session: &str, + #[case] source: &str, + #[case] unassigned: bool, + #[values("claude-code", "claude-agent-sdk")] framework: &str, +) { + let span = TraceSpansRow { + framework: framework.into(), + session_id: session.into(), + query_source: source.into(), + ..row( + "reply", + "", + "claude_code.assistant_response", + "chain", + "reader", + ) + }; + let trace = resolve_trace("trace", "ref", &[span], &[]).unwrap(); + assert_eq!(trace.spans[0].actor_id, None); + assert_eq!(trace.spans[0].actor_unassigned, unassigned); +} + +#[rstest] +fn openinference_sdk_session_keeps_recorded_root_ownership() { + let root = TraceSpansRow { + framework: "claude-agent-sdk".into(), + session_id: "sdk-session".into(), + ..row("root", "", "ClaudeAgentSDK.query", "agent", "sdk-reader") + }; + let trace = resolve_trace("trace", "ref", &[root], &[]).unwrap(); + assert_eq!(trace.summary.agent_count, 1); + assert_eq!(trace.agents[0].name, "sdk-reader"); + assert!(!trace.spans[0].actor_unassigned); +} + +#[rstest] +fn inherited_native_actor_identity_and_labels_have_bounded_size() { + let root = TraceSpansRow { + framework: "claude-code".into(), + session_id: "s".into(), + native_agent_id: "a".repeat(65_536), + query_source: format!("agent.custom.{}", "n".repeat(65_536)), + ..row("root", "", "claude_code.interaction", "agent", "assistant") + }; + let children = (0..1_024).map(|index| TraceSpansRow { + framework: "claude-code".into(), + session_id: "s".into(), + ..row( + &format!("child-{index}"), + "root", + "Read", + "tool", + "assistant", + ) + }); + let rows: Vec<_> = std::iter::once(root).chain(children).collect(); + let trace = resolve_trace("trace", "ref", &rows, &[]).unwrap(); + assert_eq!(trace.summary.agent_count, 1); + assert!(trace.spans.iter().all(|span| { + span.actor_id.as_ref().is_some_and(|id| id.len() <= 80) && span.agent.len() <= 1_024 + })); + assert!(serde_json::to_vec(&trace).unwrap().len() < 3_000_000); +} + +#[rstest] +#[case::unassigned(false)] +#[case::owned(true)] +fn deep_native_chains_resolve_owned_and_unassigned_spans(#[case] owned: bool) { + let rows: Vec<_> = (0..50_000) + .rev() + .map(|index| TraceSpansRow { + framework: "claude-code".into(), + session_id: "session".into(), + query_source: if owned && index == 0 { + "repl_main_thread".into() + } else { + String::new() + }, + ..row( + &index.to_string(), + &if index == 0 { + String::new() + } else { + (index - 1).to_string() + }, + "native", + "framework", + "assistant", + ) + }) + .collect(); + let trace = resolve_trace("trace", "ref", &rows, &[]).unwrap(); + assert_eq!(trace.spans.len(), rows.len()); + assert!( + trace + .spans + .iter() + .all(|span| span.actor_id.is_some() == owned && span.actor_unassigned != owned) + ); + assert_eq!(trace.summary.agent_count, u64::from(owned)); +} + +#[rstest] +fn capture_coverage_detects_missing_child_replies_before_content_is_paged() { + let root = TraceSpansRow { + framework: "claude-code".into(), + session_id: "s".into(), + query_source: "repl_main_thread".into(), + ..row("root", "", "claude_code.interaction", "agent", "assistant") + }; + let reply = TraceSpansRow { + has_output: true, + ..TraceSpansRow { + span_id: "reply".into(), + name: "claude_code.assistant_response".into(), + kind: litellm_traces::ObservationType::Chain, + ..root.clone() + } + }; + let child = TraceSpansRow { + span_id: "child".into(), + native_agent_id: "child-id".into(), + query_source: "agent.custom.Reader".into(), + name: "claude_code.llm_request".into(), + kind: litellm_traces::ObservationType::Llm, + ..root.clone() + }; + let unassigned = TraceSpansRow { + span_id: "unassigned".into(), + query_source: "agent".into(), + capture_warning: true, + name: "claude_code.assistant_response".into(), + kind: litellm_traces::ObservationType::Chain, + ..root.clone() + }; + let trace = resolve_trace("trace", "ref", &[root, reply, child, unassigned], &[]).unwrap(); + let capture = trace.capture.unwrap(); + assert_eq!( + ( + capture.content_events, + capture.unassigned_events, + capture.warning_events + ), + (2, 1, 1) + ); + let child = capture + .actors + .iter() + .find(|actor| actor.name == "Reader") + .unwrap(); + assert_eq!( + (child.llm_calls, child.reply_events, child.model_outputs), + (1, 0, 0) + ); + assert_eq!( + capture + .actors + .iter() + .map(|actor| actor.reply_events) + .sum::(), + 1 + ); +} + +#[rstest] +#[case::native_root(true)] +#[case::legacy_root(false)] +fn sdk_tools_inherit_only_recorded_actor_ownership(#[case] native_root: bool) { + let root = TraceSpansRow { + framework: "claude-agent-sdk".into(), + session_id: "sdk-session".into(), + ..row( + "root", + "", + if native_root { + "claude_code.interaction" + } else { + "ClaudeAgentSDK.query" + }, + "agent", + "sdk-reader", + ) + }; + let tool = TraceSpansRow { + framework: root.framework.clone(), + session_id: root.session_id.clone(), + ..row("tool", "root", "Bash", "tool", "sdk-reader") + }; + let trace = resolve_trace("trace", "ref", &[root, tool], &[]).unwrap(); + assert_eq!(trace.agents.len(), 1); + assert_eq!(trace.agents[0].tool_calls, 1); + assert_eq!(trace.spans[1].actor_id.is_some(), native_root); + assert_eq!(trace.spans[1].actor_id, trace.agents[0].actor_id); +} + +#[rstest] +fn many_native_actors_count_only_their_own_calls() { + let rows: Vec<_> = (0..50_000) + .map(|index| TraceSpansRow { + framework: "claude-code".into(), + session_id: "session".into(), + native_agent_id: index.to_string(), + ..row( + &index.to_string(), + "", + "claude_code.llm_request", + "llm", + "assistant", + ) + }) + .collect(); + let trace = resolve_trace("trace", "ref", &rows, &[]).unwrap(); + assert_eq!(trace.summary.agent_count, 50_000); + assert!( + trace + .agents + .iter() + .all(|actor| actor.llm_calls == 1 && actor.tool_calls == 0) + ); +} diff --git a/litellm/rust_bridge/trace/generated/types.py b/litellm/rust_bridge/trace/generated/types.py index 4a09e731228..050aba4f90b 100644 --- a/litellm/rust_bridge/trace/generated/types.py +++ b/litellm/rust_bridge/trace/generated/types.py @@ -48,10 +48,22 @@ class SpanErrorPage(typing_extensions.TypedDict): next_cursor: ReadOnly[str | None] +class ActorCapture(typing_extensions.TypedDict): + actor_id: ReadOnly[str] + name: ReadOnly[str] + llm_calls: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + tool_calls: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + reply_events: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + model_outputs: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + content_events: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + + SpanStatus: TypeAlias = Literal["ok", "error", "unset"] class AgentNode(typing_extensions.TypedDict): + actor_id: ReadOnly[NotRequired[str | None]] + parent_actor_id: ReadOnly[NotRequired[str | None]] name: ReadOnly[str] parent_agent: ReadOnly[str | None] invocations: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] @@ -102,6 +114,13 @@ class UIMessage(typing_extensions.TypedDict): tool_calls: ReadOnly[NotRequired[tuple[UIToolCall, ...]]] +class TraceCapture(typing_extensions.TypedDict): + actors: ReadOnly[tuple[ActorCapture, ...]] + content_events: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + unassigned_events: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + warning_events: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + + class TraceSummary(typing_extensions.TypedDict): resolution_limited: ReadOnly[NotRequired[bool]] trace_id: ReadOnly[str] @@ -133,6 +152,9 @@ class Span(typing_extensions.TypedDict): name: ReadOnly[str] type: ReadOnly[SpanType] agent: ReadOnly[str] + actor_id: ReadOnly[NotRequired[str | None]] + actor_unassigned: ReadOnly[NotRequired[bool]] + parent_actor_id: ReadOnly[NotRequired[str | None]] framework: ReadOnly[str] start_offset_ms: ReadOnly[float] duration_ms: ReadOnly[float] @@ -150,6 +172,7 @@ class Span(typing_extensions.TypedDict): class Trace(typing_extensions.TypedDict): + capture: ReadOnly[NotRequired[TraceCapture | None | None]] summary: ReadOnly[TraceSummary] agents: ReadOnly[tuple[AgentNode, ...]] spans: ReadOnly[tuple[Span, ...]] diff --git a/scripts/trace_codegen/schemas/traces/Trace.json b/scripts/trace_codegen/schemas/traces/Trace.json index ebd99880956..b0a81fafb49 100644 --- a/scripts/trace_codegen/schemas/traces/Trace.json +++ b/scripts/trace_codegen/schemas/traces/Trace.json @@ -1,8 +1,65 @@ { "$defs": { + "ActorCapture": { + "properties": { + "actor_id": { + "type": "string" + }, + "content_events": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + }, + "llm_calls": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + }, + "model_outputs": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + }, + "name": { + "type": "string" + }, + "reply_events": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + }, + "tool_calls": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "actor_id", + "name", + "llm_calls", + "tool_calls", + "reply_events", + "model_outputs", + "content_events" + ], + "type": "object" + }, "AgentNode": { "description": "One distinct agent in a trace: 200 invocations of `researcher` are one node.", "properties": { + "actor_id": { + "type": [ + "string", + "null" + ], + "x-python-optional": true + }, "duration_ms": { "format": "double", "type": "number" @@ -22,6 +79,13 @@ "name": { "type": "string" }, + "parent_actor_id": { + "type": [ + "string", + "null" + ], + "x-python-optional": true + }, "parent_agent": { "type": [ "string", @@ -49,6 +113,8 @@ } }, "required": [ + "actor_id", + "parent_actor_id", "name", "parent_agent", "invocations", @@ -62,6 +128,18 @@ }, "Span": { "properties": { + "actor_id": { + "type": [ + "string", + "null" + ], + "x-python-optional": true + }, + "actor_unassigned": { + "default": false, + "type": "boolean", + "x-python-optional": true + }, "agent": { "type": "string" }, @@ -111,6 +189,13 @@ "minimum": 0, "type": "integer" }, + "parent_actor_id": { + "type": [ + "string", + "null" + ], + "x-python-optional": true + }, "parent_span_id": { "type": [ "string", @@ -160,6 +245,9 @@ "name", "type", "agent", + "actor_id", + "actor_unassigned", + "parent_actor_id", "framework", "start_offset_ms", "duration_ms", @@ -212,6 +300,41 @@ ], "type": "string" }, + "TraceCapture": { + "properties": { + "actors": { + "items": { + "$ref": "#/$defs/ActorCapture" + }, + "type": "array" + }, + "content_events": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + }, + "unassigned_events": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + }, + "warning_events": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "actors", + "content_events", + "unassigned_events", + "warning_events" + ], + "type": "object" + }, "TraceSummary": { "properties": { "agent_count": { @@ -361,6 +484,18 @@ }, "type": "array" }, + "capture": { + "anyOf": [ + { + "$ref": "#/$defs/TraceCapture" + }, + { + "type": "null" + } + ], + "default": null, + "x-python-optional": true + }, "next_cursor": { "type": [ "string", @@ -379,6 +514,7 @@ } }, "required": [ + "capture", "summary", "agents", "spans", diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/CaptureCoverage.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/CaptureCoverage.integration.test.tsx new file mode 100644 index 00000000000..71f6432abe8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/CaptureCoverage.integration.test.tsx @@ -0,0 +1,46 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import { CaptureCoverage } from "./CaptureCoverage"; +import type { Trace } from "../../types"; + +const coverage: NonNullable = { + actors: [ + { + actor_id: "root", + name: "Root", + llm_calls: 2, + tool_calls: 1, + reply_events: 1, + model_outputs: 0, + content_events: 1, + }, + { + actor_id: "child", + name: "Reader", + llm_calls: 3, + tool_calls: 2, + reply_events: 0, + model_outputs: 0, + content_events: 0, + }, + ], + content_events: 2, + unassigned_events: 1, + warning_events: 1, +}; + +describe("CaptureCoverage", () => { + it("shows missing child capture even when the root reply was recorded", () => { + render(); + expect(screen.getByText("No reply recorded yet for Reader.")).toBeVisible(); + expect(screen.getByText(/More records may still arrive/)).toBeVisible(); + expect(screen.getByText("1 content events have unconfirmed actor ownership.")).toBeVisible(); + expect(screen.getByText(/1 records carry capture warnings/)).toBeVisible(); + }); + + it("accepts recorded model output as reply evidence", () => { + const capture = { ...coverage, actors: coverage.actors.map((actor) => ({ ...actor, model_outputs: 1 })) }; + render(); + expect(screen.queryByText(/No reply recorded yet/)).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/CaptureCoverage.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/CaptureCoverage.tsx new file mode 100644 index 00000000000..aecd75c844f --- /dev/null +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/CaptureCoverage.tsx @@ -0,0 +1,51 @@ +import { useState } from "react"; +import { Button } from "@/components/ui/button"; +import type { Trace } from "../../types"; + +export function CaptureCoverage({ capture, hasMore = false }: { capture: Trace["capture"]; hasMore?: boolean }) { + const [limit, setLimit] = useState(20); + if (!capture) return null; + const missing = capture.actors.filter((actor) => actor.llm_calls > 0 && !actor.reply_events && !actor.model_outputs); + const names = missing + .slice(0, 3) + .map((actor) => actor.name) + .join(", "); + return ( +
+

Capture coverage

+

+ {capture.content_events} content events received. Coverage for {capture.actors.length} loaded actors. More + records may still arrive. +

+ {hasMore &&

Load more steps to include coverage for later actors.

} + {missing.length > 0 && ( +

+ No reply recorded yet for {names} + {missing.length > 3 ? ` and ${missing.length - 3} more actors` : ""}. +

+ )} + {capture.unassigned_events > 0 && ( +

{capture.unassigned_events} content events have unconfirmed actor ownership.

+ )} + {capture.warning_events > 0 && ( +

{capture.warning_events} records carry capture warnings. Inspect their source spans.

+ )} +
+ View actor coverage +
    + {capture.actors.slice(0, limit).map((actor) => ( +
  • + {actor.name}: {actor.llm_calls} model calls, {actor.tool_calls}{" "} + tools, {actor.reply_events} reply events, {actor.model_outputs} recorded model outputs +
    • + ))} +
+ {capture.actors.length > limit && ( + + )} +
+
+ ); +} diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.integration.test.tsx index eb66ce65bba..c950b3a53d7 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.integration.test.tsx @@ -80,6 +80,119 @@ describe("TraceConversation", () => { expect(screen.getByRole("heading", { name: "read_file" })).toBeVisible(); }); + it("continues native child branches and merges actor coverage across trace pages", async () => { + const user = userEvent.setup(); + const spans = [ + { ...root, actor_id: "root-actor", parent_actor_id: null, agent: "Root" }, + ...Array.from({ length: 25 }, (_, index) => ({ + ...tool, + span_id: `child-${index}`, + actor_id: "child-actor", + parent_actor_id: "root-actor", + agent: "Reader", + name: `child tool ${index}`, + start_offset_ms: index + 1, + })), + { + ...tool, + span_id: "later", + actor_id: "later-actor", + parent_actor_id: "root-actor", + agent: "Later", + start_offset_ms: 30, + }, + ] as Trace["spans"]; + const coverage = (actors: string[]): NonNullable => ({ + actors: actors.map((actor_id) => ({ + actor_id, + name: actor_id, + llm_calls: 1, + tool_calls: 1, + reply_events: 1, + model_outputs: 0, + content_events: 1, + })), + content_events: 3, + unassigned_events: 0, + warning_events: 0, + }); + const first: Trace = { + ...trace, + spans: spans.slice(0, 20), + capture: coverage(["root-actor", "child-actor"]), + next_cursor: "native-next", + }; + const second: Trace = { + ...trace, + spans: spans.slice(20), + capture: coverage(["child-actor", "later-actor"]), + next_cursor: null, + }; + vi.mocked(agentTraceCall).mockImplementation(async (_token, _trace, _ref, cursor) => (cursor ? second : first)); + renderWithProviders( + , + ); + await user.click(await screen.findByRole("tab", { name: "Conversation" })); + await user.click(await screen.findByText("Subagent: Reader", { exact: true })); + const more = await screen.findByRole("button", { name: "Load next 20 entries in Reader" }); + await waitFor(() => expect(more).toBeEnabled()); + await user.click(more); + expect(await screen.findByRole("button", { name: "Expand child tool 24 tool call" })).toBeVisible(); + expect(screen.getByText(/Coverage for 3 loaded actors/)).toBeVisible(); + expect(screen.queryByText(/Load more steps to include coverage/)).not.toBeInTheDocument(); + expect(vi.mocked(agentTraceCall).mock.calls.map((call) => call[3])).toEqual([null, "native-next"]); + }); + + it("waits for an in-flight refresh before requesting another conversation page", async () => { + const user = userEvent.setup(); + const first = { ...trace, spans: [root], next_cursor: "old-page" }; + const refreshed = { ...first, next_cursor: "fresh-page" }; + const second = { ...trace, spans: [tool], next_cursor: null }; + const pending = Promise.withResolvers(); + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(first) + .mockReturnValueOnce(pending.promise) + .mockResolvedValue(second); + renderWithProviders( + , + ); + await user.click(await screen.findByRole("tab", { name: "Conversation" })); + expect(await screen.findByText("Read the release notes")).toBeVisible(); + await user.click(screen.getByRole("button", { name: "Refresh run" })); + const more = screen.getByRole("button", { name: "Load next 20 entries" }); + expect(more).toBeDisabled(); + await user.click(more); + expect(agentTraceCall).toHaveBeenCalledTimes(2); + await act(async () => pending.resolve(refreshed)); + await waitFor(() => expect(more).toBeEnabled()); + await user.click(more); + expect(await screen.findByRole("button", { name: "Expand read_file tool call" })).toBeVisible(); + expect(vi.mocked(agentTraceCall).mock.calls.map((call) => call[3])).toEqual([null, null, "fresh-page"]); + }); + + it("can load the next conversation page after a refresh fails", async () => { + const user = userEvent.setup(); + const first = { ...trace, spans: [root], next_cursor: "next-page" }; + const second = { ...trace, spans: [tool], next_cursor: null }; + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(first) + .mockRejectedValueOnce(new Error("refresh unavailable")) + .mockResolvedValue(second); + renderWithProviders( + , + ); + await user.click(await screen.findByRole("tab", { name: "Conversation" })); + expect(await screen.findByText("Read the release notes")).toBeVisible(); + await user.click(screen.getByRole("button", { name: "Refresh run" })); + expect(await screen.findByRole("alert")).toHaveTextContent("Previously received steps are still shown"); + const more = screen.getByRole("button", { name: "Load next 20 entries" }); + expect(more).toBeEnabled(); + await user.click(more); + expect(await screen.findByRole("button", { name: "Expand read_file tool call" })).toBeVisible(); + expect(vi.mocked(agentTraceCall).mock.calls.map((call) => call[3])).toEqual([null, null, "next-page"]); + expect(screen.queryByText(/Could not load more conversation entries/)).not.toBeInTheDocument(); + }); + it("renders a failed shell exchange in both views and preserves its raw result", async () => { const user = userEvent.setup(); const command = "npm test -- checkout\nprintf 'finished\\n'"; diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.tsx index 452f3306f45..b1c260f7a8c 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/conversation/TraceConversation.tsx @@ -21,6 +21,7 @@ import { ErrorBlock } from "../content/SpanError"; import { Markdown } from "../content/Markdown"; import { ToolCallBlock, ToolResultCard } from "../content/Messages"; import type { Trace, TraceMessage } from "../../types"; +import { CaptureCoverage } from "./CaptureCoverage"; import { fmtMs } from "../../utils"; import { useConversationDetails } from "./useConversationDetails"; @@ -91,6 +92,7 @@ export function TraceConversation({ return (
+ {rootErrors.map((span) => ( ))} 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 483c38000ce..5b5a0af8120 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 @@ -27,6 +27,48 @@ const detail = (span_id: string, input: unknown, output: unknown): SpanDetail => }); describe("trace conversation", () => { + it.each([false, true])("keeps native actor branches pending across details and pages (%s)", (hasMore) => { + const child = { ...root, span_id: "child", actor_id: "child-actor", parent_actor_id: "root-actor", duration_ms: 1 }; + const later = { ...root, span_id: "later", actor_id: "root-actor", start_offset_ms: 100 }; + const details = new Map([ + [later.span_id, detail(later.span_id, [], [])], + ...(hasMore ? [[child.span_id, detail(child.span_id, [], [])] as const] : []), + ]); + const pending = pendingConversationBranches([child, later], details, hasMore); + expect(pending.has("child-actor")).toBe(true); + expect(pending.has("root-actor")).toBe(true); + expect(pending.has("child")).toBe(true); + }); + + it("builds a conversation below 20,000 hidden framework spans", () => { + const chain = Array.from( + { length: 20_000 }, + (_, index): Span => ({ + ...root, + span_id: `hidden-${index}`, + parent_span_id: index ? `hidden-${index - 1}` : root.span_id, + type: "framework", + name: "internal", + start_offset_ms: 1, + }), + ); + const tool: Span = { + ...root, + span_id: "tool", + parent_span_id: "hidden-19999", + type: "tool", + name: "Read", + start_offset_ms: 2, + }; + const details = new Map([ + [root.span_id, detail(root.span_id, "Read file", "")], + [tool.span_id, detail(tool.span_id, { path: "alpha.txt" }, "value")], + ]); + const items = buildConversation([root, ...chain, tool], details, true); + expect(items.map((item) => item.span.span_id)).toEqual([root.span_id, tool.span_id]); + expect(items[1].toolResult).toBe("value"); + }); + it("keeps missing nested details pending through framework and native agent ancestors", () => { const agent = { ...root, @@ -717,3 +759,81 @@ describe("coding sessions", () => { expect(conversationWarnings(details, true)).toEqual([]); }); }); + +it.each(["claude-code", "claude-agent-sdk"])( + "keeps explicit %s actors separate across workflow branches and background resumes", + (framework) => { + const span = (id: string, actor: string, parentActor: string | null, offset: number): Span => ({ + ...root, + span_id: id, + name: "claude_code.llm_request", + type: "llm", + framework, + parent_span_id: "root", + actor_id: actor, + parent_actor_id: parentActor, + agent: "general-purpose", + start_offset_ms: offset, + }); + const spans = [ + { ...root, framework, actor_id: "root-actor", parent_actor_id: null, agent: "claude-code" }, + span("first", "actor-one", "root-actor", 1), + span("second", "actor-two", "root-actor", 2), + span("nested", "actor-three", "actor-one", 3), + { ...span("resumed", "actor-one", "root-actor", 4), parent_span_id: null }, + ]; + const details = new Map(spans.map((span) => [span.span_id, detail(span.span_id, "", `Reply ${span.span_id}`)])); + const items = buildConversation(spans, details, true); + expect(items.filter((item) => item.span.span_id === "resumed").map((item) => item.agentId)).toEqual(["actor-one"]); + const groups = groupConversation(items, spans); + const branches = groups.filter((group) => group.kind === "branch"); + expect(branches.map((group) => group.id)).toEqual(["actor-one", "actor-two"]); + expect(branches[0].children.map((group) => (group.kind === "branch" ? group.id : group.item.span.span_id))).toEqual( + ["first", "actor-three", "resumed"], + ); + expect( + items.filter((item) => ["first", "resumed"].includes(item.span.span_id)).map((item) => item.agentName), + ).toEqual(["general-purpose (1)", "general-purpose (1)"]); + }, +); + +it.each([false, true])("preserves legacy history while isolating unconfirmed actors (%s)", (unassigned) => { + const spans: Span[] = [ + { ...root, framework: "claude-code", actor_id: null, actor_unassigned: false }, + { + ...root, + span_id: "first", + type: "llm", + framework: "claude-code", + parent_span_id: root.span_id, + actor_id: null, + actor_unassigned: unassigned, + start_offset_ms: 1, + }, + { + ...root, + span_id: "second", + type: "llm", + framework: "claude-code", + parent_span_id: root.span_id, + actor_id: null, + actor_unassigned: unassigned, + start_offset_ms: 2, + }, + ]; + const user = { role: "user", content: "Read this" }; + const answer = { role: "assistant", content: "First answer" }; + const details = new Map([ + [root.span_id, detail(root.span_id, [], [])], + ["first", detail("first", [user], [answer])], + ["second", detail("second", [user, answer], [{ role: "assistant", content: "Second answer" }])], + ]); + const items = buildConversation(spans, details, true); + expect(items.map((item) => item.branchId)).toEqual( + unassigned ? ["unassigned:first", "unassigned:second"] : [root.span_id, root.span_id], + ); + expect(items[1].messages.map((message) => message.content)).toEqual( + unassigned ? ["Read this", "First answer", "Second answer"] : ["Second answer"], + ); + expect(items.every((item) => item.agentName === "Unassigned")).toBe(unassigned); +}); 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 9fc1a798a65..3d4d7ac7df6 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 @@ -27,18 +27,20 @@ export function pendingConversationBranches( const pageBoundary = spans.reduce((latest, span) => Math.max(latest, span.start_offset_ms), -Infinity); const pending = spans.filter((span) => { const missingDetail = steps.has(span.span_id) && !details.has(span.span_id); - const mayHaveLaterChildren = hasMoreSpans && span.start_offset_ms + span.duration_ms >= pageBoundary; + const mayHaveLaterChildren = + hasMoreSpans && (Boolean(span.actor_id) || span.start_offset_ms + span.duration_ms >= pageBoundary); return missingDetail || mayHaveLaterChildren; }); const initial = { - pending: new Set(pending.map((span) => span.span_id)), + pending: new Set(pending.flatMap((span) => [span.span_id, ...(span.actor_id ? [span.actor_id] : [])])), ancestors: new Map( - spans.flatMap((span) => - span.parent_span_id && byId.has(span.parent_span_id) ? [[span.span_id, span.parent_span_id] as const] : [], - ), + spans.flatMap((span) => [ + ...(span.parent_span_id && byId.has(span.parent_span_id) ? [[span.span_id, span.parent_span_id] as const] : []), + ...(span.actor_id && span.parent_actor_id ? [[span.actor_id, span.parent_actor_id] as const] : []), + ]), ), }; - const doublingPasses = Math.ceil(Math.log2(Math.max(1, spans.length))); + const doublingPasses = Math.ceil(Math.log2(Math.max(1, initial.ancestors.size + 1))); return Array.from({ length: doublingPasses }).reduce((state) => { if (!state.pending.size || !state.ancestors.size) return state; return { @@ -281,6 +283,8 @@ function isNativeAgent(span: Span): boolean { } function conversationBranch(span: Span, byId: ReadonlyMap): string { + if (span.actor_id) return span.actor_id; + if (span.actor_unassigned) return `unassigned:${span.span_id}`; if (span.type === "agent" || span.parent_span_id === null) return span.span_id; let parent = span.parent_span_id ? byId.get(span.parent_span_id) : undefined; const visited = new Set(); @@ -293,20 +297,21 @@ function conversationBranch(span: Span, byId: ReadonlyMap): string } function agentIdentity(span: Span, details: ReadonlyMap): string { - return isNativeAgent(span) ? span.span_id : details.get(span.span_id)?.attributes["gen_ai.agent.id"] || span.span_id; + return ( + span.actor_id || + (isNativeAgent(span) ? span.span_id : details.get(span.span_id)?.attributes["gen_ai.agent.id"] || span.span_id) + ); } function agentLabels(agents: readonly Span[], details: ReadonlyMap): ReadonlyMap { const identity = (agent: Span): string => agentIdentity(agent, details); - const unique = agents.filter( - (agent, index) => agents.findIndex((other) => identity(other) === identity(agent)) === index, - ); + const unique = [...new Map(agents.toReversed().map((agent) => [identity(agent), agent])).values()]; const named = unique.map((agent) => { const detail = details.get(agent.span_id); const args = detail && isNativeAgent(agent) ? toolInput(detail.input, detail.input_ui) : undefined; const description = args && typeof args === "object" && "description" in args ? args.description : undefined; const name = - agent.framework === "claude-code" && !isNativeAgent(agent) + agent.actor_id || (agent.framework === "claude-code" && !isNativeAgent(agent)) ? agent.agent || agent.name || "Agent" : agent.name || agent.agent || "Agent"; return { agent, name: typeof description === "string" && description ? description : name }; @@ -413,22 +418,31 @@ export function buildConversation( }; if (combined.length || item.showError) items.push(item); } + const byBranch = new Map( + spans + .filter((span) => span.actor_id || span.actor_unassigned) + .sort((a, b) => b.start_offset_ms - a.start_offset_ms) + .map((span) => [branch(span), span]), + ); const agents = [...new Set(conversationSteps(spans).map(branch))].flatMap((id) => { - const span = byId.get(id); + const span = byId.get(id) ?? byBranch.get(id); return span ? [span] : []; }); const labels = agentLabels(agents, details); const actor = (id: string): string => { - const span = byId.get(id); + const span = byId.get(id) ?? byBranch.get(id); + if (span?.actor_unassigned) return id; return span ? agentIdentity(span, details) : id; }; return items .map((item) => ({ ...item, agentId: actor(branch(item.span)), - agentName: labels.get(actor(branch(item.span))) || item.span.agent, + agentName: item.span.actor_unassigned ? "Unassigned" : labels.get(actor(branch(item.span))) || item.span.agent, branchId: branch(item.span), parentBranchId: (() => { + if (item.span.actor_id) return item.span.parent_actor_id || undefined; + if (item.span.actor_unassigned) return undefined; const parentId = byId.get(branch(item.span))?.parent_span_id; const parent = parentId ? byId.get(parentId) : undefined; return parent ? branch(parent) : undefined; @@ -457,6 +471,8 @@ export function groupConversation(items: readonly ConversationItem[], spans: rea const byId = new Map(spans.map((span) => [span.span_id, span])); const parentById = new Map( conversationSteps(spans).flatMap((span) => { + if (span.actor_id) return [[span.actor_id, span.parent_actor_id || undefined] as const]; + if (span.actor_unassigned) return [[conversationBranch(span, byId), undefined] as const]; const branch = byId.get(conversationBranch(span, byId)); if (!branch) return []; const parent = branch.parent_span_id ? byId.get(branch.parent_span_id) : undefined; diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunHeader.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunHeader.tsx index 33fb073d0a9..f5ec196244f 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunHeader.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunHeader.tsx @@ -1,6 +1,6 @@ "use client"; -import { ArrowLeft, Check, Copy, Link } from "lucide-react"; +import { ArrowLeft, Check, Copy, Link, RefreshCw } from "lucide-react"; import { useState } from "react"; import { useTimeout } from "usehooks-ts"; @@ -52,13 +52,14 @@ function Stat({ label, value, error = false }: { label: string; value: string; e function StatusPill({ failed }: { failed: boolean }) { return ( - {failed ? "Failed" : "Completed"} + {failed ? "Errors recorded" : "Recorded"} ); } @@ -83,10 +84,25 @@ interface RunHeaderProps { handoff: TraceHandoff; onBack: () => void; embedded: boolean; + refreshing: boolean; + onRefresh: () => void; + live: boolean; + canLive: boolean; + onLiveChange: () => void; } /** Run identity, view switch and totals in two tight rows. */ -export function RunHeader({ trace, handoff, onBack, embedded }: RunHeaderProps) { +export function RunHeader({ + trace, + handoff, + onBack, + embedded, + refreshing, + onRefresh, + live, + canLive, + onLiveChange, +}: RunHeaderProps) { const { summary } = trace; const failed = summary.status === "error"; const cost = runCost(summary); @@ -113,6 +129,20 @@ export function RunHeader({ trace, handoff, onBack, embedded }: RunHeaderProps)
+ + { ); expect(header).toHaveTextContent("Duration 40.20s"); expect(header).toHaveTextContent(`Steps ${research.summary.span_count}`); - expect(header).not.toHaveTextContent("failed"); + expect(header).toHaveTextContent("Recorded"); + expect(header).not.toHaveTextContent("Completed"); }); it("shows the agent name with the SDK logo in the run header instead of the generic agent icon", async () => { @@ -361,7 +362,7 @@ describe("RunView", () => { expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(2); }); - it("keeps a loaded snapshot on focus and reconnect", async () => { + it("refreshes a stale run on focus and reconnect without losing its selected step", async () => { testQueryClient.setQueryDefaults(["agentTrace"], { refetchOnWindowFocus: true, refetchOnReconnect: true }); vi.mocked(agentTraceCall).mockReset(); renderRun(research); @@ -374,10 +375,200 @@ describe("RunView", () => { onlineManager.setOnline(true); }); await waitFor(() => expect(testQueryClient.isFetching()).toBe(0)); - expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(1); + expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(2); + expect(screen.getByTestId("detail-pane")).toHaveAttribute("data-row-id", rootSpanId(research)); expect(screen.getByRole("tree", { name: "Spans in time order" })).toHaveTextContent(research.spans[0].name); }); + it("refreshes new spans from the header while preserving the selected step", async () => { + vi.mocked(agentTraceCall).mockReset(); + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(research) + .mockResolvedValue({ + ...research, + summary: { ...research.summary, span_count: research.summary.span_count + 1 }, + spans: [ + ...research.spans, + { + ...research.spans[0], + span_id: "new-step", + name: "newly received step", + parent_span_id: rootSpanId(research), + type: "tool", + }, + ], + }); + renderWithProviders(); + expect(await screen.findByTestId("detail-pane")).toHaveAttribute("data-row-id", rootSpanId(research)); + fireEvent.click(screen.getByRole("button", { name: "Refresh run" })); + expect(await screen.findByText("newly received step")).toBeVisible(); + expect(screen.getByRole("banner")).toHaveTextContent(`Steps ${research.summary.span_count + 1}`); + expect(screen.getByTestId("detail-pane")).toHaveAttribute("data-row-id", rootSpanId(research)); + }); + + it("receives a growing run on the polling interval", async () => { + vi.useFakeTimers(); + try { + vi.mocked(agentTraceCall).mockReset(); + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(research) + .mockResolvedValue({ + ...research, + summary: { ...research.summary, span_count: research.summary.span_count + 7 }, + }); + renderWithProviders(); + await act(async () => { + await vi.advanceTimersByTimeAsync(100); + }); + expect(screen.getByRole("banner")).toHaveTextContent(`Steps ${research.summary.span_count}`); + await act(async () => { + await vi.advanceTimersByTimeAsync(30_100); + }); + expect(screen.getByRole("banner")).toHaveTextContent(`Steps ${research.summary.span_count + 7}`); + expect(screen.getByTestId("detail-pane")).toHaveAttribute("data-row-id", rootSpanId(research)); + } finally { + vi.useRealTimers(); + } + }); + + it("keeps fixed snapshot sources out of live mode", async () => { + vi.mocked(agentTraceCall).mockResolvedValue(research); + renderWithProviders( + + + , + ); + expect(await screen.findByRole("button", { name: "Live updates" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Live updates" })).toHaveAttribute("aria-pressed", "false"); + expect(screen.getByRole("button", { name: "Refresh run" })).toBeEnabled(); + }); + + it("can pause live updates without disabling manual refresh", async () => { + vi.mocked(agentTraceCall).mockReset(); + renderRun(research); + await screen.findByTestId("detail-pane"); + fireEvent.click(screen.getByRole("button", { name: "Live updates" })); + expect(screen.getByRole("button", { name: "Live updates" })).toHaveAttribute("aria-pressed", "false"); + await testQueryClient.invalidateQueries({ queryKey: ["agentTrace"], refetchType: "none" }); + await act(async () => { + focusManager.setFocused(false); + onlineManager.setOnline(false); + focusManager.setFocused(true); + onlineManager.setOnline(true); + }); + expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(1); + fireEvent.click(screen.getByRole("button", { name: "Refresh run" })); + await waitFor(() => expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(2)); + }); + + it("cancels an in-flight refresh and stops interval requests when live updates pause", async () => { + vi.useFakeTimers(); + try { + const updated = { ...research, summary: { ...research.summary, span_count: research.summary.span_count + 7 } }; + vi.mocked(agentTraceCall).mockReset(); + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(research) + .mockImplementationOnce( + () => + new Promise((resolve) => { + setTimeout(() => resolve(updated), 1_000); + }), + ); + renderWithProviders(); + await act(async () => { + await vi.advanceTimersByTimeAsync(30_100); + }); + expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(2); + fireEvent.click(screen.getByRole("button", { name: "Live updates" })); + await act(async () => { + await vi.advanceTimersByTimeAsync(60_100); + }); + expect(screen.getByRole("banner")).toHaveTextContent(`Steps ${research.summary.span_count}`); + expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(2); + } finally { + vi.useRealTimers(); + } + }); + + it.each(["Refresh run", "Load more steps"])( + "keeps an in-flight manual %s when live updates pause", + async (action) => { + vi.useFakeTimers(); + try { + const first: Trace = { ...research, next_cursor: "next-page" }; + const updated: Trace = { + ...research, + summary: { ...research.summary, span_count: research.summary.span_count + 7 }, + next_cursor: null, + }; + vi.mocked(agentTraceCall).mockReset(); + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(first) + .mockImplementationOnce( + () => + new Promise((resolve) => { + setTimeout(() => resolve(updated), 1_000); + }), + ); + renderWithProviders( + , + ); + await act(async () => { + await vi.advanceTimersByTimeAsync(100); + }); + fireEvent.click(screen.getByRole("button", { name: action })); + fireEvent.click(screen.getByRole("button", { name: "Live updates" })); + await act(async () => { + await vi.advanceTimersByTimeAsync(60_100); + }); + expect(screen.queryByRole("button", { name: "Load more steps" })).not.toBeInTheDocument(); + expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(2); + if (action === "Refresh run") { + expect(screen.getByRole("banner")).toHaveTextContent(`Steps ${updated.summary.span_count}`); + } + } finally { + vi.useRealTimers(); + } + }, + ); + + it("retries a failed refresh without calling the next-page operation", async () => { + vi.mocked(agentTraceCall).mockReset(); + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(research) + .mockRejectedValueOnce(new Error("refresh unavailable")) + .mockResolvedValue(research); + renderWithProviders(); + await screen.findByTestId("detail-pane"); + fireEvent.click(screen.getByRole("button", { name: "Refresh run" })); + expect(await screen.findByRole("alert")).toHaveTextContent("Previously received steps are still shown"); + expect(screen.queryByText(/could not load more/i)).not.toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Retry refresh" })); + await waitFor(() => expect(screen.queryByRole("alert")).not.toBeInTheDocument()); + expect(vi.mocked(agentTraceCall).mock.calls.map((call) => call[3])).toEqual([null, null, null]); + }); + + it("refreshes paginated runs using the newly returned cursor", async () => { + const first = { ...research, spans: research.spans.slice(0, 1), next_cursor: "old-page" }; + const second = { ...research, spans: research.spans.slice(1), next_cursor: null }; + vi.mocked(agentTraceCall).mockReset(); + vi.mocked(agentTraceCall) + .mockResolvedValueOnce(first) + .mockResolvedValueOnce(second) + .mockResolvedValueOnce({ ...first, next_cursor: "new-page" }) + .mockImplementationOnce(async (_token, _id, _ref, cursor) => { + if (cursor !== "new-page") throw new Error("stale cursor"); + return second; + }); + renderWithProviders(); + fireEvent.click(await screen.findByRole("button", { name: "Load more steps" })); + await waitFor(() => expect(screen.queryByRole("button", { name: "Load more steps" })).not.toBeInTheDocument()); + fireEvent.click(screen.getByRole("button", { name: "Refresh run" })); + await waitFor(() => expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(4)); + expect(vi.mocked(agentTraceCall).mock.calls.map((call) => call[3])).toEqual([null, "old-page", null, "new-page"]); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + it("keeps a way back to the runs table when a run fails to load", async () => { const user = userEvent.setup(); const onBack = vi.fn(); @@ -457,7 +648,7 @@ describe("RunView", () => { it("distinguishes a completed run with recovered step errors from a failed run", async () => { renderRun({ ...research, summary: { ...research.summary, status: "ok", error_count: 2 } }); const header = await screen.findByRole("banner"); - expect(header).toHaveTextContent("Completed"); + expect(header).toHaveTextContent("Recorded"); expect(header).toHaveTextContent("Step errors 2"); expect(header).not.toHaveTextContent("Failed"); }); diff --git a/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunView.tsx b/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunView.tsx index d00481498f1..46377b44ed1 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunView.tsx +++ b/ui/litellm-dashboard/src/components/lens/traces/detail/run/RunView.tsx @@ -2,7 +2,7 @@ import { QueryErrorResetBoundary, useQueryClient, useSuspenseInfiniteQuery } from "@tanstack/react-query"; import { ArrowLeft } from "lucide-react"; -import { Suspense, useDeferredValue, useEffect, useMemo } from "react"; +import { Suspense, useDeferredValue, useEffect, useMemo, useState } from "react"; import { ErrorBoundary } from "react-error-boundary"; import { LoadingState } from "@/components/shared/LoadingState"; @@ -91,6 +91,8 @@ function LoadedRun({ }: RunViewProps & { switching: boolean }) { const traces = useTracesApi(accessToken); const queryClient = useQueryClient(); + const [live, setLive] = useState(traces.live); + const [manualRead, setManualRead] = useState(false); const queryKey = ["agentTrace", traceId, traceRef, accessToken]; const traceQueryOptions = { queryKey, @@ -98,29 +100,61 @@ function LoadedRun({ initialPageParam: null as string | null, getNextPageParam: (lastPage: Trace) => lastPage.next_cursor ?? undefined, staleTime: 30_000, - refetchOnWindowFocus: false, - refetchOnReconnect: false, - refetchOnMount: false, + refetchOnWindowFocus: live, + refetchOnReconnect: live, + refetchOnMount: true, + refetchInterval: live ? 30_000 : (false as const), retry: traceReadRetry, retryDelay: traceReadRetryDelay, }; const traceQuery = useSuspenseInfiniteQuery(traceQueryOptions); const refreshTrace = () => queryClient.resetQueries({ queryKey, exact: true }); - const failure = traceQuery.error ? classifyTraceReadFailure(traceQuery.error) : null; + const failure = traceQuery.isFetchNextPageError ? classifyTraceReadFailure(traceQuery.error) : null; + const readManually = (read: () => Promise) => { + setManualRead(true); + void read().finally(() => setManualRead(false)); + }; + const refreshRun = () => readManually(() => traceQuery.refetch()); + const toggleLive = () => { + if (live && !manualRead && !traceQuery.isFetchingNextPage) { + void queryClient.cancelQueries({ queryKey, exact: true }); + } + setLive((enabled) => !enabled); + }; const trace = useMemo(() => { const pages = traceQuery.data.pages; + const capture = pages[0].capture; return { ...pages[0], + agents: [ + ...new Map( + pages.flatMap((page) => + page.agents.map( + (agent) => [agent.actor_id ?? JSON.stringify([agent.name, agent.parent_agent]), agent] as const, + ), + ), + ).values(), + ], + capture: capture + ? { + ...capture, + actors: [ + ...new Map( + pages.flatMap((page) => page.capture?.actors.map((actor) => [actor.actor_id, actor] as const) ?? []), + ).values(), + ], + } + : capture, spans: pages.flatMap((page) => page.spans), next_cursor: pages[pages.length - 1].next_cursor, }; }, [traceQuery.data]); const seekingSpan = !switching && selectedSpanMissing(trace, selection.spanId); - const { hasNextPage, isFetching, isError, fetchNextPage } = traceQuery; + const { hasNextPage, isFetching, isFetchNextPageError, fetchNextPage } = traceQuery; const canSeek = seekingSpan && hasNextPage; useEffect(() => { - if (canSeek && !isFetching && !isError) void fetchNextPage(); - }, [canSeek, isFetching, isError, fetchNextPage]); + if (canSeek && !isFetching && !isFetchNextPageError) void fetchNextPage(); + }, [canSeek, isFetching, isFetchNextPageError, fetchNextPage]); return ( + {traceQuery.isRefetchError && ( +
+ Could not refresh this run. Previously received steps are still shown. + +
+ )} {(traceQuery.hasNextPage || failure) && ( void traceQuery.fetchNextPage()} - onRefresh={() => void refreshTrace()} + onLoadMore={() => readManually(() => traceQuery.fetchNextPage())} + onRefresh={() => readManually(refreshTrace)} /> )} diff --git a/ui/litellm-dashboard/src/components/lens/traces/utils.ts b/ui/litellm-dashboard/src/components/lens/traces/utils.ts index 63ae58486cc..21f90428703 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/utils.ts +++ b/ui/litellm-dashboard/src/components/lens/traces/utils.ts @@ -113,7 +113,8 @@ export function errorSource(span: Span): ErrorSource | null { type GroupOrSpan = Span | { group: Span[] }; -const groupKey = (span: Pick): string => `${span.agent}|${span.name}|${span.type}`; +const groupKey = (span: Pick): string => + `${span.actor_id ?? span.agent}|${span.name}|${span.type}`; function groupChildren(children: readonly Span[]): GroupOrSpan[] { const runs: Span[][] = []; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8fd8b4449be..11ec7caebdf 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25469,6 +25469,23 @@ export interface components { */ team_id: string; }; + /** ActorCapture */ + ActorCapture: { + /** Actor Id */ + actor_id: string; + /** Content Events */ + content_events: number; + /** Llm Calls */ + llm_calls: number; + /** Model Outputs */ + model_outputs: number; + /** Name */ + name: string; + /** Reply Events */ + reply_events: number; + /** Tool Calls */ + tool_calls: number; + }; /** AdaptiveRouterWeights */ AdaptiveRouterWeights: { /** @@ -25826,6 +25843,8 @@ export interface components { }; /** AgentNode */ AgentNode: { + /** Actor Id */ + actor_id?: string | null; /** Duration Ms */ duration_ms: number; /** Invocations */ @@ -25834,6 +25853,8 @@ export interface components { llm_calls: number; /** Name */ name: string; + /** Parent Actor Id */ + parent_actor_id?: string | null; /** Parent Agent */ parent_agent: string | null; /** Priced Calls */ @@ -45261,6 +45282,10 @@ export interface components { }; /** Span */ Span: { + /** Actor Id */ + actor_id?: string | null; + /** Actor Unassigned */ + actor_unassigned?: boolean; /** Agent */ agent: string; /** Duration Ms */ @@ -45283,6 +45308,8 @@ export interface components { name: string; /** Output Tokens */ output_tokens: number; + /** Parent Actor Id */ + parent_actor_id?: string | null; /** Parent Span Id */ parent_span_id: string | null; /** Span Id */ @@ -47189,12 +47216,24 @@ export interface components { Trace: { /** Agents */ agents: components["schemas"]["AgentNode"][]; + capture?: components["schemas"]["TraceCapture"] | null; /** Next Cursor */ next_cursor?: string | null; /** Spans */ spans: components["schemas"]["Span"][]; summary: components["schemas"]["TraceSummary"]; }; + /** TraceCapture */ + TraceCapture: { + /** Actors */ + actors: components["schemas"]["ActorCapture"][]; + /** Content Events */ + content_events: number; + /** Unassigned Events */ + unassigned_events: number; + /** Warning Events */ + warning_events: number; + }; /** TraceFindingCount */ TraceFindingCount: { /** Finding Count */