feat(lens): own trace reads behind a cached TraceStore port

Move storage-independent trace reads into litellm-traces-cache behind a
TraceStore port that ClickHouse implements. One keyset pager drives the
span, list span and spend reads, and a run list batch reads spend once.

Trace opens, pages and list summaries share one resolved read per trace
in an in-process cache with single-flight loading. Live traces and reads
with unknown spend expire after 5s, quiet traces after 10 minutes, failed
reads are never cached, and the accepted list page size is remembered per
scope.

Trace read failures map to their own status and code (400, 409, 413, 503
with Retry-After), and the trace drawer retries temporary failures while
offering only a refresh for changed or oversized traces.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

litellm-rust/Cargo.lock

@@ -4156,6 +4156,7 @@ dependencies = [
  "litellm-storage-clickhouse",
  "litellm-token-counter",
  "litellm-traces",
+ "litellm-traces-cache",
  "litellm-traces-clickhouse",
  "litellm-tracing",
  "prost",
@@ -4474,13 +4475,17 @@ dependencies = [
 name = "litellm-traces-cache"
 version = "0.1.0"
 dependencies = [
+ "base64 0.22.1",
  "litellm-traces",
  "moka",
  "rstest",
+ "serde",
  "serde_json",
  "sha2 0.10.9",
  "thiserror 2.0.19",
+ "time",
  "tokio",
+ "tracing",
 ]
 
 [[package]]
@@ -4488,11 +4493,9 @@ name = "litellm-traces-clickhouse"
 version = "0.1.0"
 dependencies = [
  "askama",
- "base64 0.22.1",
  "flate2",
  "futures-util",
  "hmac 0.12.1",
- "itertools 0.14.0",
  "jsonschema",
  "litellm-http",
  "litellm-migrate",
@@ -4511,7 +4514,6 @@ dependencies = [
  "thiserror 2.0.19",
  "time",
  "tokio",
- "tracing",
  "url",
  "wiremock",
 ]

litellm-rust/crates/python-bridge/Cargo.toml

@@ -22,6 +22,7 @@ tiktoken = ["litellm-token-counter/tiktoken"]
 fancy-regex.workspace = true
 litellm-tracing.workspace = true
 litellm-traces.workspace = true
+litellm-traces-cache.workspace = true
 litellm-traces-clickhouse.workspace = true
 litellm-storage-clickhouse.workspace = true
 litellm-host.workspace = true

litellm-rust/crates/python-bridge/src/routes/traces.rs

@@ -1,8 +1,11 @@
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, sync::Arc};
 
 use litellm_http::ClientVariant;
 use litellm_traces::{QueryScope, ReadQuery, Tenant, query::named::ReadAccessParams};
-use litellm_traces_clickhouse::{Config, Error, InsertTable, Parameter, QueryReaders};
+use litellm_traces_cache::{ReadError, TraceReader};
+use litellm_traces_clickhouse::{
+    ClickHouseTraces, Config, Error, InsertTable, Parameter, QueryReaders,
+};
 use prost::Message;
 use pyo3::{
     exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
@@ -10,6 +13,8 @@ use pyo3::{
     types::PyBytes,
 };
 
+pyo3::import_exception!(litellm.rust_bridge.trace.errors, TraceChanged);
+
 #[derive(Message)]
 struct OtlpErrorStatus {
     #[prost(int32, tag = "1")]
@@ -35,15 +40,12 @@ fn map_error_ref(error: &Error) -> PyErr {
     use litellm_storage_clickhouse::Error as StorageError;
 
     match error {
-        Error::Decode(litellm_traces::Error::TooLarge)
-        | Error::InsertTooLarge
-        | Error::ReadTooLarge => PyOverflowError::new_err(error.to_string()),
+        Error::Decode(litellm_traces::Error::TooLarge) | Error::InsertTooLarge => {
+            PyOverflowError::new_err(error.to_string())
+        }
         Error::InvalidRow
         | Error::InvalidLimit(_)
         | Error::InvalidTable
-        | Error::InvalidCursor(_)
-        | Error::AmbiguousTrace
-        | Error::TraceChanged
         | Error::Decode(_)
         | Error::InvalidSchema
         | Error::InvalidQuery
@@ -78,6 +80,18 @@ fn map_error_ref(error: &Error) -> PyErr {
     }
 }
 
+fn map_read_error(error: ReadError<Error>) -> PyErr {
+    match error {
+        error @ (ReadError::InvalidParameters
+        | ReadError::InvalidCursor(_)
+        | ReadError::AmbiguousTrace) => PyValueError::new_err(error.to_string()),
+        error @ ReadError::TraceChanged => TraceChanged::new_err(error.to_string()),
+        error @ ReadError::TooLarge => PyOverflowError::new_err(error.to_string()),
+        error @ ReadError::Encode(_) => PyRuntimeError::new_err(error.to_string()),
+        ReadError::Store(error) => map_error_ref(&error),
+    }
+}
+
 fn map_sql_error(error: Error) -> PyErr {
     match error {
         Error::Storage(litellm_storage_clickhouse::Error::QueryFailed(400 | 404)) => {
@@ -112,6 +126,7 @@ impl NativeTraceConfig {
 pub struct NativeTraceStorage {
     config: Config,
     query_readers: QueryReaders,
+    reader: Arc<TraceReader>,
 }
 
 #[pymethods]
@@ -123,6 +138,9 @@ impl NativeTraceStorage {
                 config.inner.storage().writer().clone(),
                 config.inner.storage().database().to_owned(),
             ),
+            reader: Arc::new(TraceReader::new(
+                litellm_storage_clickhouse::READ_LIMITS.response_bytes,
+            )),
             config: config.inner.clone(),
         })
     }
@@ -218,21 +236,16 @@ impl NativeTraceStorage {
     ) -> PyResult<Bound<'py, PyAny>> {
         let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
         let connection = self.config.storage().reader().clone();
+        let reader = Arc::clone(&self.reader);
         crate::execution::run_async(
             py,
             async move {
-                litellm_traces_clickhouse::list_traces(
-                    &client,
-                    &connection,
-                    &scope,
-                    start_ms,
-                    end_ms,
-                    cursor.as_deref(),
-                    limit,
-                )
-                .await
+                let store = ClickHouseTraces::new(client, connection);
+                reader
+                    .list_traces(&store, &scope, start_ms, end_ms, cursor.as_deref(), limit)
+                    .await
             },
-            map_error,
+            map_read_error,
         )
     }
 
@@ -248,34 +261,31 @@ impl NativeTraceStorage {
     ) -> PyResult<Bound<'py, PyAny>> {
         let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
         let connection = self.config.storage().reader().clone();
+        let reader = Arc::clone(&self.reader);
         crate::execution::run_async(
             py,
             async move {
+                let store = ClickHouseTraces::new(client, connection);
                 if let Some(page_size) = page_size {
-                    litellm_traces_clickhouse::get_trace_page(
-                        &client,
-                        &connection,
-                        &scope,
-                        &trace_id,
-                        &trace_ref,
-                        cursor.as_deref(),
-                        page_size,
-                    )
-                    .await
+                    reader
+                        .get_trace_page(
+                            &store,
+                            &scope,
+                            &trace_id,
+                            &trace_ref,
+                            cursor.as_deref(),
+                            page_size,
+                        )
+                        .await
                 } else if cursor.is_some() {
-                    Err(Error::InvalidParameters)
+                    Err(ReadError::InvalidParameters)
                 } else {
-                    litellm_traces_clickhouse::get_trace(
-                        &client,
-                        &connection,
-                        &scope,
-                        &trace_id,
-                        &trace_ref,
-                    )
-                    .await
+                    reader
+                        .get_trace(&store, &scope, &trace_id, &trace_ref)
+                        .await
                 }
             },
-            map_error,
+            map_read_error,
         )
     }
 
@@ -289,20 +299,16 @@ impl NativeTraceStorage {
     ) -> PyResult<Bound<'py, PyAny>> {
         let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
         let connection = self.config.storage().reader().clone();
+        let reader = Arc::clone(&self.reader);
         crate::execution::run_async(
             py,
             async move {
-                litellm_traces_clickhouse::get_span(
-                    &client,
-                    &connection,
-                    &scope,
-                    &trace_id,
-                    &span_id,
-                    &trace_ref,
-                )
-                .await
+                let store = ClickHouseTraces::new(client, connection);
+                reader
+                    .get_span(&store, &scope, &trace_id, &span_id, &trace_ref)
+                    .await
             },
-            map_error,
+            map_read_error,
         )
     }
 
@@ -318,21 +324,23 @@ impl NativeTraceStorage {
     ) -> PyResult<Bound<'py, PyAny>> {
         let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
         let connection = self.config.storage().reader().clone();
+        let reader = Arc::clone(&self.reader);
         crate::execution::run_async(
             py,
             async move {
-                litellm_traces_clickhouse::get_span_error(
-                    &client,
-                    &connection,
-                    &scope,
-                    &trace_id,
-                    &span_id,
-                    &trace_ref,
-                    cursor.as_deref(),
-                )
-                .await
+                let store = ClickHouseTraces::new(client, connection);
+                reader
+                    .get_span_error(
+                        &store,
+                        &scope,
+                        &trace_id,
+                        &span_id,
+                        &trace_ref,
+                        cursor.as_deref(),
+                    )
+                    .await
             },
-            map_error,
+            map_read_error,
         )
     }
 
@@ -495,11 +503,7 @@ mod tests {
         Error::Decode(litellm_traces::Error::InvalidLimit("OTLP_MAX_SPANS")),
         "ValueError"
     )]
-    #[case::cursor(Error::InvalidCursor("trace"), "ValueError")]
-    #[case::ambiguous(Error::AmbiguousTrace, "ValueError")]
-    #[case::changed_snapshot(Error::TraceChanged, "ValueError")]
-    #[case::read_budget(Error::ReadTooLarge, "OverflowError")]
-    fn trace_read_and_ingest_failures_preserve_public_exception_types(
+    fn trace_ingest_failures_preserve_public_exception_types(
         #[case] error: Error,
         #[case] exception_name: &str,
     ) {
@@ -511,4 +515,43 @@ mod tests {
             );
         });
     }
+
+    #[rstest]
+    #[case::invalid_parameters(ReadError::InvalidParameters, "ValueError")]
+    #[case::invalid_cursor(ReadError::InvalidCursor("trace"), "ValueError")]
+    #[case::ambiguous(ReadError::AmbiguousTrace, "ValueError")]
+    #[case::changed_snapshot(ReadError::TraceChanged, "TraceChanged")]
+    #[case::read_budget(ReadError::TooLarge, "OverflowError")]
+    #[case::encode(
+        ReadError::Encode(Arc::new(serde_json::Error::io(std::io::Error::other("invalid")))),
+        "RuntimeError"
+    )]
+    #[case::store(ReadError::Store(Arc::new(Error::InvalidScope)), "ValueError")]
+    fn trace_read_failures_preserve_public_exception_types(
+        #[case] error: ReadError<Error>,
+        #[case] exception_name: &str,
+    ) {
+        Python::initialize();
+        Python::attach(|py| {
+            let repository = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
+                .ancestors()
+                .nth(3)
+                .unwrap()
+                .to_str()
+                .unwrap();
+            pyo3::types::PyModule::import(py, "sys")
+                .unwrap()
+                .getattr("path")
+                .unwrap()
+                .call_method1("insert", (0, repository))
+                .unwrap();
+            let message = error.to_string();
+            let exception = map_read_error(error);
+            assert_eq!(exception.get_type(py).name().unwrap(), exception_name);
+            assert_eq!(
+                exception.value(py).str().unwrap().to_str().unwrap(),
+                message
+            );
+        });
+    }
 }

litellm-rust/crates/traces-cache/AGENTS.md

@@ -1,5 +1,5 @@
-Own resolved-trace snapshot storage, cache identity, weighting, and expiry
-Depend on trace domain types, never storage, HTTP, or Python
+Own storage-independent trace reads over `TraceStore`: the in-process read cache (identity, single-flight, freshness expiry, weighting), cursor formats, paging, response splitting, spend windows and run batching
+Depend on trace domain types, never storage, HTTP or Python
 Preserve the full source and authorization scope in every cache key
 Keep snapshots immutable and expose borrowed data
-Keep cursor formats and database reads in their existing owners
+Storage adapters implement `TraceStore`; keep SQL and row encoding there

litellm-rust/crates/traces-cache/Cargo.toml

@@ -6,11 +6,15 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
+base64.workspace = true
 litellm-traces.workspace = true
 moka.workspace = true
+serde.workspace = true
 serde_json.workspace = true
 sha2.workspace = true
+time.workspace = true
 thiserror.workspace = true
+tracing.workspace = true
 
 [dev-dependencies]
 rstest.workspace = true

litellm-rust/crates/traces-cache/src/cache.rs

