collab: Add `is_staff` to spans (#16389)

Marshall Bowers created 1 year ago

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

Change summary

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

Detailed changes

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

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(