fix(lens): preserve numeric tags and report OTLP error codes (#45206)

litellm-rust/crates/lens/src/ingest.rs

@@ -46,6 +46,15 @@ pub fn response(content_type: Option<&str>, outcome: Result<(), Error>) -> Respo
         .map(|_| StatusCode::OK)
         .unwrap_or_else(|error| error.status());
     let message = status.canonical_reason().unwrap_or("Trace request failed");
+    let rpc_code = match status {
+        StatusCode::OK => 0,
+        StatusCode::BAD_REQUEST => 3,
+        StatusCode::UNAUTHORIZED => 16,
+        StatusCode::PAYLOAD_TOO_LARGE | StatusCode::TOO_MANY_REQUESTS => 8,
+        StatusCode::CONFLICT => 10,
+        StatusCode::SERVICE_UNAVAILABLE => 14,
+        _ => 2,
+    };
     let protobuf = content_type.is_some_and(|value| {
         value
             .split(';')
@@ -58,7 +67,7 @@ pub fn response(content_type: Option<&str>, outcome: Result<(), Error>) -> Respo
                 Vec::new()
             } else {
                 OtlpError {
-                    code: 0,
+                    code: rpc_code,
                     message: message.into(),
                 }
                 .encode_to_vec()
@@ -70,7 +79,7 @@ pub fn response(content_type: Option<&str>, outcome: Result<(), Error>) -> Respo
             if outcome.is_ok() {
                 b"{}".to_vec()
             } else {
-                serde_json::json!({"code": 0, "message": message})
+                serde_json::json!({"code": rpc_code, "message": message})
                     .to_string()
                     .into_bytes()
             },
@@ -169,3 +178,67 @@ async fn store(
     .await?;
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{OtlpError, response};
+    use crate::Error;
+    use axum::body::to_bytes;
+    use prost::Message;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case::invalid_request(Error::InvalidRequest, 400, 3, false)]
+    #[case::unauthenticated(Error::Unauthorized, 401, 16, false)]
+    #[case::payload_too_large(Error::TooLarge, 413, 8, false)]
+    #[case::credentials_pending(Error::CredentialsPending, 429, 8, true)]
+    #[case::storage_unavailable(Error::Unavailable, 503, 14, true)]
+    #[case::conflict(Error::TraceChanged, 409, 10, false)]
+    #[tokio::test]
+    async fn rejected_batches_have_matching_http_and_rpc_errors(
+        #[case] error: Error,
+        #[case] http_status: u16,
+        #[case] rpc_code: i32,
+        #[case] retryable: bool,
+        #[values("application/json", "application/x-protobuf")] content_type: &str,
+    ) {
+        let reply = response(Some(content_type), Err(error));
+        assert_eq!(reply.status().as_u16(), http_status);
+        assert_eq!(reply.headers()["content-type"], content_type);
+        assert_eq!(
+            reply
+                .headers()
+                .get("retry-after")
+                .map(|v| v.to_str().unwrap()),
+            retryable.then_some("5")
+        );
+        let message = reply.status().canonical_reason().unwrap();
+        let body = to_bytes(reply.into_body(), 1024).await.unwrap();
+        if content_type == "application/x-protobuf" {
+            let status = OtlpError::decode(body).unwrap();
+            assert_eq!(status.code, rpc_code);
+            assert_eq!(status.message, message);
+        } else {
+            let status: serde_json::Value = serde_json::from_slice(&body).unwrap();
+            assert_eq!(
+                status,
+                serde_json::json!({"code": rpc_code, "message": message})
+            );
+        }
+    }
+
+    #[rstest]
+    #[case::json("application/json", b"{}")]
+    #[case::protobuf("application/x-protobuf", b"")]
+    #[tokio::test]
+    async fn accepted_batches_keep_the_empty_export_response(
+        #[case] content_type: &str,
+        #[case] expected: &[u8],
+    ) {
+        let reply = response(Some(content_type), Ok(()));
+        assert_eq!(reply.status(), 200);
+        assert_eq!(reply.headers()["content-type"], content_type);
+        assert!(!reply.headers().contains_key("retry-after"));
+        assert_eq!(to_bytes(reply.into_body(), 1024).await.unwrap(), expected);
+    }
+}

litellm/tracing/types.py

@@ -49,7 +49,7 @@ class SpendLogPayload(TypedDict, total=False):
     cache_hit: ReadOnly[bool | None]
     session_id: ReadOnly[str | None]
     trace_id: ReadOnly[str | None]
-    request_tags: ReadOnly[Sequence[str] | None]
+    request_tags: ReadOnly[Sequence[object] | None]
     messages: ReadOnly[object]
     response: ReadOnly[object]
 

tests/unit/tracing/test_exporter.py

@@ -9,6 +9,57 @@ import pytest
 from litellm.tracing.exporter import MAX_BUFFER_EVENTS, MAX_EVENT_BYTES, ExportFailure, LensExporter, encode_record
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failed", (False, True), ids=("success-event", "failure-event"))
+@pytest.mark.parametrize(
+    ("tags", "expected"),
+    (
+        pytest.param((7,), ("7",), id="numeric"),
+        pytest.param(("env:prod", 0, -7, 2.5, True, None), ("env:prod", "0", "-7", "2.5", "True", "None"), id="mixed"),
+        pytest.param(("env:prod", "agent:research"), ("env:prod", "agent:research"), id="strings"),
+    ),
+)
+async def test_request_tags_are_normalized_without_dropping_the_record(
+    failed: bool, tags: tuple[object, ...], expected: tuple[str, ...]
+) -> None:
+    requests: Final = asyncio.Queue[httpx.Request]()
+
+    def accept(request: httpx.Request) -> httpx.Response:
+        requests.put_nowait(request)
+        return httpx.Response(204)
+
+    async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(accept)) as client:
+        exporter: Final = LensExporter(client)
+        exporter.start()
+        callback: Final = exporter.async_log_failure_event if failed else exporter.async_log_success_event
+        await callback(
+            {
+                "response_cost": 0.12,
+                "standard_logging_object": {
+                    "id": "tagged-request",
+                    "status": "failure" if failed else "success",
+                    "response_cost": 0.12,
+                    "request_tags": list(tags),
+                },
+            },
+            None,
+            None,
+            None,
+        )
+        await exporter.aclose()
+    assert exporter.rows_written == 1
+    assert exporter.rows_dropped == 0
+    assert requests.qsize() == 1
+    request: Final = requests.get_nowait()
+    rows: Final = json.loads(request.content)
+    assert request.url.path == "/internal/spend"
+    assert len(rows) == 1
+    assert rows[0]["request_id"] == "tagged-request"
+    assert rows[0]["status"] == ("failure" if failed else "success")
+    assert rows[0]["spend"] == 0.12
+    assert rows[0]["request_tags"] == list(expected)
+
+
 @pytest.mark.asyncio
 async def test_request_export_ignores_unrelated_model_metadata_and_preserves_billing() -> None:
     received: Final = asyncio.Future[httpx.Request]()