From 9233418cb85ee73f9ece42e3a1ed97791a60edf9 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Thu, 15 Aug 2024 17:06:20 -0400
Subject: [PATCH] collab: Attach GitHub login to LLM spans (#16316)

This PR updates the LLM service to include the GitHub login on its
spans.

We need to pass this information through on the LLM token, so it will
temporarily be `None` until this change is deployed and new tokens have
been issued.

Release Notes:

- N/A
---
 crates/collab/src/llm.rs       | 1 +
 crates/collab/src/llm/token.rs | 8 ++++++++
 crates/collab/src/main.rs      | 1 +
 crates/collab/src/rpc.rs       | 1 +
 4 files changed, 11 insertions(+)

diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs
index 2bffc39947..151b6a247c 100644
--- a/crates/collab/src/llm.rs
+++ b/crates/collab/src/llm.rs
@@ -140,6 +140,7 @@ async fn validate_api_token<B>(mut req: Request<B>, next: Next<B>) -> impl IntoR
 
             tracing::Span::current()
                 .record("user_id", claims.user_id)
+                .record("login", claims.github_user_login.clone())
                 .record("authn.jti", &claims.jti);
 
             req.extensions_mut().insert(claims);
diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs
index e2350a853a..f789e5c220 100644
--- a/crates/collab/src/llm/token.rs
+++ b/crates/collab/src/llm/token.rs
@@ -13,6 +13,12 @@ pub struct LlmTokenClaims {
     pub exp: u64,
     pub jti: String,
     pub user_id: u64,
+    // This field is temporarily optional so it can be added
+    // in a backwards-compatible way. We can make it required
+    // once all of the LLM tokens have cycled (~1 hour after
+    // this change has been deployed).
+    #[serde(default)]
+    pub github_user_login: Option<String>,
     pub is_staff: bool,
     pub plan: rpc::proto::Plan,
 }
@@ -22,6 +28,7 @@ const LLM_TOKEN_LIFETIME: Duration = Duration::from_secs(60 * 60);
 impl LlmTokenClaims {
     pub fn create(
         user_id: UserId,
+        github_user_login: String,
         is_staff: bool,
         plan: rpc::proto::Plan,
         config: &Config,
@@ -37,6 +44,7 @@ impl LlmTokenClaims {
             exp: (now + LLM_TOKEN_LIFETIME).timestamp() as u64,
             jti: uuid::Uuid::new_v4().to_string(),
             user_id: user_id.to_proto(),
+            github_user_login: Some(github_user_login),
             is_staff,
             plan,
         };
diff --git a/crates/collab/src/main.rs b/crates/collab/src/main.rs
index d070036e64..5d4fc2abe0 100644
--- a/crates/collab/src/main.rs
+++ b/crates/collab/src/main.rs
@@ -151,6 +151,7 @@ async fn main() -> Result<()> {
                             method = ?request.method(),
                             matched_path,
                             user_id = tracing::field::Empty,
+                            login = tracing::field::Empty,
                             authn.jti = tracing::field::Empty
                         )
                     })
diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs
index c84f818635..b3dc904df1 100644
--- a/crates/collab/src/rpc.rs
+++ b/crates/collab/src/rpc.rs
@@ -4941,6 +4941,7 @@ async fn get_llm_api_token(
     }
     let token = LlmTokenClaims::create(
         user.id,
+        user.github_login.clone(),
         session.is_staff(),
         session.current_plan(db).await?,
         &session.app_state.config,