From 5e05821d1875f8963ab957bc9d359289a2fee1bb Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Thu, 15 Aug 2024 15:49:12 -0400
Subject: [PATCH] collab: Attach `user_id` to LLM spans (#16311)

This PR updates the LLM service to attach the user ID to the spans.

Release Notes:

- N/A
---
 crates/collab/src/llm.rs  | 4 +++-
 crates/collab/src/main.rs | 1 +
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs
index 609610f15c..2bffc39947 100644
--- a/crates/collab/src/llm.rs
+++ b/crates/collab/src/llm.rs
@@ -138,7 +138,9 @@ async fn validate_api_token<B>(mut req: Request<B>, next: Next<B>) -> impl IntoR
                 ));
             }
 
-            tracing::Span::current().record("authn.jti", &claims.jti);
+            tracing::Span::current()
+                .record("user_id", claims.user_id)
+                .record("authn.jti", &claims.jti);
 
             req.extensions_mut().insert(claims);
             Ok::<_, Error>(next.run(req).await.into_response())
diff --git a/crates/collab/src/main.rs b/crates/collab/src/main.rs
index 6994109443..d070036e64 100644
--- a/crates/collab/src/main.rs
+++ b/crates/collab/src/main.rs
@@ -150,6 +150,7 @@ async fn main() -> Result<()> {
                             "http_request",
                             method = ?request.method(),
                             matched_path,
+                            user_id = tracing::field::Empty,
                             authn.jti = tracing::field::Empty
                         )
                     })