From 21ecd0af557597d7fc2cad48a3aae2e8c2f1094b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 3 Oct 2026 23:43:10 +0000 Subject: [PATCH] fix(traces): reject conflicting spend aliases and unrelated HTTP siblings (#44456) Co-authored-by: Yujong Lee --- .../crates/traces-clickhouse/src/span_row.rs | 2 +- .../tests/queries/support.rs | 3 +- .../normalize/instrumentation/http_client.rs | 41 ++++--- .../crates/traces/src/normalize/mod.rs | 3 + litellm-rust/crates/traces/src/otlp/span.rs | 4 +- .../crates/traces/src/resolve/resolution.rs | 16 ++- .../crates/traces/src/resolve/spend.rs | 50 ++++++++- litellm-rust/crates/traces/tests/captures.rs | 12 +- .../traces/tests/normalization_formats.rs | 3 +- litellm-rust/crates/traces/tests/normalize.rs | 3 + litellm-rust/crates/traces/tests/resolve.rs | 106 ++++++++++++++++-- 11 files changed, 204 insertions(+), 39 deletions(-) diff --git a/litellm-rust/crates/traces-clickhouse/src/span_row.rs b/litellm-rust/crates/traces-clickhouse/src/span_row.rs index c3b652129a4..0189dbff2c8 100644 --- a/litellm-rust/crates/traces-clickhouse/src/span_row.rs +++ b/litellm-rust/crates/traces-clickhouse/src/span_row.rs @@ -211,7 +211,7 @@ fn request_id(evidence: &CallEvidence) -> &str { .flatten() .find_map(|key| match key { CallKey::ProviderResponse(id) => Some(id.as_str()), - CallKey::LiteLlmRequest(_) | CallKey::Transport => None, + CallKey::LiteLlmRequest(_) | CallKey::Transport | CallKey::GatewayAttempt => None, }) .unwrap_or_default() } diff --git a/litellm-rust/crates/traces-clickhouse/tests/queries/support.rs b/litellm-rust/crates/traces-clickhouse/tests/queries/support.rs index a9f42119fb4..371424c63d6 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/queries/support.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/queries/support.rs @@ -158,7 +158,8 @@ fn span_row(span: &DecodedSpan, team: &str, key: &str) -> BTreeMap Some(id.as_str()), - litellm_traces::CallKey::Transport => None, + litellm_traces::CallKey::Transport + | litellm_traces::CallKey::GatewayAttempt => None, }) .unwrap_or_default() ), diff --git a/litellm-rust/crates/traces/src/normalize/instrumentation/http_client.rs b/litellm-rust/crates/traces/src/normalize/instrumentation/http_client.rs index f7b563535c1..1ee6c24b9d4 100644 --- a/litellm-rust/crates/traces/src/normalize/instrumentation/http_client.rs +++ b/litellm-rust/crates/traces/src/normalize/instrumentation/http_client.rs @@ -12,23 +12,30 @@ const SCOPES: [&str; 7] = [ ]; pub(super) fn matches(context: &SpanContext<'_>) -> bool { - SCOPES.contains(&context.scope) - || (context.scope == "litellm.gateway.client" - && context.name == "gateway.request" - && context - .attributes - .get("litellm.gateway.attempt") - .is_some_and(|value| value == "true") - && context - .attributes - .get("http.request.method") - .is_some_and(|value| value == "POST")) + SCOPES.contains(&context.scope) || matches_gateway_attempt(context) } -pub(super) fn adjust(facts: SpanFacts) -> SpanFacts { +fn matches_gateway_attempt(context: &SpanContext<'_>) -> bool { + context.scope == "litellm.gateway.client" + && context.name == "gateway.request" + && context + .attributes + .get("litellm.gateway.attempt") + .is_some_and(|value| value == "true") + && context + .attributes + .get("http.request.method") + .is_some_and(|value| value == "POST") +} + +pub(super) fn adjust(context: &SpanContext<'_>, facts: SpanFacts) -> SpanFacts { SpanFacts { role: Some(RoleEvidence::Declared(ObservationType::Framework)), - calls: CallEvidence::complete(CallKey::Transport), + calls: CallEvidence::complete(if matches_gateway_attempt(context) { + CallKey::GatewayAttempt + } else { + CallKey::Transport + }), ..facts } } @@ -42,7 +49,11 @@ impl Rule for HttpClient { fn integration(&self, _: &SpanContext<'_>) -> Option { None } - fn adjust(&self, _: &SpanContext<'_>, extraction: super::Extraction) -> super::Extraction { - extraction.map_facts(adjust) + fn adjust( + &self, + context: &SpanContext<'_>, + extraction: super::Extraction, + ) -> super::Extraction { + extraction.map_facts(|facts| adjust(context, facts)) } } diff --git a/litellm-rust/crates/traces/src/normalize/mod.rs b/litellm-rust/crates/traces/src/normalize/mod.rs index 7d17f53daa8..c3ad76fde02 100644 --- a/litellm-rust/crates/traces/src/normalize/mod.rs +++ b/litellm-rust/crates/traces/src/normalize/mod.rs @@ -53,6 +53,7 @@ pub enum CallKey { ProviderResponse(String), /// The span is the HTTP request itself; LiteLLM logs its `traceparent` span id. Transport, + GatewayAttempt, } impl fmt::Display for CallKey { @@ -61,6 +62,7 @@ impl fmt::Display for CallKey { Self::LiteLlmRequest(id) => write!(formatter, "litellm_request:{id}"), Self::ProviderResponse(id) => write!(formatter, "provider_response:{id}"), Self::Transport => formatter.write_str("transport:"), + Self::GatewayAttempt => formatter.write_str("gateway_attempt:"), } } } @@ -77,6 +79,7 @@ impl FromStr for CallKey { Ok(Self::LiteLlmRequest(id.to_owned())) } Some(("transport", "")) => Ok(Self::Transport), + Some(("gateway_attempt", "")) => Ok(Self::GatewayAttempt), _ => Err(crate::InvalidCallKey), } } diff --git a/litellm-rust/crates/traces/src/otlp/span.rs b/litellm-rust/crates/traces/src/otlp/span.rs index 58aba3c68b9..0261e4aeb07 100644 --- a/litellm-rust/crates/traces/src/otlp/span.rs +++ b/litellm-rust/crates/traces/src/otlp/span.rs @@ -171,7 +171,9 @@ fn decoded_span( crate::CallKey::LiteLlmRequest(id) | crate::CallKey::ProviderResponse(id) => { id.len() + size_of::() } - crate::CallKey::Transport => size_of::(), + crate::CallKey::Transport | crate::CallKey::GatewayAttempt => { + size_of::() + } }) .sum::() + normalized.model.as_ref().map_or(0, String::len) diff --git a/litellm-rust/crates/traces/src/resolve/resolution.rs b/litellm-rust/crates/traces/src/resolve/resolution.rs index 1b8ffffa812..b469fb95019 100644 --- a/litellm-rust/crates/traces/src/resolve/resolution.rs +++ b/litellm-rust/crates/traces/src/resolve/resolution.rs @@ -134,11 +134,13 @@ impl<'a> Resolution<'a> { ) } - /// The request attempts a model call made: its transport descendants, or, for bridges that - /// emit the request beside the call instead of under it, transport siblings inside the call's - /// time window when the call is the only model call under that parent. fn transports(&self, call: usize) -> Vec { - let is_transport = |index: &usize| self.row(*index).call_keys.contains(&CallKey::Transport); + let is_transport = |index: &usize| { + self.row(*index) + .call_keys + .iter() + .any(|key| matches!(key, CallKey::Transport | CallKey::GatewayAttempt)) + }; let nested: Vec = self .graph .descendants(call) @@ -162,7 +164,11 @@ impl<'a> Resolution<'a> { let call_end_ns = call_start_ns + i128::from(call_row.duration_ns); siblings .into_iter() - .filter(is_transport) + .filter(|sibling| { + self.row(*sibling) + .call_keys + .contains(&CallKey::GatewayAttempt) + }) .filter(|sibling| { let transport = self.row(*sibling); let transport_start_ns = i128::from(transport.start_ns); diff --git a/litellm-rust/crates/traces/src/resolve/spend.rs b/litellm-rust/crates/traces/src/resolve/spend.rs index 34c90a7d12c..ce23362b7f8 100644 --- a/litellm-rust/crates/traces/src/resolve/spend.rs +++ b/litellm-rust/crates/traces/src/resolve/spend.rs @@ -48,7 +48,9 @@ impl SpendLookup { trace_ids: sorted( keys() .filter_map(|(row, key)| match key { - CallKey::Transport if !row.trace_id.is_empty() => { + CallKey::Transport | CallKey::GatewayAttempt + if !row.trace_id.is_empty() => + { Some(row.trace_id.clone()) } _ => None, @@ -80,6 +82,21 @@ impl Ownership<'_> { pub(super) type Requests<'a> = Vec<&'a SpendRow>; +#[derive(Clone, Copy, Eq, Ord, PartialEq, PartialOrd)] +enum KeyFamily { + GatewayCall, + ProviderResponse, + Transport, +} + +fn key_family(key: &CallKey) -> KeyFamily { + match key { + CallKey::LiteLlmRequest(_) => KeyFamily::GatewayCall, + CallKey::ProviderResponse(_) => KeyFamily::ProviderResponse, + CallKey::Transport | CallKey::GatewayAttempt => KeyFamily::Transport, + } +} + pub(super) enum KeyMatch<'a> { Missing, Unique(&'a SpendRow), @@ -174,7 +191,7 @@ fn matches<'a>( && (spend.litellm_call_id == *id || (spend.litellm_call_id.is_empty() && spend.request_id == *id)) } - CallKey::Transport => { + CallKey::Transport | CallKey::GatewayAttempt => { !row.trace_id.is_empty() && !row.span_id.is_empty() && spend.trace_id == row.trace_id @@ -216,12 +233,37 @@ pub(super) fn requests<'a>( && anchored .iter() .all(|request| request.litellm_call_id.is_empty()); - let matches = keyed + let aliases: Vec<_> = keyed .into_iter() .filter(|(key, requests)| { !(legacy_rows && requests.is_empty() && matches!(key, CallKey::LiteLlmRequest(_))) }) - .map(|(_, requests)| KeyMatch::new(requests)) + .collect(); + let families: BTreeSet<_> = aliases.iter().map(|(key, _)| key_family(key)).collect(); + let compatible_rows: Vec> = families + .into_iter() + .map(|family| { + aliases + .iter() + .filter(|(key, _)| key_family(key) == family) + .flat_map(|(_, requests)| requests.iter().map(|request| request.identity())) + .collect() + }) + .collect(); + let matches = aliases + .into_iter() + .map(|(_, requests)| { + KeyMatch::new( + requests + .into_iter() + .filter(|request| { + compatible_rows + .iter() + .all(|family| family.contains(&request.identity())) + }) + .collect(), + ) + }) .collect(); match evidence.kind() { CallEvidenceKind::Complete => SpendEvidence::Complete(matches), diff --git a/litellm-rust/crates/traces/tests/captures.rs b/litellm-rust/crates/traces/tests/captures.rs index 30f84086a9b..fc3527f6253 100644 --- a/litellm-rust/crates/traces/tests/captures.rs +++ b/litellm-rust/crates/traces/tests/captures.rs @@ -172,7 +172,7 @@ fn trace_span(span: DecodedSpan) -> TraceSpansRow { .flatten() .find_map(|key| match key { CallKey::ProviderResponse(id) => Some(id.clone()), - CallKey::LiteLlmRequest(_) | CallKey::Transport => None, + CallKey::LiteLlmRequest(_) | CallKey::Transport | CallKey::GatewayAttempt => None, }) .unwrap_or_default(); TraceSpansRow { @@ -359,7 +359,11 @@ fn unrelated_sibling_transport_leaves_cost_unchanged( let calls: Vec<_> = rows .iter() .filter(|row| { - row.kind == ObservationType::Llm && !row.call_keys.contains(&CallKey::Transport) + row.kind == ObservationType::Llm + && !row + .call_keys + .iter() + .any(|key| matches!(key, CallKey::Transport | CallKey::GatewayAttempt)) }) .cloned() .collect(); @@ -408,7 +412,9 @@ fn redundant_genai_response_id_keeps_call_evidence( .iter() .filter_map(|key| match key { CallKey::ProviderResponse(id) => Some(id.clone()), - CallKey::LiteLlmRequest(_) | CallKey::Transport => None, + CallKey::LiteLlmRequest(_) | CallKey::Transport | CallKey::GatewayAttempt => { + None + } }) .collect(); (!response_ids.is_empty()).then(|| { diff --git a/litellm-rust/crates/traces/tests/normalization_formats.rs b/litellm-rust/crates/traces/tests/normalization_formats.rs index dfa5753978d..0cb2152ab19 100644 --- a/litellm-rust/crates/traces/tests/normalization_formats.rs +++ b/litellm-rust/crates/traces/tests/normalization_formats.rs @@ -746,6 +746,7 @@ fn transport_contract_keeps_independent_call_ids(span: Span) { #[case::unrelated_scope("custom", "gateway.request", "true", "POST", false)] #[case::unrelated_span("litellm.gateway.client", "step", "true", "POST", false)] #[case::missing_contract("litellm.gateway.client", "gateway.request", "", "POST", false)] +#[case::disabled_contract("litellm.gateway.client", "gateway.request", "false", "POST", false)] #[case::unrelated_method("litellm.gateway.client", "gateway.request", "true", "GET", false)] fn gateway_attempt_contract_requires_recorded_request_boundary( span: Span, @@ -774,7 +775,7 @@ fn gateway_attempt_contract_requires_recorded_request_boundary( decoded.normalized.calls, if complete { CallEvidence::Complete(std::collections::BTreeSet::from([ - CallKey::Transport, + CallKey::GatewayAttempt, gateway, ])) } else { diff --git a/litellm-rust/crates/traces/tests/normalize.rs b/litellm-rust/crates/traces/tests/normalize.rs index cab1bb8d08f..244b814cdd6 100644 --- a/litellm-rust/crates/traces/tests/normalize.rs +++ b/litellm-rust/crates/traces/tests/normalize.rs @@ -255,6 +255,7 @@ fn llamaindex_wrapped_responses_keep_provider_call_keys(#[case] body: &[u8]) { #[case::request(litellm_traces::CallKey::LiteLlmRequest("request:with:colons".to_owned()))] #[case::response(litellm_traces::CallKey::ProviderResponse("response:with:colons".to_owned()))] #[case::transport(litellm_traces::CallKey::Transport)] +#[case::gateway_attempt(litellm_traces::CallKey::GatewayAttempt)] fn call_keys_round_trip_through_storage(#[case] key: litellm_traces::CallKey) { assert_eq!( key.to_string().parse::().unwrap(), @@ -272,6 +273,8 @@ fn call_keys_round_trip_through_storage(#[case] key: litellm_traces::CallKey) { #[case::missing_response("provider_response:")] #[case::missing_request("litellm_request:")] #[case::transport_id("transport:unexpected")] +#[case::gateway_attempt_separator("gateway_attempt")] +#[case::gateway_attempt_id("gateway_attempt:unexpected")] #[case::unknown("unknown:id")] fn malformed_call_keys_are_rejected_at_the_boundary(#[case] encoded: &str) { assert!(encoded.parse::().is_err()); diff --git a/litellm-rust/crates/traces/tests/resolve.rs b/litellm-rust/crates/traces/tests/resolve.rs index 009c9b8f929..6e987cd4970 100644 --- a/litellm-rust/crates/traces/tests/resolve.rs +++ b/litellm-rust/crates/traces/tests/resolve.rs @@ -579,7 +579,7 @@ fn sibling_transports_belong_to_the_only_model_call_under_their_parent( 10, ); transport.trace_id = "trace".into(); - transport.call_keys = vec!["transport:".parse().unwrap()]; + transport.call_keys = vec![litellm_traces::CallKey::GatewayAttempt]; transport.call_evidence = Some(litellm_traces::CallEvidenceKind::Complete); let mut rows = vec![ owned( @@ -605,11 +605,15 @@ fn sibling_transports_belong_to_the_only_model_call_under_their_parent( } #[rstest] -#[case::without_tool_http_sibling(None, Some(0.5))] -#[case::after_call(Some((200, 10)), Some(0.5))] -#[case::inside_call_without_spend(Some((10, 10)), None)] +#[case::without_tool_http_sibling(None, false, litellm_traces::CallKey::Transport, Some(0.5))] +#[case::after_call(Some((200, 10)), false, litellm_traces::CallKey::Transport, Some(0.5))] +#[case::inside_call_without_spend(Some((10, 10)), false, litellm_traces::CallKey::Transport, Some(0.5))] +#[case::inside_call_with_unrelated_spend(Some((10, 10)), true, litellm_traces::CallKey::Transport, Some(0.5))] +#[case::missing_gateway_attempt(Some((10, 10)), false, litellm_traces::CallKey::GatewayAttempt, None)] fn sibling_transport_does_not_lose_model_call_spend( #[case] transport_timing: Option<(i64, u64)>, + #[case] unrelated_spend: bool, + #[case] key: litellm_traces::CallKey, #[case] expected: Option, ) { let call = owned( @@ -637,24 +641,110 @@ fn sibling_transport_does_not_lose_model_call_spend( ]; let rows: Vec<_> = base_rows .into_iter() - .chain(transport_timing.into_iter().map(|(start, duration)| { + .chain(transport_timing.map(|(start, duration)| { let mut transport = at( row("tool-http", "step", "GET", "framework", ""), start, duration, ); transport.trace_id = "trace".into(); - transport.call_keys = vec![litellm_traces::CallKey::Transport]; + transport.call_keys = vec![key]; transport.call_evidence = Some(litellm_traces::CallEvidenceKind::Complete); owned(transport, "team", "", "key") })) .collect(); - let logged = spend("chatcmpl-1", "chatcmpl-1", "team", "", "key", 0.5); - let trace = resolve_trace("trace", "ref", &rows, &[logged]).unwrap(); + let logs: Vec<_> = std::iter::once(spend("chatcmpl-1", "chatcmpl-1", "team", "", "key", 0.5)) + .chain(unrelated_spend.then(|| SpendByResponseIdsRow { + trace_id: "trace".into(), + span_id: "tool-http".into(), + ..spend("unrelated", "unrelated", "team", "", "key", 0.75) + })) + .collect(); + let trace = resolve_trace("trace", "ref", &rows, &logs).unwrap(); assert_eq!(trace.summary.spend, expected); assert_eq!(trace.agents[0].spend, expected); } +#[rstest] +#[case::agreeing_ids( + litellm_traces::CallKey::Transport, + "call-a", + Some("response-a"), + Some(0.25) +)] +#[case::conflicting_gateway_id(litellm_traces::CallKey::Transport, "call-b", None, None)] +#[case::conflicting_response_id( + litellm_traces::CallKey::Transport, + "call-a", + Some("response-b"), + None +)] +#[case::conflicting_gateway_and_response( + litellm_traces::CallKey::Transport, + "call-b", + Some("response-b"), + None +)] +#[case::agreeing_gateway_attempt( + litellm_traces::CallKey::GatewayAttempt, + "call-a", + Some("response-a"), + Some(0.25) +)] +#[case::conflicting_gateway_attempt(litellm_traces::CallKey::GatewayAttempt, "call-b", None, None)] +fn gateway_attempt_identifiers_must_match_one_spend_row( + #[case] transport: litellm_traces::CallKey, + #[case] call_id: &str, + #[case] response_id: Option<&str>, + #[case] expected: Option, +) { + let keys = [ + transport, + litellm_traces::CallKey::LiteLlmRequest(call_id.into()), + ] + .into_iter() + .chain(response_id.map(|id| litellm_traces::CallKey::ProviderResponse(id.into()))) + .collect(); + let rows = [ + owned( + row("agent", "", "agent", "agent", "agent"), + "team", + "", + "key", + ), + owned(llm("call", "agent", "agent", ""), "team", "", "key"), + owned( + TraceSpansRow { + trace_id: "trace".into(), + call_keys: keys, + call_evidence: Some(litellm_traces::CallEvidenceKind::Complete), + ..row("attempt", "call", "gateway.request", "framework", "") + }, + "team", + "", + "key", + ), + ]; + let logs = [ + SpendByResponseIdsRow { + litellm_call_id: "call-a".into(), + trace_id: "trace".into(), + span_id: "attempt".into(), + ..spend("request-a", "response-a", "team", "", "key", 0.25) + }, + SpendByResponseIdsRow { + litellm_call_id: "call-b".into(), + trace_id: "trace".into(), + span_id: "other-attempt".into(), + ..spend("request-b", "response-b", "team", "", "key", 0.5) + }, + ]; + let trace = resolve_trace("trace", "ref", &rows, &logs).unwrap(); + assert_eq!(trace.summary.spend, expected); + assert_eq!(trace.agents[0].spend, expected); + assert_eq!(trace.spans[2].spend, expected); +} + #[rstest] #[case::legacy_row("", Some(0.5))] #[case::other_call("other-call", None)]