diff --git a/litellm-rust/crates/traces-cache/src/reader.rs b/litellm-rust/crates/traces-cache/src/reader.rs index 96a05969d9e..2cdef7bed8f 100644 --- a/litellm-rust/crates/traces-cache/src/reader.rs +++ b/litellm-rust/crates/traces-cache/src/reader.rs @@ -10,7 +10,7 @@ use crate::{ spend::spend, }; use litellm_traces::{ - SpanDetail, SpanErrorPage, Trace, TracePage, + ObservationType, SpanDetail, SpanErrorPage, Trace, TracePage, query::named::{ ListTracesParams, ReadAccessParams, SpanDetailParams, SpanDetailsParams, SpanErrorParams, TraceIdentityParams, TraceSpansParams, @@ -350,6 +350,19 @@ fn page( .iter() .filter_map(|span| span.actor_id.as_deref()) .collect(); + let legacy_agents: HashSet<&str> = page_spans + .iter() + .filter(|span| span.actor_id.is_none() && !span.actor_unassigned) + .filter_map(|span| { + if !span.agent.is_empty() { + Some(span.agent.as_str()) + } else if span.kind == ObservationType::Agent { + Some(span.name.as_str()) + } else { + None + } + }) + .collect(); Trace { capture: snapshot.trace().capture.as_ref().map(|capture| { litellm_traces::TraceCapture { @@ -369,11 +382,9 @@ fn page( .trace() .agents .iter() - .filter(|agent| { - agent - .actor_id - .as_deref() - .is_none_or(|id| actors.contains(id)) + .filter(|agent| match agent.actor_id.as_deref() { + Some(id) => actors.contains(id), + None => legacy_agents.contains(agent.name.as_str()), }) .cloned() .collect(), diff --git a/litellm-rust/crates/traces-cache/tests/read.rs b/litellm-rust/crates/traces-cache/tests/read.rs index 50fd08bef38..01587466246 100644 --- a/litellm-rust/crates/traces-cache/tests/read.rs +++ b/litellm-rust/crates/traces-cache/tests/read.rs @@ -571,6 +571,54 @@ async fn actor_metadata_pages_with_spans_without_losing_trace_totals( assert_eq!(store.calls(Operation::TraceSpans), 1); } +#[rstest] +#[case::explicit_names(4096, true)] +#[case::span_name_fallback(4096, false)] +#[case::larger_pages(8192, true)] +#[tokio::test] +async fn legacy_agent_metadata_fits_each_page_without_losing_agents( + #[case] response_bytes: usize, + #[case] explicit_names: bool, +) { + let rows = (0..120) + .map(|index| TraceSpansRow { + kind: ObservationType::Agent, + name: format!("legacy-{index}"), + agent: if explicit_names { + format!("legacy-{index}") + } else { + String::new() + }, + parent_span_id: String::new(), + ..span(index) + }) + .collect(); + let store = FakeStore::with_spans("ref", rows); + let reader = TraceReader::new(response_bytes); + let mut cursor = None; + let mut names = 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, 120); + assert_eq!(page.agents.len(), page.spans.len()); + for agent in page.agents { + assert!(page.spans.iter().any(|span| span.name == agent.name)); + assert!(names.insert(agent.name)); + } + cursor = page.next_cursor; + if cursor.is_none() { + break; + } + } + assert_eq!(names.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() {