From 3d997e5fd6f899f424d7507345c597406e731adf Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Fri, 16 Aug 2024 18:42:44 -0400
Subject: [PATCH] collab: Add `is_staff` to spans (#16389)

This PR adds the `is_staff` field to our LLM spans so that we can
distinguish between staff and non-staff traffic.

Release Notes:

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

diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs
index 4ec8b70ac0cb11..cf254a42eeb41b 100644
--- a/crates/collab/src/llm.rs
+++ b/crates/collab/src/llm.rs
@@ -141,7 +141,8 @@ 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);
+                .record("authn.jti", &claims.jti)
+                .record("is_staff", &claims.is_staff);
 
             req.extensions_mut().insert(claims);
             Ok::<_, Error>(next.run(req).await.into_response())
@@ -558,6 +559,7 @@ impl<S> Drop for TokenCountingStream<S> {
                     user_id = claims.user_id,
                     login = claims.github_user_login,
                     authn.jti = claims.jti,
+                    is_staff = claims.is_staff,
                     requests_this_minute = usage.requests_this_minute,
                     tokens_this_minute = usage.tokens_this_minute,
                 );
diff --git a/crates/collab/src/main.rs b/crates/collab/src/main.rs
index 35a80b702e0c24..a946d1662019b7 100644
--- a/crates/collab/src/main.rs
+++ b/crates/collab/src/main.rs
@@ -154,7 +154,8 @@ async fn main() -> Result<()> {
                             matched_path,
                             user_id = tracing::field::Empty,
                             login = tracing::field::Empty,
-                            authn.jti = tracing::field::Empty
+                            authn.jti = tracing::field::Empty,
+                            is_staff = tracing::field::Empty
                         )
                     })
                     .on_response(