cloud_api_client: Key the cached LlmApiToken by organization (#58915)

The tokens are organization-specific (the organization id in the JWT
claims). Before this commit, the `LlmApiToken` cache was an
`Option<String>`, and `cached()` ignored the `organization_id` argument
on cache hit. Correctness relied on `RefreshLlmTokenListener` clearing
and refreshing the token when it observed an `OrganizationChanged`
event.

That leaves two windows where a token minted for the wrong organization
could be served from the cache:

- the refresh runs in a spawned task, so between the organization switch
and the task acquiring the write lock, any request calling `cached()`
would get the previous organization's token.
- `Client::authenticated_llm_request` snapshots the organization id at
call time; if the server demanded a token refresh after the user had
switched organizations mid-request, the retry path would mint a fresh
token for the old organization and write it into the shared cache,
poisoning it for all subsequent requests until the next refresh event.

The fix is to store the organization id corresponding to the token next
to it in the cache, and to check that the cached token is for the
correct organization in `cached()`.

Release Notes:

- Fixed a race where LLM and edit prediction requests made immediately
after switching organizations could be attributed to the previously
selected organization.

crates/cloud_api_client/src/llm_token.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{fmt, sync::Arc};
 
 use async_lock::{RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard};
 use cloud_api_types::OrganizationId;
@@ -6,11 +6,27 @@ use cloud_api_types::OrganizationId;
 use crate::{ClientApiError, CloudApiClient};
 
 #[derive(Clone, Default)]
-pub struct LlmApiToken(Arc<RwLock<Option<String>>>);
+pub struct LlmApiToken(Arc<RwLock<Option<CachedLlmApiToken>>>);
+
+struct CachedLlmApiToken {
+    /// The organization ID the token was minted for.
+    organization_id: OrganizationId,
+    token: String,
+}
+
+impl fmt::Debug for CachedLlmApiToken {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("CachedLlmApiToken")
+            .field("organization_id", &self.organization_id)
+            .field("token", &"<redacted>")
+            .finish()
+    }
+}
 
 impl LlmApiToken {
-    /// Returns the cached LLM token, fetching a fresh one only if none has
-    /// been cached yet. The returned token is not validated; callers must
+    /// Returns the cached LLM token, fetching a fresh one if none has been
+    /// cached yet or if the cached token was minted for a different
+    /// organization. The returned token is not validated; callers must
     /// be prepared to refresh it (via [`LlmApiToken::refresh`]) if the
     /// server rejects it.
     pub async fn cached(
@@ -20,7 +36,12 @@ impl LlmApiToken {
         organization_id: OrganizationId,
     ) -> Result<String, ClientApiError> {
         let lock = self.0.upgradable_read().await;
-        if let Some(token) = lock.as_ref() {
+        if let Some(CachedLlmApiToken {
+            organization_id: cached_organization_id,
+            token,
+        }) = lock.as_ref()
+            && *cached_organization_id == organization_id
+        {
             Ok(token.to_string())
         } else {
             Self::fetch(
@@ -62,15 +83,21 @@ impl LlmApiToken {
     }
 
     async fn fetch(
-        mut lock: RwLockWriteGuard<'_, Option<String>>,
+        mut lock: RwLockWriteGuard<'_, Option<CachedLlmApiToken>>,
         client: &CloudApiClient,
         system_id: Option<String>,
         organization_id: OrganizationId,
     ) -> Result<String, ClientApiError> {
-        let result = client.create_llm_token(system_id, organization_id).await;
+        let result = client
+            .create_llm_token(system_id, organization_id.clone())
+            .await;
         match result {
             Ok(response) => {
-                *lock = Some(response.token.0.clone());
+                *lock = Some(CachedLlmApiToken {
+                    organization_id,
+                    token: response.token.0.clone(),
+                });
+
                 Ok(response.token.0)
             }
             Err(err) => {