@@ -0,0 +1,395 @@
+use std::{future::Future, sync::Arc, time::Duration};
+
+use litellm_traces::{
+    Trace, TraceSummary,
+    query::named::{ReadAccessParams, TraceSpansRow},
+};
+use moka::{Expiry, future::Cache};
+use serde::Serialize;
+use sha2::{Digest, Sha256};
+
+use crate::Error;
+
+pub const LIVE_TTL: Duration = Duration::from_secs(5);
+pub const SETTLED_TTL: Duration = Duration::from_secs(10 * 60);
+const SETTLED_AFTER_MS: u64 = 5 * 60 * 1000;
+const MAX_INDEX_ENTRIES: u64 = 100_000;
+
+#[derive(Clone, Eq, Hash, PartialEq)]
+pub struct SnapshotKey(String);
+
+impl SnapshotKey {
+    fn digest(fields: &impl Serialize) -> Result<Self, Error> {
+        Ok(Self(format!(
+            "{:x}",
+            Sha256::digest(serde_json::to_vec(fields)?)
+        )))
+    }
+
+    pub fn new(
+        source: &str,
+        access: &ReadAccessParams,
+        trace_id: &str,
+        trace_ref: &str,
+        snapshot_ms: u64,
+    ) -> Result<Self, Error> {
+        Self::digest(&(source, access, trace_id, trace_ref, snapshot_ms))
+    }
+
+    pub fn latest(
+        source: &str,
+        access: &ReadAccessParams,
+        trace_id: &str,
+        trace_ref: &str,
+    ) -> Result<Self, Error> {
+        Self::digest(&(source, access, trace_id, trace_ref))
+    }
+
+    pub(crate) fn run(
+        source: &str,
+        access: &ReadAccessParams,
+        run: (&str, &str, &str, &str),
+    ) -> Result<Self, Error> {
+        Self::digest(&("run", source, access, run))
+    }
+
+    pub(crate) fn scope(source: &str, access: &ReadAccessParams) -> Result<Self, Error> {
+        Self::digest(&("scope", source, access))
+    }
+}
+
+/// 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`.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum Freshness {
+    Live,
+    Settled,
+}
+
+impl Freshness {
+    pub fn of(rows: &[TraceSpansRow], spend_known: bool, snapshot_ms: u64) -> Self {
+        let last_end_ms = rows
+            .iter()
+            .map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns) / 1_000_000)
+            .max()
+            .unwrap_or(i64::MAX);
+        let quiet_ms = i64::try_from(snapshot_ms)
+            .unwrap_or(i64::MAX)
+            .saturating_sub(last_end_ms);
+        if spend_known && quiet_ms >= SETTLED_AFTER_MS as i64 {
+            Self::Settled
+        } else {
+            Self::Live
+        }
+    }
+
+    fn ttl(self) -> Duration {
+        match self {
+            Self::Live => LIVE_TTL,
+            Self::Settled => SETTLED_TTL,
+        }
+    }
+}
+
+trait Fresh {
+    fn freshness(&self) -> Freshness;
+}
+
+struct ByFreshness;
+
+impl<K, V: Fresh> Expiry<K, V> for ByFreshness {
+    fn expire_after_create(&self, _: &K, value: &V, _: std::time::Instant) -> Option<Duration> {
+        Some(value.freshness().ttl())
+    }
+}
+
+pub struct Snapshot {
+    trace: Trace,
+    version: String,
+    snapshot_ms: u64,
+    freshness: Freshness,
+    weight: u32,
+}
+
+impl Snapshot {
+    pub fn trace(&self) -> &Trace {
+        &self.trace
+    }
+
+    pub fn version(&self) -> &str {
+        &self.version
+    }
+
+    pub fn snapshot_ms(&self) -> u64 {
+        self.snapshot_ms
+    }
+
+    pub fn freshness(&self) -> Freshness {
+        self.freshness
+    }
+}
+
+#[derive(Clone, Copy)]
+struct Latest {
+    snapshot_ms: u64,
+    freshness: Freshness,
+}
+
+impl Fresh for Latest {
+    fn freshness(&self) -> Freshness {
+        self.freshness
+    }
+}
+
+/// Resolved trace snapshots pinned by `snapshot_ms` for paging, plus which snapshot each trace
+/// currently serves so repeated opens reuse one read until its freshness expires.
+pub struct SnapshotCache {
+    pinned: Cache<SnapshotKey, Arc<Snapshot>>,
+    latest: Cache<SnapshotKey, Latest>,
+    max_graph_bytes: usize,
+}
+
+impl SnapshotCache {
+    pub fn new(max_graph_bytes: usize, idle: Duration) -> Self {
+        Self {
+            pinned: Cache::builder()
+                .max_capacity((max_graph_bytes as u64).saturating_mul(2))
+                .weigher(|_: &SnapshotKey, snapshot: &Arc<Snapshot>| snapshot.weight)
+                .time_to_idle(idle)
+                .build(),
+            latest: Cache::builder()
+                .max_capacity(MAX_INDEX_ENTRIES)
+                .expire_after(ByFreshness)
+                .build(),
+            max_graph_bytes,
+        }
+    }
+
+    pub async fn get(&self, key: &SnapshotKey) -> Option<Arc<Snapshot>> {
+        self.pinned.get(key).await
+    }
+
+    /// Returns the snapshot pinned at `key`, running `load` once for all concurrent callers on a
+    /// miss. A failed load is not cached.
+    pub async fn pinned_or_load<E, F>(
+        &self,
+        key: SnapshotKey,
+        snapshot_ms: u64,
+        load: F,
+    ) -> Result<Arc<Snapshot>, Arc<E>>
+    where
+        E: From<Error> + Send + Sync + 'static,
+        F: Future<Output = Result<(Trace, Freshness), E>>,
+    {
+        self.pinned
+            .try_get_with(key, async {
+                let (trace, freshness) = load.await?;
+                Ok(Arc::new(self.snapshot(trace, snapshot_ms, freshness)?))
+            })
+            .await
+    }
+
+    /// Returns the snapshot `latest` currently serves. On a miss, `load_at(now_ms)` runs once for
+    /// all concurrent callers and its snapshot is served until its freshness expires.
+    pub async fn latest_or_load<E, F, Fut>(
+        &self,
+        latest: SnapshotKey,
+        now_ms: u64,
+        load_at: F,
+    ) -> Result<Arc<Snapshot>, Arc<E>>
+    where
+        E: Send + Sync + 'static,
+        F: Fn(u64) -> Fut,
+        Fut: Future<Output = Result<Arc<Snapshot>, Arc<E>>>,
+    {
+        let entry = self
+            .latest
+            .try_get_with(latest, async {
+                let snapshot = load_at(now_ms).await?;
+                Ok::<_, Arc<E>>(Latest {
+                    snapshot_ms: snapshot.snapshot_ms,
+                    freshness: snapshot.freshness,
+                })
+            })
+            .await
+            .map_err(|error| Arc::clone(&*error))?;
+        load_at(entry.snapshot_ms).await
+    }
+
+    fn snapshot(
+        &self,
+        trace: Trace,
+        snapshot_ms: u64,
+        freshness: Freshness,
+    ) -> Result<Snapshot, Error> {
+        let encoded = serde_json::to_vec(&trace)?;
+        if encoded.len() > self.max_graph_bytes {
+            return Err(Error::ReadTooLarge);
+        }
+        let span_ids: Vec<&str> = trace
+            .spans
+            .iter()
+            .map(|span| span.span_id.as_str())
+            .collect();
+        let version = format!("{:x}", Sha256::digest(serde_json::to_vec(&span_ids)?));
+        Ok(Snapshot {
+            trace,
+            version,
+            snapshot_ms,
+            freshness,
+            weight: u32::try_from(encoded.len().saturating_mul(2)).unwrap_or(u32::MAX),
+        })
+    }
+
+    #[cfg(test)]
+    async fn weighted_size(&self) -> u64 {
+        self.pinned.run_pending_tasks().await;
+        self.pinned.weighted_size()
+    }
+}
+
+#[derive(Clone)]
+pub(crate) enum ListedRun {
+    Resolved(Box<TraceSummary>, Freshness),
+    Limited,
+}
+
+impl Fresh for ListedRun {
+    fn freshness(&self) -> Freshness {
+        match self {
+            Self::Resolved(_, freshness) => *freshness,
+            Self::Limited => Freshness::Settled,
+        }
+    }
+}
+
+pub(crate) struct ListCache {
+    pub(crate) runs: Cache<SnapshotKey, ListedRun>,
+    pub(crate) limits: Cache<SnapshotKey, u32>,
+}
+
+impl ListCache {
+    pub(crate) fn new() -> Self {
+        Self {
+            runs: Cache::builder()
+                .max_capacity(MAX_INDEX_ENTRIES)
+                .expire_after(ByFreshness)
+                .build(),
+            limits: Cache::builder()
+                .max_capacity(MAX_INDEX_ENTRIES)
+                .time_to_live(SETTLED_TTL)
+                .build(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use litellm_traces::{
+        SpanStatus,
+        query::named::{SpendByResponseIdsRow, TraceSpansRow},
+        resolve_trace,
+    };
+    use rstest::rstest;
+
+    use super::*;
+
+    fn row(span_id: &str) -> TraceSpansRow {
+        TraceSpansRow {
+            trace_id: String::new(),
+            span_id: span_id.into(),
+            parent_span_id: String::new(),
+            name: "run".into(),
+            kind: litellm_traces::ObservationType::Agent,
+            wrapper_candidate: false,
+            agent: "agent".into(),
+            framework: String::new(),
+            status: SpanStatus::Ok,
+            status_message: String::new(),
+            error_truncated: false,
+            start_ns: 1_790_742_989_000_000_000,
+            duration_ns: 10_000_000,
+            service: "agent-demo".into(),
+            input_preview: format!("input of {span_id}"),
+            model: String::new(),
+            input_tokens: 0,
+            output_tokens: 0,
+            litellm_request_id: String::new(),
+            call_keys: Vec::new(),
+            call_evidence: None,
+            tool_call_id: String::new(),
+            team_id: String::new(),
+            api_key_hash: String::new(),
+            user_id: String::new(),
+        }
+    }
+
+    fn trace(span_id: &str) -> Trace {
+        resolve_trace(
+            "trace",
+            "ref",
+            &[row(span_id)],
+            &[] as &[SpendByResponseIdsRow],
+        )
+        .expect("fixture should resolve")
+    }
+
+    fn key(suffix: &str) -> SnapshotKey {
+        SnapshotKey::new(
+            "source",
+            &ReadAccessParams {
+                all_teams: false,
+                user_id: String::new(),
+                team_ids: vec!["team".into()],
+            },
+            suffix,
+            "ref",
+            100,
+        )
+        .unwrap()
+    }
+
+    #[tokio::test]
+    async fn weighted_capacity_bounds_retained_snapshots() {
+        let limit = ["first", "second", "third"]
+            .iter()
+            .map(|span_id| serde_json::to_vec(&trace(span_id)).unwrap().len())
+            .max()
+            .unwrap();
+        let cache = SnapshotCache::new(limit, Duration::from_secs(120));
+
+        for (key, span_id) in [
+            (key("a"), "first"),
+            (key("b"), "second"),
+            (key("c"), "third"),
+        ] {
+            cache
+                .pinned_or_load(key, 100, async {
+                    Ok::<_, Error>((trace(span_id), Freshness::Settled))
+                })
+                .await
+                .unwrap();
+        }
+
+        assert!(cache.weighted_size().await <= (limit as u64) * 2);
+    }
+
+    const LAST_END_MS: u64 = 1_790_742_989_010;
+
+    #[rstest]
+    #[case::just_ended(LAST_END_MS, true, Freshness::Live)]
+    #[case::quiet_just_under(LAST_END_MS + SETTLED_AFTER_MS - 1, true, Freshness::Live)]
+    #[case::quiet_long_enough(LAST_END_MS + SETTLED_AFTER_MS, true, Freshness::Settled)]
+    #[case::spend_unknown(LAST_END_MS + SETTLED_AFTER_MS, false, Freshness::Live)]
+    fn freshness_settles_once_spans_stop_and_spend_is_known(
+        #[case] snapshot_ms: u64,
+        #[case] spend_known: bool,
+        #[case] expected: Freshness,
+    ) {
+        assert_eq!(
+            Freshness::of(&[row("root")], spend_known, snapshot_ms),
+            expected
+        );
+    }
+}

litellm-rust/crates/traces-cache/src/cursor.rs

@@ -0,0 +1,124 @@
+use base64::{Engine, engine::general_purpose::URL_SAFE};
+use serde::{Deserialize, Serialize};
+
+use crate::ReadError;
+
+pub(super) fn encode_cursor<T: Serialize>(position: &T) -> String {
+    URL_SAFE.encode(serde_json::to_vec(position).unwrap_or_default())
+}
+
+pub(super) fn decode_cursor<T: for<'de> Deserialize<'de>, E>(
+    cursor: &str,
+    kind: &'static str,
+) -> Result<T, ReadError<E>> {
+    URL_SAFE
+        .decode(cursor)
+        .ok()
+        .and_then(|json| serde_json::from_slice(&json).ok())
+        .ok_or(ReadError::InvalidCursor(kind))
+}
+
+pub(super) fn trace_position<E>(cursor: Option<&str>) -> Result<(i64, String), ReadError<E>> {
+    let Some(cursor) = cursor.filter(|cursor| !cursor.is_empty()) else {
+        return Ok((0, String::new()));
+    };
+    match decode_cursor::<(i64, String), E>(cursor, "trace")? {
+        (start_ms, trace_ref) if start_ms > 0 && !trace_ref.is_empty() => Ok((start_ms, trace_ref)),
+        _ => Err(ReadError::InvalidCursor("trace")),
+    }
+}
+
+#[derive(Deserialize, Serialize)]
+pub(super) struct ErrorPosition {
+    pub(super) offset: u64,
+    pub(super) version: String,
+}
+
+pub(super) fn error_position<E>(
+    cursor: Option<&str>,
+) -> Result<Option<ErrorPosition>, ReadError<E>> {
+    let Some(cursor) = cursor else {
+        return Ok(None);
+    };
+    let position: ErrorPosition = decode_cursor(cursor, "diagnostic")?;
+    let valid_version = position.version.len() == 64
+        && position
+            .version
+            .bytes()
+            .all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte));
+    if i64::try_from(position.offset).is_err() || !valid_version {
+        return Err(ReadError::InvalidCursor("diagnostic"));
+    }
+    Ok(Some(position))
+}
+
+#[derive(Deserialize, Serialize)]
+pub(super) struct SpanPosition {
+    pub(super) trace_ref: String,
+    pub(super) snapshot_ms: u64,
+    pub(super) offset: usize,
+    pub(super) version: String,
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    #[rstest]
+    fn trace_cursor_round_trips_the_last_listed_run() {
+        let cursor = encode_cursor(&(1_790_742_989_377_i64, "4bad42b84e9de3ba46fc870185f8f023"));
+        assert_eq!(
+            trace_position::<std::io::Error>(Some(&cursor)).unwrap(),
+            (
+                1_790_742_989_377,
+                "4bad42b84e9de3ba46fc870185f8f023".to_owned()
+            )
+        );
+        assert_eq!(
+            trace_position::<std::io::Error>(None).unwrap(),
+            (0, String::new())
+        );
+        assert_eq!(
+            trace_position::<std::io::Error>(Some("")).unwrap(),
+            (0, String::new())
+        );
+    }
+
+    #[rstest]
+    #[case::not_base64("abc")]
+    #[case::not_json("bm90LWpzb24=")]
+    #[case::numeric_reference("WzEsIDJd")]
+    #[case::zero_start("WzAsICJ0Il0=")]
+    fn malformed_trace_cursors_are_rejected(#[case] cursor: &str) {
+        let result: Result<(i64, String), ReadError<std::io::Error>> = trace_position(Some(cursor));
+        assert!(matches!(result, Err(ReadError::InvalidCursor("trace"))));
+    }
+
+    #[rstest]
+    #[case::not_base64("garbage")]
+    #[case::missing_fields("e30=")]
+    #[case::not_an_object("WzEsMl0=")]
+    fn malformed_diagnostic_cursors_are_rejected(#[case] cursor: &str) {
+        let result: Result<Option<ErrorPosition>, ReadError<std::io::Error>> =
+            error_position(Some(cursor));
+        assert!(matches!(
+            result,
+            Err(ReadError::InvalidCursor("diagnostic"))
+        ));
+    }
+
+    #[rstest]
+    #[case::lowercase_version("a".repeat(64))]
+    #[case::short_version("A".repeat(63))]
+    fn diagnostic_cursor_requires_a_content_version(#[case] version: String) {
+        let cursor = encode_cursor(&ErrorPosition { offset: 1, version });
+        let result: Result<Option<ErrorPosition>, ReadError<std::io::Error>> =
+            error_position(Some(&cursor));
+        assert!(matches!(
+            result,
+            Err(ReadError::InvalidCursor("diagnostic"))
+        ));
+    }
+}

litellm-rust/crates/traces-cache/src/error.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
     #[error("trace snapshot serialization failed")]
@@ -5,3 +7,45 @@ pub enum Error {
     #[error("trace snapshot exceeds the size limit")]
     ReadTooLarge,
 }
+
+/// Cheap to clone so one failed single-flight read can be returned to every waiting caller.
+#[derive(Debug, thiserror::Error)]
+pub enum ReadError<E> {
+    #[error("invalid trace read parameters")]
+    InvalidParameters,
+    #[error("Invalid {0} cursor")]
+    InvalidCursor(&'static str),
+    #[error("Multiple traces have this ID; provide trace_ref")]
+    AmbiguousTrace,
+    #[error("Trace changed while paging; refresh the trace to continue")]
+    TraceChanged,
+    #[error("Trace exceeds the interactive read budget; use a filtered trace query")]
+    TooLarge,
+    #[error("trace could not be encoded")]
+    Encode(#[source] Arc<serde_json::Error>),
+    #[error(transparent)]
+    Store(Arc<E>),
+}
+
+impl<E> Clone for ReadError<E> {
+    fn clone(&self) -> Self {
+        match self {
+            Self::InvalidParameters => Self::InvalidParameters,
+            Self::InvalidCursor(kind) => Self::InvalidCursor(kind),
+            Self::AmbiguousTrace => Self::AmbiguousTrace,
+            Self::TraceChanged => Self::TraceChanged,
+            Self::TooLarge => Self::TooLarge,
+            Self::Encode(error) => Self::Encode(Arc::clone(error)),
+            Self::Store(error) => Self::Store(Arc::clone(error)),
+        }
+    }
+}
+
+impl<E> From<Error> for ReadError<E> {
+    fn from(error: Error) -> Self {
+        match error {
+            Error::ReadTooLarge => Self::TooLarge,
+            Error::Serialization(error) => Self::Encode(Arc::new(error)),
+        }
+    }
+}

litellm-rust/crates/traces-cache/src/lib.rs

@@ -1,166 +1,12 @@
-use std::{sync::Arc, time::Duration};
-
-use litellm_traces::{Trace, query::named::ReadAccessParams};
-use moka::future::Cache;
-use sha2::{Digest, Sha256};
-
+mod cache;
+mod cursor;
 mod error;
+mod list;
+mod reader;
+mod spend;
+mod store;
 
-pub use error::Error;
-
-#[derive(Clone, Eq, Hash, PartialEq)]
-pub struct SnapshotKey(String);
-
-impl SnapshotKey {
-    pub fn new(
-        source: &str,
-        access: &ReadAccessParams,
-        trace_id: &str,
-        trace_ref: &str,
-        snapshot_ms: u64,
-    ) -> Result<Self, Error> {
-        let encoded = serde_json::to_vec(&(source, access, trace_id, trace_ref, snapshot_ms))?;
-        Ok(Self(format!("{:x}", Sha256::digest(encoded))))
-    }
-}
-
-pub struct Snapshot {
-    trace: Trace,
-    version: String,
-    weight: u32,
-}
-
-impl Snapshot {
-    pub fn trace(&self) -> &Trace {
-        &self.trace
-    }
-
-    pub fn version(&self) -> &str {
-        &self.version
-    }
-}
-
-pub struct SnapshotCache {
-    entries: Cache<SnapshotKey, Arc<Snapshot>>,
-    max_graph_bytes: usize,
-}
-
-impl SnapshotCache {
-    pub fn new(max_graph_bytes: usize, ttl: Duration) -> Self {
-        Self {
-            entries: Cache::builder()
-                .max_capacity((max_graph_bytes as u64).saturating_mul(2))
-                .weigher(|_: &SnapshotKey, snapshot: &Arc<Snapshot>| snapshot.weight)
-                .time_to_live(ttl)
-                .build(),
-            max_graph_bytes,
-        }
-    }
-
-    pub async fn get(&self, key: &SnapshotKey) -> Option<Arc<Snapshot>> {
-        self.entries.get(key).await
-    }
-
-    pub async fn insert(&self, key: SnapshotKey, trace: Trace) -> Result<Arc<Snapshot>, Error> {
-        let encoded = serde_json::to_vec(&trace)?;
-        if encoded.len() > self.max_graph_bytes {
-            return Err(Error::ReadTooLarge);
-        }
-
-        let span_ids: Vec<&str> = trace
-            .spans
-            .iter()
-            .map(|span| span.span_id.as_str())
-            .collect();
-
-        let version = format!("{:x}", Sha256::digest(serde_json::to_vec(&span_ids)?));
-
-        let snapshot = Arc::new(Snapshot {
-            trace,
-            version,
-            weight: u32::try_from(encoded.len().saturating_mul(2)).unwrap_or(u32::MAX),
-        });
-
-        self.entries.insert(key, Arc::clone(&snapshot)).await;
-        Ok(snapshot)
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use litellm_traces::{
-        SpanStatus,
-        query::named::{SpendByResponseIdsRow, TraceSpansRow},
-        resolve_trace,
-    };
-
-    use super::*;
-
-    fn trace(span_id: &str) -> Trace {
-        let rows = [TraceSpansRow {
-            trace_id: String::new(),
-            span_id: span_id.into(),
-            parent_span_id: String::new(),
-            name: "run".into(),
-            kind: litellm_traces::ObservationType::Agent,
-            wrapper_candidate: false,
-            agent: "agent".into(),
-            framework: String::new(),
-            status: SpanStatus::Ok,
-            status_message: String::new(),
-            error_truncated: false,
-            start_ns: 1_790_742_989_000_000_000,
-            duration_ns: 10_000_000,
-            service: "agent-demo".into(),
-            input_preview: format!("input of {span_id}"),
-            model: String::new(),
-            input_tokens: 0,
-            output_tokens: 0,
-            litellm_request_id: String::new(),
-            call_keys: Vec::new(),
-            call_evidence: None,
-            tool_call_id: String::new(),
-            team_id: String::new(),
-            api_key_hash: String::new(),
-            user_id: String::new(),
-        }];
-        resolve_trace("trace", "ref", &rows, &[] as &[SpendByResponseIdsRow])
-            .expect("fixture should resolve")
-    }
-
-    fn key(suffix: &str) -> SnapshotKey {
-        SnapshotKey::new(
-            "source",
-            &ReadAccessParams {
-                all_teams: false,
-                user_id: String::new(),
-                team_ids: vec!["team".into()],
-            },
-            suffix,
-            "ref",
-            100,
-        )
-        .unwrap()
-    }
-
-    #[tokio::test]
-    async fn weighted_capacity_bounds_retained_snapshots() {
-        let limit = ["first", "second", "third"]
-            .iter()
-            .map(|span_id| serde_json::to_vec(&trace(span_id)).unwrap().len())
-            .max()
-            .unwrap();
-        let cache = SnapshotCache::new(limit, Duration::from_secs(120));
-
-        for (key, span_id) in [
-            (key("a"), "first"),
-            (key("b"), "second"),
-            (key("c"), "third"),
-        ] {
-            cache.insert(key, trace(span_id)).await.unwrap();
-        }
-
-        cache.entries.run_pending_tasks().await;
-        assert!(cache.entries.weighted_size() <= (limit as u64) * 2);
-    }
-}
+pub use cache::{Freshness, LIVE_TTL, SETTLED_TTL, Snapshot, SnapshotCache, SnapshotKey};
+pub use error::{Error, ReadError};
+pub use reader::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, TraceReader};
+pub use store::{StoreError, TraceStore};

litellm-rust/crates/traces-cache/src/list.rs

@@ -0,0 +1,191 @@
+use std::collections::HashMap;
+
+use crate::{
+    ReadError, SnapshotKey, TraceReader, TraceStore,
+    cache::{Freshness, ListedRun},
+    reader::{map_store_error, now_ms},
+    spend::{spend, spend_window, spend_within},
+    store::StoreError,
+};
+use litellm_traces::{
+    TraceSummary, listed_summary,
+    query::named::{ListTracesRow, ReadAccessParams, TracePageSpansParams, TraceSpansRow},
+    resolve_trace,
+};
+
+const RUNS_PER_SPAN_READ: usize = 16;
+
+fn run_key(team_id: &str, api_key_hash: &str, trace_id: &str) -> (String, String, String) {
+    (
+        team_id.to_owned(),
+        api_key_hash.to_owned(),
+        trace_id.to_owned(),
+    )
+}
+
+fn cache_key<E>(
+    source: &str,
+    access: &ReadAccessParams,
+    row: &ListTracesRow,
+) -> Result<SnapshotKey, ReadError<E>> {
+    Ok(SnapshotKey::run(
+        source,
+        access,
+        (
+            &row.team_id,
+            &row.api_key_hash,
+            &row.trace_id,
+            &row.trace_ref,
+        ),
+    )?)
+}
+
+fn summary(row: &ListTracesRow, listed: Option<&ListedRun>) -> TraceSummary {
+    match listed {
+        Some(ListedRun::Resolved(summary, _)) => (**summary).clone(),
+        Some(ListedRun::Limited) | None => listed_summary(row),
+    }
+}
+
+/// Summaries for one batch of listed runs. Runs resolved within their freshness window come
+/// from the cache; only the rest are read from storage, with one span and one spend read.
+pub(super) async fn list_summaries<S: TraceStore>(
+    reader: &TraceReader,
+    store: &S,
+    access: &ReadAccessParams,
+    runs: &[ListTracesRow],
+) -> Result<Vec<TraceSummary>, ReadError<S::Error>> {
+    let mut keys = Vec::with_capacity(runs.len());
+    let mut listed = Vec::with_capacity(runs.len());
+    for row in runs {
+        let key = cache_key(store.source(), access, row)?;
+        listed.push(reader.lists.runs.get(&key).await);
+        keys.push(key);
+    }
+    let misses: Vec<&ListTracesRow> = runs
+        .iter()
+        .zip(&listed)
+        .filter_map(|(row, listed)| listed.is_none().then_some(row))
+        .collect();
+    let mut resolved = resolve_runs(reader, store, access, &misses)
+        .await?
+        .into_iter();
+    let mut summaries = Vec::with_capacity(runs.len());
+    for ((row, key), cached) in runs.iter().zip(keys).zip(listed) {
+        let listed = match cached {
+            Some(listed) => Some(listed),
+            None => {
+                let listed = resolved.next().flatten();
+                if let Some(listed) = &listed {
+                    reader.lists.runs.insert(key, listed.clone()).await;
+                }
+                listed
+            }
+        };
+        summaries.push(summary(row, listed.as_ref()));
+    }
+    Ok(summaries)
+}
+
+/// One entry per run; `None` means the run could not be resolved and keeps its listed summary
+/// without being cached.
+async fn resolve_runs<S: TraceStore>(
+    reader: &TraceReader,
+    store: &S,
+    access: &ReadAccessParams,
+    runs: &[&ListTracesRow],
+) -> Result<Vec<Option<ListedRun>>, ReadError<S::Error>> {
+    let (Some(start_ms), Some(end_ms)) = (
+        runs.iter().map(|row| row.start_ms).min(),
+        runs.iter()
+            .map(|row| row.start_ms.saturating_add(row.duration_ms))
+            .max(),
+    ) else {
+        return Ok(Vec::new());
+    };
+    let params = TracePageSpansParams {
+        access: access.clone(),
+        trace_refs: runs.iter().map(|row| row.trace_ref.clone()).collect(),
+        start_ms,
+        end_ms: end_ms.saturating_add(1),
+    };
+    let snapshot_ms = now_ms();
+    let spans = match store.run_spans(&params, snapshot_ms).await {
+        Ok(spans) => spans,
+        Err(StoreError::TooLarge) => {
+            let mut resolved = Vec::with_capacity(runs.len());
+            for row in runs {
+                resolved.push(resolve_run(reader, store, access, row).await?);
+            }
+            return Ok(resolved);
+        }
+        Err(error) => return Err(map_store_error(error)),
+    };
+    let spend_rows = spend(store, access, &spans).await;
+    let spend_known = spend_rows.is_some();
+    let spend_rows = spend_rows.unwrap_or_default();
+    let mut spans = spans;
+    spans.sort_by(|left, right| {
+        run_key(&left.team_id, &left.api_key_hash, &left.trace_id)
+            .cmp(&run_key(
+                &right.team_id,
+                &right.api_key_hash,
+                &right.trace_id,
+            ))
+            .then(left.start_ns.cmp(&right.start_ns))
+    });
+    let by_run: HashMap<_, &[TraceSpansRow]> = spans
+        .chunk_by(|left, right| {
+            (&left.team_id, &left.api_key_hash, &left.trace_id)
+                == (&right.team_id, &right.api_key_hash, &right.trace_id)
+        })
+        .map(|run| {
+            (
+                run_key(&run[0].team_id, &run[0].api_key_hash, &run[0].trace_id),
+                run,
+            )
+        })
+        .collect();
+    Ok(runs
+        .iter()
+        .map(|row| {
+            let spans = by_run
+                .get(&run_key(&row.team_id, &row.api_key_hash, &row.trace_id))
+                .copied()
+                .unwrap_or_default();
+            let spend =
+                spend_window(spans).map_or(&[][..], |window| spend_within(&spend_rows, window));
+            resolve_trace(&row.trace_id, &row.trace_ref, spans, spend).map(|trace| {
+                ListedRun::Resolved(
+                    Box::new(trace.summary),
+                    Freshness::of(spans, spend_known, snapshot_ms),
+                )
+            })
+        })
+        .collect())
+}
+
+async fn resolve_run<S: TraceStore>(
+    reader: &TraceReader,
+    store: &S,
+    access: &ReadAccessParams,
+    row: &ListTracesRow,
+) -> Result<Option<ListedRun>, ReadError<S::Error>> {
+    match reader
+        .current(store, access, &row.trace_id, &row.trace_ref)
+        .await
+    {
+        Ok(snapshot) => Ok(snapshot.map(|snapshot| {
+            ListedRun::Resolved(
+                Box::new(snapshot.trace().summary.clone()),
+                snapshot.freshness(),
+            )
+        })),
+        Err(ReadError::TooLarge) => Ok(Some(ListedRun::Limited)),
+        Err(error) => Err(error),
+    }
+}
+
+pub(super) fn run_batches<T>(runs: &[T]) -> impl Iterator<Item = &[T]> + '_ {
+    runs.chunks(RUNS_PER_SPAN_READ)
+}

litellm-rust/crates/traces-cache/src/reader.rs

@@ -0,0 +1,365 @@
+use std::{sync::Arc, time::Duration};
+
+use crate::{
+    ReadError, Snapshot, SnapshotCache, SnapshotKey, StoreError, TraceStore,
+    cache::{Freshness, ListCache},
+    cursor::{
+        ErrorPosition, SpanPosition, decode_cursor, encode_cursor, error_position, trace_position,
+    },
+    list::{list_summaries, run_batches},
+    spend::spend,
+};
+use litellm_traces::{
+    SpanDetail, SpanErrorPage, Trace, TracePage,
+    query::named::{
+        ListTracesParams, ReadAccessParams, SpanDetailParams, SpanErrorParams, TraceIdentityParams,
+        TraceSpansParams,
+    },
+    resolve_trace, to_ui_content,
+};
+
+pub const MAX_GRAPH_BYTES: usize = 64 * 1024 * 1024;
+pub const MAX_GRAPH_SPANS: usize = 100_000;
+
+const SNAPSHOT_IDLE: Duration = Duration::from_secs(120);
+
+/// A read that found no trace, kept apart from failures so single-flight waiters share it
+/// without it being cached.
+pub(super) enum Miss<E> {
+    Absent,
+    Read(ReadError<E>),
+}
+
+impl<E> From<crate::Error> for Miss<E> {
+    fn from(error: crate::Error) -> Self {
+        Self::Read(error.into())
+    }
+}
+
+fn settle<T, E>(result: Result<T, Arc<Miss<E>>>) -> Result<Option<T>, ReadError<E>> {
+    match result {
+        Ok(value) => Ok(Some(value)),
+        Err(miss) => match &*miss {
+            Miss::Absent => Ok(None),
+            Miss::Read(error) => Err(error.clone()),
+        },
+    }
+}
+
+pub struct TraceReader {
+    snapshots: SnapshotCache,
+    pub(super) lists: ListCache,
+    response_bytes: usize,
+}
+
+impl TraceReader {
+    pub fn new(response_bytes: usize) -> Self {
+        Self {
+            snapshots: SnapshotCache::new(MAX_GRAPH_BYTES, SNAPSHOT_IDLE),
+            lists: ListCache::new(),
+            response_bytes,
+        }
+    }
+
+    pub async fn list_traces<S: TraceStore>(
+        &self,
+        store: &S,
+        access: &ReadAccessParams,
+        start_ms: i64,
+        end_ms: i64,
+        cursor: Option<&str>,
+        limit: u32,
+    ) -> Result<TracePage, ReadError<S::Error>> {
+        if limit == 0 {
+            return Err(ReadError::InvalidParameters);
+        }
+        let (cursor_ms, cursor_trace_id) = trace_position(cursor)?;
+        let scope = SnapshotKey::scope(store.source(), access)?;
+        let accepted = self.lists.limits.get(&scope).await.unwrap_or(u32::MAX);
+        let mut params = ListTracesParams {
+            access: access.clone(),
+            start_ms,
+            end_ms,
+            cursor_ms,
+            cursor_trace_id,
+            limit: limit.min(500).min(accepted),
+        };
+        let page = loop {
+            match store.list_runs(&params).await {
+                Err(StoreError::TooLarge) if params.limit > 1 => {
+                    params.limit /= 2;
+                    self.lists.limits.insert(scope.clone(), params.limit).await;
+                }
+                Err(StoreError::TooLarge) => return Err(ReadError::TooLarge),
+                result => break result.map_err(map_store_error)?,
+            }
+        };
+        let next_cursor = page
+            .last()
+            .filter(|_| page.len() == params.limit as usize)
+            .map(|last| encode_cursor(&(last.start_ms, &last.trace_ref)));
+        let data = {
+            let mut summaries = Vec::with_capacity(page.len());
+            for batch in run_batches(&page) {
+                summaries.extend(list_summaries(self, store, access, batch).await?);
+            }
+            summaries
+        };
+        Ok(TracePage { data, next_cursor })
+    }
+
+    pub async fn get_trace<S: TraceStore>(
+        &self,
+        store: &S,
+        access: &ReadAccessParams,
+        trace_id: &str,
+        trace_ref: &str,
+    ) -> Result<Option<Trace>, ReadError<S::Error>> {
+        let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else {
+            return Ok(None);
+        };
+        Ok(self
+            .current(store, access, trace_id, &trace_ref)
+            .await?
+            .map(|snapshot| snapshot.trace().clone()))
+    }
+
+    pub async fn get_trace_page<S: TraceStore>(
+        &self,
+        store: &S,
+        access: &ReadAccessParams,
+        trace_id: &str,
+        trace_ref: &str,
+        cursor: Option<&str>,
+        page_size: u32,
+    ) -> Result<Option<Trace>, ReadError<S::Error>> {
+        if !(1..=500).contains(&page_size) {
+            return Err(ReadError::InvalidParameters);
+        }
+        let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else {
+            return Ok(None);
+        };
+        let Some(cursor) = cursor else {
+            let Some(snapshot) = self.current(store, access, trace_id, &trace_ref).await? else {
+                return Ok(None);
+            };
+            let position = SpanPosition {
+                trace_ref,
+                snapshot_ms: snapshot.snapshot_ms(),
+                offset: 0,
+                version: snapshot.version().to_owned(),
+            };
+            return page(&snapshot, &position, page_size, self.response_bytes).map(Some);
+        };
+        let position: SpanPosition = decode_cursor(cursor, "span")?;
+        if position.trace_ref != trace_ref || position.snapshot_ms == 0 {
+            return Err(ReadError::InvalidCursor("span"));
+        }
+        let Some(snapshot) = settle(
+            self.pinned(store, access, trace_id, &trace_ref, position.snapshot_ms)
+                .await,
+        )?
+        else {
+            return Ok(None);
+        };
+        if position.version != snapshot.version() {
+            return Err(ReadError::TraceChanged);
+        }
+        if position.offset > snapshot.trace().spans.len() {
+            return Err(ReadError::InvalidCursor("span"));
+        }
+        page(&snapshot, &position, page_size, self.response_bytes).map(Some)
+    }
+
+    pub(super) async fn current<S: TraceStore>(
+        &self,
+        store: &S,
+        access: &ReadAccessParams,
+        trace_id: &str,
+        trace_ref: &str,
+    ) -> Result<Option<Arc<Snapshot>>, ReadError<S::Error>> {
+        let latest = SnapshotKey::latest(store.source(), access, trace_id, trace_ref)?;
+        settle(
+            self.snapshots
+                .latest_or_load(latest, now_ms(), |snapshot_ms| {
+                    self.pinned(store, access, trace_id, trace_ref, snapshot_ms)
+                })
+                .await,
+        )
+    }
+
+    async fn pinned<S: TraceStore>(
+        &self,
+        store: &S,
+        access: &ReadAccessParams,
+        trace_id: &str,
+        trace_ref: &str,
+        snapshot_ms: u64,
+    ) -> Result<Arc<Snapshot>, Arc<Miss<S::Error>>> {
+        let key = SnapshotKey::new(store.source(), access, trace_id, trace_ref, snapshot_ms)
+            .map_err(|error| Arc::new(error.into()))?;
+        self.snapshots
+            .pinned_or_load(key, snapshot_ms, async {
+                let params = TraceSpansParams {
+                    access: access.clone(),
+                    trace_id: trace_id.to_owned(),
+                    trace_ref: trace_ref.to_owned(),
+                };
+                let rows = store
+                    .trace_spans(&params, snapshot_ms)
+                    .await
+                    .map_err(|error| Miss::Read(map_store_error(error)))?;
+                let spend_rows = spend(store, access, &rows).await;
+                let freshness = Freshness::of(&rows, spend_rows.is_some(), snapshot_ms);
+                resolve_trace(
+                    trace_id,
+                    trace_ref,
+                    &rows,
+                    spend_rows.as_deref().unwrap_or_default(),
+                )
+                .map(|trace| (trace, freshness))
+                .ok_or(Miss::Absent)
+            })
+            .await
+    }
+
+    pub async fn get_span<S: TraceStore>(
+        &self,
+        store: &S,
+        access: &ReadAccessParams,
+        trace_id: &str,
+        span_id: &str,
+        trace_ref: &str,
+    ) -> Result<Option<SpanDetail>, ReadError<S::Error>> {
+        let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else {
+            return Ok(None);
+        };
+        let params = SpanDetailParams {
+            access: access.clone(),
+            trace_id: trace_id.to_owned(),
+            trace_ref,
+            span_id: span_id.to_owned(),
+        };
+        let row = store.span_detail(&params).await.map_err(map_store_error)?;
+        Ok(row.map(|row| SpanDetail {
+            input_ui: to_ui_content(&row.input),
+            output_ui: to_ui_content(&row.output),
+            span_id: row.span_id,
+            input: row.input,
+            output: row.output,
+            attributes: row.attributes,
+        }))
+    }
+
+    pub async fn get_span_error<S: TraceStore>(
+        &self,
+        store: &S,
+        access: &ReadAccessParams,
+        trace_id: &str,
+        span_id: &str,
+        trace_ref: &str,
+        cursor: Option<&str>,
+    ) -> Result<Option<SpanErrorPage>, ReadError<S::Error>> {
+        let position = error_position(cursor)?;
+        let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else {
+            return Ok(None);
+        };
+        let offset = position.as_ref().map_or(0, |position| position.offset);
+        let params = SpanErrorParams {
+            access: access.clone(),
+            trace_id: trace_id.to_owned(),
+            trace_ref,
+            span_id: span_id.to_owned(),
+            error_offset: offset,
+            error_version: position
+                .map(|position| position.version)
+                .unwrap_or_default(),
+        };
+        let Some(row) = store.span_error(&params).await.map_err(map_store_error)? else {
+            return Ok(None);
+        };
+        let next_offset = offset + row.message.chars().count() as u64;
+        let next_cursor = (next_offset < row.total_chars).then(|| {
+            encode_cursor(&ErrorPosition {
+                offset: next_offset,
+                version: row.version,
+            })
+        });
+        Ok(Some(SpanErrorPage {
+            span_id: row.span_id,
+            message: row.message,
+            total_chars: row.total_chars,
+            next_cursor,
+        }))
+    }
+}
+
+async fn reference<S: TraceStore>(
+    store: &S,
+    access: &ReadAccessParams,
+    trace_id: &str,
+    trace_ref: &str,
+) -> Result<Option<String>, ReadError<S::Error>> {
+    if !trace_ref.is_empty() {
+        return Ok(Some(trace_ref.to_owned()));
+    }
+    let params = TraceIdentityParams {
+        access: access.clone(),
+        trace_id: trace_id.to_owned(),
+    };
+    let identities = store.trace_refs(&params).await.map_err(map_store_error)?;
+    if identities.len() > 1 {
+        return Err(ReadError::AmbiguousTrace);
+    }
+    Ok(identities.into_iter().next())
+}
+
+fn page<E>(
+    snapshot: &Snapshot,
+    position: &SpanPosition,
+    page_size: u32,
+    response_bytes: usize,
+) -> Result<Trace, ReadError<E>> {
+    let spans = &snapshot.trace().spans;
+    let create_page = |count: usize| {
+        let end = position.offset.saturating_add(count).min(spans.len());
+        Trace {
+            summary: snapshot.trace().summary.clone(),
+            agents: snapshot.trace().agents.clone(),
+            spans: spans[position.offset..end].to_vec(),
+            next_cursor: (end < spans.len()).then(|| {
+                encode_cursor(&SpanPosition {
+                    trace_ref: position.trace_ref.clone(),
+                    snapshot_ms: position.snapshot_ms,
+                    offset: end,
+                    version: snapshot.version().to_owned(),
+                })
+            }),
+        }
+    };
+    let mut trace = create_page(page_size as usize);
+    loop {
+        if serde_json::to_vec(&trace)
+            .map_err(|error| ReadError::Encode(Arc::new(error)))?
+            .len()
+            <= response_bytes
+        {
+            return Ok(trace);
+        }
+        if trace.spans.len() <= 1 {
+            return Err(ReadError::TooLarge);
+        }
+        trace = create_page(trace.spans.len() / 2);
+    }
+}
+
+pub(super) fn now_ms() -> u64 {
+    (time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) as u64
+}
+
+pub(super) fn map_store_error<E>(error: StoreError<E>) -> ReadError<E> {
+    match error {
+        StoreError::TooLarge => ReadError::TooLarge,
+        StoreError::Failed(error) => ReadError::Store(Arc::new(error)),
+    }
+}

litellm-rust/crates/traces-cache/src/spend.rs

@@ -0,0 +1,67 @@
+use std::ops::Range;
+
+use crate::TraceStore;
+use litellm_traces::{
+    SpendLookup,
+    query::named::{
+        ReadAccessParams, SpendByResponseIdsParams, SpendByResponseIdsRow, TraceSpansRow,
+    },
+};
+
+const NANOS_PER_MS: i64 = 1_000_000;
+const SPEND_WINDOW_MS: i64 = 30 * 60 * 1000;
+
+pub(super) fn spend_window(rows: &[TraceSpansRow]) -> Option<Range<i64>> {
+    let start_ns = rows.iter().map(|row| row.start_ns).min()?;
+    let end_ns = rows
+        .iter()
+        .map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns))
+        .max()?;
+    Some(
+        start_ns.div_euclid(NANOS_PER_MS) - SPEND_WINDOW_MS
+            ..end_ns.div_euclid(NANOS_PER_MS) + SPEND_WINDOW_MS,
+    )
+}
+
+pub(super) fn spend_within(
+    spend: &[SpendByResponseIdsRow],
+    window: Range<i64>,
+) -> &[SpendByResponseIdsRow] {
+    let first = spend.partition_point(|row| row.start_ms < window.start);
+    let end = spend.partition_point(|row| row.start_ms < window.end);
+    &spend[first..end.max(first)]
+}
+
+/// Spend rows sorted by `start_ms`, or `None` when the lookup failed and spend is unknown.
+pub(super) async fn spend<S: TraceStore>(
+    store: &S,
+    access: &ReadAccessParams,
+    rows: &[TraceSpansRow],
+) -> Option<Vec<SpendByResponseIdsRow>> {
+    let lookup = SpendLookup::new(rows);
+    let Some(window) = spend_window(rows) else {
+        return Some(Vec::new());
+    };
+    if lookup.is_empty() {
+        return Some(Vec::new());
+    }
+    let params = SpendByResponseIdsParams {
+        access: access.clone(),
+        response_ids: lookup.response_ids,
+        request_ids: lookup.request_ids,
+        trace_ids: lookup.trace_ids,
+        start_ms: window.start,
+        end_ms: window.end,
+    };
+    match store.spend(&params).await {
+        Ok(rows) => {
+            let mut rows = rows;
+            rows.sort_by_key(|row| row.start_ms);
+            Some(rows)
+        }
+        Err(error) => {
+            tracing::warn!(%error, "trace spend lookup unavailable");
+            None
+        }
+    }
+}

