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
This commit is contained in:
Marshall Bowers 2024-08-15 17:06:20 -04:00 committed by GitHub
parent df20bae80e
commit 9233418cb8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 11 additions and 0 deletions

View File

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

View File

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

View File

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

View File

@ -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,