Add support for refreshing outdated LLM tokens (#47512)

Marshall Bowers created 2 weeks ago

This PR adds support for refreshing LLM tokens that are "outdated"—that
is, that are missing some required claims.

Release Notes:

- Fixed some instances of authentication errors with the Zed API that
could be resolved automatically by refreshing the token.

Change summary

crates/cloud_llm_client/src/cloud_llm_client.rs | 8 ++++++++
crates/language_model/src/model/cloud_model.rs  | 3 ++-
2 files changed, 10 insertions(+), 1 deletion(-)

Detailed changes

crates/cloud_llm_client/src/cloud_llm_client.rs

@@ -17,6 +17,14 @@ pub const ZED_VERSION_HEADER_NAME: &str = "x-zed-version";
 /// The client may use this as a signal to refresh the token.
 pub const EXPIRED_LLM_TOKEN_HEADER_NAME: &str = "x-zed-expired-token";
 
+/// The name of the header used to indicate when a request failed due to an outdated LLM token.
+///
+/// A token is considered "outdated" when we can't parse the claims (e.g., after adding a new required claim).
+///
+/// This is distinct from [`EXPIRED_LLM_TOKEN_HEADER_NAME`] which indicates the token's time-based validity has passed.
+/// An outdated token means the token's structure is incompatible with the current server expectations.
+pub const OUTDATED_LLM_TOKEN_HEADER_NAME: &str = "x-zed-outdated-token";
+
 /// The name of the header used to indicate the usage limit for edit predictions.
 pub const EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME: &str = "x-zed-edit-predictions-usage-limit";
 

crates/language_model/src/model/cloud_model.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 use anyhow::Result;
 use client::Client;
 use cloud_api_types::websocket_protocol::MessageToClient;
-use cloud_llm_client::EXPIRED_LLM_TOKEN_HEADER_NAME;
+use cloud_llm_client::{EXPIRED_LLM_TOKEN_HEADER_NAME, OUTDATED_LLM_TOKEN_HEADER_NAME};
 use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Global, ReadGlobal as _};
 use smol::lock::{RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard};
 use thiserror::Error;
@@ -61,6 +61,7 @@ pub trait NeedsLlmTokenRefresh {
 impl NeedsLlmTokenRefresh for http_client::Response<http_client::AsyncBody> {
     fn needs_llm_token_refresh(&self) -> bool {
         self.headers().get(EXPIRED_LLM_TOKEN_HEADER_NAME).is_some()
+            || self.headers().get(OUTDATED_LLM_TOKEN_HEADER_NAME).is_some()
     }
 }