litellm-rust/crates/traces-cache/src/store.rs

@@ -0,0 +1,62 @@
+use std::future::Future;
+
+use litellm_traces::query::named::{
+    ListTracesParams, ListTracesRow, SpanDetailParams, SpanDetailRow, SpanErrorParams,
+    SpanErrorRow, SpendByResponseIdsParams, SpendByResponseIdsRow, TraceIdentityParams,
+    TracePageSpansParams, TraceSpansParams, TraceSpansRow,
+};
+
+#[derive(Debug, thiserror::Error)]
+pub enum StoreError<E> {
+    #[error("trace read exceeds the storage read budget")]
+    TooLarge,
+    #[error(transparent)]
+    Failed(E),
+}
+
+pub trait TraceStore: Sync {
+    type Error: std::error::Error + Send + Sync + 'static;
+
+    /// Identifies the backing storage for snapshot cache keys.
+    fn source(&self) -> &str;
+
+    fn trace_refs(
+        &self,
+        params: &TraceIdentityParams,
+    ) -> impl Future<Output = Result<Vec<String>, StoreError<Self::Error>>> + Send;
+
+    /// Returns `TooLarge` when the response exceeds the storage limit so the reader can halve `limit`.
+    fn list_runs(
+        &self,
+        params: &ListTracesParams,
+    ) -> impl Future<Output = Result<Vec<ListTracesRow>, StoreError<Self::Error>>> + Send;
+
+    /// Returns spans visible at `snapshot_ms`, sorted by `start_ns`, or `TooLarge` past `MAX_GRAPH_BYTES`/`MAX_GRAPH_SPANS`.
+    fn trace_spans(
+        &self,
+        params: &TraceSpansParams,
+        snapshot_ms: u64,
+    ) -> impl Future<Output = Result<Vec<TraceSpansRow>, StoreError<Self::Error>>> + Send;
+
+    /// Returns spans visible at `snapshot_ms`, sorted by `start_ns`, or `TooLarge` past `MAX_GRAPH_BYTES`/`MAX_GRAPH_SPANS`.
+    fn run_spans(
+        &self,
+        params: &TracePageSpansParams,
+        snapshot_ms: u64,
+    ) -> impl Future<Output = Result<Vec<TraceSpansRow>, StoreError<Self::Error>>> + Send;
+
+    fn spend(
+        &self,
+        params: &SpendByResponseIdsParams,
+    ) -> impl Future<Output = Result<Vec<SpendByResponseIdsRow>, StoreError<Self::Error>>> + Send;
+
+    fn span_detail(
+        &self,
+        params: &SpanDetailParams,
+    ) -> impl Future<Output = Result<Option<SpanDetailRow>, StoreError<Self::Error>>> + Send;
+
+    fn span_error(
+        &self,
+        params: &SpanErrorParams,
+    ) -> impl Future<Output = Result<Option<SpanErrorRow>, StoreError<Self::Error>>> + Send;
+}

litellm-rust/crates/traces-cache/tests/read.rs

