collab: Attach GitHub login to LLM spans (#16316)

Marshall Bowers created 1 year ago

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

Change summary

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(+)

Detailed changes

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);

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,
         };

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
                         )
                     })

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,