fix(lens): include bounded legacy agent metadata

This commit is contained in:
moe-berri 2026-10-06 16:19:53 -07:00
commit 779da12cae
2 changed files with 65 additions and 6 deletions

View file

@ -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<E>(
.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<E>(
.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(),

View file

@ -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() {