diff --git a/litellm-rust/.agents/skills/rust-tracing/SKILL.md b/litellm-rust/.agents/skills/rust-tracing/SKILL.md new file mode 100644 index 00000000000..c1c7f43c456 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-tracing/SKILL.md @@ -0,0 +1,22 @@ +--- +name: rust-tracing +description: Add or change Rust diagnostic tracing in litellm-rust, including route spans, subscriber layers, and Python logger delivery +--- + +# Rust tracing + +Use upstream `tracing` throughout Rust, including `#[tracing::instrument]`, events, and span propagation. Centralize collection and delivery infrastructure in `crates/tracing`. Direct upstream imports still reach our configured subscriber; re-exporting macros does not control delivery. Do not introduce Rust `log` or `pyo3-log` for this path + +`litellm-tracing` owns shared subscriber layers, span field collection, and diagnostic processing. Keep adapters composable as `tracing_subscriber::Layer`s, with `Logger` providing host setup. Runtime-specific delivery belongs in the host bridge. The Python bridge delivers directly to the existing Python SDK logger, preserving its handlers, filtering, redaction, and request correlation. Keep Python dependencies out of `crates/tracing` + +Hosts configure subscribers. Keep Python execution scoped to its captured dispatch rather than installing a process-wide subscriber. Propagate both span context and dispatch across spawned work and returned streams + +In core, instrument execution shared by native calls and hosted machines. Use consistent route, model, provider, streaming, and outcome fields. Put status recording at shared provider boundaries instead of scattering basic logging through handlers. Keep upstream HTTP status separate from route success + +Use `skip_all` and explicitly selected fields. Basic tracing excludes bodies, credentials, headers, and raw error strings. Avoid automatic `ret` or `err` capture of sensitive values. Keep payload diagnostics separate and subject to existing redaction + +A returned stream retains its route span until exhaustion, error, or drop, with exactly one terminal outcome. Builder construction does not start a trace. Never hold a span entry guard across an await. Diagnostic tracing remains separate from lifecycle callbacks and `CustomLogger` dispatch + +Use `litellm_tracing::sink_layer` to compose a sink with other subscriber layers. It inherits span fields into events and emits span-close summaries with elapsed time. Test observable records, concurrent isolation, dynamic filtering, sensitive-field exclusion, and stream cancellation when changing this behavior + +Consult the [tracing API](https://docs.rs/tracing/latest/tracing/) and [subscriber layers](https://docs.rs/tracing-subscriber/latest/tracing_subscriber/layer/index.html) for implementation details diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index b1dc35d3698..fe0aac56f0c 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,5 +1,7 @@ # Rust workspace rules +For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md) + ## Test placement - Never create a `tests.rs` (or `test.rs`) file under `src/`, and never `#[path = "tests.rs"] mod tests;` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 19458cf34b1..26d36580cf0 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3395,6 +3395,7 @@ dependencies = [ "time", "tokio", "tokio-tungstenite", + "tracing", "url", "veil", "wiremock", @@ -3596,6 +3597,7 @@ dependencies = [ "thiserror 2.0.19", "tokio", "tokio-tungstenite", + "tracing", "veil", "webpki-roots", ] @@ -3698,6 +3700,7 @@ dependencies = [ "thiserror 2.0.19", "tokio", "tokio-tungstenite", + "tracing", "url", "veil", "wiremock", diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 1ecf6dce8ea..365f806142f 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -18,6 +18,7 @@ litellm-auth-aws.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-tracing.workspace = true +tracing.workspace = true moka.workspace = true mime_guess = "2.0.5" rand.workspace = true diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index e387ba2c6df..57d6a3003f2 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -31,9 +31,24 @@ impl AudioTranscriptionRoute { } } + #[tracing::instrument(name = "litellm.route", skip_all, fields( + route = "audio_transcription", + model = request.model, + provider, + resolved_model, + stream = false, + outcome + ))] pub async fn execute(&self, request: AudioTranscriptionRequest<'_>) -> Result { - let request = - prepare_audio_transcription_provider_call(request, self.secrets.as_ref()).await?; - execute_audio_transcription_provider_call(&self.http, &self.auth, request).await + crate::diagnostic::unary(async { + let request = + prepare_audio_transcription_provider_call(request, self.secrets.as_ref()).await?; + crate::diagnostic::provider(&request.model, &request.custom_llm_provider); + let execute: futures_util::future::BoxFuture<'_, Result> = Box::pin( + execute_audio_transcription_provider_call(&self.http, &self.auth, request), + ); + execute.await + }) + .await } } diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 4c951352db2..3c2be913b0b 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -30,6 +30,7 @@ fn provider_config(provider: LlmProviders) -> Option<&'static dyn BaseAudioTrans } } +#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)] pub async fn prepare_audio_transcription_provider_call( request: AudioTranscriptionRequest<'_>, secrets: &dyn SecretSource, diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index c6b43ff9eea..8f593689988 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -68,17 +68,33 @@ impl ChatCompletionsRoute { litellm_host::lifecycle::observe_unary(hooks.observer(), self.run(request, hooks)).await } + #[tracing::instrument(name = "litellm.route", skip_all, fields( + route = "chat_completions", + model = %request.model, + provider, + resolved_model, + stream = false, + outcome + ))] async fn run( &self, request: ChatCompletionsRequest<'_>, hooks: &impl litellm_host::hooks::RouteHooks, ) -> Result { - let resolved = resolve_request(request)?; - let snapshot = self - .secrets - .resolve(&resolved.config.secret_names()) - .await?; - let prepared = prepare_provider_request(resolved, snapshot)?; - handler::execute(&self.http, &self.auth, prepared, hooks).await + crate::diagnostic::unary(async { + let resolved = resolve_request(request)?; + let snapshot = self + .secrets + .resolve(&resolved.config.secret_names()) + .await?; + let prepared = prepare_provider_request(resolved, snapshot)?; + crate::diagnostic::provider(&prepared.model, &prepared.custom_llm_provider); + let execute: futures_util::future::BoxFuture< + '_, + Result, + > = Box::pin(handler::execute(&self.http, &self.auth, prepared, hooks)); + execute.await + }) + .await } } diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index 6b3246d44b5..c85ed8fed75 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -96,6 +96,7 @@ fn validate_environment( }) } +#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)] pub(super) fn prepare_provider_request( request: ResolvedChatCompletionsRequest<'_>, secrets: Secrets, diff --git a/litellm-rust/crates/core/src/diagnostic.rs b/litellm-rust/crates/core/src/diagnostic.rs new file mode 100644 index 00000000000..c1ac4451c81 --- /dev/null +++ b/litellm-rust/crates/core/src/diagnostic.rs @@ -0,0 +1,324 @@ +use std::{ + future::Future, + pin::Pin, + task::{Context, Poll}, +}; + +use futures_util::{Stream, stream::BoxStream}; +use litellm_host::call::CallOutput; +use litellm_tracing::Logger; +use tracing::Span; + +struct Completion { + span: Span, + outcome: &'static str, +} + +impl Completion { + fn new(name: &str) -> Self { + let current = Span::current(); + Self { + span: if current + .metadata() + .is_some_and(|metadata| metadata.name() == name) + { + current + } else { + Span::none() + }, + outcome: "cancelled", + } + } + + fn finish(mut self, outcome: &'static str) { + self.outcome = outcome; + } +} + +impl Drop for Completion { + fn drop(&mut self) { + self.span.record("outcome", self.outcome); + } +} + +pub(crate) fn provider(model: &str, provider: &str) { + let span = Span::current(); + span.record("resolved_model", model); + span.record("provider", provider); +} + +pub(crate) async fn unary(execute: impl Future>) -> Result { + operation("litellm.route", execute).await +} + +pub(crate) async fn operation( + name: &str, + execute: impl Future>, +) -> Result { + let completion = Completion::new(name); + let result = execute.await; + completion.finish(if result.is_ok() { "success" } else { "failure" }); + result +} + +pub(crate) async fn call( + execute: impl Future, E>>, +) -> Result, E> +where + C: Send + 'static, + E: Send + 'static, +{ + let completion = Completion::new("litellm.route"); + match execute.await { + Err(error) => { + completion.finish("failure"); + Err(error) + } + Ok(CallOutput::Complete(response)) => { + completion.span.record("stream", false); + completion.finish("success"); + Ok(CallOutput::Complete(response)) + } + Ok(CallOutput::Stream { head, chunks }) => { + completion.span.record("stream", true); + Ok(CallOutput::Stream { + head, + chunks: Box::pin(TracedStream { + state: Some(StreamState { + chunks, + completion, + logger: Logger::current(), + }), + }), + }) + } + } +} + +struct StreamState { + chunks: BoxStream<'static, Result>, + completion: Completion, + logger: Logger, +} + +impl StreamState { + fn close(self, outcome: &'static str) { + self.logger + .scope(|| self.completion.span.in_scope(|| drop(self.chunks))); + self.completion.finish(outcome); + } +} + +struct TracedStream { + state: Option>, +} + +impl Stream for TracedStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + let Some(state) = self.state.as_mut() else { + return Poll::Ready(None); + }; + let next = state.logger.scope(|| { + state + .completion + .span + .in_scope(|| state.chunks.as_mut().poll_next(context)) + }); + let outcome = match &next { + Poll::Ready(None) => "success", + Poll::Ready(Some(Err(_))) => "failure", + _ => return next, + }; + if let Some(state) = self.state.take() { + state.close(outcome); + } + next + } +} + +impl Drop for TracedStream { + fn drop(&mut self) { + if let Some(state) = self.state.take() { + state.close("cancelled"); + } + } +} + +#[cfg(test)] +mod tests { + use std::{sync::mpsc, task::Context}; + + use futures_util::{StreamExt, task::noop_waker_ref}; + use litellm_tracing::{Metadata, Record, Sink}; + use rstest::{fixture, rstest}; + use serde_json::{Value, json}; + + use super::*; + + struct Capture(mpsc::Sender); + + impl Sink for Capture { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + *metadata.level() <= tracing::Level::INFO + } + fn emit(&self, record: &Record) { + self.0.send(Value::Object(record.fields.clone())).unwrap(); + } + } + + #[fixture] + fn logger() -> (Logger, mpsc::Receiver) { + let (sender, receiver) = mpsc::channel(); + (Logger::new(Capture(sender)), receiver) + } + + struct Chunks(std::vec::IntoIter>); + + impl Stream for Chunks { + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + tracing::info!(event = "poll"); + Poll::Ready(self.0.next()) + } + } + + impl Drop for Chunks { + fn drop(&mut self) { + tracing::info!(event = "drop"); + } + } + + #[tracing::instrument( + name = "litellm.route", + skip_all, + fields(route = "fixture", stream, outcome) + )] + async fn streamed() -> Result, &'static str> { + call(async { + Ok(CallOutput::Stream { + head: (), + chunks: Box::pin(Chunks(vec![Ok(1), Err("broken"), Ok(2)].into_iter())), + }) + }) + .await + } + + #[rstest] + #[tokio::test] + async fn stream_errors_finish_once_and_poll_and_drop_use_the_captured_context( + logger: (Logger, mpsc::Receiver), + ) { + let (logger, records) = logger; + let CallOutput::Stream { mut chunks, .. } = logger.instrument(streamed()).await.unwrap() + else { + panic!() + }; + assert!(records.try_recv().is_err()); + tokio::spawn(async move { + assert_eq!(chunks.next().await, Some(Ok(1))); + assert_eq!(chunks.next().await, Some(Err("broken"))); + assert_eq!(chunks.next().await, None); + let emitted = records.try_iter().collect::>(); + assert_eq!(emitted.len(), 4); + assert!(emitted.iter().all(|record| record["route"] == "fixture")); + assert_eq!(emitted[0]["event"], "poll"); + assert_eq!(emitted[1]["event"], "poll"); + assert_eq!(emitted[2]["event"], "drop"); + assert_eq!(emitted[3]["outcome"], "failure"); + drop(chunks); + assert!(records.try_recv().is_err()); + }) + .await + .unwrap(); + } + + #[tracing::instrument(name = "litellm.route", skip_all, fields(route = "waiting", outcome))] + async fn waiting(streaming: bool) { + if streaming { + let _: Result, ()> = call(std::future::pending()).await; + } else { + let _: Result<(), ()> = unary(std::future::pending()).await; + } + } + + #[rstest] + #[case::unary(false)] + #[case::streaming(true)] + fn cancellation_before_headers_closes_the_span( + logger: (Logger, mpsc::Receiver), + #[case] streaming: bool, + ) { + let (logger, records) = logger; + let mut future = Box::pin(logger.instrument(waiting(streaming))); + assert!( + future + .as_mut() + .poll(&mut Context::from_waker(noop_waker_ref())) + .is_pending() + ); + assert!(records.try_recv().is_err()); + drop(future); + let summary = records.try_recv().unwrap(); + assert_eq!(summary["outcome"], "cancelled"); + assert_eq!(summary["route"], "waiting"); + assert!(records.try_recv().is_err()); + } + + #[rstest] + #[tokio::test] + async fn dropped_stream_teardown_uses_its_original_logger( + logger: (Logger, mpsc::Receiver), + ) { + let (logger, records) = logger; + let output = logger.instrument(streamed()).await.unwrap(); + Logger::default().scope(|| drop(output)); + let emitted = records.try_iter().collect::>(); + assert_eq!(emitted.len(), 2); + assert_eq!( + emitted[0], + json!({"route":"fixture", "stream":true, "event":"drop"}) + ); + assert_eq!(emitted[1]["outcome"], "cancelled"); + } + + #[tracing::instrument(name = "disabled", level = "debug", skip_all, fields(outcome))] + async fn disabled_child() { + let _: Result<(), ()> = operation("disabled", async { Err(()) }).await; + } + + #[rstest] + #[tokio::test] + async fn a_filtered_operation_does_not_overwrite_its_parent_outcome( + logger: (Logger, mpsc::Receiver), + ) { + let (logger, records) = logger; + logger + .instrument(async { + let parent = tracing::info_span!("parent", outcome = "original"); + tracing::Instrument::instrument(disabled_child(), parent).await; + }) + .await; + let summary = records.try_recv().unwrap(); + assert_eq!(summary["span_name"], "parent"); + assert_eq!(summary["outcome"], "original"); + assert!(records.try_recv().is_err()); + } + + #[tracing::instrument(name = "litellm.route", skip_all, fields(stream = true, outcome))] + async fn completed() -> Result, ()> { + call(async { Ok(CallOutput::Complete(())) }).await + } + + #[rstest] + #[tokio::test] + async fn streaming_mode_reflects_the_returned_output(logger: (Logger, mpsc::Receiver)) { + let (logger, records) = logger; + logger.instrument(completed()).await.unwrap(); + let summary = records.try_recv().unwrap(); + assert_eq!(summary["stream"], false); + assert_eq!(summary["outcome"], "success"); + assert!(records.try_recv().is_err()); + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 66643eb65ff..a0ae582210e 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,3 +1,5 @@ +mod diagnostic; + pub mod audio_transcription; pub mod chat_completions; pub mod constants; diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 604a9abaca7..d015c238043 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -199,20 +199,20 @@ fn decoded_chunks( } fn log_request_body(provider: &str, stream: bool, body: &serde_json::Value) { - litellm_tracing::debug!(provider, stream, body = %body, "provider request"); + tracing::debug!(provider, stream, body = %body, "provider request"); } fn log_response_body(body: &str) { - litellm_tracing::debug!(body, "provider response body"); + tracing::debug!(body, "provider response body"); } fn log_error_body(status: u16, body: &str) { - litellm_tracing::debug!(status, body, "provider error body"); + tracing::debug!(status, body, "provider error body"); } fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) { let chunk = ByteChunk::new(data); - litellm_tracing::debug!(provider, stage, encoding = chunk.encoding(), chunk = %chunk, "stream chunk"); + tracing::debug!(provider, stage, encoding = chunk.encoding(), chunk = %chunk, "stream chunk"); } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 6a08df5b411..f74f2c54771 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -39,12 +39,26 @@ impl MessagesRoute { litellm_host::lifecycle::observe_call(hooks.observer(), self.run(call, hooks)).await } + #[tracing::instrument(name = "litellm.route", skip_all, fields( + route = "messages", + model = %call.body.model, + provider, + resolved_model, + stream = call.body.params.stream == Some(true), + outcome + ))] async fn run( &self, call: MessagesCall, hooks: &impl litellm_host::hooks::RouteHooks, ) -> Result { - let request = prepare::prepare(call, self.secrets.as_ref()).await?; - handler::execute(&self.http, &self.auth, request, hooks).await + crate::diagnostic::call(async { + let request = prepare::prepare(call, self.secrets.as_ref()).await?; + crate::diagnostic::provider(&request.body.model, request.provider.as_str()); + let execute: futures_util::future::BoxFuture<'_, Result> = + Box::pin(handler::execute(&self.http, &self.auth, request, hooks)); + execute.await + }) + .await } } diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 0e4c72e4270..13e77d4649b 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -34,6 +34,7 @@ pub(super) struct ProviderMessagesRequest { pub(super) api_key: Option, } +#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)] pub(super) async fn prepare( call: MessagesCall, secrets: &dyn SecretSource, diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 4ab11f6692e..f39a8379a3c 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -28,31 +28,47 @@ impl OcrRoute { litellm_host::lifecycle::observe_unary(hooks.observer(), self.run(request, hooks)).await } + #[tracing::instrument(name = "litellm.route", skip_all, fields( + route = "ocr", + model = %request.model, + resolved_model = %request.model, + provider = <&str>::from(request.config.provider()), + stream = false, + outcome + ))] pub(super) async fn run( &self, request: LiteLLMOcrRequest, hooks: &impl RouteHooks, ) -> Result { - let caller_document = matches!(&request.document, OcrDocumentInput::Document(_)); - let prepared = prepare_request_document(request).await?; - let execute: futures_util::future::BoxFuture<'_, Result> = - Box::pin(perform_ocr_request( - &self.client, - prepared, - hooks, - caller_document, - )); - execute.await + crate::diagnostic::unary(async { + let caller_document = matches!(&request.document, OcrDocumentInput::Document(_)); + let prepared = prepare_request_document(request).await?; + let execute: futures_util::future::BoxFuture<'_, Result> = + Box::pin(perform_ocr_request( + &self.client, + prepared, + hooks, + caller_document, + )); + execute.await + }) + .await } } +#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)] async fn prepare_request_document( request: LiteLLMOcrRequest, ) -> Result { if let OcrDocumentInput::Document(_) = &request.document { return request.map_document(super::document::prepare_document); } - tokio::task::spawn_blocking(move || request.map_document(super::document::prepare_document)) - .await - .map_err(|error| Error::DocumentTask(Arc::new(error)))? + let logger = litellm_tracing::Logger::current(); + let span = tracing::Span::current(); + tokio::task::spawn_blocking(move || { + logger.scope(|| span.in_scope(|| request.map_document(super::document::prepare_document))) + }) + .await + .map_err(|error| Error::DocumentTask(Arc::new(error)))? } diff --git a/litellm-rust/crates/core/src/outbound.rs b/litellm-rust/crates/core/src/outbound.rs index 6d16b2d88af..4d038e0eb00 100644 --- a/litellm-rust/crates/core/src/outbound.rs +++ b/litellm-rust/crates/core/src/outbound.rs @@ -4,11 +4,19 @@ use litellm_http::outbound::OutboundRequest; use litellm_llms::base_llm::auth::Authenticated; use serde_json::Value; +#[tracing::instrument( + name = "litellm.provider.send", + level = "debug", + skip_all, + fields(status) +)] pub(crate) async fn send( request: OutboundRequest, client: &litellm_http::Client, ) -> Result { - request.send(client).await + request.send(client).await.inspect(|response| { + tracing::Span::current().record("status", response.status().as_u16()); + }) } /// Header credentials are already in `headers`; SigV4 is applied here, over the diff --git a/litellm-rust/crates/core/src/responses/mod.rs b/litellm-rust/crates/core/src/responses/mod.rs index 35c79a7e801..cba9313d0ad 100644 --- a/litellm-rust/crates/core/src/responses/mod.rs +++ b/litellm-rust/crates/core/src/responses/mod.rs @@ -41,12 +41,29 @@ impl ResponsesRoute { litellm_host::lifecycle::observe_call(hooks.observer(), self.run(call, hooks)).await } + #[tracing::instrument(name = "litellm.route", skip_all, fields( + route = "responses", + model = %call.model, + provider, + resolved_model, + stream = call.optional_params.get("stream").and_then(serde_json::Value::as_bool).unwrap_or(false), + outcome + ))] async fn run( &self, call: ResponsesCall, hooks: &impl RouteHooks, ) -> Result { - let request = prepare::prepare(call, self.secrets.as_ref()).await?; - handler::execute(&self.http, &self.auth, request, hooks).await + crate::diagnostic::call(async { + let request = prepare::prepare(call, self.secrets.as_ref()).await?; + crate::diagnostic::provider( + &request.context.model, + &request.context.custom_llm_provider, + ); + let execute: futures_util::future::BoxFuture<'_, Result> = + Box::pin(handler::execute(&self.http, &self.auth, request, hooks)); + execute.await + }) + .await } } diff --git a/litellm-rust/crates/core/src/responses/prepare.rs b/litellm-rust/crates/core/src/responses/prepare.rs index 6e00ecd577b..cf7eb11813d 100644 --- a/litellm-rust/crates/core/src/responses/prepare.rs +++ b/litellm-rust/crates/core/src/responses/prepare.rs @@ -11,6 +11,7 @@ use super::{ types::{ProviderResponsesRequest, ResponsesCall}, }; +#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)] pub(super) async fn prepare( call: ResponsesCall, secrets: &dyn SecretSource, diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index b7a101a8ddf..69165186d25 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -29,83 +29,121 @@ pub struct ResponsesWebSocketConnection { } impl ResponsesWebSocketConnection { + #[tracing::instrument( + name = "litellm.websocket.connect_url", + level = "debug", + skip_all, + fields(outcome) + )] pub async fn connect_url( url: &str, headers: &HashMap, timeout: Option, ) -> Result { - let mut request = url.into_client_request().map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - })?; - for (name, value) in headers { - let header_name = name - .parse::() - .map_err(|error| Error::InvalidRequest(error.to_string().into()))?; - let header_value = HeaderValue::from_str(value) - .map_err(|error| Error::InvalidRequest(error.to_string().into()))?; - request.headers_mut().insert(header_name, header_value); - } - let connect = connect_upstream(request); - let result = match timeout { - Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { - Error::Transport(litellm_http::transport::Error::Network( - "Responses WebSocket connection timed out".into(), - )) - })?, - None => connect.await, - }; - let (socket, _) = result.map_err(|error| match *error { - tokio_tungstenite::tungstenite::Error::Http(response) => { - Error::Transport(litellm_http::transport::Error::Http { - status: response.status().as_u16(), - body: String::new(), - }) - } - other => Error::Transport(litellm_http::transport::Error::Network(other.to_string())), - })?; - Ok(Self { - socket: Arc::new(Mutex::new(Some(socket))), - }) - } - - pub async fn send_text(&self, text: String) -> Result<(), Error> { - let mut socket = self.socket.lock().await; - let Some(socket) = socket.as_mut() else { - return Err(Error::Transport(litellm_http::transport::Error::Network( - "Responses WebSocket is closed".into(), - ))); - }; - socket.send(Message::Text(text)).await.map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - }) - } - - pub async fn recv_text(&self) -> Result, Error> { - let mut socket = self.socket.lock().await; - let Some(socket) = socket.as_mut() else { - return Ok(None); - }; - match socket.next().await { - Some(Ok(Message::Text(text))) => Ok(Some(text)), - Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec()) - .map(Some) - .map_err(|error| Error::InvalidResponse(error.to_string().into())), - Some(Ok(Message::Close(_))) | None => Ok(None), - Some(Ok(_)) => Ok(None), - Some(Err(error)) => Err(Error::Transport(litellm_http::transport::Error::Network( - error.to_string(), - ))), - } - } - - pub async fn close(&self) -> Result<(), Error> { - let mut socket = self.socket.lock().await; - if let Some(socket) = socket.as_mut() { - socket.close(None).await.map_err(|error| { + crate::diagnostic::operation("litellm.websocket.connect_url", async { + let mut request = url.into_client_request().map_err(|error| { Error::Transport(litellm_http::transport::Error::Network(error.to_string())) })?; - } - *socket = None; - Ok(()) + for (name, value) in headers { + let header_name = name + .parse::() + .map_err(|error| Error::InvalidRequest(error.to_string().into()))?; + let header_value = HeaderValue::from_str(value) + .map_err(|error| Error::InvalidRequest(error.to_string().into()))?; + request.headers_mut().insert(header_name, header_value); + } + let connect = connect_upstream(request); + let result = match timeout { + Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { + Error::Transport(litellm_http::transport::Error::Network( + "Responses WebSocket connection timed out".into(), + )) + })?, + None => connect.await, + }; + let (socket, _) = result.map_err(|error| match *error { + tokio_tungstenite::tungstenite::Error::Http(response) => { + Error::Transport(litellm_http::transport::Error::Http { + status: response.status().as_u16(), + body: String::new(), + }) + } + other => { + Error::Transport(litellm_http::transport::Error::Network(other.to_string())) + } + })?; + Ok(Self { + socket: Arc::new(Mutex::new(Some(socket))), + }) + }) + .await + } + + #[tracing::instrument( + name = "litellm.websocket.send_text", + level = "debug", + skip_all, + fields(outcome) + )] + pub async fn send_text(&self, text: String) -> Result<(), Error> { + crate::diagnostic::operation("litellm.websocket.send_text", async { + let mut socket = self.socket.lock().await; + let Some(socket) = socket.as_mut() else { + return Err(Error::Transport(litellm_http::transport::Error::Network( + "Responses WebSocket is closed".into(), + ))); + }; + socket.send(Message::Text(text)).await.map_err(|error| { + Error::Transport(litellm_http::transport::Error::Network(error.to_string())) + }) + }) + .await + } + + #[tracing::instrument( + name = "litellm.websocket.recv_text", + level = "debug", + skip_all, + fields(outcome) + )] + pub async fn recv_text(&self) -> Result, Error> { + crate::diagnostic::operation("litellm.websocket.recv_text", async { + let mut socket = self.socket.lock().await; + let Some(socket) = socket.as_mut() else { + return Ok(None); + }; + match socket.next().await { + Some(Ok(Message::Text(text))) => Ok(Some(text)), + Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec()) + .map(Some) + .map_err(|error| Error::InvalidResponse(error.to_string().into())), + Some(Ok(Message::Close(_))) | None => Ok(None), + Some(Ok(_)) => Ok(None), + Some(Err(error)) => Err(Error::Transport(litellm_http::transport::Error::Network( + error.to_string(), + ))), + } + }) + .await + } + + #[tracing::instrument( + name = "litellm.websocket.close", + level = "debug", + skip_all, + fields(outcome) + )] + pub async fn close(&self) -> Result<(), Error> { + crate::diagnostic::operation("litellm.websocket.close", async { + let mut socket = self.socket.lock().await; + if let Some(socket) = socket.as_mut() { + socket.close(None).await.map_err(|error| { + Error::Transport(litellm_http::transport::Error::Network(error.to_string())) + })?; + } + *socket = None; + Ok(()) + }) + .await } } diff --git a/litellm-rust/crates/core/tests/audio_transcription.rs b/litellm-rust/crates/core/tests/audio_transcription.rs index 1e41df75168..02dbbc50bac 100644 --- a/litellm-rust/crates/core/tests/audio_transcription.rs +++ b/litellm-rust/crates/core/tests/audio_transcription.rs @@ -250,3 +250,31 @@ async fn an_unreadable_success_body_is_an_invalid_response( assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); } + +#[rstest] +#[tokio::test] +async fn transcription_records_route_and_resolved_provider( + request: AudioTranscriptionRequest<'static>, + traces: TraceCapture, +) { + let upstream = upstream([transcript_response("hello")]).await; + let base = upstream.uri(); + let model = request.model; + traces + .logger() + .instrument(transcribe(AudioTranscriptionRequest { + api_base: Some(&base), + ..request + })) + .await + .unwrap(); + let summaries = traces.summaries("litellm.route"); + assert_eq!(summaries.len(), 1); + assert_eq!(summaries[0]["route"], "audio_transcription"); + assert_eq!(summaries[0]["model"], model); + assert_eq!(summaries[0]["resolved_model"], model); + assert_eq!(summaries[0]["provider"], "bedrock"); + assert_eq!(summaries[0]["outcome"], "success"); + assert_eq!(summaries[0]["stream"], false); + assert!(!format!("{:?}", traces.records()).contains("secret-key")); +} diff --git a/litellm-rust/crates/core/tests/chat_completions.rs b/litellm-rust/crates/core/tests/chat_completions.rs index 1b0cba12ce1..28e0a65589c 100644 --- a/litellm-rust/crates/core/tests/chat_completions.rs +++ b/litellm-rust/crates/core/tests/chat_completions.rs @@ -434,3 +434,33 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry( assert_eq!(*source, Error::InvalidRequest("callback rejected".into())); assert_eq!(received(&upstream).await.len(), 1); } + +#[rstest] +#[tokio::test] +async fn completed_chat_records_route_and_resolved_provider( + request: ChatCompletionsRequest<'static>, + traces: TraceCapture, +) { + let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; + let base = upstream.uri(); + let model = request.model; + traces + .logger() + .instrument(complete(ChatCompletionsRequest { + api_base: Some(&base), + ..request + })) + .await + .unwrap(); + let summaries = traces.summaries("litellm.route"); + assert_eq!(summaries.len(), 1); + assert_eq!(summaries[0]["route"], "chat_completions"); + assert_eq!(summaries[0]["model"], model); + assert_eq!(summaries[0]["provider"], "anthropic"); + assert_eq!( + summaries[0]["resolved_model"], + only_request(&upstream).await.json()["model"] + ); + assert_eq!(summaries[0]["outcome"], "success"); + assert_eq!(summaries[0]["stream"], false); +} diff --git a/litellm-rust/crates/core/tests/messages/response.rs b/litellm-rust/crates/core/tests/messages/response.rs index 47505aa86a4..87cfb1bb651 100644 --- a/litellm-rust/crates/core/tests/messages/response.rs +++ b/litellm-rust/crates/core/tests/messages/response.rs @@ -273,3 +273,34 @@ fn a_body_that_does_not_parse_is_an_invalid_request(#[case] raw: Value) { "{error:?}" ); } + +#[rstest] +#[tokio::test] +async fn message_route_summary_excludes_payload_diagnostics( + call: MessagesCall, + traces: TraceCapture, +) { + let upstream = upstream([message_response()]).await; + let model = call.body.model.clone(); + traces + .logger() + .instrument(run_message(MessagesCall { + api_key: Some("private-key-sentinel".into()), + api_base: Some(upstream.uri()), + ..call + })) + .await; + let summaries = traces.summaries("litellm.route"); + assert_eq!(summaries.len(), 1); + assert_eq!(summaries[0]["route"], "messages"); + assert_eq!(summaries[0]["model"], model); + assert_eq!( + summaries[0]["resolved_model"], + only_request(&upstream).await.json()["model"] + ); + assert_eq!(summaries[0]["provider"], "anthropic"); + assert_eq!(summaries[0]["outcome"], "success"); + assert_eq!(summaries[0]["stream"], false); + assert!(summaries[0].get("body").is_none()); + assert!(!format!("{:?}", traces.records()).contains("private-key-sentinel")); +} diff --git a/litellm-rust/crates/core/tests/messages/stream.rs b/litellm-rust/crates/core/tests/messages/stream.rs index fa4594525c9..8b51704d968 100644 --- a/litellm-rust/crates/core/tests/messages/stream.rs +++ b/litellm-rust/crates/core/tests/messages/stream.rs @@ -1,4 +1,4 @@ -use std::sync::Mutex; +use std::sync::{Mutex, mpsc}; use bytes::Bytes; use futures_util::{StreamExt, TryStreamExt}; @@ -7,6 +7,7 @@ use litellm_core::messages::{ route::{Messages, MessagesStreamHead}, }; use litellm_host::protocol::Demand; +use litellm_tracing::{Logger, Metadata, Record, Sink}; use rstest::rstest; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, @@ -28,6 +29,20 @@ enum Seen { Deliver(Bytes), } +struct TraceSink(mpsc::Sender<(String, Value)>); + +impl Sink for TraceSink { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + metadata.target().starts_with("litellm_core::messages") + } + + fn emit(&self, record: &Record) { + self.0 + .send((record.message.clone(), Value::Object(record.fields.clone()))) + .unwrap(); + } +} + /// Projects like `LocalMessagesHost`, records every stream op in the order the route /// performs it, and detaches after `detach_after` ops. struct RecordingStreamHost { @@ -161,6 +176,37 @@ async fn upstream_headers_are_on_the_stream_head_before_the_first_chunk(call: Me assert_eq!(delivered, SSE_BODY.as_bytes()); } +#[rstest] +#[tokio::test] +async fn debug_trace_keeps_provider_input_and_every_stream_chunk(call: MessagesCall) { + let upstream = upstream([sse_response()]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); + let (sender, receiver) = mpsc::channel(); + + Logger::new(TraceSink(sender)) + .instrument(stream_through(&host)) + .await + .unwrap(); + + let records: Vec<(String, Value)> = receiver.try_iter().collect(); + let request = records + .iter() + .find(|(message, _)| message == "provider request") + .unwrap(); + let body: Value = serde_json::from_str(request.1["body"].as_str().unwrap()).unwrap(); + assert_eq!(body["messages"][0]["content"], "hi"); + assert_eq!(request.1["stream"], true); + let chunks: String = records + .iter() + .filter(|(message, fields)| { + message == "stream chunk" && fields["stage"] == "provider_response" + }) + .map(|(_, fields)| fields["chunk"].as_str().unwrap()) + .collect(); + assert_eq!(chunks, SSE_BODY); + assert!(!format!("{records:?}").contains("sk-ant")); +} + #[rstest] #[case::at_open(1)] #[case::after_the_first_chunk(2)] diff --git a/litellm-rust/crates/core/tests/ocr/lifecycle.rs b/litellm-rust/crates/core/tests/ocr/lifecycle.rs index 7febef105c6..b138fafff17 100644 --- a/litellm-rust/crates/core/tests/ocr/lifecycle.rs +++ b/litellm-rust/crates/core/tests/ocr/lifecycle.rs @@ -370,3 +370,34 @@ async fn direct_execution_uses_hooks_without_a_machine() { ] )); } + +#[rstest] +#[case::native(false)] +#[case::hosted(true)] +#[tokio::test] +async fn ocr_records_one_route_summary_across_both_execution_paths( + traces: TraceCapture, + #[case] hosted: bool, +) { + let upstream = upstream([pages_response()]).await; + let request = ocr_request("mistral/model", &upstream.uri(), json!({})); + let model = request.model.clone(); + traces + .logger() + .instrument(async { + if hosted { + perform_with(LocalOcrHost::new(request)).await + } else { + perform(request).await + } + }) + .await + .unwrap(); + let summaries = traces.summaries("litellm.route"); + assert_eq!(summaries.len(), 1); + assert_eq!(summaries[0]["route"], "ocr"); + assert_eq!(summaries[0]["model"], model); + assert_eq!(summaries[0]["provider"], "mistral"); + assert_eq!(summaries[0]["outcome"], "success"); + assert_eq!(summaries[0]["stream"], false); +} diff --git a/litellm-rust/crates/core/tests/responses.rs b/litellm-rust/crates/core/tests/responses.rs index a2652fdd610..84a0a7d14f3 100644 --- a/litellm-rust/crates/core/tests/responses.rs +++ b/litellm-rust/crates/core/tests/responses.rs @@ -231,3 +231,212 @@ async fn unsupported_providers_fail_before_secrets_or_transport( assert!(secrets.requested().is_empty()); assert!(received(&upstream).await.is_empty()); } + +#[rstest] +#[case::success(200, json!({"id":"response-1", "model":"test-model", "output":[]}), "success")] +#[case::invalid_response(200, json!("private-response-sentinel"), "failure")] +#[case::upstream_error(429, json!({"error":"private-response-sentinel"}), "failure")] +#[tokio::test] +async fn route_tracing_covers_native_and_hosted_outcomes( + call: ResponsesCall, + traces: TraceCapture, + #[case] status: u16, + #[case] body: serde_json::Value, + #[case] outcome: &str, + #[values(false, true)] hosted: bool, +) { + let upstream = upstream([status_response(status, body)]).await; + let host = RecordingCall::::new(ResponsesCall { + api_base: Some(upstream.uri()), + input: json!("private-prompt-sentinel"), + api_key: Some("private-key-sentinel".into()), + ..call + }); + let route = responses_route(no_secrets()); + let result = traces + .logger() + .instrument(async { + if hosted { + litellm_host::in_process::run_hosted( + route.clone().machine(host.request().unwrap()), + host.runtime(), + ) + .await + .map(|_| ()) + } else { + route + .execute(host.request().unwrap(), &()) + .await + .map(|_| ()) + } + }) + .await; + assert_eq!(result.is_ok(), outcome == "success"); + let summaries = traces.summaries("litellm.route"); + let [summary] = summaries.as_slice() else { + panic!("expected one route summary: {summaries:?}") + }; + assert_eq!(summary["route"], "responses"); + assert_eq!(summary["model"], "openai/test-model"); + assert_eq!(summary["resolved_model"], "test-model"); + assert_eq!(summary["provider"], "openai"); + assert_eq!(summary["stream"], false); + assert_eq!(summary["outcome"], outcome); + assert!(summary["duration_ms"].as_f64().unwrap() >= 0.0); + let sends = traces.summaries("litellm.provider.send"); + assert_eq!(sends.len(), 1); + assert_eq!(sends[0]["status"], status); + assert!(!format!("{:?}", traces.records()).contains("private-")); +} + +#[rstest] +#[case::exhausted(true, "success")] +#[case::dropped(false, "cancelled")] +#[tokio::test] +async fn stream_trace_survives_handoff_and_closes_before_the_stream_object_is_dropped( + call: ResponsesCall, + traces: TraceCapture, + #[case] exhaust: bool, + #[case] outcome: &str, +) { + let upstream = upstream([ResponseTemplate::new(200).set_body_string("stream-bytes")]).await; + let output = traces + .logger() + .instrument(async { + responses_route(no_secrets()) + .execute( + ResponsesCall { + api_base: Some(upstream.uri()), + optional_params: json!({"stream":true}).as_object().unwrap().clone(), + ..call + }, + &(), + ) + .await + }) + .await + .unwrap(); + assert!(traces.summaries("litellm.route").is_empty()); + let ResponsesOutput::Stream { mut chunks, .. } = output else { + panic!("expected stream") + }; + let captured = traces.clone(); + tokio::spawn(async move { + if exhaust { + while chunks.try_next().await.unwrap().is_some() {} + assert_eq!(captured.summaries("litellm.route").len(), 1); + assert!(chunks.try_next().await.unwrap().is_none()); + } + drop(chunks); + }) + .await + .unwrap(); + let summaries = traces.summaries("litellm.route"); + assert_eq!(summaries.len(), 1); + assert_eq!(summaries[0]["stream"], true); + assert_eq!(summaries[0]["outcome"], outcome); +} + +#[rstest] +#[tokio::test] +async fn preparation_failure_is_traced_but_unpolled_builders_are_not( + call: ResponsesCall, + traces: TraceCapture, +) { + traces.logger().scope(|| { + drop(responses_route(no_secrets()).execute( + ResponsesCall { + model: "unknown/model".into(), + ..call + }, + &(), + )) + }); + assert!(traces.records().is_empty()); + let result = traces + .logger() + .instrument(async { + responses_route(no_secrets()) + .execute( + ResponsesCall { + model: "unknown/model".into(), + ..self::call() + }, + &(), + ) + .await + }) + .await; + assert!(result.is_err()); + let summaries = traces.summaries("litellm.route"); + assert_eq!(summaries.len(), 1); + assert_eq!(summaries[0]["outcome"], "failure"); + assert!(traces.summaries("litellm.provider.send").is_empty()); +} + +#[rstest] +#[tokio::test] +async fn websocket_operations_trace_outcomes_without_capturing_frames_or_credentials( + traces: TraceCapture, +) { + use futures_util::{SinkExt, StreamExt}; + use litellm_core::responses::websocket::ResponsesWebSocketConnection; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let message = socket.next().await.unwrap().unwrap(); + socket.send(message).await.unwrap(); + let _ = socket.next().await; + }); + traces + .logger() + .instrument(async { + let connection = ResponsesWebSocketConnection::connect_url( + &format!("ws://{address}/responses"), + &std::collections::HashMap::from([( + "authorization".into(), + "private-key-sentinel".into(), + )]), + None, + ) + .await + .unwrap(); + connection + .send_text("private-frame-sentinel".into()) + .await + .unwrap(); + assert_eq!( + connection.recv_text().await.unwrap().as_deref(), + Some("private-frame-sentinel") + ); + connection.close().await.unwrap(); + assert!( + connection + .send_text("private-frame-sentinel".into()) + .await + .is_err() + ); + }) + .await; + server.await.unwrap(); + let sends = traces.summaries("litellm.websocket.send_text"); + assert_eq!(sends.len(), 2); + assert_eq!(sends[0]["outcome"], "success"); + assert_eq!(sends[1]["outcome"], "failure"); + assert_eq!( + traces.summaries("litellm.websocket.connect_url")[0]["outcome"], + "success" + ); + assert_eq!( + traces.summaries("litellm.websocket.recv_text")[0]["outcome"], + "success" + ); + assert_eq!( + traces.summaries("litellm.websocket.close")[0]["outcome"], + "success" + ); + assert!(!format!("{:?}", traces.records()).contains("private-")); +} diff --git a/litellm-rust/crates/core/tests/support/mod.rs b/litellm-rust/crates/core/tests/support/mod.rs index baf21d2970a..82663fbb418 100644 --- a/litellm-rust/crates/core/tests/support/mod.rs +++ b/litellm-rust/crates/core/tests/support/mod.rs @@ -350,3 +350,41 @@ where self.events.0.lock().unwrap().push(event.clone()); } } + +#[derive(Clone, Default)] +pub struct TraceCapture(Arc>>); + +impl TraceCapture { + pub fn logger(&self) -> litellm_tracing::Logger { + litellm_tracing::Logger::new(self.clone()) + } + + pub fn records(&self) -> Vec { + self.0.lock().unwrap().clone() + } + + pub fn summaries(&self, name: &str) -> Vec { + self.records() + .into_iter() + .filter(|record| record["span_name"] == name) + .collect() + } +} + +impl litellm_tracing::Sink for TraceCapture { + fn enabled(&self, metadata: &litellm_tracing::Metadata<'_>) -> bool { + metadata.target().starts_with("litellm_core") + } + + fn emit(&self, record: &litellm_tracing::Record) { + self.0 + .lock() + .unwrap() + .push(Value::Object(record.fields.clone())); + } +} + +#[rstest::fixture] +pub fn traces() -> TraceCapture { + TraceCapture::default() +} diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index 046ad9df77c..a27597e0bfe 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -20,6 +20,7 @@ serde.workspace = true serde_json.workspace = true thiserror.workspace = true tokio.workspace = true +tracing.workspace = true veil.workspace = true webpki-roots.workspace = true diff --git a/litellm-rust/crates/http/src/websocket.rs b/litellm-rust/crates/http/src/websocket.rs index 3f108ea9703..e904392152a 100644 --- a/litellm-rust/crates/http/src/websocket.rs +++ b/litellm-rust/crates/http/src/websocket.rs @@ -44,6 +44,7 @@ fn tls_config() -> Result, Box( request: R, ) -> Result<(UpstreamWebSocket, Response), Box> diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 7cfb3f207d4..329fb63c8e7 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -61,6 +61,7 @@ tokio = { workspace = true, features = ["rt", "sync"] } url.workspace = true [dev-dependencies] +tracing.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-secrets-aws.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/logger/tests.rs b/litellm-rust/crates/python-bridge/src/logger/tests.rs index a9b7ef41c08..21ebf432c99 100644 --- a/litellm-rust/crates/python-bridge/src/logger/tests.rs +++ b/litellm-rust/crates/python-bridge/src/logger/tests.rs @@ -66,6 +66,19 @@ fn warning(py: Python<'_>) { }); } +#[tracing::instrument(name = "litellm.route", skip_all, fields(route = "fixture", outcome))] +async fn traced_operation(_secret: &str) -> PyResult<()> { + tokio::task::yield_now().await; + tracing::info!("inside route"); + tracing::Span::current().record("outcome", "success"); + Ok(()) +} + +#[pyfunction] +fn span_warning(py: Python<'_>) -> PyResult> { + super::run_async_value(py, traced_operation("private-key-sentinel")) +} + #[pyfunction] fn levels(py: Python<'_>) { super::capture(py).scope(|| { @@ -109,7 +122,7 @@ fn http_warning(py: Python<'_>) -> PyResult<()> { crate::http::call_config(py, &PyDict::new(py), false).map(|_| ()) } -#[test] +#[rstest::rstest] fn native_events_reach_python_with_levels_context_reentry_and_http_deduplication() { if std::env::var_os("LITELLM_LOGGER_TEST_PROCESS").is_none() { let output = Command::new(std::env::current_exe().unwrap()) @@ -156,6 +169,9 @@ fn native_events_reach_python_with_levels_context_reentry_and_http_deduplication locals .set_item("warning", wrap_pyfunction!(warning, py).unwrap()) .unwrap(); + locals + .set_item("span_warning", wrap_pyfunction!(span_warning, py).unwrap()) + .unwrap(); locals .set_item( "asynchronous_warning", @@ -263,6 +279,25 @@ try: ('warn', logging.WARNING), ('error', logging.ERROR), ] + before_spans = len(capture.records) + async def traced_request(): + session = session_id_var.set('span-session') + trace = trace_id_var.set('span-trace') + try: + await span_warning() + finally: + trace_id_var.reset(trace) + session_id_var.reset(session) + asyncio.run(traced_request()) + span_records = capture.records[before_spans:] + assert [r.getMessage() for r in span_records] == ['inside route', 'span closed'] + assert all(r.session_id == 'span-session' and r.trace_id == 'span-trace' for r in span_records) + assert all(r.rust_fields['route'] == 'fixture' for r in span_records) + assert span_records[1].rust_fields['outcome'] == 'success' + assert span_records[1].rust_fields['span_name'] == 'litellm.route' + assert span_records[1].rust_fields['duration_ms'] >= 0 + assert 'private-key-sentinel' not in repr([r.rust_fields for r in span_records]) + before = len(capture.records) litellm.ssl_ecdh_curve = 'logger-test-unsupported-curve' http_warning() diff --git a/litellm-rust/crates/tracing/README.md b/litellm-rust/crates/tracing/README.md index 7f086067f3e..16591ce5dfd 100644 --- a/litellm-rust/crates/tracing/README.md +++ b/litellm-rust/crates/tracing/README.md @@ -1,8 +1,8 @@ # Native diagnostic tracing -`litellm-tracing` connects standard `tracing` events to a host-provided `Sink`. It has no Python dependency and does not install a global subscriber +`litellm-tracing` connects standard `tracing` spans and events to a host-provided `Sink`. It has no Python dependency. Hosts choose scoped dispatch or explicitly install a global subscriber -Use the exported `debug!`, `info!`, `warn!`, and `error!` macros in native code. A host creates a `Logger` with its sink, uses `scope` for synchronous operations, and wraps futures with `instrument`. Instrument spawned futures explicitly because thread-local subscribers do not automatically follow spawned work +Use upstream `tracing` macros and `#[tracing::instrument(skip_all, fields(...))]` in native code. A host creates a `Logger` with its sink, uses `scope` for synchronous operations, and wraps futures with `instrument`. Propagate both spans and dispatch into spawned work and returned streams. `Logger::current()` captures the current dispatch. Existing event macro re-exports remain available Bindings implement `litellm_tracing::Sink` to connect events to their host runtime: @@ -17,6 +17,8 @@ Pass the implementation to `litellm_tracing::Logger::new(sink)`, then call `logg Records retain event metadata, the message, and typed event fields. Sink filtering runs for each event so runtime level changes take effect. Logging from inside a sink is suppressed to prevent recursion +`sink_layer(sink)` exposes the same adapter as a composable `tracing_subscriber::Layer`, with an independent dynamic filter. It records span field updates and inherits fields from outer to inner spans, with event fields taking precedence. Closing a span emits `span closed` at the span's level with `span_name` and monotonic `duration_ms`. Delivery rechecks the sink filter. Core records route outcomes and retains route spans until a stream ends, fails, or is dropped + The Python bridge scopes native execution to a sink that uses LiteLLM's existing Python logger. It preserves request correlation, redacts before delivering to handlers, maps Rust trace events to Python debug, and reports handler failures through `sys.unraisablehook`. It accepts LiteLLM targets only, keeping dependency wire diagnostics out of the application logger Python consumers continue using `litellm._logging` and its existing loggers, filters, formatters, and context setters. Catalog dispatch selects the processing backend for both Python and native diagnostics. The pure `Processor` takes explicit settings and never emits events diff --git a/litellm-rust/crates/tracing/src/layer.rs b/litellm-rust/crates/tracing/src/layer.rs new file mode 100644 index 00000000000..150cffa4d7f --- /dev/null +++ b/litellm-rust/crates/tracing/src/layer.rs @@ -0,0 +1,143 @@ +use std::{sync::Arc, time::Instant}; + +use serde_json::Map; +use tracing::{ + Dispatch, Event, Metadata, Subscriber, + span::{Attributes, Id, Record as SpanRecord}, +}; +use tracing_subscriber::{ + Layer, Registry, filter::dynamic_filter_fn, layer::Context, prelude::*, registry::LookupSpan, +}; + +use crate::{Emitting, Record, Sink}; + +pub fn sink_layer(sink: S) -> impl Layer +where + S: Sink, + R: Subscriber + for<'a> LookupSpan<'a>, +{ + let sink = Arc::new(sink); + let filter_sink = sink.clone(); + Output(sink).with_filter(dynamic_filter_fn(move |metadata, _| { + enabled(filter_sink.as_ref(), metadata) + })) +} + +pub(crate) fn dispatch(sink: impl Sink) -> Dispatch { + let sink = Arc::new(sink); + let filter_sink = sink.clone(); + Dispatch::new( + Registry::default() + .with(Output(sink)) + .with(dynamic_filter_fn(move |metadata, _| { + enabled(filter_sink.as_ref(), metadata) + })), + ) +} + +fn enabled(sink: &impl Sink, metadata: &Metadata<'_>) -> bool { + let Some(_guard) = Emitting::enter() else { + return false; + }; + sink.enabled(metadata) +} + +struct Output(Arc); + +struct SpanData { + record: Record, + started: Instant, +} + +impl Layer for Output +where + S: Sink, + R: Subscriber + for<'a> LookupSpan<'a>, +{ + fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, R>) { + let Some(_guard) = Emitting::enter() else { + return; + }; + let Some(span) = context.span(id) else { + return; + }; + let mut extensions = span.extensions_mut(); + if extensions.get_mut::().is_some() { + return; + } + let mut record = Record { + metadata: attributes.metadata(), + message: String::new(), + fields: Map::new(), + }; + attributes.record(&mut record); + extensions.insert(SpanData { + record, + started: Instant::now(), + }); + } + + fn on_record(&self, id: &Id, values: &SpanRecord<'_>, context: Context<'_, R>) { + let Some(_guard) = Emitting::enter() else { + return; + }; + let Some(span) = context.span(id) else { + return; + }; + if let Some(data) = span.extensions_mut().get_mut::() { + values.record(&mut data.record); + } + } + + fn on_event(&self, event: &Event<'_>, context: Context<'_, R>) { + let Some(_guard) = Emitting::enter() else { + return; + }; + let mut record = Record { + metadata: event.metadata(), + message: String::new(), + fields: Map::new(), + }; + if let Some(scope) = context.event_scope(event) { + for span in scope.from_root() { + if let Some(data) = span.extensions().get::() { + record.fields.extend(data.record.fields.clone()); + } + } + } + event.record(&mut record); + self.0.emit(&record); + } + + fn on_close(&self, id: Id, context: Context<'_, R>) { + let Some(_guard) = Emitting::enter() else { + return; + }; + let Some(span) = context.span(&id) else { + return; + }; + if !self.0.enabled(span.metadata()) { + return; + } + let mut record = Record { + metadata: span.metadata(), + message: "span closed".into(), + fields: Map::new(), + }; + for ancestor in span.scope().from_root() { + if let Some(data) = ancestor.extensions().get::() { + record.fields.extend(data.record.fields.clone()); + } + } + let extensions = span.extensions(); + let Some(data) = extensions.get::() else { + return; + }; + record.fields.insert("span_name".into(), span.name().into()); + record.fields.insert( + "duration_ms".into(), + (data.started.elapsed().as_secs_f64() * 1000.0).into(), + ); + self.0.emit(&record); + } +} diff --git a/litellm-rust/crates/tracing/src/lib.rs b/litellm-rust/crates/tracing/src/lib.rs index 4c6ec104f3a..16d299f1608 100644 --- a/litellm-rust/crates/tracing/src/lib.rs +++ b/litellm-rust/crates/tracing/src/lib.rs @@ -8,15 +8,15 @@ use std::{ use base64::{Engine, engine::general_purpose::STANDARD}; use serde_json::{Map, Value}; use tracing::{ - Dispatch, Event, Subscriber, + Dispatch, field::{Field, Visit}, - subscriber::Interest, }; -use tracing_subscriber::{Layer, Registry, layer::Context, prelude::*}; +mod layer; mod processing; mod redaction; +pub use layer::sink_layer; pub use processing::{DiagnosticInput, DiagnosticOutput, Policy, Processor}; pub use redaction::{REDACTED, SecretRedactor}; pub use tracing::{Level, Metadata, debug, error, info, trace, warn}; @@ -66,7 +66,13 @@ pub struct Logger { impl Logger { pub fn new(sink: impl Sink) -> Self { Self { - dispatch: Dispatch::new(Registry::default().with(Output(sink))), + dispatch: layer::dispatch(sink), + } + } + + pub fn current() -> Self { + Self { + dispatch: tracing::dispatcher::get_default(Clone::clone), } } @@ -108,37 +114,9 @@ impl Drop for Emitting { } } -struct Output(S); - -impl Layer for Output { - fn register_callsite(&self, _: &'static Metadata<'static>) -> Interest { - Interest::sometimes() - } - - fn enabled(&self, metadata: &Metadata<'_>, _: Context<'_, R>) -> bool { - let Some(_guard) = Emitting::enter() else { - return false; - }; - self.0.enabled(metadata) - } - - fn on_event(&self, event: &Event<'_>, _: Context<'_, R>) { - let Some(_guard) = Emitting::enter() else { - return; - }; - let mut record = Record { - metadata: event.metadata(), - message: String::new(), - fields: Map::new(), - }; - event.record(&mut record); - self.0.emit(&record); - } -} - impl Record { fn field(&mut self, field: &Field, value: Value) { - if field.name() == "message" { + if field.name() == "message" && self.metadata.is_event() { self.message = match value { Value::String(message) => message, value => value.to_string(), diff --git a/litellm-rust/crates/tracing/tests/logging.rs b/litellm-rust/crates/tracing/tests/logging.rs index 3387f259f8b..249ac021ca3 100644 --- a/litellm-rust/crates/tracing/tests/logging.rs +++ b/litellm-rust/crates/tracing/tests/logging.rs @@ -137,3 +137,145 @@ fn byte_chunk_logging_preserves_exact_bytes(#[case] bytes: &[u8], #[case] encodi }; assert_eq!(recovered, bytes); } + +#[rstest] +fn spans_inherit_fields_accept_updates_and_close_once() { + let (sender, receiver) = mpsc::channel(); + let logger = Logger::new(Output { + enabled: Arc::new(AtomicBool::new(true)), + sender, + }); + logger.scope(|| { + let route = tracing::info_span!( + "route", + route = "messages", + model = "requested", + message = "span message", + outcome = tracing::field::Empty + ); + route.in_scope(|| { + let child = tracing::debug_span!( + "provider", + model = "resolved", + status = tracing::field::Empty + ); + child.record("status", 201_u64); + child.in_scope(|| info!(model = "event", "response")); + route.record("outcome", "success"); + }); + }); + let records = receiver.try_iter().collect::>(); + assert_eq!(records.len(), 3); + assert_eq!(records[0].0, "response"); + assert_eq!( + records[0].1, + json!({"route": "messages", "model": "event", "status": 201, "message": "span message"}) + ); + assert_eq!(records[1].1["span_name"], "provider"); + assert_eq!(records[1].1["model"], "resolved"); + assert_eq!(records[1].1["status"], 201); + assert_eq!(records[2].1["span_name"], "route"); + assert_eq!(records[2].1["model"], "requested"); + assert_eq!(records[2].1["outcome"], "success"); + assert!( + records[2].1["duration_ms"].as_f64().unwrap() + >= records[1].1["duration_ms"].as_f64().unwrap() + ); +} + +#[rstest] +fn sink_layers_filter_independently_and_recheck_span_delivery() { + use tracing_subscriber::prelude::*; + + let (first_sender, first) = mpsc::channel(); + let (second_sender, second) = mpsc::channel(); + let enabled = Arc::new(AtomicBool::new(false)); + let subscriber = tracing_subscriber::Registry::default() + .with(litellm_tracing::sink_layer(Output { + enabled: enabled.clone(), + sender: first_sender, + })) + .with(litellm_tracing::sink_layer(Output { + enabled: Arc::new(AtomicBool::new(true)), + sender: second_sender, + })); + tracing::subscriber::with_default(subscriber, || { + info!("second only"); + enabled.store(true, Ordering::Relaxed); + let span = tracing::info_span!("shared", route = "responses"); + span.in_scope(|| info!("both")); + enabled.store(false, Ordering::Relaxed); + drop(span); + }); + assert_eq!( + first.try_iter().map(|record| record.0).collect::>(), + ["both"] + ); + let records = second.try_iter().collect::>(); + assert_eq!( + records + .iter() + .map(|record| record.0.as_str()) + .collect::>(), + ["second only", "both", "span closed"] + ); + assert_eq!(records[1].1["route"], "responses"); + assert_eq!(records[2].1["route"], "responses"); +} + +#[rstest] +fn disabled_spans_do_not_evaluate_fields() { + let (sender, receiver) = mpsc::channel(); + let logger = Logger::new(Output { + enabled: Arc::new(AtomicBool::new(false)), + sender, + }); + let evaluated = AtomicBool::new(false); + logger.scope(|| { + let _span = tracing::info_span!( + "disabled", + field = { + evaluated.store(true, Ordering::Relaxed); + "value" + } + ); + assert!(_span.is_disabled()); + }); + assert!(!evaluated.load(Ordering::Relaxed)); + assert!(receiver.try_recv().is_err()); +} + +#[tracing::instrument(name = "worker", skip_all, fields(id))] +async fn traced_worker(id: u64) { + tracing::Span::current().record("id", id); + tokio::task::yield_now().await; + info!("inside worker"); +} + +#[rstest] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn instrumented_futures_keep_span_fields_in_their_captured_sink() { + let tasks = (0..2) + .map(|id| { + let (sender, receiver) = mpsc::channel(); + let logger = Logger::new(Output { + enabled: Arc::new(AtomicBool::new(true)), + sender, + }); + ( + id, + tokio::spawn(logger.instrument(traced_worker(id))), + receiver, + ) + }) + .collect::>(); + for (id, task, receiver) in tasks { + task.await.unwrap(); + let records = receiver.try_iter().collect::>(); + assert_eq!(records.len(), 2); + assert_eq!(records[0].0, "inside worker"); + assert_eq!(records[0].1["id"], id); + assert_eq!(records[1].1["id"], id); + assert_eq!(records[1].1["span_name"], "worker"); + } +}