From 8a85d6ef96cc9d38e1c0e7e420765cd4b31b6954 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Mon, 9 Dec 2024 17:58:14 -0500
Subject: [PATCH] collab: Make `metrics_id` required in `LlmTokenClaims`
 (#21771)

This PR makes the `metrics_id` field on the `LlmTokenClaims` required,
as we always have one in practice.

Release Notes:

- N/A
---
 crates/collab/src/api/billing.rs | 2 +-
 crates/collab/src/api/events.rs  | 4 ++--
 crates/collab/src/llm/token.rs   | 6 ++----
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/crates/collab/src/api/billing.rs b/crates/collab/src/api/billing.rs
index f6c2d88105..2a866879a9 100644
--- a/crates/collab/src/api/billing.rs
+++ b/crates/collab/src/api/billing.rs
@@ -130,7 +130,7 @@ async fn update_billing_preferences(
 
     SnowflakeRow::new(
         "Spend Limit Updated",
-        Some(user.metrics_id),
+        user.metrics_id,
         user.admin,
         None,
         json!({
diff --git a/crates/collab/src/api/events.rs b/crates/collab/src/api/events.rs
index 1dc036ca86..15085db2c2 100644
--- a/crates/collab/src/api/events.rs
+++ b/crates/collab/src/api/events.rs
@@ -1593,7 +1593,7 @@ pub struct SnowflakeRow {
 impl SnowflakeRow {
     pub fn new(
         event_type: impl Into<String>,
-        metrics_id: Option<Uuid>,
+        metrics_id: Uuid,
         is_staff: bool,
         system_id: Option<String>,
         event_properties: serde_json::Value,
@@ -1602,7 +1602,7 @@ impl SnowflakeRow {
             time: chrono::Utc::now(),
             event_type: event_type.into(),
             device_id: system_id,
-            user_id: metrics_id.map(|id| id.to_string()),
+            user_id: Some(metrics_id.to_string()),
             insert_id: Some(uuid::Uuid::new_v4().to_string()),
             event_properties,
             user_properties: Some(json!({"is_staff": is_staff})),
diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs
index 7e0706e2d5..bc0a14cda1 100644
--- a/crates/collab/src/llm/token.rs
+++ b/crates/collab/src/llm/token.rs
@@ -17,10 +17,8 @@ pub struct LlmTokenClaims {
     pub exp: u64,
     pub jti: String,
     pub user_id: u64,
-    #[serde(default)]
     pub system_id: Option<String>,
-    #[serde(default)]
-    pub metrics_id: Option<Uuid>,
+    pub metrics_id: Uuid,
     pub github_user_login: String,
     pub is_staff: bool,
     pub has_llm_closed_beta_feature_flag: bool,
@@ -56,7 +54,7 @@ impl LlmTokenClaims {
             jti: uuid::Uuid::new_v4().to_string(),
             user_id: user.id.to_proto(),
             system_id,
-            metrics_id: Some(user.metrics_id),
+            metrics_id: user.metrics_id,
             github_user_login: user.github_login.clone(),
             is_staff,
             has_llm_closed_beta_feature_flag,