@@ -0,0 +1,687 @@
+use std::{
+    collections::{HashMap, HashSet},
+    sync::{
+        Mutex,
+        atomic::{AtomicUsize, Ordering},
+    },
+    time::Duration,
+};
+
+use litellm_traces::{
+    CallEvidenceKind, CallKey, ObservationType, SpanStatus,
+    query::named::{
+        ListTracesParams, ListTracesRow, ReadAccessParams, SpanDetailParams, SpanDetailRow,
+        SpanErrorParams, SpanErrorRow, SpendByResponseIdsParams, SpendByResponseIdsRow,
+        TraceIdentityParams, TracePageSpansParams, TraceSpansParams, TraceSpansRow,
+    },
+};
+use litellm_traces_cache::{LIVE_TTL, ReadError, StoreError, TraceReader, TraceStore};
+use rstest::rstest;
+
+const START_NS: i64 = 1_790_742_989_000_000_000;
+
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
+enum Operation {
+    TraceRefs,
+    ListRuns,
+    TraceSpans,
+    RunSpans,
+    Spend,
+    SpanDetail,
+    SpanError,
+}
+
+#[derive(Clone, Copy)]
+enum Failure {
+    TooLarge,
+    Failed,
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("fake trace store failed")]
+struct FakeError;
+
+#[derive(Default)]
+struct State {
+    failures: HashMap<Operation, Failure>,
+    trace_refs: Vec<String>,
+    list_runs: Vec<ListTracesRow>,
+    trace_spans: HashMap<String, Vec<TraceSpansRow>>,
+    run_spans: Vec<TraceSpansRow>,
+    spend: Vec<SpendByResponseIdsRow>,
+    span_detail: Option<SpanDetailRow>,
+    span_error: Option<SpanErrorRow>,
+    list_runs_too_large_above: Option<u32>,
+    trace_too_large_refs: HashSet<String>,
+}
+
+#[derive(Default)]
+struct Calls {
+    trace_refs: AtomicUsize,
+    list_runs: AtomicUsize,
+    trace_spans: AtomicUsize,
+    run_spans: AtomicUsize,
+    spend: AtomicUsize,
+    span_detail: AtomicUsize,
+    span_error: AtomicUsize,
+}
+
+#[derive(Default)]
+struct FakeStore {
+    state: Mutex<State>,
+    calls: Calls,
+}
+
+impl FakeStore {
+    fn with_spans(trace_ref: &str, spans: Vec<TraceSpansRow>) -> Self {
+        Self {
+            state: Mutex::new(State {
+                trace_spans: HashMap::from([(trace_ref.to_owned(), spans)]),
+                ..State::default()
+            }),
+            calls: Calls::default(),
+        }
+    }
+
+    fn set_failure(&self, operation: Operation, failure: Failure) {
+        self.state
+            .lock()
+            .unwrap()
+            .failures
+            .insert(operation, failure);
+    }
+
+    fn set_trace_refs(&self, trace_refs: Vec<String>) {
+        self.state.lock().unwrap().trace_refs = trace_refs;
+    }
+
+    fn set_list_runs(&self, rows: Vec<ListTracesRow>) {
+        self.state.lock().unwrap().list_runs = rows;
+    }
+
+    fn set_list_runs_too_large_above(&self, limit: u32) {
+        self.state.lock().unwrap().list_runs_too_large_above = Some(limit);
+    }
+
+    fn set_run_spans(&self, rows: Vec<TraceSpansRow>) {
+        self.state.lock().unwrap().run_spans = rows;
+    }
+
+    fn set_trace_spans_too_large(&self, trace_ref: &str) {
+        self.state
+            .lock()
+            .unwrap()
+            .trace_too_large_refs
+            .insert(trace_ref.to_owned());
+    }
+
+    fn calls(&self, operation: Operation) -> usize {
+        match operation {
+            Operation::TraceRefs => self.calls.trace_refs.load(Ordering::SeqCst),
+            Operation::ListRuns => self.calls.list_runs.load(Ordering::SeqCst),
+            Operation::TraceSpans => self.calls.trace_spans.load(Ordering::SeqCst),
+            Operation::RunSpans => self.calls.run_spans.load(Ordering::SeqCst),
+            Operation::Spend => self.calls.spend.load(Ordering::SeqCst),
+            Operation::SpanDetail => self.calls.span_detail.load(Ordering::SeqCst),
+            Operation::SpanError => self.calls.span_error.load(Ordering::SeqCst),
+        }
+    }
+
+    fn failure(state: &State, operation: Operation) -> Result<(), StoreError<FakeError>> {
+        match state.failures.get(&operation) {
+            Some(Failure::TooLarge) => Err(StoreError::TooLarge),
+            Some(Failure::Failed) => Err(StoreError::Failed(FakeError)),
+            None => Ok(()),
+        }
+    }
+}
+
+impl TraceStore for FakeStore {
+    type Error = FakeError;
+
+    fn source(&self) -> &str {
+        "fake"
+    }
+
+    async fn trace_refs(
+        &self,
+        _: &TraceIdentityParams,
+    ) -> Result<Vec<String>, StoreError<Self::Error>> {
+        self.calls.trace_refs.fetch_add(1, Ordering::SeqCst);
+        let state = self.state.lock().unwrap();
+        Self::failure(&state, Operation::TraceRefs)?;
+        Ok(state.trace_refs.clone())
+    }
+
+    async fn list_runs(
+        &self,
+        params: &ListTracesParams,
+    ) -> Result<Vec<ListTracesRow>, StoreError<Self::Error>> {
+        self.calls.list_runs.fetch_add(1, Ordering::SeqCst);
+        let state = self.state.lock().unwrap();
+        Self::failure(&state, Operation::ListRuns)?;
+        if state
+            .list_runs_too_large_above
+            .is_some_and(|limit| params.limit > limit)
+        {
+            return Err(StoreError::TooLarge);
+        }
+        Ok(state
+            .list_runs
+            .iter()
+            .take(params.limit as usize)
+            .cloned()
+            .collect())
+    }
+
+    async fn trace_spans(
+        &self,
+        params: &TraceSpansParams,
+        _: u64,
+    ) -> Result<Vec<TraceSpansRow>, StoreError<Self::Error>> {
+        self.calls.trace_spans.fetch_add(1, Ordering::SeqCst);
+        tokio::task::yield_now().await;
+        let state = self.state.lock().unwrap();
+        Self::failure(&state, Operation::TraceSpans)?;
+        if state.trace_too_large_refs.contains(&params.trace_ref) {
+            return Err(StoreError::TooLarge);
+        }
+        Ok(state
+            .trace_spans
+            .get(&params.trace_ref)
+            .cloned()
+            .unwrap_or_default())
+    }
+
+    async fn run_spans(
+        &self,
+        _: &TracePageSpansParams,
+        _: u64,
+    ) -> Result<Vec<TraceSpansRow>, StoreError<Self::Error>> {
+        self.calls.run_spans.fetch_add(1, Ordering::SeqCst);
+        tokio::task::yield_now().await;
+        let state = self.state.lock().unwrap();
+        Self::failure(&state, Operation::RunSpans)?;
+        Ok(state.run_spans.clone())
+    }
+
+    async fn spend(
+        &self,
+        _: &SpendByResponseIdsParams,
+    ) -> Result<Vec<SpendByResponseIdsRow>, StoreError<Self::Error>> {
+        self.calls.spend.fetch_add(1, Ordering::SeqCst);
+        let state = self.state.lock().unwrap();
+        Self::failure(&state, Operation::Spend)?;
+        Ok(state.spend.clone())
+    }
+
+    async fn span_detail(
+        &self,
+        _: &SpanDetailParams,
+    ) -> Result<Option<SpanDetailRow>, StoreError<Self::Error>> {
+        self.calls.span_detail.fetch_add(1, Ordering::SeqCst);
+        let state = self.state.lock().unwrap();
+        Self::failure(&state, Operation::SpanDetail)?;
+        Ok(state.span_detail.clone())
+    }
+
+    async fn span_error(
+        &self,
+        _: &SpanErrorParams,
+    ) -> Result<Option<SpanErrorRow>, StoreError<Self::Error>> {
+        self.calls.span_error.fetch_add(1, Ordering::SeqCst);
+        let state = self.state.lock().unwrap();
+        Self::failure(&state, Operation::SpanError)?;
+        Ok(state.span_error.clone())
+    }
+}
+
+fn access() -> ReadAccessParams {
+    ReadAccessParams {
+        all_teams: true,
+        user_id: String::new(),
+        team_ids: Vec::new(),
+    }
+}
+
+fn span(index: usize) -> TraceSpansRow {
+    TraceSpansRow {
+        trace_id: "trace".into(),
+        span_id: format!("span-{index}"),
+        parent_span_id: if index == 0 {
+            String::new()
+        } else {
+            "span-0".into()
+        },
+        name: "agent".into(),
+        kind: ObservationType::Agent,
+        wrapper_candidate: false,
+        agent: "agent".into(),
+        framework: String::new(),
+        status: SpanStatus::Ok,
+        status_message: String::new(),
+        error_truncated: false,
+        start_ns: START_NS + index as i64 * 1_000_000,
+        duration_ns: 10_000_000,
+        service: "test".into(),
+        input_preview: format!("span input {index}"),
+        model: String::new(),
+        input_tokens: 0,
+        output_tokens: 0,
+        litellm_request_id: String::new(),
+        call_keys: Vec::new(),
+        call_evidence: None,
+        tool_call_id: String::new(),
+        team_id: "team".into(),
+        api_key_hash: "key".into(),
+        user_id: "user".into(),
+    }
+}
+
+fn run(trace_id: &str, trace_ref: &str) -> ListTracesRow {
+    ListTracesRow {
+        trace_id: trace_id.into(),
+        trace_ref: trace_ref.into(),
+        team_id: "team".into(),
+        api_key_hash: "key".into(),
+        user_id: "user".into(),
+        name: "listed".into(),
+        service: "test".into(),
+        input_preview: String::new(),
+        status: SpanStatus::Ok,
+        start_ms: 1_790_742_989_000,
+        duration_ms: 10,
+        span_count: 1,
+        agent_count: 1,
+        agent_invocations: 1,
+        agent_names: vec!["agent".into()],
+        frameworks: Vec::new(),
+        llm_calls: 0,
+        tool_calls: 0,
+        input_tokens: 0,
+        output_tokens: 0,
+        models: Vec::new(),
+        error_count: 0,
+        request_ids: Vec::new(),
+    }
+}
+
+#[rstest]
+#[tokio::test]
+async fn pages_reuse_one_trace_snapshot_and_concatenate_in_order() {
+    let store = FakeStore::with_spans("ref", (0..5).map(span).collect());
+    let reader = TraceReader::new(usize::MAX);
+    let access = access();
+    let first = reader
+        .get_trace_page(&store, &access, "trace", "ref", None, 2)
+        .await
+        .unwrap()
+        .unwrap();
+    let second = reader
+        .get_trace_page(
+            &store,
+            &access,
+            "trace",
+            "ref",
+            first.next_cursor.as_deref(),
+            2,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    let third = reader
+        .get_trace_page(
+            &store,
+            &access,
+            "trace",
+            "ref",
+            second.next_cursor.as_deref(),
+            2,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    let ids: Vec<_> = first
+        .spans
+        .iter()
+        .chain(&second.spans)
+        .chain(&third.spans)
+        .map(|span| span.span_id.as_str())
+        .collect();
+    assert_eq!(ids, ["span-0", "span-1", "span-2", "span-3", "span-4"]);
+    assert!(third.next_cursor.is_none());
+    assert_eq!(store.calls(Operation::TraceSpans), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn snapshot_versions_are_stable_across_readers_and_detect_changes() {
+    let access = access();
+    let original = FakeStore::with_spans("ref", vec![span(0), span(1)]);
+    let reader_a = TraceReader::new(usize::MAX);
+    let first = reader_a
+        .get_trace_page(&original, &access, "trace", "ref", None, 1)
+        .await
+        .unwrap()
+        .unwrap();
+    let cursor = first.next_cursor.unwrap();
+
+    let changed = FakeStore::with_spans("ref", vec![span(0), span(1), span(2)]);
+    let reader_b = TraceReader::new(usize::MAX);
+    let result = reader_b
+        .get_trace_page(&changed, &access, "trace", "ref", Some(&cursor), 1)
+        .await;
+    assert!(matches!(result, Err(ReadError::TraceChanged)));
+
+    let unchanged = FakeStore::with_spans("ref", vec![span(0), span(1)]);
+    let reader_c = TraceReader::new(usize::MAX);
+    let next = reader_c
+        .get_trace_page(&unchanged, &access, "trace", "ref", Some(&cursor), 1)
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(next.spans[0].span_id, "span-1");
+}
+
+#[rstest]
+#[tokio::test]
+async fn response_size_splits_pages_and_rejects_a_single_oversized_span() {
+    let spans: Vec<_> = (0..4)
+        .map(|index| {
+            let mut row = span(index);
+            row.input_preview = "x".repeat(256);
+            row
+        })
+        .collect();
+    let access = access();
+    let full_budget_reader = TraceReader::new(usize::MAX);
+    let one_span = full_budget_reader
+        .get_trace_page(
+            &FakeStore::with_spans("ref", spans.clone()),
+            &access,
+            "trace",
+            "ref",
+            None,
+            1,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    let response_bytes = serde_json::to_vec(&one_span).unwrap().len() + 128;
+    let reader = TraceReader::new(response_bytes);
+    let store = FakeStore::with_spans("ref", spans.clone());
+    let page = reader
+        .get_trace_page(&store, &access, "trace", "ref", None, 4)
+        .await
+        .unwrap()
+        .unwrap();
+    assert!(!page.spans.is_empty());
+    assert!(page.spans.len() < 4);
+    assert!(page.next_cursor.is_some());
+    let continued = reader
+        .get_trace_page(
+            &store,
+            &access,
+            "trace",
+            "ref",
+            page.next_cursor.as_deref(),
+            4,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    assert!(!continued.spans.is_empty());
+    assert!(matches!(
+        TraceReader::new(1)
+            .get_trace_page(&store, &access, "trace", "ref", None, 1)
+            .await,
+        Err(ReadError::TooLarge)
+    ));
+}
+
+#[rstest]
+#[tokio::test]
+async fn list_run_budget_halves_the_limit_and_cursor_requires_a_full_page() {
+    let store = FakeStore::default();
+    store.set_list_runs(
+        (0..3)
+            .map(|index| run(&format!("trace-{index}"), &format!("ref-{index}")))
+            .collect(),
+    );
+    store.set_list_runs_too_large_above(2);
+    let reader = TraceReader::new(usize::MAX);
+    let access = access();
+    let page = reader
+        .list_traces(&store, &access, 0, i64::MAX, None, 8)
+        .await
+        .unwrap();
+    assert_eq!(page.data.len(), 2);
+    assert!(page.next_cursor.is_some());
+    assert_eq!(store.calls(Operation::ListRuns), 3);
+
+    let shorter = FakeStore::default();
+    shorter.set_list_runs(vec![run("only", "ref-only")]);
+    shorter.set_list_runs_too_large_above(2);
+    let page = reader
+        .list_traces(&shorter, &access, 0, i64::MAX, None, 8)
+        .await
+        .unwrap();
+    assert_eq!(page.data.len(), 1);
+    assert!(page.next_cursor.is_none());
+    assert_eq!(shorter.calls(Operation::ListRuns), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn oversized_run_batch_falls_back_to_each_run_and_keeps_listed_summaries() {
+    let store = FakeStore::with_spans("ref-good", vec![span(0)]);
+    store.set_list_runs(vec![
+        run("trace-large", "ref-large"),
+        run("trace-good", "ref-good"),
+    ]);
+    store.set_trace_spans_too_large("ref-large");
+    store.set_run_spans(Vec::new());
+    store.set_failure(Operation::RunSpans, Failure::TooLarge);
+    let reader = TraceReader::new(usize::MAX);
+    let page = reader
+        .list_traces(&store, &access(), 0, i64::MAX, None, 2)
+        .await
+        .unwrap();
+    assert_eq!(page.data.len(), 2);
+    assert!(page.data[0].resolution_limited);
+    assert_eq!(page.data[0].trace_ref, "ref-large");
+    assert!(!page.data[1].resolution_limited);
+    assert_eq!(page.data[1].trace_ref, "ref-good");
+    assert_eq!(store.calls(Operation::RunSpans), 1);
+    assert_eq!(store.calls(Operation::TraceSpans), 2);
+
+    let again = reader
+        .list_traces(&store, &access(), 0, i64::MAX, None, 2)
+        .await
+        .unwrap();
+    assert_eq!(again.data, page.data);
+    assert_eq!(store.calls(Operation::RunSpans), 1);
+    assert_eq!(store.calls(Operation::TraceSpans), 2);
+}
+
+#[rstest]
+#[tokio::test]
+async fn failed_spend_lookup_preserves_the_trace_with_unknown_spend() {
+    let mut row = span(0);
+    row.litellm_request_id = "response".into();
+    row.call_keys = vec![CallKey::ProviderResponse("response".into())];
+    row.call_evidence = Some(CallEvidenceKind::Complete);
+    let store = FakeStore::with_spans("ref", vec![row]);
+    store.set_failure(Operation::Spend, Failure::Failed);
+    let trace = TraceReader::new(usize::MAX)
+        .get_trace(&store, &access(), "trace", "ref")
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(trace.summary.spend, None);
+    assert_eq!(trace.spans[0].spend, None);
+    assert_eq!(store.calls(Operation::Spend), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn ambiguous_trace_references_fail_and_a_single_reference_is_resolved() {
+    let reader = TraceReader::new(usize::MAX);
+    let access = access();
+    let ambiguous = FakeStore::default();
+    ambiguous.set_trace_refs(vec!["ref-a".into(), "ref-b".into()]);
+    assert!(matches!(
+        reader.get_trace(&ambiguous, &access, "trace", "").await,
+        Err(ReadError::AmbiguousTrace)
+    ));
+
+    let unique = FakeStore::with_spans("ref-only", vec![span(0)]);
+    unique.set_trace_refs(vec!["ref-only".into()]);
+    let trace = reader
+        .get_trace(&unique, &access, "trace", "")
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(trace.summary.trace_ref, "ref-only");
+}
+
+#[rstest]
+#[case::zero(0)]
+#[case::above_max(501)]
+#[tokio::test]
+async fn invalid_page_sizes_are_rejected(#[case] page_size: u32) {
+    let store = FakeStore::default();
+    let reader = TraceReader::new(usize::MAX);
+    let access = access();
+    assert!(matches!(
+        reader
+            .get_trace_page(&store, &access, "trace", "ref", None, page_size)
+            .await,
+        Err(ReadError::InvalidParameters)
+    ));
+}
+
+#[rstest]
+#[tokio::test]
+async fn zero_list_limit_is_rejected() {
+    let store = FakeStore::default();
+    let reader = TraceReader::new(usize::MAX);
+    let access = access();
+    assert!(matches!(
+        reader
+            .list_traces(&store, &access, 0, i64::MAX, None, 0)
+            .await,
+        Err(ReadError::InvalidParameters)
+    ));
+}
+
+fn now_ns() -> i64 {
+    time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64
+}
+
+#[rstest]
+#[tokio::test]
+async fn concurrent_and_repeated_opens_share_one_storage_read() {
+    let store = FakeStore::with_spans("ref", (0..3).map(span).collect());
+    let reader = TraceReader::new(usize::MAX);
+    let access = access();
+    let (first, second) = tokio::join!(
+        reader.get_trace(&store, &access, "trace", "ref"),
+        reader.get_trace_page(&store, &access, "trace", "ref", None, 2),
+    );
+    let first = first.unwrap().unwrap();
+    let second = second.unwrap().unwrap();
+    let reopened = reader
+        .get_trace_page(&store, &access, "trace", "ref", None, 2)
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(first.spans.len(), 3);
+    assert_eq!(second.spans, first.spans[..2]);
+    assert_eq!(reopened.next_cursor, second.next_cursor);
+    assert_eq!(store.calls(Operation::TraceSpans), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn failed_reads_are_not_cached() {
+    let store = FakeStore::with_spans("ref", vec![span(0)]);
+    store.set_failure(Operation::TraceSpans, Failure::Failed);
+    let reader = TraceReader::new(usize::MAX);
+    let access = access();
+    assert!(matches!(
+        reader.get_trace(&store, &access, "trace", "ref").await,
+        Err(ReadError::Store(_))
+    ));
+    store.state.lock().unwrap().failures.clear();
+    let trace = reader
+        .get_trace(&store, &access, "trace", "ref")
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(trace.spans.len(), 1);
+    assert_eq!(store.calls(Operation::TraceSpans), 2);
+}
+
+#[rstest]
+#[tokio::test]
+async fn listed_runs_are_read_once_until_a_live_run_expires() {
+    let live = TraceSpansRow {
+        trace_id: "trace-live".into(),
+        start_ns: now_ns(),
+        ..span(0)
+    };
+    let settled = TraceSpansRow {
+        trace_id: "trace-settled".into(),
+        ..span(0)
+    };
+    let store = FakeStore::default();
+    store.set_list_runs(vec![
+        run("trace-live", "ref-live"),
+        run("trace-settled", "ref-settled"),
+    ]);
+    store.set_run_spans(vec![live, settled]);
+    let reader = TraceReader::new(usize::MAX);
+    let access = access();
+    let list = || reader.list_traces(&store, &access, 0, i64::MAX, None, 2);
+
+    let first = list().await.unwrap();
+    assert!(first.data.iter().all(|summary| summary.name == "agent"));
+    list().await.unwrap();
+    assert_eq!(store.calls(Operation::RunSpans), 1);
+
+    store.set_run_spans(Vec::new());
+    tokio::time::sleep(LIVE_TTL + Duration::from_millis(200)).await;
+    let after = list().await.unwrap();
+    assert_eq!(store.calls(Operation::RunSpans), 2);
+    assert_eq!(after.data[0].name, "listed");
+    assert_eq!(after.data[1], first.data[1]);
+}
+
+#[rstest]
+#[tokio::test]
+async fn concurrent_pages_of_an_evicted_snapshot_share_one_storage_read() {
+    let access = access();
+    let first = TraceReader::new(usize::MAX)
+        .get_trace_page(
+            &FakeStore::with_spans("ref", (0..3).map(span).collect()),
+            &access,
+            "trace",
+            "ref",
+            None,
+            1,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    let cursor = first.next_cursor.as_deref();
+    let store = FakeStore::with_spans("ref", (0..3).map(span).collect());
+    let reader = TraceReader::new(usize::MAX);
+    let (left, right) = tokio::join!(
+        reader.get_trace_page(&store, &access, "trace", "ref", cursor, 1),
+        reader.get_trace_page(&store, &access, "trace", "ref", cursor, 1),
+    );
+    assert_eq!(left.unwrap().unwrap().spans[0].span_id, "span-1");
+    assert_eq!(right.unwrap().unwrap().spans[0].span_id, "span-1");
+    assert_eq!(store.calls(Operation::TraceSpans), 1);
+}

litellm-rust/crates/traces-cache/tests/snapshots.rs

@@ -5,7 +5,9 @@ use litellm_traces::{
     query::named::{ReadAccessParams, SpendByResponseIdsRow, TraceSpansRow},
     resolve_trace,
 };
-use litellm_traces_cache::{Error, SnapshotCache, SnapshotKey};
+use std::sync::Arc;
+
+use litellm_traces_cache::{Error, Freshness, Snapshot, SnapshotCache, SnapshotKey};
 use rstest::{fixture, rstest};
 
 const T0: i64 = 1_790_742_989_000_000_000;
@@ -60,6 +62,18 @@ fn key(
     SnapshotKey::new(source, access, trace_id, trace_ref, ms).unwrap()
 }
 
+async fn insert(
+    cache: &SnapshotCache,
+    key: SnapshotKey,
+    trace: Trace,
+) -> Result<Arc<Snapshot>, Arc<Error>> {
+    cache
+        .pinned_or_load(key, 100, async {
+            Ok::<_, Error>((trace, Freshness::Settled))
+        })
+        .await
+}
+
 #[fixture]
 fn trace() -> Trace {
     resolve_trace(
@@ -85,7 +99,7 @@ async fn cached_trace_is_isolated_by_access_scope(
     let cache = SnapshotCache::new(1024 * 1024, TTL);
     let stored = key("source", &access(), "trace", "ref", 100);
 
-    cache.insert(stored.clone(), trace.clone()).await.unwrap();
+    insert(&cache, stored.clone(), trace.clone()).await.unwrap();
 
     let other_access = ReadAccessParams {
         all_teams,
@@ -115,7 +129,7 @@ async fn cached_trace_is_isolated_by_key_fields(
     let cache = SnapshotCache::new(1024 * 1024, TTL);
     let stored = key("source", &access(), "trace", "ref", 100);
 
-    cache.insert(stored.clone(), trace.clone()).await.unwrap();
+    insert(&cache, stored.clone(), trace.clone()).await.unwrap();
 
     let other = key(source, &access(), trace_id, trace_ref, snapshot_ms);
     assert!(cache.get(&other).await.is_none());
@@ -129,7 +143,7 @@ async fn snapshot_at_the_size_limit_is_accepted(trace: Trace) {
     let cache = SnapshotCache::new(size, TTL);
     let stored = key("source", &access(), "trace", "ref", 100);
 
-    cache.insert(stored.clone(), trace).await.unwrap();
+    insert(&cache, stored.clone(), trace).await.unwrap();
     assert!(cache.get(&stored).await.is_some());
 }
 
@@ -141,8 +155,8 @@ async fn snapshot_one_byte_over_the_size_limit_is_rejected(trace: Trace) {
     let stored = key("source", &access(), "trace", "ref", 100);
 
     assert!(matches!(
-        cache.insert(stored.clone(), trace).await,
-        Err(Error::ReadTooLarge)
+        insert(&cache, stored.clone(), trace).await,
+        Err(error) if matches!(*error, Error::ReadTooLarge)
     ));
     assert!(cache.get(&stored).await.is_none());
 }
@@ -166,25 +180,31 @@ async fn snapshot_version_tracks_the_ordered_span_ids(
     };
     let cache = SnapshotCache::new(1024 * 1024, TTL);
 
-    let first = cache
-        .insert(key("source", &access(), "a", "ref", 100), build(first_ids))
-        .await
-        .unwrap();
-    let second = cache
-        .insert(key("source", &access(), "b", "ref", 100), build(second_ids))
-        .await
-        .unwrap();
+    let first = insert(
+        &cache,
+        key("source", &access(), "a", "ref", 100),
+        build(first_ids),
+    )
+    .await
+    .unwrap();
+    let second = insert(
+        &cache,
+        key("source", &access(), "b", "ref", 100),
+        build(second_ids),
+    )
+    .await
+    .unwrap();
 
     assert_eq!(first.version() == second.version(), equal);
 }
 
 #[rstest]
 #[tokio::test]
-async fn snapshots_expire_after_the_ttl(trace: Trace) {
+async fn snapshots_expire_when_idle(trace: Trace) {
     let cache = SnapshotCache::new(1024 * 1024, Duration::from_millis(50));
     let stored = key("source", &access(), "trace", "ref", 100);
 
-    cache.insert(stored.clone(), trace).await.unwrap();
+    insert(&cache, stored.clone(), trace).await.unwrap();
     tokio::time::sleep(Duration::from_millis(200)).await;
 
     assert!(cache.get(&stored).await.is_none());

litellm-rust/crates/traces-clickhouse/Cargo.toml

@@ -12,11 +12,9 @@ schema = ["dep:schemars", "litellm-traces/schema"]
 macro_rules_attribute.workspace = true
 schemars = { workspace = true, optional = true }
 askama.workspace = true
-base64.workspace = true
 flate2.workspace = true
 futures-util.workspace = true
 hmac = "0.12.1"
-itertools = "0.14.0"
 litellm-http.workspace = true
 litellm-migrate.workspace = true
 litellm-storage-clickhouse.workspace = true
@@ -30,7 +28,6 @@ strum.workspace = true
 thiserror.workspace = true
 time = { workspace = true, features = ["formatting"] }
 tokio.workspace = true
-tracing.workspace = true
 url.workspace = true
 
 [dev-dependencies]

litellm-rust/crates/traces-clickhouse/src/error.rs

@@ -16,8 +16,6 @@ pub enum Error {
     InvalidResponse,
     #[error("ClickHouse insert exceeds the encoded size limit")]
     InsertTooLarge,
-    #[error("Trace exceeds the interactive read budget; use a filtered trace query")]
-    ReadTooLarge,
     #[error("ClickHouse schema setup failed with HTTP status {0}")]
     SchemaFailed(u16),
     #[error("ClickHouse schema setup transport failed")]
@@ -34,12 +32,6 @@ pub enum Error {
     ProvisionFailed(u16),
     #[error("ClickHouse reader provisioning transport failed")]
     ProvisionTransport,
-    #[error("Invalid {0} cursor")]
-    InvalidCursor(&'static str),
-    #[error("Multiple traces have this ID; provide trace_ref")]
-    AmbiguousTrace,
-    #[error("Trace changed while paging; refresh the trace to continue")]
-    TraceChanged,
     #[error(transparent)]
     Decode(#[from] litellm_traces::Error),
     #[error("trace ingestion task failed")]
@@ -49,12 +41,3 @@ pub enum Error {
     #[error(transparent)]
     Cached(#[from] std::sync::Arc<Error>),
 }
-
-impl From<litellm_traces_cache::Error> for Error {
-    fn from(error: litellm_traces_cache::Error) -> Self {
-        match error {
-            litellm_traces_cache::Error::Serialization(_) => Self::InvalidResponse,
-            litellm_traces_cache::Error::ReadTooLarge => Self::ReadTooLarge,
-        }
-    }
-}

litellm-rust/crates/traces-clickhouse/src/lib.rs

@@ -31,7 +31,7 @@ pub use litellm_storage_clickhouse::{Connection, Parameter};
 pub use litellm_traces::{QueryScope, ReadQuery};
 pub use query::{QueryHelp, execute_read, query_help, query_sql};
 pub use query_access::QueryReaders;
-pub use reads::{get_span, get_span_error, get_trace, get_trace_page, list_traces};
+pub use reads::ClickHouseTraces;
 pub use schema::{
     NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, ensure_schema, schema_statements,
 };

litellm-rust/crates/traces-clickhouse/src/reads.rs

@@ -1,26 +1,14 @@
-//! Scoped trace reads: the trace list, one trace resolved with its spend, and span payloads.
-
-use std::sync::LazyLock;
-use std::time::Duration;
-
-use base64::{Engine, engine::general_purpose::URL_SAFE};
-use futures_util::{StreamExt, TryStreamExt, stream};
-use itertools::Itertools;
 use litellm_http::Client;
+use litellm_storage_clickhouse::Error as StorageError;
 use litellm_storage_clickhouse::{Query, fetch};
-use litellm_traces::{
-    SpanDetail, SpanErrorPage, SpendLookup, Trace, TracePage, listed_summary,
-    query::named as contracts, resolve_trace, to_ui_content,
-};
-use litellm_traces_cache::{SnapshotCache, SnapshotKey};
-use serde::{Deserialize, Serialize};
+use litellm_traces::query::named as contracts;
+use litellm_traces_cache::{StoreError, TraceStore};
 
 use crate::{
     Connection, Error,
     query::named::{
-        ListTracesParams, ListTracesRow, ReadAccessParams, SpanDetail as SpanDetailQuery,
-        SpanDetailParams, SpanError, SpanErrorParams, SpendByResponseIdsParams, TraceIdentity,
-        TraceIdentityParams, TracePageSpansParams, TraceSpansParams,
+        ListTracesParams, ListTracesRow, SpanDetail as SpanDetailQuery, SpanError, SpanErrorParams,
+        SpendByResponseIdsParams, TraceIdentity, TracePageSpansParams,
     },
 };
 
@@ -36,501 +24,103 @@ impl Query for RunCandidates {
     );
 }
 
-// Cursor pages share a bounded snapshot so advancing does not resolve the whole graph again.
-static TRACE_SNAPSHOTS: LazyLock<SnapshotCache> = LazyLock::new(|| {
-    SnapshotCache::new(
-        crate::span_batches::MAX_GRAPH_BYTES,
-        Duration::from_secs(120),
-    )
-});
-
-const NANOS_PER_MS: i64 = 1_000_000;
-const SPEND_WINDOW_MS: i64 = 30 * 60 * 1000;
-const SPEND_CONCURRENCY: usize = 4;
-
-fn encode_cursor<T: Serialize>(position: &T) -> String {
-    URL_SAFE.encode(serde_json::to_vec(position).unwrap_or_default())
+pub struct ClickHouseTraces {
+    client: Client,
+    connection: Connection,
 }
 
-fn decode_cursor<T: for<'de> Deserialize<'de>>(
-    cursor: &str,
-    kind: &'static str,
-) -> Result<T, Error> {
-    URL_SAFE
-        .decode(cursor)
-        .ok()
-        .and_then(|json| serde_json::from_slice(&json).ok())
-        .ok_or(Error::InvalidCursor(kind))
-}
-
-fn trace_position(cursor: Option<&str>) -> Result<(i64, String), Error> {
-    let Some(cursor) = cursor.filter(|cursor| !cursor.is_empty()) else {
-        return Ok((0, String::new()));
-    };
-    match decode_cursor::<(i64, String)>(cursor, "trace")? {
-        (start_ms, trace_ref) if start_ms > 0 && !trace_ref.is_empty() => Ok((start_ms, trace_ref)),
-        _ => Err(Error::InvalidCursor("trace")),
+impl ClickHouseTraces {
+    pub fn new(client: Client, connection: Connection) -> Self {
+        Self { client, connection }
     }
 }
 
-#[derive(Deserialize, Serialize)]
-struct ErrorPosition {
-    offset: u64,
-    version: String,
-}
+impl TraceStore for ClickHouseTraces {
+    type Error = Error;
 
-fn error_position(cursor: Option<&str>) -> Result<Option<ErrorPosition>, Error> {
-    let Some(cursor) = cursor else {
-        return Ok(None);
-    };
-    let position = decode_cursor::<ErrorPosition>(cursor, "diagnostic")?;
-    let valid_version = position.version.len() == 64
-        && position
-            .version
-            .bytes()
-            .all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte));
-    if i64::try_from(position.offset).is_err() || !valid_version {
-        return Err(Error::InvalidCursor("diagnostic"));
+    fn source(&self) -> &str {
+        self.connection.url().as_str()
     }
-    Ok(Some(position))
-}
 
-/// The stored run a trace id names for this caller; ids can repeat across tenants and runs.
-async fn reference(
-    client: &Client,
-    connection: &Connection,
-    access: &ReadAccessParams,
-    trace_id: &str,
-    trace_ref: &str,
-) -> Result<Option<String>, Error> {
-    if !trace_ref.is_empty() {
-        return Ok(Some(trace_ref.to_owned()));
+    async fn trace_refs(
+        &self,
+        params: &contracts::TraceIdentityParams,
+    ) -> Result<Vec<String>, StoreError<Self::Error>> {
+        fetch::<TraceIdentity>(&self.client, &self.connection, params)
+            .await
+            .map(|rows| rows.into_iter().map(|row| row.trace_ref).collect())
+            .map_err(failed)
     }
-    let params = TraceIdentityParams {
-        access: access.clone(),
-        trace_id: trace_id.to_owned(),
-    };
-    let mut identities = fetch::<TraceIdentity>(client, connection, &params).await?;
-    if identities.len() > 1 {
-        return Err(Error::AmbiguousTrace);
-    }
-    Ok(identities.pop().map(|identity| identity.trace_ref))
-}
 
-/// Spend records behind the spans' calls. A failed lookup leaves cost unknown instead of failing
-/// the read.
-async fn spend(
-    client: &Client,
-    connection: &Connection,
-    access: &ReadAccessParams,
-    rows: &[contracts::TraceSpansRow],
-) -> Vec<contracts::SpendByResponseIdsRow> {
-    let lookup = SpendLookup::new(rows);
-    let (Some(start_ns), Some(end_ns)) = (
-        rows.iter().map(|row| row.start_ns).min(),
-        rows.iter()
-            .map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns))
-            .max(),
-    ) else {
-        return Vec::new();
-    };
-    if lookup.is_empty() {
-        return Vec::new();
-    }
-    let params = SpendByResponseIdsParams::from(contracts::SpendByResponseIdsParams {
-        access: access.clone(),
-        response_ids: lookup.response_ids,
-        request_ids: lookup.request_ids,
-        trace_ids: lookup.trace_ids,
-        start_ms: start_ns.div_euclid(NANOS_PER_MS) - SPEND_WINDOW_MS,
-        end_ms: end_ns.div_euclid(NANOS_PER_MS) + SPEND_WINDOW_MS,
-    });
-    match crate::span_batches::read_spend(client, connection, params).await {
-        Ok(rows) => rows,
-        Err(error) => {
-            tracing::warn!(%error, "trace spend lookup unavailable");
-            Vec::new()
+    async fn list_runs(
+        &self,
+        params: &contracts::ListTracesParams,
+    ) -> Result<Vec<contracts::ListTracesRow>, StoreError<Self::Error>> {
+        let storage_params = ListTracesParams::from(params.clone());
+        match fetch::<RunCandidates>(&self.client, &self.connection, &storage_params).await {
+            Ok(rows) => Ok(rows.into_iter().map(|row| row.0).collect()),
+            Err(StorageError::ResponseTooLarge) => Err(StoreError::TooLarge),
+            Err(error) => Err(failed(error)),
         }
     }
-}
 
-pub async fn list_traces(
-    client: &Client,
-    connection: &Connection,
-    access: &ReadAccessParams,
-    start_ms: i64,
-    end_ms: i64,
-    cursor: Option<&str>,
-    limit: u32,
-) -> Result<TracePage, Error> {
-    if limit == 0 {
-        return Err(Error::InvalidParameters);
+    async fn trace_spans(
+        &self,
+        params: &contracts::TraceSpansParams,
+        snapshot_ms: u64,
+    ) -> Result<Vec<contracts::TraceSpansRow>, StoreError<Self::Error>> {
+        crate::span_batches::read_spans(&self.client, &self.connection, params.clone(), snapshot_ms)
+            .await
     }
-    let (cursor_ms, cursor_trace_id) = trace_position(cursor)?;
-    let mut params = ListTracesParams::from(contracts::ListTracesParams {
-        access: access.clone(),
-        start_ms,
-        end_ms,
-        cursor_ms,
-        cursor_trace_id,
-        limit: limit.min(500),
-    });
-    let page: Vec<contracts::ListTracesRow> = loop {
-        match fetch::<RunCandidates>(client, connection, &params).await {
-            Err(litellm_storage_clickhouse::Error::ResponseTooLarge) if params.0.limit > 1 => {
-                params.0.limit /= 2;
-            }
-            Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
-                return Err(Error::ReadTooLarge);
-            }
-            result => break result?.into_iter().map(|row| row.0).collect(),
-        }
-    };
-    let next_cursor = page
-        .last()
-        .filter(|_| page.len() == params.0.limit as usize)
-        .map(|last| encode_cursor(&(last.start_ms, &last.trace_ref)));
-    let data = stream::iter(page.chunks(16))
-        .then(|batch| list_summaries(client, connection, access, batch))
-        .try_collect::<Vec<_>>()
-        .await?
-        .into_iter()
-        .flatten()
-        .collect();
-    Ok(TracePage { data, next_cursor })
-}
 
-async fn list_summaries(
-    client: &Client,
-    connection: &Connection,
-    access: &ReadAccessParams,
-    runs: &[contracts::ListTracesRow],
-) -> Result<Vec<litellm_traces::TraceSummary>, Error> {
-    let (Some(start_ms), Some(end_ms)) = (
-        runs.iter().map(|row| row.start_ms).min(),
-        runs.iter()
-            .map(|row| row.start_ms.saturating_add(row.duration_ms))
-            .max(),
-    ) else {
-        return Ok(Vec::new());
-    };
-    let params = TracePageSpansParams::from(contracts::TracePageSpansParams {
-        access: access.clone(),
-        trace_refs: runs.iter().map(|row| row.trace_ref.clone()).collect(),
-        start_ms,
-        end_ms: end_ms.saturating_add(1),
-    });
-    let spans = match crate::span_batches::read_list_spans(client, connection, params).await {
-        Ok(spans) => spans,
-        Err(Error::ReadTooLarge) => {
-            return stream::iter(runs)
-                .then(|row| async move {
-                    match get_trace(client, connection, access, &row.trace_id, &row.trace_ref).await
-                    {
-                        Ok(trace) => {
-                            Ok(trace.map_or_else(|| listed_summary(row), |trace| trace.summary))
-                        }
-                        Err(Error::ReadTooLarge) => Ok(listed_summary(row)),
-                        Err(error) => Err(error),
-                    }
-                })
-                .try_collect()
-                .await;
-        }
-        Err(error) => return Err(error),
-    };
-    let by_trace = spans.into_iter().into_group_map_by(|span| {
-        (
-            span.team_id.clone(),
-            span.api_key_hash.clone(),
-            span.trace_id.clone(),
+    async fn run_spans(
+        &self,
+        params: &contracts::TracePageSpansParams,
+        snapshot_ms: u64,
+    ) -> Result<Vec<contracts::TraceSpansRow>, StoreError<Self::Error>> {
+        crate::span_batches::read_list_spans(
+            &self.client,
+            &self.connection,
+            TracePageSpansParams::from(params.clone()),
+            snapshot_ms,
         )
-    });
-    let summaries = runs
-        .iter()
-        .map(|row| {
-            let spans = by_trace
-                .get(&(
-                    row.team_id.clone(),
-                    row.api_key_hash.clone(),
-                    row.trace_id.clone(),
-                ))
-                .map(Vec::as_slice)
-                .unwrap_or_default();
-            async move {
-                let spend_rows = spend(client, connection, access, spans).await;
-                resolve_trace(&row.trace_id, &row.trace_ref, spans, &spend_rows)
-                    .map_or_else(|| listed_summary(row), |trace| trace.summary)
-            }
-        })
-        .collect::<Vec<_>>();
-    Ok(stream::iter(summaries)
-        .buffered(SPEND_CONCURRENCY)
-        .collect()
-        .await)
-}
-
-pub async fn get_trace(
-    client: &Client,
-    connection: &Connection,
-    access: &ReadAccessParams,
-    trace_id: &str,
-    trace_ref: &str,
-) -> Result<Option<Trace>, Error> {
-    let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else {
-        return Ok(None);
-    };
-    let params = TraceSpansParams {
-        access: access.clone(),
-        trace_id: trace_id.to_owned(),
-        trace_ref: trace_ref.clone(),
-    };
-    let rows = crate::span_batches::read_spans(client, connection, params, u64::MAX).await?;
-    if rows.is_empty() {
-        return Ok(None);
+        .await
     }
-    let spend_rows = spend(client, connection, access, &rows).await;
-    Ok(resolve_trace(trace_id, &trace_ref, &rows, &spend_rows))
-}
 
-#[derive(Deserialize, Serialize)]
-struct SpanPosition {
-    trace_ref: String,
-    snapshot_ms: u64,
-    offset: usize,
-    version: String,
-}
-
-pub async fn get_trace_page(
-    client: &Client,
-    connection: &Connection,
-    access: &ReadAccessParams,
-    trace_id: &str,
-    trace_ref: &str,
-    cursor: Option<&str>,
-    page_size: u32,
-) -> Result<Option<Trace>, Error> {
-    if !(1..=500).contains(&page_size) {
-        return Err(Error::InvalidParameters);
+    async fn spend(
+        &self,
+        params: &contracts::SpendByResponseIdsParams,
+    ) -> Result<Vec<contracts::SpendByResponseIdsRow>, StoreError<Self::Error>> {
+        crate::span_batches::read_spend(
+            &self.client,
+            &self.connection,
+            SpendByResponseIdsParams::from(params.clone()),
+        )
+        .await
     }
-    let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else {
-        return Ok(None);
-    };
-    let position = match cursor {
-        Some(cursor) => {
-            let position: SpanPosition = decode_cursor(cursor, "span")?;
-            if position.trace_ref != trace_ref || position.snapshot_ms == 0 {
-                return Err(Error::InvalidCursor("span"));
-            }
-            position
+
+    async fn span_detail(
+        &self,
+        params: &contracts::SpanDetailParams,
+    ) -> Result<Option<contracts::SpanDetailRow>, StoreError<Self::Error>> {
+        match fetch::<SpanDetailQuery>(&self.client, &self.connection, params).await {
+            Ok(rows) => Ok(rows.into_iter().next()),
+            Err(error) => Err(failed(error)),
         }
-        None => SpanPosition {
-            trace_ref: trace_ref.clone(),
-            snapshot_ms: (time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000)
-                as u64,
-            offset: 0,
-            version: String::new(),
-        },
-    };
-    let key = SnapshotKey::new(
-        connection.url().as_str(),
-        access,
-        trace_id,
-        &trace_ref,
-        position.snapshot_ms,
-    )
-    .map_err(|_| Error::InvalidParameters)?;
-    let snapshot = match TRACE_SNAPSHOTS.get(&key).await {
-        Some(snapshot) => snapshot,
-        None => {
-            let params = TraceSpansParams {
-                access: access.clone(),
-                trace_id: trace_id.to_owned(),
-                trace_ref: trace_ref.clone(),
-            };
-            let rows =
-                crate::span_batches::read_spans(client, connection, params, position.snapshot_ms)
-                    .await?;
-            let spend_rows = spend(client, connection, access, &rows).await;
-            let Some(trace) = resolve_trace(trace_id, &trace_ref, &rows, &spend_rows) else {
-                return Ok(None);
-            };
-            TRACE_SNAPSHOTS.insert(key, trace).await?
+    }
+
+    async fn span_error(
+        &self,
+        params: &contracts::SpanErrorParams,
+    ) -> Result<Option<contracts::SpanErrorRow>, StoreError<Self::Error>> {
+        let storage_params = SpanErrorParams::from(params.clone());
+        match fetch::<SpanError>(&self.client, &self.connection, &storage_params).await {
+            Ok(rows) => Ok(rows.into_iter().next().map(|row| row.0)),
+            Err(error) => Err(failed(error)),
         }
-    };
-    let spans = &snapshot.trace().spans;
-    if cursor.is_some() && position.version != snapshot.version() {
-        return Err(Error::TraceChanged);
-    }
-    let mut trace = Trace {
-        summary: snapshot.trace().summary.clone(),
-        agents: snapshot.trace().agents.clone(),
-        spans: Vec::new(),
-        next_cursor: None,
-    };
-    if position.offset > spans.len() {
-        return Err(Error::InvalidCursor("span"));
-    }
-    let end = position
-        .offset
-        .saturating_add(page_size as usize)
-        .min(spans.len());
-    trace.next_cursor = (end < spans.len()).then(|| {
-        encode_cursor(&SpanPosition {
-            offset: end,
-            version: snapshot.version().to_owned(),
-            ..position
-        })
-    });
-    trace.spans = spans[position.offset..end].to_vec();
-    while serde_json::to_vec(&trace)
-        .map_err(|_| Error::InvalidResponse)?
-        .len()
-        > litellm_storage_clickhouse::READ_LIMITS.response_bytes
-    {
-        if trace.spans.len() <= 1 {
-            return Err(Error::ReadTooLarge);
-        }
-        trace.spans.truncate(trace.spans.len() / 2);
-        trace.next_cursor = Some(encode_cursor(&SpanPosition {
-            trace_ref: trace_ref.clone(),
-            snapshot_ms: position.snapshot_ms,
-            offset: position.offset + trace.spans.len(),
-            version: snapshot.version().to_owned(),
-        }));
-    }
-    Ok(Some(trace))
-}
-
-pub async fn get_span(
-    client: &Client,
-    connection: &Connection,
-    access: &ReadAccessParams,
-    trace_id: &str,
-    span_id: &str,
-    trace_ref: &str,
-) -> Result<Option<SpanDetail>, Error> {
-    let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else {
-        return Ok(None);
-    };
-    let params = SpanDetailParams {
-        access: access.clone(),
-        trace_id: trace_id.to_owned(),
-        trace_ref,
-        span_id: span_id.to_owned(),
-    };
-    let row = fetch::<SpanDetailQuery>(client, connection, &params)
-        .await?
-        .into_iter()
-        .next();
-    Ok(row.map(|row| SpanDetail {
-        input_ui: to_ui_content(&row.input),
-        output_ui: to_ui_content(&row.output),
-        span_id: row.span_id,
-        input: row.input,
-        output: row.output,
-        attributes: row.attributes,
-    }))
-}
-
-pub async fn get_span_error(
-    client: &Client,
-    connection: &Connection,
-    access: &ReadAccessParams,
-    trace_id: &str,
-    span_id: &str,
-    trace_ref: &str,
-    cursor: Option<&str>,
-) -> Result<Option<SpanErrorPage>, Error> {
-    let position = error_position(cursor)?;
-    let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else {
-        return Ok(None);
-    };
-    let offset = position.as_ref().map_or(0, |position| position.offset);
-    let params = SpanErrorParams::from(contracts::SpanErrorParams {
-        access: access.clone(),
-        trace_id: trace_id.to_owned(),
-        trace_ref,
-        span_id: span_id.to_owned(),
-        error_offset: offset,
-        error_version: position
-            .map(|position| position.version)
-            .unwrap_or_default(),
-    });
-    let Some(row) = fetch::<SpanError>(client, connection, &params)
-        .await?
-        .into_iter()
-        .next()
-    else {
-        return Ok(None);
-    };
-    let row = row.0;
-    let next_offset = offset + row.message.chars().count() as u64;
-    let next_cursor = (next_offset < row.total_chars).then(|| {
-        encode_cursor(&ErrorPosition {
-            offset: next_offset,
-            version: row.version,
-        })
-    });
-    Ok(Some(SpanErrorPage {
-        span_id: row.span_id,
-        message: row.message,
-        total_chars: row.total_chars,
-        next_cursor,
-    }))
-}
-
-#[cfg(test)]
-mod tests {
-    use rstest::rstest;
-
-    use super::*;
-
-    #[rstest]
-    fn trace_cursor_round_trips_the_last_listed_run() {
-        let cursor = encode_cursor(&(1_790_742_989_377_i64, "4bad42b84e9de3ba46fc870185f8f023"));
-        assert_eq!(
-            trace_position(Some(&cursor)).unwrap(),
-            (
-                1_790_742_989_377,
-                "4bad42b84e9de3ba46fc870185f8f023".to_owned()
-            )
-        );
-        assert_eq!(trace_position(None).unwrap(), (0, String::new()));
-        assert_eq!(trace_position(Some("")).unwrap(), (0, String::new()));
-    }
-
-    #[rstest]
-    #[case::not_base64("abc")]
-    #[case::not_json("bm90LWpzb24=")]
-    #[case::numeric_reference("WzEsIDJd")]
-    #[case::zero_start("WzAsICJ0Il0=")]
-    fn malformed_trace_cursors_are_rejected(#[case] cursor: &str) {
-        assert!(matches!(
-            trace_position(Some(cursor)),
-            Err(Error::InvalidCursor("trace"))
-        ));
-    }
-
-    #[rstest]
-    #[case::not_base64("garbage")]
-    #[case::missing_fields("e30=")]
-    #[case::not_an_object("WzEsMl0=")]
-    fn malformed_diagnostic_cursors_are_rejected(#[case] cursor: &str) {
-        assert!(matches!(
-            error_position(Some(cursor)),
-            Err(Error::InvalidCursor("diagnostic"))
-        ));
-    }
-
-    #[rstest]
-    #[case::lowercase_version("a".repeat(64))]
-    #[case::short_version("A".repeat(63))]
-    fn diagnostic_cursor_requires_a_content_version(#[case] version: String) {
-        let cursor = encode_cursor(&ErrorPosition { offset: 1, version });
-        assert!(matches!(
-            error_position(Some(&cursor)),
-            Err(Error::InvalidCursor("diagnostic"))
-        ));
     }
 }
+
+fn failed(error: StorageError) -> StoreError<Error> {
+    StoreError::Failed(Error::Storage(error))
+}

litellm-rust/crates/traces-clickhouse/src/span_batches.rs

@@ -1,15 +1,20 @@
-use futures_util::{TryStreamExt, stream};
-use itertools::Itertools;
+//! Keyset-paged reads that shrink their page when ClickHouse rejects a response as too large and
+//! stop accumulating once a graph exceeds the interactive budget.
+
+use std::{future::Future, marker::PhantomData};
+
 use litellm_http::Client;
 use litellm_storage_clickhouse::{Query, fetch};
 use litellm_traces::query::named as contracts;
-use serde::Serialize;
+use litellm_traces_cache::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, StoreError};
+use serde::{Serialize, de::DeserializeOwned};
 
-use crate::{Connection, Error, query::named::TraceSpansRow};
+use crate::{
+    Connection, Error,
+    query::named::{SpendByResponseIdsParams, SpendByResponseIdsRow, TraceSpansRow},
+};
 
 const PAGE_SIZE: u32 = 256;
-pub(crate) const MAX_GRAPH_BYTES: usize = 64 * 1024 * 1024;
-const MAX_GRAPH_SPANS: usize = 100_000;
 
 #[derive(Default)]
 struct ReadBudget {
@@ -18,40 +23,138 @@ struct ReadBudget {
 }
 
 impl ReadBudget {
-    fn checked_add(&self, bytes: usize, rows: usize) -> Result<Self, Error> {
-        let next = Self {
-            bytes: self.bytes.saturating_add(bytes),
-            rows: self.rows.saturating_add(rows),
-        };
-        if next.bytes > MAX_GRAPH_BYTES || next.rows > MAX_GRAPH_SPANS {
-            return Err(Error::ReadTooLarge);
+    fn reserve(&mut self, bytes: usize) -> Result<(), StoreError<Error>> {
+        self.bytes = self.bytes.saturating_add(bytes);
+        if self.bytes > MAX_GRAPH_BYTES || self.rows == MAX_GRAPH_SPANS {
+            return Err(StoreError::TooLarge);
         }
-        Ok(next)
-    }
-
-    fn record(&mut self, row: &impl Serialize) -> Result<(), Error> {
-        let bytes = serde_json::to_vec(row).map_err(|_| Error::InvalidResponse)?;
-        *self = self.checked_add(bytes.len(), 1)?;
+        self.rows += 1;
         Ok(())
     }
+
+    fn record(&mut self, row: &impl Serialize) -> Result<(), StoreError<Error>> {
+        let bytes =
+            serde_json::to_vec(row).map_err(|_| StoreError::Failed(Error::InvalidResponse))?;
+        self.reserve(bytes.len())
+    }
+}
+
+/// One keyset position in a paged query: the SQL reads the cursor fields of `Self` plus the
+/// `page_size` that [`Batch`] adds.
+trait Keyset: Serialize + Sized + Send + Sync {
+    type Row: Serialize + DeserializeOwned + Send;
+    const SQL: &'static str;
+
+    /// The position just after `last`.
+    fn after(self, last: &Self::Row) -> Self;
 }
 
 #[derive(Serialize)]
-struct Parameters {
+struct Batch<K> {
+    #[serde(flatten)]
+    keyset: K,
+    page_size: u32,
+}
+
+trait PageSource<K: Keyset> {
+    fn page(
+        &self,
+        batch: &Batch<K>,
+    ) -> impl Future<Output = Result<Vec<K::Row>, litellm_storage_clickhouse::Error>> + Send;
+}
+
+struct Paged<K>(PhantomData<K>);
+
+impl<K: Keyset> Query for Paged<K> {
+    type Params = Batch<K>;
+    type Row = K::Row;
+    const SQL: &'static str = K::SQL;
+}
+
+/// Reads every row after `keyset`. A page ClickHouse rejects as too large is retried at half the
+/// size, and the smaller page is kept for the rest of the read because row sizes within one graph
+/// rarely shrink again. Halving a one-row page means a single row exceeds the response limit.
+async fn read_all<K: Keyset, S: PageSource<K>>(
+    source: &S,
+    keyset: K,
+) -> Result<Vec<K::Row>, StoreError<Error>> {
+    let mut batch = Batch {
+        keyset,
+        page_size: PAGE_SIZE,
+    };
+    let mut rows = Vec::new();
+    let mut budget = ReadBudget::default();
+    loop {
+        let page = match source.page(&batch).await {
+            Err(litellm_storage_clickhouse::Error::ResponseTooLarge) if batch.page_size > 1 => {
+                batch.page_size /= 2;
+                continue;
+            }
+            Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
+                return Err(StoreError::TooLarge);
+            }
+            result => result.map_err(|error| StoreError::Failed(Error::Storage(error)))?,
+        };
+        let complete = page.len() < batch.page_size as usize;
+        for row in &page {
+            budget.record(row)?;
+        }
+        if let Some(last) = page.last() {
+            batch.keyset = batch.keyset.after(last);
+        }
+        rows.extend(page);
+        if complete {
+            return Ok(rows);
+        }
+    }
+}
+
+struct ClickHouse<'a> {
+    client: &'a Client,
+    connection: &'a Connection,
+}
+
+impl<K: Keyset> PageSource<K> for ClickHouse<'_> {
+    fn page(
+        &self,
+        batch: &Batch<K>,
+    ) -> impl Future<Output = Result<Vec<K::Row>, litellm_storage_clickhouse::Error>> + Send {
+        fetch::<Paged<K>>(self.client, self.connection, batch)
+    }
+}
+
+async fn read_paged<K: Keyset>(
+    client: &Client,
+    connection: &Connection,
+    keyset: K,
+) -> Result<Vec<K::Row>, StoreError<Error>> {
+    let source = ClickHouse { client, connection };
+    read_all(&source, keyset).await
+}
+
+fn by_start(mut rows: Vec<contracts::TraceSpansRow>) -> Vec<contracts::TraceSpansRow> {
+    rows.sort_by_key(|row| row.start_ns);
+    rows
+}
+
+#[derive(Serialize)]
+struct SpanKeyset {
     #[serde(flatten)]
     trace: contracts::TraceSpansParams,
     after_span_id: String,
-    page_size: u32,
     snapshot_ms: u64,
 }
 
-struct SpanBatch;
-
-impl Query for SpanBatch {
-    type Params = Parameters;
+impl Keyset for SpanKeyset {
     type Row = TraceSpansRow;
-
     const SQL: &'static str = include_str!("../query/trace_span_batch.sql");
+
+    fn after(self, last: &TraceSpansRow) -> Self {
+        Self {
+            after_span_id: last.0.span_id.clone(),
+            ..self
+        }
+    }
 }
 
 pub(crate) async fn read_spans(
@@ -59,199 +162,109 @@ pub(crate) async fn read_spans(
     connection: &Connection,
     trace: contracts::TraceSpansParams,
     snapshot_ms: u64,
-) -> Result<Vec<contracts::TraceSpansRow>, Error> {
-    let mut parameters = Parameters {
+) -> Result<Vec<contracts::TraceSpansRow>, StoreError<Error>> {
+    let keyset = SpanKeyset {
         trace,
         after_span_id: String::new(),
-        page_size: PAGE_SIZE,
         snapshot_ms,
     };
-    let mut spans = Vec::new();
-    let mut budget = ReadBudget::default();
-    loop {
-        let page = match fetch::<SpanBatch>(client, connection, &parameters).await {
-            Err(litellm_storage_clickhouse::Error::ResponseTooLarge)
-                if parameters.page_size > 1 =>
-            {
-                parameters.page_size /= 2;
-                continue;
-            }
-            Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
-                return Err(Error::ReadTooLarge);
-            }
-            result => result?,
-        };
-        let complete = page.len() < parameters.page_size as usize;
-        if let Some(last) = page.last() {
-            parameters.after_span_id.clone_from(&last.0.span_id);
-        }
-        for row in page {
-            budget.record(&row)?;
-            spans.push(row.0);
-        }
-        if complete {
-            spans.sort_by_key(|row| row.start_ns);
-            return Ok(spans);
-        }
-        parameters.page_size = (parameters.page_size * 2).min(PAGE_SIZE);
-    }
+    let rows = read_paged(client, connection, keyset).await?;
+    Ok(by_start(rows.into_iter().map(|row| row.0).collect()))
 }
 
 #[derive(Serialize)]
-struct ListParameters {
+struct ListSpanKeyset {
     #[serde(flatten)]
     runs: crate::query::named::TracePageSpansParams,
     after_team: String,
     after_key: String,
     after_trace: String,
     after_span: String,
-    page_size: u32,
     snapshot_ms: u64,
 }
 
-struct ListSpanBatch;
-
-impl Query for ListSpanBatch {
-    type Params = ListParameters;
+impl Keyset for ListSpanKeyset {
     type Row = TraceSpansRow;
-
     const SQL: &'static str = include_str!("../query/trace_list_span_batch.sql");
+
+    fn after(self, last: &TraceSpansRow) -> Self {
+        Self {
+            after_team: last.0.team_id.clone(),
+            after_key: last.0.api_key_hash.clone(),
+            after_trace: last.0.trace_id.clone(),
+            after_span: last.0.span_id.clone(),
+            ..self
+        }
+    }
 }
 
 pub(crate) async fn read_list_spans(
     client: &Client,
     connection: &Connection,
     runs: crate::query::named::TracePageSpansParams,
-) -> Result<Vec<contracts::TraceSpansRow>, Error> {
-    let parameters = ListParameters {
+    snapshot_ms: u64,
+) -> Result<Vec<contracts::TraceSpansRow>, StoreError<Error>> {
+    let keyset = ListSpanKeyset {
         runs,
         after_team: String::new(),
         after_key: String::new(),
         after_trace: String::new(),
         after_span: String::new(),
-        page_size: PAGE_SIZE,
-        snapshot_ms: (time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) as u64,
+        snapshot_ms,
     };
-    let pages = stream::try_unfold(
-        (Some(parameters), ReadBudget::default()),
-        |(parameters, budget)| async move {
-            let Some(parameters) = parameters else {
-                return Ok(None);
-            };
-            let page = match fetch::<ListSpanBatch>(client, connection, &parameters).await {
-                Err(litellm_storage_clickhouse::Error::ResponseTooLarge)
-                    if parameters.page_size > 1 =>
-                {
-                    let retry = ListParameters {
-                        page_size: parameters.page_size / 2,
-                        ..parameters
-                    };
-                    return Ok(Some((Vec::new(), (Some(retry), budget))));
-                }
-                Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
-                    return Err(Error::ReadTooLarge);
-                }
-                result => result?,
-            };
-            let next = page
-                .last()
-                .filter(|_| page.len() == parameters.page_size as usize)
-                .map(|last| ListParameters {
-                    after_team: last.0.team_id.clone(),
-                    after_key: last.0.api_key_hash.clone(),
-                    after_trace: last.0.trace_id.clone(),
-                    after_span: last.0.span_id.clone(),
-                    page_size: (parameters.page_size * 2).min(PAGE_SIZE),
-                    ..parameters
-                });
-            let next_budget = page.iter().try_fold(budget, |budget, row| {
-                let bytes = serde_json::to_vec(row).map_err(|_| Error::InvalidResponse)?;
-                budget.checked_add(bytes.len(), 1)
-            })?;
-            Ok(Some((page, (next, next_budget))))
-        },
-    )
-    .try_collect::<Vec<_>>()
-    .await?;
-    Ok(pages
-        .into_iter()
-        .flatten()
-        .map(|row| row.0)
-        .sorted_by_key(|row| row.start_ns)
-        .collect())
+    let rows = read_paged(client, connection, keyset).await?;
+    Ok(by_start(rows.into_iter().map(|row| row.0).collect()))
 }
 
 #[derive(Serialize)]
-struct SpendParameters {
+struct SpendKeyset {
     #[serde(flatten)]
-    lookup: crate::query::named::SpendByResponseIdsParams,
+    lookup: SpendByResponseIdsParams,
     has_cursor: u8,
     after_team: String,
     after_ms: i64,
     after_id: String,
-    page_size: u32,
 }
 
-struct SpendBatch;
-
-impl Query for SpendBatch {
-    type Params = SpendParameters;
-    type Row = crate::query::named::SpendByResponseIdsRow;
-
+impl Keyset for SpendKeyset {
+    type Row = SpendByResponseIdsRow;
     const SQL: &'static str = include_str!("../query/spend_batch.sql");
+
+    fn after(self, last: &SpendByResponseIdsRow) -> Self {
+        Self {
+            has_cursor: 1,
+            after_team: last.0.team_id.clone(),
+            after_ms: last.0.start_ms,
+            after_id: last.0.request_id.clone(),
+            ..self
+        }
+    }
 }
 
 pub(crate) async fn read_spend(
     client: &Client,
     connection: &Connection,
-    lookup: crate::query::named::SpendByResponseIdsParams,
-) -> Result<Vec<contracts::SpendByResponseIdsRow>, Error> {
-    let mut parameters = SpendParameters {
+    lookup: SpendByResponseIdsParams,
+) -> Result<Vec<contracts::SpendByResponseIdsRow>, StoreError<Error>> {
+    let keyset = SpendKeyset {
         lookup,
         has_cursor: 0,
         after_team: String::new(),
         after_ms: 0,
         after_id: String::new(),
-        page_size: PAGE_SIZE,
     };
-    let mut rows = Vec::new();
-    let mut budget = ReadBudget::default();
-    loop {
-        let page = match fetch::<SpendBatch>(client, connection, &parameters).await {
-            Err(litellm_storage_clickhouse::Error::ResponseTooLarge)
-                if parameters.page_size > 1 =>
-            {
-                parameters.page_size /= 2;
-                continue;
-            }
-            Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
-                return Err(Error::ReadTooLarge);
-            }
-            result => result?,
-        };
-        let complete = page.len() < parameters.page_size as usize;
-        if let Some(last) = page.last() {
-            parameters.has_cursor = 1;
-            parameters.after_team.clone_from(&last.0.team_id);
-            parameters.after_ms = last.0.start_ms;
-            parameters.after_id.clone_from(&last.0.request_id);
-        }
-        for row in page {
-            budget.record(&row)?;
-            rows.push(row.0);
-        }
-        if complete {
-            return Ok(rows);
-        }
-        parameters.page_size = (parameters.page_size * 2).min(PAGE_SIZE);
-    }
+    let rows = read_paged(client, connection, keyset).await?;
+    Ok(rows.into_iter().map(|row| row.0).collect())
 }
 
 #[cfg(test)]
 mod tests {
-    use super::*;
+    use std::sync::Mutex;
+
     use rstest::rstest;
 
+    use super::*;
+
     #[rstest]
     #[case::byte_boundary(MAX_GRAPH_BYTES - 1, 0, 1, false)]
     #[case::byte_overflow(MAX_GRAPH_BYTES - 1, 0, 2, true)]
@@ -264,7 +277,77 @@ mod tests {
         #[case] next: usize,
         #[case] rejected: bool,
     ) {
-        let budget = ReadBudget { bytes, rows };
-        assert_eq!(budget.checked_add(next, 1).is_err(), rejected);
+        let mut budget = ReadBudget { bytes, rows };
+        assert_eq!(budget.reserve(next).is_err(), rejected);
+    }
+
+    #[derive(Serialize)]
+    struct Numbers {
+        after: u32,
+    }
+
+    impl Keyset for Numbers {
+        type Row = u32;
+        const SQL: &'static str = "";
+
+        fn after(self, last: &u32) -> Self {
+            Self { after: *last }
+        }
+    }
+
+    /// A table of `total` rows whose transport rejects any page larger than `largest_page`.
+    struct Table {
+        total: u32,
+        largest_page: u32,
+        requests: Mutex<Vec<u32>>,
+    }
+
+    impl PageSource<Numbers> for Table {
+        async fn page(
+            &self,
+            batch: &Batch<Numbers>,
+        ) -> Result<Vec<u32>, litellm_storage_clickhouse::Error> {
+            self.requests.lock().unwrap().push(batch.page_size);
+            if batch.page_size > self.largest_page {
+                return Err(litellm_storage_clickhouse::Error::ResponseTooLarge);
+            }
+            let end = (batch.keyset.after + batch.page_size).min(self.total);
+            Ok((batch.keyset.after + 1..=end).collect())
+        }
+    }
+
+    #[rstest]
+    #[case::fits(1000, PAGE_SIZE, &[256, 256, 256, 256])]
+    #[case::uniform_large_rows(1000, 100, &[256, 128, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64])]
+    #[tokio::test]
+    async fn a_rejected_page_size_is_not_retried(
+        #[case] total: u32,
+        #[case] largest_page: u32,
+        #[case] requests: &[u32],
+    ) {
+        let table = Table {
+            total,
+            largest_page,
+            requests: Mutex::new(Vec::new()),
+        };
+        let rows = read_all(&table, Numbers { after: 0 }).await.unwrap();
+        assert_eq!(rows, (1..=total).collect::<Vec<_>>());
+        assert_eq!(table.requests.lock().unwrap().as_slice(), requests);
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn a_single_oversized_row_fails_the_read() {
+        let table = Table {
+            total: 10,
+            largest_page: 0,
+            requests: Mutex::new(Vec::new()),
+        };
+        let result = read_all(&table, Numbers { after: 0 }).await;
+        assert!(matches!(result, Err(StoreError::TooLarge)), "{result:?}");
+        assert_eq!(
+            table.requests.lock().unwrap().as_slice(),
+            &[256, 128, 64, 32, 16, 8, 4, 2, 1]
+        );
     }
 }

litellm-rust/crates/traces-clickhouse/tests/reads.rs

@@ -1,8 +1,10 @@
 use std::collections::BTreeMap;
 
+use litellm_http::Client;
 use litellm_traces::query::named::ReadAccessParams;
+use litellm_traces_cache::{ReadError, TraceReader};
 use litellm_traces_clickhouse::{
-    Connection, InsertTable, QueryScope, get_trace, get_trace_page, insert_rows, list_traces,
+    ClickHouseTraces, Connection, InsertTable, QueryScope, insert_rows,
 };
 use rstest::rstest;
 use serde_json::json;
@@ -14,6 +16,13 @@ mod support;
 use fixtures::{DATABASE, SeededDatabase, migrated_database, seeded_database};
 use support::TestResult;
 
+fn make_reader(client: &Client, connection: Connection) -> (TraceReader, ClickHouseTraces) {
+    (
+        TraceReader::new(litellm_storage_clickhouse::READ_LIMITS.response_bytes),
+        ClickHouseTraces::new(client.clone(), connection),
+    )
+}
+
 #[rstest]
 #[case::api_key("key-a", "")]
 #[case::user("", "user-a")]
@@ -74,16 +83,19 @@ async fn list_costs_match_each_run_when_response_ids_are_reused(
             .collect(),
     )
     .await?;
-    let reader = fixture
+    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: user_id.into(),
         team_ids: vec!["team-a".into()],
     };
-    let page = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 50).await?;
+    let page = reader
+        .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50)
+        .await?;
     assert_eq!(page.data.len(), runs.len());
     for (trace_id, _, cost) in runs {
         let summary = page
@@ -91,7 +103,8 @@ async fn list_costs_match_each_run_when_response_ids_are_reused(
             .iter()
             .find(|summary| summary.trace_id == trace_id)
             .ok_or("missing run")?;
-        let detail = get_trace(client, &reader, &access, trace_id, &summary.trace_ref)
+        let detail = reader
+            .get_trace(&store, &access, trace_id, &summary.trace_ref)
             .await?
             .ok_or("missing trace")?;
         assert_eq!(detail.summary.spend, Some(cost));
@@ -199,16 +212,19 @@ async fn large_runs_remain_complete_under_default_reader_limits(
             .collect::<Vec<_>>();
         insert_rows(client, &writer, DATABASE, InsertTable::SpendLogs, costs).await?;
     }
-    let reader = fixture
+    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 page = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 500).await?;
+    let page = reader
+        .list_traces(&store, &access, 0, 2_000_000_000_000, None, 500)
+        .await?;
     assert_eq!(page.data.len(), runs);
     assert!(
         page.data
@@ -222,42 +238,17 @@ async fn large_runs_remain_complete_under_default_reader_limits(
             .send()
             .await?
             .error_for_status()?;
-        let read_queries = client.post(writer.url().clone()).body(format!(
-            "SELECT count() FROM system.query_log WHERE type = 'QueryFinish' AND current_database = '{DATABASE}' AND query LIKE '%FROM otel_traces AS o%' AND query NOT LIKE '%system.query_log%'"
-        )).send().await?.error_for_status()?.text().await?;
-        let read_queries = read_queries.trim().parse::<usize>()?;
-        assert!(
-            read_queries > 0 && read_queries < runs,
-            "{read_queries} span queries for {runs} runs"
-        );
-        if costed {
-            let overlapping = client
-                .post(writer.url().clone())
-                .body(format!(
-                    "WITH spend_reads AS (
-                        SELECT query_start_time_microseconds AS started, event_time_microseconds AS finished
-                        FROM system.query_log
-                        WHERE type = 'QueryFinish' AND current_database = '{DATABASE}'
-                          AND query LIKE '%FROM spend_logs FINAL%' AND query NOT LIKE '%system.query_log%'
-                    ), events AS (
-                        SELECT started AS at, 1 AS delta FROM spend_reads
-                        UNION ALL SELECT finished AS at, -1 AS delta FROM spend_reads
-                    )
-                    SELECT max(active) FROM (
-                        SELECT sum(delta) OVER (ORDER BY at, delta ROWS UNBOUNDED PRECEDING) AS active
-                        FROM events
-                    )"
-                ))
-                .send()
-                .await?
-                .error_for_status()?
-                .text()
-                .await?
-                .trim()
-                .parse::<usize>()?;
+        for table in ["otel_traces AS o", "spend_logs FINAL"]
+            .into_iter()
+            .take(if costed { 2 } else { 1 })
+        {
+            let read_queries = client.post(writer.url().clone()).body(format!(
+                "SELECT count() FROM system.query_log WHERE type = 'QueryFinish' AND current_database = '{DATABASE}' AND query LIKE '%FROM {table}%' AND query NOT LIKE '%system.query_log%'"
+            )).send().await?.error_for_status()?.text().await?;
+            let read_queries = read_queries.trim().parse::<usize>()?;
             assert!(
-                (2..=4).contains(&overlapping),
-                "{overlapping} simultaneous spend reads for {runs} runs"
+                read_queries > 0 && read_queries < runs,
+                "{read_queries} {table} queries for {runs} runs"
             );
         }
     }
@@ -281,7 +272,8 @@ async fn large_runs_remain_complete_under_default_reader_limits(
         .find(|run| run.trace_id == "trace-0000")
         .ok_or("missing run")?
         .trace_ref;
-    let detail = get_trace(client, &reader, &access, "trace-0000", trace_ref)
+    let detail = reader
+        .get_trace(&store, &access, "trace-0000", trace_ref)
         .await?
         .ok_or("missing trace")?;
     assert_eq!(detail.spans.len(), steps);
@@ -303,24 +295,25 @@ async fn large_runs_remain_complete_under_default_reader_limits(
         ..access.clone()
     };
     assert!(
-        get_trace(client, &reader, &denied, "trace-0000", trace_ref)
+        reader
+            .get_trace(&store, &denied, "trace-0000", trace_ref)
             .await?
             .is_none()
     );
     let mut cursor = None;
     let mut ids = Vec::new();
     loop {
-        let page = get_trace_page(
-            client,
-            &reader,
-            &access,
-            "trace-0000",
-            trace_ref,
-            cursor.as_deref(),
-            200,
-        )
-        .await?
-        .ok_or("missing page")?;
+        let page = reader
+            .get_trace_page(
+                &store,
+                &access,
+                "trace-0000",
+                trace_ref,
+                cursor.as_deref(),
+                200,
+            )
+            .await?
+            .ok_or("missing page")?;
         assert_eq!(page.summary, detail.summary);
         assert!(page.spans.len() <= 200);
         assert!(
@@ -329,17 +322,17 @@ async fn large_runs_remain_complete_under_default_reader_limits(
         );
         if ids.is_empty() {
             assert!(
-                get_trace_page(
-                    client,
-                    &reader,
-                    &denied,
-                    "trace-0000",
-                    trace_ref,
-                    page.next_cursor.as_deref(),
-                    200,
-                )
-                .await?
-                .is_none()
+                reader
+                    .get_trace_page(
+                        &store,
+                        &denied,
+                        "trace-0000",
+                        trace_ref,
+                        page.next_cursor.as_deref(),
+                        200,
+                    )
+                    .await?
+                    .is_none()
             );
             client
                 .post(writer.url().clone())
@@ -372,45 +365,43 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive(
 ) -> TestResult {
     let fixture = seeded_database?;
     let client = &fixture.database.client;
-    let reader = fixture
+    let connection = fixture
         .readers
         .connection(client, &QueryScope::All, "fixture-secret")
         .await?;
+    let (reader, store) = make_reader(client, connection.clone());
     let access = ReadAccessParams {
         all_teams: true,
         user_id: String::new(),
         team_ids: Vec::new(),
     };
-    let listed = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 10).await?;
+    let listed = reader
+        .list_traces(&store, &access, 0, 2_000_000_000_000, None, 10)
+        .await?;
     let summary = listed
         .data
         .iter()
         .find(|summary| summary.span_count == 3)
         .ok_or("missing fixture")?;
-    let first = get_trace_page(
-        client,
-        &reader,
-        &access,
-        &summary.trace_id,
-        &summary.trace_ref,
-        None,
-        1,
-    )
-    .await?
-    .ok_or("missing first page")?;
-    let original_ids = get_trace(
-        client,
-        &reader,
-        &access,
-        &summary.trace_id,
-        &summary.trace_ref,
-    )
-    .await?
-    .ok_or("missing trace")?
-    .spans
-    .into_iter()
-    .map(|span| span.span_id)
-    .collect::<Vec<_>>();
+    let first = reader
+        .get_trace_page(
+            &store,
+            &access,
+            &summary.trace_id,
+            &summary.trace_ref,
+            None,
+            1,
+        )
+        .await?
+        .ok_or("missing first page")?;
+    let original_ids = reader
+        .get_trace(&store, &access, &summary.trace_id, &summary.trace_ref)
+        .await?
+        .ok_or("missing trace")?
+        .spans
+        .into_iter()
+        .map(|span| span.span_id)
+        .collect::<Vec<_>>();
     let writer = Connection::writer(&fixture.database.url)?;
     insert_rows(
         client,
@@ -434,17 +425,17 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive(
         team_ids: vec!["not-this-team".into()],
     };
     assert!(
-        get_trace_page(
-            client,
-            &reader,
-            &denied,
-            &summary.trace_id,
-            &summary.trace_ref,
-            first.next_cursor.as_deref(),
-            1
-        )
-        .await?
-        .is_none()
+        reader
+            .get_trace_page(
+                &store,
+                &denied,
+                &summary.trace_id,
+                &summary.trace_ref,
+                first.next_cursor.as_deref(),
+                1
+            )
+            .await?
+            .is_none()
     );
     let first_cursor = first.next_cursor.clone();
     let mut cursor = first.next_cursor;
@@ -454,44 +445,45 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive(
         .map(|span| span.span_id)
         .collect::<Vec<_>>();
     while let Some(current) = cursor {
-        let next = get_trace_page(
-            client,
-            &reader,
-            &access,
-            &summary.trace_id,
-            &summary.trace_ref,
-            Some(&current),
-            1,
-        )
-        .await?
-        .ok_or("missing next page")?;
+        let next = reader
+            .get_trace_page(
+                &store,
+                &access,
+                &summary.trace_id,
+                &summary.trace_ref,
+                Some(&current),
+                1,
+            )
+            .await?
+            .ok_or("missing next page")?;
         assert_eq!(next.summary.span_count, 3);
         ids.extend(next.spans.into_iter().map(|span| span.span_id));
         cursor = next.next_cursor;
     }
     assert_eq!(ids, original_ids);
-    let refreshed = get_trace(
-        client,
-        &reader,
-        &access,
-        &summary.trace_id,
-        &summary.trace_ref,
-    )
-    .await?
-    .ok_or("missing refreshed trace")?;
+    let cached = reader
+        .get_trace(&store, &access, &summary.trace_id, &summary.trace_ref)
+        .await?
+        .ok_or("missing cached trace")?;
+    assert_eq!(cached.spans.len(), 3);
+    let (fresh_reader, fresh_store) = make_reader(client, connection);
+    let refreshed = fresh_reader
+        .get_trace(&fresh_store, &access, &summary.trace_id, &summary.trace_ref)
+        .await?
+        .ok_or("missing refreshed trace")?;
     assert_eq!(refreshed.spans.len(), 4);
     assert!(matches!(
-        get_trace_page(
-            client,
-            &reader,
-            &access,
-            &summary.trace_id,
-            &summary.trace_ref,
-            Some("invalid"),
-            1
-        )
-        .await,
-        Err(litellm_traces_clickhouse::Error::InvalidCursor("span"))
+        reader
+            .get_trace_page(
+                &store,
+                &access,
+                &summary.trace_id,
+                &summary.trace_ref,
+                Some("invalid"),
+                1
+            )
+            .await,
+        Err(ReadError::InvalidCursor("span"))
     ));
     let backdated = json!({
         "Timestamp": "2026-09-01 00:00:00.000000000",
@@ -509,20 +501,21 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive(
         .send()
         .await?
         .error_for_status()?;
-    let uncached_reader =
+    let uncached_connection =
         Connection::reader(&format!("{}?max_threads=1", fixture.database.url), DATABASE)?;
-    let changed = get_trace_page(
-        client,
-        &uncached_reader,
-        &access,
-        &summary.trace_id,
-        &summary.trace_ref,
-        first_cursor.as_deref(),
-        1,
-    )
-    .await;
+    let (uncached_reader, uncached_store) = make_reader(client, uncached_connection);
+    let changed = uncached_reader
+        .get_trace_page(
+            &uncached_store,
+            &access,
+            &summary.trace_id,
+            &summary.trace_ref,
+            first_cursor.as_deref(),
+            1,
+        )
+        .await;
     assert!(
-        matches!(changed, Err(litellm_traces_clickhouse::Error::TraceChanged)),
+        matches!(changed, Err(ReadError::TraceChanged)),
         "{changed:?}"
     );
     Ok(())
@@ -535,16 +528,19 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals(
 ) -> TestResult {
     let fixture = seeded_database?;
     let client = &fixture.database.client;
-    let reader = fixture
+    let connection = fixture
         .readers
         .connection(client, &QueryScope::All, "fixture-secret")
         .await?;
+    let (reader, store) = make_reader(client, connection.clone());
     let access = ReadAccessParams {
         all_teams: true,
         user_id: String::new(),
         team_ids: Vec::new(),
     };
-    let before = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 50).await?;
+    let before = reader
+        .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50)
+        .await?;
     let run = before
         .data
         .iter()
@@ -571,7 +567,14 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals(
         ])],
     )
     .await?;
-    let after = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 50).await?;
+    let cached = reader
+        .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50)
+        .await?;
+    assert_eq!(cached.data, before.data);
+    let (reader, store) = make_reader(client, connection);
+    let after = reader
+        .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50)
+        .await?;
     assert_eq!(after.data.len(), before.data.len());
     let limited = after
         .data
@@ -588,17 +591,10 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals(
             .all(|item| !item.resolution_limited)
     );
     assert!(matches!(
-        get_trace_page(
-            client,
-            &reader,
-            &access,
-            &run.trace_id,
-            &run.trace_ref,
-            None,
-            200
-        )
-        .await,
-        Err(litellm_traces_clickhouse::Error::ReadTooLarge)
+        reader
+            .get_trace_page(&store, &access, &run.trace_id, &run.trace_ref, None, 200)
+            .await,
+        Err(ReadError::TooLarge)
     ));
     Ok(())
 }
@@ -685,16 +681,19 @@ async fn gateway_ids_resolve_through_detail_and_batch_reads_with_legacy_fallback
             .collect(),
     )
     .await?;
-    let reader = fixture
+    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 page = list_traces(client, &reader, &access, 0, 2_000_000_000_000, None, 50).await?;
+    let page = reader
+        .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50)
+        .await?;
     assert_eq!(page.data.len(), cases.len());
     for (id, _, _, _, _, expected) in cases {
         let summary = page
@@ -702,7 +701,8 @@ async fn gateway_ids_resolve_through_detail_and_batch_reads_with_legacy_fallback
             .iter()
             .find(|summary| summary.trace_id == id)
             .ok_or("missing run")?;
-        let detail = get_trace(client, &reader, &access, id, &summary.trace_ref)
+        let detail = reader
+            .get_trace(&store, &access, id, &summary.trace_ref)
             .await?
             .ok_or("missing trace")?;
         assert_eq!(detail.summary.spend, expected, "{id}");

litellm-rust/crates/traces/src/query/named.rs

@@ -15,7 +15,7 @@ pub struct ReadAccessParams {
     pub team_ids: Vec<String>,
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct ListTracesParams {
     #[serde(flatten)]
     pub access: ReadAccessParams,
@@ -26,7 +26,7 @@ pub struct ListTracesParams {
     pub limit: u32,
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct ListTracesRow {
     pub trace_id: String,
     pub trace_ref: String,
@@ -56,7 +56,7 @@ pub struct ListTracesRow {
     pub request_ids: Vec<String>,
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct TraceSpansParams {
     #[serde(flatten)]
     pub access: ReadAccessParams,
@@ -113,7 +113,7 @@ pub struct TraceSpansRow {
     pub user_id: String,
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct TracePageSpansParams {
     #[serde(flatten)]
     pub access: ReadAccessParams,
@@ -131,7 +131,7 @@ pub struct SpanDetailParams {
     pub span_id: String,
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct SpanDetailRow {
     pub span_id: String,
     pub input: String,
@@ -139,7 +139,7 @@ pub struct SpanDetailRow {
     pub attributes: BTreeMap<String, String>,
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct SpanErrorParams {
     #[serde(flatten)]
     pub access: ReadAccessParams,
@@ -150,7 +150,7 @@ pub struct SpanErrorParams {
     pub error_version: String,
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct SpanErrorRow {
     pub span_id: String,
     pub message: String,
@@ -158,7 +158,7 @@ pub struct SpanErrorRow {
     pub version: String,
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct SpendByResponseIdsParams {
     #[serde(flatten)]
     pub access: ReadAccessParams,
@@ -169,7 +169,7 @@ pub struct SpendByResponseIdsParams {
     pub end_ms: i64,
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct SpendByResponseIdsRow {
     pub request_id: String,
     pub litellm_call_id: String,

litellm/constants.py

@@ -55,6 +55,7 @@ DEFAULT_AGENT_TRACING_RETENTION_DAYS: Final = 14
 OTLP_MAX_BODY_BYTES: Final = get_env_int("OTLP_MAX_BODY_BYTES", 16 * 1024 * 1024)
 OTLP_MAX_ATTRIBUTE_VALUE_BYTES: Final = get_env_int("OTLP_MAX_ATTRIBUTE_VALUE_BYTES", 64 * 1024)
 OTLP_RETRY_AFTER_SECONDS: Final = get_env_int("OTLP_RETRY_AFTER_SECONDS", 2)
+TRACE_READ_RETRY_AFTER_SECONDS: Final = get_env_int("TRACE_READ_RETRY_AFTER_SECONDS", 2)
 OTLP_MAX_CONCURRENT_INGESTS: Final = get_env_int("OTLP_MAX_CONCURRENT_INGESTS", 2)
 AGENT_TRACING_INPUT_PREVIEW_CHARS: Final = get_env_int("AGENT_TRACING_INPUT_PREVIEW_CHARS", 240)
 AGENT_TRACING_LIST_PAGE_SIZE: Final = get_env_int("AGENT_TRACING_LIST_PAGE_SIZE", 50)

litellm/proxy/tracing_endpoints.py

@@ -13,19 +13,21 @@ from dataclasses import dataclass
 from functools import partial
 from http.client import responses
 from types import MappingProxyType
-from typing import Annotated, Final
+from typing import Annotated, Final, Literal
 
 from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response
 from pydantic import BaseModel, ConfigDict
+from typing_extensions import assert_never
 
 from litellm._logging import verbose_proxy_logger
-from litellm.constants import OTLP_RETRY_AFTER_SECONDS
+from litellm.constants import OTLP_RETRY_AFTER_SECONDS, TRACE_READ_RETRY_AFTER_SECONDS
 from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.auth.authorization import AllRows, ReadScope, resolve_trace_read_scope
 from litellm.proxy.auth.authorization_dependencies import LogTeamLookupDependency
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_utils.http_parsing_utils import is_otlp_trace_request
 from litellm.proxy.tracing_runtime import provide_receiver, require_receiver
+from litellm.rust_bridge.trace.errors import TraceChanged
 from litellm.rust_bridge.trace.generated.models import TraceQueryHelp
 from litellm.rust_bridge.trace.generated.types import (
     AllQueryScope,
@@ -135,6 +137,49 @@ async def ingest_otlp_traces(
     return Response(content=body, media_type=media_type)
 
 
+class TraceReadFailure(BaseModel):
+    """The body of every failed trace read. Clients branch on `code`, never on `message`."""
+
+    model_config = ConfigDict(frozen=True)
+
+    code: Literal["invalid_request", "trace_changed", "too_large", "unavailable"]
+    message: str
+
+
+def read_failure(error: TraceChanged | ValueError | OverflowError | RuntimeError) -> HTTPException:
+    """One status per failure kind, so a client can tell a bad cursor (400, fix the request) from a
+    traversal it must restart (409), a result it cannot page through (413), and an outage it should
+    retry after `Retry-After` (503)."""
+    match error:
+        case TraceChanged():
+            return HTTPException(
+                status_code=409,
+                detail=TraceReadFailure(code="trace_changed", message=str(error)).model_dump(),
+            )
+        case ValueError():
+            return HTTPException(
+                status_code=400, detail=TraceReadFailure(code="invalid_request", message=str(error)).model_dump()
+            )
+        case OverflowError():
+            return HTTPException(
+                status_code=413,
+                detail=TraceReadFailure(
+                    code="too_large", message="Trace is too large for this view. Use a filtered trace query."
+                ).model_dump(),
+            )
+        case RuntimeError():
+            verbose_proxy_logger.warning("Trace read unavailable: %s", error)
+            return HTTPException(
+                status_code=503,
+                detail=TraceReadFailure(
+                    code="unavailable", message="Traces are temporarily unavailable. Please try again."
+                ).model_dump(),
+                headers={"Retry-After": str(TRACE_READ_RETRY_AFTER_SECONDS)},
+            )
+        case _:
+            assert_never(error)
+
+
 @router.get("/v1/traces", response_model=TracePage)
 async def list_agent_traces(
     context: Annotated[TraceAccessContext, Depends(provide_trace_access)],
@@ -151,15 +196,8 @@ async def list_agent_traces(
             end_ms=end_ms if end_ms is not None else now_ms,
             cursor=cursor,
         )
-    except ValueError as error:
-        raise HTTPException(status_code=400, detail=str(error)) from error
-    except OverflowError as error:
-        raise HTTPException(
-            status_code=413, detail="Trace is too large for this view. Use a filtered trace query."
-        ) from error
-    except RuntimeError as error:
-        verbose_proxy_logger.warning("Trace read unavailable: %s", error)
-        raise HTTPException(status_code=503, detail="Traces are temporarily unavailable. Please try again.") from error
+    except (TraceChanged, ValueError, OverflowError, RuntimeError) as error:
+        raise read_failure(error) from error
 
 
 class TraceQueryRequest(BaseModel):
@@ -241,15 +279,8 @@ async def get_agent_trace(
     tracing, scope = context.reader()
     try:
         trace: Final = await tracing.get_trace(trace_id, scope, trace_ref, cursor, page_size)
-    except ValueError as error:
-        raise HTTPException(status_code=400, detail=str(error)) from error
-    except OverflowError as error:
-        raise HTTPException(
-            status_code=413, detail="Trace is too large for this view. Use a filtered trace query."
-        ) from error
-    except RuntimeError as error:
-        verbose_proxy_logger.warning("Trace read unavailable: %s", error)
-        raise HTTPException(status_code=503, detail="Traces are temporarily unavailable. Please try again.") from error
+    except (TraceChanged, ValueError, OverflowError, RuntimeError) as error:
+        raise read_failure(error) from error
     if trace is None:
         raise HTTPException(status_code=404, detail=f"Trace {trace_id} not found")
     return trace
@@ -265,15 +296,8 @@ async def get_agent_trace_span(
     tracing, scope = context.reader()
     try:
         span: Final = await tracing.get_span(trace_id, span_id, scope, trace_ref)
-    except ValueError as error:
-        raise HTTPException(status_code=400, detail=str(error)) from error
-    except OverflowError as error:
-        raise HTTPException(
-            status_code=413, detail="Trace is too large for this view. Use a filtered trace query."
-        ) from error
-    except RuntimeError as error:
-        verbose_proxy_logger.warning("Trace read unavailable: %s", error)
-        raise HTTPException(status_code=503, detail="Traces are temporarily unavailable. Please try again.") from error
+    except (TraceChanged, ValueError, OverflowError, RuntimeError) as error:
+        raise read_failure(error) from error
     if span is None:
         raise HTTPException(status_code=404, detail=f"Span {span_id} not found")
     return span
@@ -290,15 +314,8 @@ async def get_agent_trace_span_error(
     try:
         tracing, scope = context.reader()
         page: Final = await tracing.get_span_error(trace_id, span_id, scope, trace_ref, cursor)
-    except ValueError as error:
-        raise HTTPException(status_code=400, detail=str(error)) from error
-    except OverflowError as error:
-        raise HTTPException(
-            status_code=413, detail="Trace is too large for this view. Use a filtered trace query."
-        ) from error
-    except RuntimeError as error:
-        verbose_proxy_logger.warning("Trace read unavailable: %s", error)
-        raise HTTPException(status_code=503, detail="Traces are temporarily unavailable. Please try again.") from error
+    except (TraceChanged, ValueError, OverflowError, RuntimeError) as error:
+        raise read_failure(error) from error
     if page is None:
         raise HTTPException(status_code=404, detail="Span diagnostic not found or no longer available")
     return page

litellm/rust_bridge/trace/errors.py

@@ -0,0 +1,5 @@
+class TraceChanged(Exception):
+    """The paging snapshot no longer matches the stored trace, so the client must start a new traversal.
+
+    Raised by the Rust trace reader when a cursor's snapshot version differs from the graph it rebuilt.
+    """

tests/unit/proxy/test_tracing_endpoints.py

@@ -12,6 +12,7 @@ import pytest
 from fastapi import FastAPI, HTTPException
 from fastapi.testclient import TestClient
 
+from litellm.constants import TRACE_READ_RETRY_AFTER_SECONDS
 from litellm.proxy import tracing_endpoints
 from litellm.proxy._types import LitellmUserRoles, ProxyLifespanState, UserAPIKeyAuth
 from litellm.proxy.auth.authorization import OwnedRows, ReadScope
@@ -19,6 +20,7 @@ from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.tracing_runtime import manage_tracing, provide_storage
 from litellm.rust_bridge import loader
+from litellm.rust_bridge.trace.errors import TraceChanged
 from litellm.rust_bridge.trace.generated.models import TraceQueryHelp
 from litellm.rust_bridge.trace.generated.types import AllQueryScope, TraceScope
 from litellm.rust_bridge.trace.queries import TraceSQLResponse
@@ -297,19 +299,45 @@ def test_trace_detail_passes_scoped_reference(client, receiver, suffix, cursor,
     ),
 )
 @pytest.mark.parametrize(
-    "error,status,message",
+    "error,status,code,message",
     (
-        (RuntimeError("private database details"), 503, "Traces are temporarily unavailable. Please try again."),
-        (OverflowError("private query details"), 413, "Trace is too large for this view. Use a filtered trace query."),
+        (
+            RuntimeError("private database details"),
+            503,
+            "unavailable",
+            "Traces are temporarily unavailable. Please try again.",
+        ),
+        (
+            OverflowError("private query details"),
+            413,
+            "too_large",
+            "Trace is too large for this view. Use a filtered trace query.",
+        ),
+        (
+            TraceChanged("Trace changed while paging; refresh the trace to continue"),
+            409,
+            "trace_changed",
+            "Trace changed while paging; refresh the trace to continue",
+        ),
+        (ValueError("Invalid span cursor"), 400, "invalid_request", "Invalid span cursor"),
     ),
 )
-def test_read_failures_are_actionable_without_exposing_database_details(
-    client: TestClient, receiver: MagicMock, path: str, method: str, error: Exception, status: int, message: str
+def test_read_failures_carry_a_code_per_kind_without_exposing_database_details(
+    client: TestClient,
+    receiver: MagicMock,
+    path: str,
+    method: str,
+    error: Exception,
+    status: int,
+    code: str,
+    message: str,
 ) -> None:
     getattr(receiver, method).side_effect = error
     response: Final = client.get(path)
     assert response.status_code == status
-    assert response.json() == {"detail": message}
+    assert response.json() == {"detail": {"code": code, "message": message}}
+    retry_after: Final = response.headers.get("Retry-After")
+    assert (retry_after == str(TRACE_READ_RETRY_AFTER_SECONDS)) == (status == 503), retry_after
 
 
 @pytest.mark.parametrize("query", ("page_size=0", "page_size=501", "cursor=" + "x" * 513))

ui/litellm-dashboard/src/components/lens/traces/detail/DetailContent.tsx

@@ -10,6 +10,12 @@ import { cn } from "@/lib/cva.config";
 
 import { type KeyValue, KeyValueRows, objectEntries } from "./KeyValueRows";
 import { Card, MessageCard, Section, ToolResultCard } from "./MessageCard";
+import {
+  classifyTraceReadFailure,
+  isRetryableTraceRead,
+  traceReadRetry,
+  traceReadRetryDelay,
+} from "../list/traceReadFailure";
 import type { ErrorSource } from "../tree";
 import type { Span, SpanDetail, SpanErrorPage, TraceMessage, UIContent, UIMessage } from "../types";
 import { errorSource, parseJson, parseMessages, prettyPayload } from "../utils";
@@ -155,9 +161,11 @@ function DiagnosticContent({ accessToken, traceId, traceRef, span }: DetailConte
     enabled: opened,
     staleTime: Infinity,
     gcTime: 0,
-    retry: false,
+    retry: traceReadRetry,
+    retryDelay: traceReadRetryDelay,
   };
   const query = useQuery(queryOptions);
+  const failure = query.error ? classifyTraceReadFailure(query.error) : null;
   return (
     <section aria-label="Stored diagnostic" className="mx-3 mb-3 space-y-2">
       {span.error_truncated && <p className="text-xs text-muted-foreground">Error preview truncated</p>}
@@ -167,12 +175,18 @@ function DiagnosticContent({ accessToken, traceId, traceRef, span }: DetailConte
         </Button>
       )}
       {opened && query.isPending && <p role="status">Loading diagnostic…</p>}
-      {opened && query.isError && (
+      {failure && (
         <div role="alert">
-          Could not load diagnostic: {query.error.message}
-          <Button variant="outline" size="sm" onClick={() => query.refetch()}>
-            Retry
-          </Button>
+          Could not load diagnostic: {failure.message}
+          {isRetryableTraceRead(failure) ? (
+            <Button variant="outline" size="sm" onClick={() => query.refetch()}>
+              Retry
+            </Button>
+          ) : (
+            <Button variant="outline" size="sm" onClick={() => setCursor(null)}>
+              Back to beginning
+            </Button>
+          )}
         </div>
       )}
       {opened && query.data && (

ui/litellm-dashboard/src/components/lens/traces/detail/TraceDrawer.test.tsx

@@ -3,6 +3,8 @@ import { focusManager, onlineManager } from "@tanstack/react-query";
 import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
+import { ApiError } from "@/lib/http/client";
+
 import { renderWithProviders, testQueryClient } from "../../../../../tests/test-utils";
 import researchTrace from "../__fixtures__/research_trace.json";
 import swarmTrace from "../__fixtures__/swarm_trace.json";
@@ -260,14 +262,21 @@ describe("RunView", () => {
     vi.mocked(agentTraceCall)
       .mockResolvedValueOnce(first)
       .mockResolvedValueOnce(second)
-      .mockRejectedValueOnce(new Error("Trace changed while paging; refresh the trace"))
+      .mockRejectedValueOnce(
+        new ApiError("Trace changed while paging", 409, {
+          detail: { code: "trace_changed", message: "Trace changed while paging" },
+        }),
+      )
       .mockResolvedValueOnce(fresh)
       .mockResolvedValueOnce(freshSecond);
     renderWithProviders(<RoutedRunView traceId={research.summary.trace_id} accessToken="sk-test" onBack={vi.fn()} />);
     await user.click(await screen.findByRole("button", { name: "Load more steps" }));
     expect(await screen.findByText("old-snapshot-tool")).toBeVisible();
     await user.click(screen.getByRole("button", { name: "Load more steps" }));
-    await user.click(await screen.findByRole("button", { name: "Refresh trace" }));
+    expect(await screen.findByText(/This trace changed while you were browsing/)).toBeVisible();
+    expect(screen.queryByRole("button", { name: "Retry" })).not.toBeInTheDocument();
+    expect(screen.getByText("old-snapshot-tool")).toBeVisible();
+    await user.click(screen.getByRole("button", { name: "Refresh trace" }));
     expect(await screen.findByText("Showing 1 of 2 steps")).toBeVisible();
     expect(screen.queryByText("old-snapshot-tool")).not.toBeInTheDocument();
     await user.click(screen.getByRole("button", { name: "Load more steps" }));
@@ -282,6 +291,55 @@ describe("RunView", () => {
     ]);
   });
 
+  it("retries an unavailable page by itself after Retry-After and keeps loaded steps", async () => {
+    const user = userEvent.setup();
+    const summary = { ...research.summary, span_count: 2 };
+    const first: Trace = { ...research, summary, spans: research.spans.slice(0, 1), next_cursor: "next-page" };
+    const second: Trace = {
+      ...research,
+      summary,
+      spans: [
+        { ...research.spans[1], type: "tool", name: "later-page-tool", parent_span_id: research.spans[0].span_id },
+      ],
+      next_cursor: null,
+    };
+    const outage = new ApiError(
+      "Traces are temporarily unavailable",
+      503,
+      { detail: { code: "unavailable", message: "Traces are temporarily unavailable" } },
+      0,
+    );
+    vi.mocked(agentTraceCall).mockReset();
+    vi.mocked(agentTraceCall).mockResolvedValueOnce(first).mockRejectedValueOnce(outage).mockResolvedValueOnce(second);
+    renderWithProviders(<RoutedRunView traceId={research.summary.trace_id} accessToken="sk-test" onBack={vi.fn()} />);
+
+    await user.click(await screen.findByRole("button", { name: "Load more steps" }));
+    expect(await screen.findByText("later-page-tool")).toBeVisible();
+    expect(screen.getAllByRole("treeitem")).toHaveLength(2);
+    expect(screen.queryByRole("button", { name: "Refresh trace" })).not.toBeInTheDocument();
+    expect(vi.mocked(agentTraceCall).mock.calls.map((call) => call[3])).toEqual([null, "next-page", "next-page"]);
+  });
+
+  it("offers no retry for a page that is too large, only a refresh", async () => {
+    const user = userEvent.setup();
+    const summary = { ...research.summary, span_count: 2 };
+    const first: Trace = { ...research, summary, spans: research.spans.slice(0, 1), next_cursor: "next-page" };
+    vi.mocked(agentTraceCall).mockReset();
+    vi.mocked(agentTraceCall)
+      .mockResolvedValueOnce(first)
+      .mockRejectedValueOnce(
+        new ApiError("Trace is too large", 413, { detail: { code: "too_large", message: "Trace is too large" } }),
+      );
+    renderWithProviders(<RoutedRunView traceId={research.summary.trace_id} accessToken="sk-test" onBack={vi.fn()} />);
+
+    await user.click(await screen.findByRole("button", { name: "Load more steps" }));
+    expect(await screen.findByText(/too large to load here/)).toBeVisible();
+    expect(screen.queryByRole("button", { name: "Retry" })).not.toBeInTheDocument();
+    expect(screen.getByRole("button", { name: "Refresh trace" })).toBeVisible();
+    expect(screen.getByRole("tree", { name: "Spans in time order" })).toHaveTextContent(research.spans[0].name);
+    expect(vi.mocked(agentTraceCall)).toHaveBeenCalledTimes(2);
+  });
+
   it("keeps a loaded snapshot on focus and reconnect", async () => {
     testQueryClient.setQueryDefaults(["agentTrace"], { refetchOnWindowFocus: true, refetchOnReconnect: true });
     vi.mocked(agentTraceCall).mockReset();

ui/litellm-dashboard/src/components/lens/traces/detail/TraceDrawer.tsx

@@ -23,6 +23,13 @@ import { TraceConversation } from "./TraceConversation";
 import { FrameworkLogo, traceFramework } from "../ui/TraceFramework";
 import { type RunSelection, traceKey } from "../routing";
 import type { SpanTreeState, TreeRow } from "../tree";
+import {
+  classifyTraceReadFailure,
+  isRetryableTraceRead,
+  traceReadRetry,
+  traceReadRetryDelay,
+} from "../list/traceReadFailure";
+import type { TraceReadFailure } from "../list/traceReadFailure";
 import type { Trace } from "../types";
 import {
   buildTreeRows,
@@ -340,6 +347,60 @@ interface RunViewProps {
   embedded?: boolean;
 }
 
+interface PagingBannerProps {
+  readonly loaded: number;
+  readonly total: number;
+  readonly failure: TraceReadFailure | null;
+  readonly busy: boolean;
+  readonly onLoadMore: () => void;
+  readonly onRefresh: () => void;
+}
+
+/**
+ * Progress through the steps, or what the loaded steps are still worth after a later page failed.
+ * Retry keeps the loaded pages and asks for the same page again; Refresh starts a new traversal.
+ * Only a temporary failure gets a Retry, because a changed, invalid, or oversized read fails the
+ * same way every time.
+ */
+function PagingBanner({ loaded, total, failure, busy, onLoadMore, onRefresh }: PagingBannerProps) {
+  const retryable = failure !== null && isRetryableTraceRead(failure);
+  const nextLabel = failure ? "Retry" : "Load more steps";
+  return (
+    <div className="flex items-center justify-between gap-3 border-b px-3 py-2 text-xs" role="status">
+      <span>
+        {failure
+          ? pageFailureMessage(failure)
+          : `Showing ${loaded.toLocaleString()} of ${total.toLocaleString()} steps`}
+      </span>
+      {failure && (
+        <Button size="xs" variant={retryable ? "ghost" : "outline"} disabled={busy} onClick={onRefresh}>
+          Refresh trace
+        </Button>
+      )}
+      {(!failure || retryable) && (
+        <Button size="xs" variant="outline" disabled={busy} onClick={onLoadMore}>
+          {busy ? "Loading…" : nextLabel}
+        </Button>
+      )}
+    </div>
+  );
+}
+
+/** What the loaded steps are still worth after a later page fails, and what the user can do about it. */
+export function pageFailureMessage(failure: TraceReadFailure): string {
+  switch (failure.kind) {
+    case "changed":
+      return "This trace changed while you were browsing. Refresh to continue from its latest steps.";
+    case "invalid":
+      return "This page could not be continued. Refresh the trace to start over.";
+    case "too_large":
+      return "A step on the next page is too large to load here. Your loaded steps are still available.";
+    case "unavailable":
+    case "unknown":
+      return "Could not load more steps. Your loaded steps are still available.";
+  }
+}
+
 function selectedSpanMissing(trace: Trace, spanId: string | null): boolean {
   return Boolean(spanId && !trace.spans.some((span) => span.span_id === spanId));
 }
@@ -427,10 +488,12 @@ function LoadedRun({
     refetchOnWindowFocus: false,
     refetchOnReconnect: false,
     refetchOnMount: false,
-    retry: false,
+    retry: traceReadRetry,
+    retryDelay: traceReadRetryDelay,
   };
   const traceQuery = useSuspenseInfiniteQuery(traceQueryOptions);
   const refreshTrace = () => queryClient.resetQueries({ queryKey, exact: true });
+  const failure = traceQuery.error ? classifyTraceReadFailure(traceQuery.error) : null;
   const trace = useMemo(() => {
     const [first, ...rest] = traceQuery.data.pages;
     return { ...first, spans: [first, ...rest].flatMap((page) => page.spans) };
@@ -441,7 +504,6 @@ function LoadedRun({
   useEffect(() => {
     if (canSeek && !isFetching && !isError) void fetchNextPage();
   }, [canSeek, isFetching, isError, fetchNextPage]);
-  const pageAction = isError ? "Retry" : "Load more steps";
 
   return (
     <Tabs
@@ -462,27 +524,15 @@ function LoadedRun({
         onBack={onBack}
         embedded={embedded}
       />
-      {(traceQuery.hasNextPage || traceQuery.isError) && (
-        <div className="flex items-center justify-between gap-3 border-b px-3 py-2 text-xs" role="status">
-          <span>
-            {traceQuery.isError
-              ? "Could not load more steps. Your loaded steps are still available."
-              : `Showing ${trace.spans.length.toLocaleString()} of ${trace.summary.span_count.toLocaleString()} steps`}
-          </span>
-          {traceQuery.isError && (
-            <Button size="xs" variant="ghost" disabled={traceQuery.isFetching} onClick={() => void refreshTrace()}>
-              Refresh trace
-            </Button>
-          )}
-          <Button
-            size="xs"
-            variant="outline"
-            disabled={traceQuery.isFetching}
-            onClick={() => void (traceQuery.hasNextPage ? traceQuery.fetchNextPage() : refreshTrace())}
-          >
-            {traceQuery.isFetching ? "Loading…" : pageAction}
-          </Button>
-        </div>
+      {(traceQuery.hasNextPage || failure) && (
+        <PagingBanner
+          loaded={trace.spans.length}
+          total={trace.summary.span_count}
+          failure={failure}
+          busy={traceQuery.isFetching}
+          onLoadMore={() => void traceQuery.fetchNextPage()}
+          onRefresh={() => void refreshTrace()}
+        />
       )}
       <RunBody
         key={seekingSpan ? "seeking" : "loaded"}

ui/litellm-dashboard/src/components/lens/traces/list/traceReadFailure.test.ts

@@ -0,0 +1,61 @@
+import { describe, expect, it } from "vitest";
+
+import { ApiError } from "@/lib/http/client";
+
+import {
+  classifyTraceReadFailure,
+  isRetryableTraceRead,
+  TRACE_READ_AUTO_RETRIES,
+  traceReadRetry,
+  traceReadRetryDelay,
+} from "./traceReadFailure";
+
+const failure = (status: number, code?: string, retryAfterMs: number | null = null) =>
+  new ApiError("boom", status, { detail: code ? { code, message: "boom" } : "boom" }, retryAfterMs);
+
+describe("classifyTraceReadFailure", () => {
+  it.each([
+    ["invalid_request", "invalid"],
+    ["trace_changed", "changed"],
+    ["too_large", "too_large"],
+    ["unavailable", "unavailable"],
+  ] as const)("reads the %s code ahead of the status", (code, kind) => {
+    expect(classifyTraceReadFailure(failure(500, code)).kind).toBe(kind);
+  });
+
+  it.each([
+    [400, "invalid"],
+    [409, "changed"],
+    [413, "too_large"],
+    [503, "unavailable"],
+    [500, "unknown"],
+  ] as const)("falls back to status %i without a code", (status, kind) => {
+    expect(classifyTraceReadFailure(failure(status)).kind).toBe(kind);
+  });
+
+  it("treats a non-HTTP error as unknown and keeps its message", () => {
+    expect(classifyTraceReadFailure(new Error("offline"))).toEqual({
+      kind: "unknown",
+      message: "offline",
+      retryAfterMs: null,
+    });
+  });
+});
+
+describe("retry policy", () => {
+  it("retries an outage automatically after the server's Retry-After and stops at the cap", () => {
+    const outage = failure(503, "unavailable", 2_000);
+    expect(traceReadRetry(0, outage)).toBe(true);
+    expect(traceReadRetry(TRACE_READ_AUTO_RETRIES, outage)).toBe(false);
+    expect(traceReadRetryDelay(0, outage)).toBe(2_000);
+    expect(traceReadRetryDelay(0, failure(503, "unavailable"))).toBe(1_000);
+  });
+
+  it("never retries a changed, invalid, or oversized read", () => {
+    for (const code of ["trace_changed", "invalid_request", "too_large"]) {
+      expect(traceReadRetry(0, failure(500, code))).toBe(false);
+      expect(isRetryableTraceRead(classifyTraceReadFailure(failure(500, code)))).toBe(false);
+    }
+    expect(isRetryableTraceRead(classifyTraceReadFailure(new Error("offline")))).toBe(true);
+  });
+});

ui/litellm-dashboard/src/components/lens/traces/list/traceReadFailure.ts

@@ -0,0 +1,58 @@
+import { ApiError } from "@/lib/http/client";
+
+/**
+ * Why a trace read failed, as the proxy reports it in the `code` of a failed response. Each kind has
+ * one recovery: `invalid` and `changed` need a new traversal, `too_large` cannot be paged through at
+ * all, `unavailable` is worth retrying after the server's `Retry-After`, and `unknown` is treated like
+ * `unavailable` because the loaded pages are still good.
+ */
+export type TraceReadFailureKind = "invalid" | "changed" | "too_large" | "unavailable" | "unknown";
+
+export interface TraceReadFailure {
+  readonly kind: TraceReadFailureKind;
+  readonly message: string;
+  readonly retryAfterMs: number | null;
+}
+
+const KIND_BY_CODE: Readonly<Record<string, TraceReadFailureKind>> = {
+  invalid_request: "invalid",
+  trace_changed: "changed",
+  too_large: "too_large",
+  unavailable: "unavailable",
+};
+
+const KIND_BY_STATUS: Readonly<Record<number, TraceReadFailureKind>> = {
+  400: "invalid",
+  409: "changed",
+  413: "too_large",
+  503: "unavailable",
+};
+
+const failureCode = (body: unknown): string | undefined => {
+  if (typeof body !== "object" || body === null || !("detail" in body)) return undefined;
+  const detail = body.detail;
+  if (typeof detail !== "object" || detail === null || !("code" in detail)) return undefined;
+  return typeof detail.code === "string" ? detail.code : undefined;
+};
+
+export function classifyTraceReadFailure(error: unknown): TraceReadFailure {
+  if (!(error instanceof ApiError)) {
+    return { kind: "unknown", message: error instanceof Error ? error.message : String(error), retryAfterMs: null };
+  }
+  const code = failureCode(error.body);
+  const kind = (code && KIND_BY_CODE[code]) || KIND_BY_STATUS[error.status] || "unknown";
+  return { kind, message: error.message, retryAfterMs: error.retryAfterMs };
+}
+
+/** Only an outage is worth retrying without the user; a changed or invalid cursor fails the same way every time. */
+export const isRetryableTraceRead = (failure: TraceReadFailure): boolean =>
+  failure.kind === "unavailable" || failure.kind === "unknown";
+
+export const TRACE_READ_AUTO_RETRIES = 2;
+const DEFAULT_RETRY_DELAY_MS = 1_000;
+
+export const traceReadRetry = (failureCount: number, error: unknown): boolean =>
+  failureCount < TRACE_READ_AUTO_RETRIES && classifyTraceReadFailure(error).kind === "unavailable";
+
+export const traceReadRetryDelay = (_attempt: number, error: unknown): number =>
+  classifyTraceReadFailure(error).retryAfterMs ?? DEFAULT_RETRY_DELAY_MS;

ui/litellm-dashboard/src/lib/http/api.ts

@@ -1,7 +1,7 @@
 import createFetchClient, { type Middleware } from "openapi-fetch";
 import createQueryClient from "openapi-react-query";
 import type { paths } from "./schema";
-import { ApiError, deriveErrorMessage } from "./client";
+import { ApiError, deriveErrorMessage, retryAfterMs } from "./client";
 import { getAuthHeaderName, getAuthToken, getRequestBaseUrl, reportError } from "./runtime";
 import { resolveRequestUrl } from "./resolveApiBase";
 
@@ -32,7 +32,7 @@ const middleware: Middleware = {
       message = raw || `HTTP ${response.status}`;
     }
     reportError(message);
-    throw new ApiError(message, response.status, body);
+    throw new ApiError(message, response.status, body, retryAfterMs(response.headers));
   },
 };
 

ui/litellm-dashboard/src/lib/http/client.ts

@@ -32,15 +32,28 @@ export interface RequestOptions {
 export class ApiError extends Error {
   readonly status: number;
   readonly body: unknown;
+  /** The server's `Retry-After` delay, when it sent one. */
+  readonly retryAfterMs: number | null;
 
-  constructor(message: string, status: number, body: unknown) {
+  constructor(message: string, status: number, body: unknown, retryAfterMs: number | null = null) {
     super(message);
     this.name = "ApiError";
     this.status = status;
     this.body = body;
+    this.retryAfterMs = retryAfterMs;
   }
 }
 
+/** `Retry-After` as milliseconds; the header is whole seconds or an HTTP date. */
+export const retryAfterMs = (headers?: Headers): number | null => {
+  const header = headers?.get("retry-after");
+  if (header === null || header === undefined) return null;
+  const seconds = Number(header);
+  if (Number.isFinite(seconds)) return Math.max(0, seconds * 1000);
+  const at = Date.parse(header);
+  return Number.isNaN(at) ? null : Math.max(0, at - Date.now());
+};
+
 /**
  * Best-effort extraction of a human-readable message from a proxy error body.
  * Lives here because error parsing is the client's job; networking.tsx re-exports
@@ -175,7 +188,7 @@ export function createApiClient(config: ApiClientConfig): ApiClient {
         message = raw || `HTTP ${response.status}`;
       }
       onError?.(message);
-      throw new ApiError(message, response.status, errorBody);
+      throw new ApiError(message, response.status, errorBody, retryAfterMs(response.headers));
     }
 
     return response;