yenru0/zed
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
c8518b331739f072b6bbfc7cc8a4cb3bc6e990d9
zed/crates/language_models_cloud/src/language_models_cloud.rs
T
Conrad Irwin c8518b3317 Preserve language model provider rejection details (#63095) 
Language model providers can reject a request inside an otherwise
successful response stream. Those failures were represented by
overlapping completion-error variants and provider-specific code enums,
and OpenAI Responses errors such as `cyber_policy` fell through to
`Other`. Agent sessions consequently retried permanent provider
rejections and user interfaces lost the provider's actionable message.

This gives every provider-originated failure one shared representation.
`ProviderRejection` preserves the provider, exact message, retry delay,
raw provider code, and an HTTP status only when one was actually
reported. `ProviderErrorCategory` supplies the normalized meaning used
by shared UI and policy code. Provider adapters map their native wire
errors into that category explicitly, so typed SDK exceptions and errors
embedded in successful streams no longer fabricate HTTP statuses.

Retry selection is independent from error construction.
`LanguageModelCompletionError::retry_delay` owns transient
classification and a bounded schedule of 5, 10, 20, then 40 seconds. A
provider-supplied `retry_after` overrides that schedule. The scheduling
layer adds positive jitter of up to 10%, while interactive requests and
evaluation runs retain their own attempt budgets. Permanent status-less
errors such as `cyber_policy` do not retry. OpenRouter also keeps its
numeric body code distinct from an optional HTTP status, including for
errors embedded in a stream.

Testing performed:

- `corgi test -p language_model_core`
- `corgi test -p anthropic` (33 tests)
- `corgi test -p open_ai` (77 tests)
- `cargo nextest run -p open_router -p copilot_chat -p
language_models_cloud -p language_models` (125 tests)
- `cargo nextest run -p agent -p agent_ui` (1,153 tests; 11 skipped)
- `cargo check -p open_router -p copilot_chat -p language_models_cloud
-p language_models -p agent -p agent_ui`
- `./script/clippy -p language_model_core -p anthropic -p open_ai -p
open_router -p copilot_chat -p language_models_cloud -p language_models
-p agent -p agent_ui` completed the release Clippy phase; Cargo Shear
then reported the repository's existing warnings
- `cargo fmt --all --check`

Suggested .rules additions

None.

Release Notes:

- Fixed permanent language model provider rejections being retried and
displayed without their actionable provider message.
2026-08-24 19:36:04 +00:00

1908 lines
72 KiB
Rust
Raw Blame History

use anthropic::AnthropicModelMode;
use anyhow::{Context as _, Result};
use cloud_llm_client::{
CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, CLIENT_SUPPORTS_STATUS_STREAM_ENDED_HEADER_NAME,
CLIENT_SUPPORTS_X_AI_HEADER_NAME, CompletionBody, CompletionEvent, CompletionRequestStatus,
EXPIRED_LLM_TOKEN_HEADER_NAME, ListModelsResponse, OUTDATED_LLM_TOKEN_HEADER_NAME,
SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, ZED_VERSION_HEADER_NAME,
};
use futures::{
AsyncBufReadExt, AsyncReadExt as _, FutureExt, Stream, StreamExt,
future::BoxFuture,
io::BufReader,
stream::{self, BoxStream},
};
use google_ai::GoogleModelMode;
use gpui::{AppContext, AsyncApp, Context, Task};
use http_client::http::{HeaderMap, HeaderValue};
use http_client::{
AsyncBody, HttpClient, HttpClientWithUrl, HttpRequestExt, Method, Response, StatusCode,
};
use language_model::{
ANTHROPIC_PROVIDER_ID, ANTHROPIC_PROVIDER_NAME, BASETEN_PROVIDER_ID, BASETEN_PROVIDER_NAME,
CompactionResult, DisabledReason, GOOGLE_PROVIDER_ID, GOOGLE_PROVIDER_NAME, LanguageModel,
LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelEffortLevel,
LanguageModelId, LanguageModelName, LanguageModelProviderId, LanguageModelProviderName,
LanguageModelRequest, LanguageModelToolChoice, LanguageModelToolSchemaFormat,
OPEN_AI_PROVIDER_ID, OPEN_AI_PROVIDER_NAME, ProviderErrorCategory, RateLimiter,
X_AI_PROVIDER_ID, X_AI_PROVIDER_NAME, ZED_CLOUD_PROVIDER_ID, ZED_CLOUD_PROVIDER_NAME,
};
use schemars::JsonSchema;
use semver::Version;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use std::collections::VecDeque;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::Arc;
use std::task::Poll;
use std::time::Duration;
use thiserror::Error;
use anthropic::completion::{
AnthropicEventMapper, AnthropicPromptCacheMode, collect_compaction_result, into_anthropic,
};
use google_ai::completion::{GoogleEventMapper, into_google};
use open_ai::completion::{
ChatCompletionMaxTokensParameter, OpenAiEventMapper, OpenAiResponseEventMapper, into_open_ai,
into_open_ai_response, token_usage_from_response_usage,
};
const PROVIDER_ID: LanguageModelProviderId = ZED_CLOUD_PROVIDER_ID;
const PROVIDER_NAME: LanguageModelProviderName = ZED_CLOUD_PROVIDER_NAME;
/// Trait for acquiring and refreshing LLM authentication tokens.
pub trait CloudLlmTokenProvider: Send + Sync {
type AuthContext: Clone + Send + 'static;
fn auth_context(&self, cx: &impl AppContext) -> Self::AuthContext;
fn cached_token(&self, auth_context: Self::AuthContext) -> BoxFuture<'static, Result<String>>;
fn refresh_token(&self, auth_context: Self::AuthContext) -> BoxFuture<'static, Result<String>>;
/// Whether the user has consented to upstream providers retaining
/// inference logs for models that require it (see
/// [`LanguageModel::requires_data_retention`]).
fn has_data_retention_consent(&self, cx: &impl AppContext) -> bool;
}
/// Sends an authenticated request to the Zed LLM service, retrying once with
/// a refreshed token if the server signals that the cached LLM token is
/// expired or otherwise rejected. Returns the raw response so callers can
/// inspect headers and stream the body.
pub async fn authenticated_llm_request<TP: CloudLlmTokenProvider>(
http_client: &HttpClientWithUrl,
token_provider: &TP,
auth_context: TP::AuthContext,
build_request: impl Fn(&str) -> Result<http_client::Request<AsyncBody>>,
) -> Result<Response<AsyncBody>> {
let token = token_provider.cached_token(auth_context.clone()).await?;
let response = http_client.send(build_request(&token)?).await?;
if !needs_llm_token_refresh(&response) && response.status() != StatusCode::UNAUTHORIZED {
return Ok(response);
}
log::info!("LLM token rejected; refreshing and retrying request");
let token = token_provider.refresh_token(auth_context).await?;
http_client.send(build_request(&token)?).await
}
#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum ModelMode {
#[default]
Default,
Thinking {
/// The maximum number of tokens to use for reasoning. Must be lower than the model's `max_output_tokens`.
budget_tokens: Option<u32>,
},
}
impl From<ModelMode> for AnthropicModelMode {
fn from(value: ModelMode) -> Self {
match value {
ModelMode::Default => AnthropicModelMode::Default,
ModelMode::Thinking { budget_tokens } => AnthropicModelMode::Thinking { budget_tokens },
}
}
}
pub struct CloudLanguageModel<TP: CloudLlmTokenProvider> {
pub id: LanguageModelId,
pub model: Arc<cloud_llm_client::LanguageModel>,
pub token_provider: Arc<TP>,
pub http_client: Arc<HttpClientWithUrl>,
pub app_version: Option<Version>,
pub request_limiter: RateLimiter,
}
pub struct PerformLlmCompletionResponse {
pub response: Response<AsyncBody>,
pub includes_status_messages: bool,
}
impl<TP: CloudLlmTokenProvider> CloudLanguageModel<TP> {
pub async fn perform_llm_completion(
http_client: &HttpClientWithUrl,
token_provider: &TP,
auth_context: TP::AuthContext,
app_version: Option<Version>,
body: CompletionBody,
) -> Result<PerformLlmCompletionResponse, LanguageModelCompletionError> {
Self::perform_llm_request(
"/completions",
true,
http_client,
token_provider,
auth_context,
app_version,
body,
)
.await
}
async fn perform_llm_compaction(
http_client: &HttpClientWithUrl,
token_provider: &TP,
auth_context: TP::AuthContext,
app_version: Option<Version>,
body: CompletionBody,
) -> Result<PerformLlmCompletionResponse, LanguageModelCompletionError> {
Self::perform_llm_request(
"/completions/compact",
false,
http_client,
token_provider,
auth_context,
app_version,
body,
)
.await
}
async fn perform_llm_request(
path: &str,
request_status_messages: bool,
http_client: &HttpClientWithUrl,
token_provider: &TP,
auth_context: TP::AuthContext,
app_version: Option<Version>,
body: CompletionBody,
) -> Result<PerformLlmCompletionResponse, LanguageModelCompletionError> {
let url = http_client
.build_zed_llm_url(path, &[])
.map_err(LanguageModelCompletionError::Other)?;
let host = url.host_str().unwrap_or(url.as_str()).to_owned();
let body = serde_json::to_string(&body).map_err(|error| {
LanguageModelCompletionError::SerializeRequest {
provider: PROVIDER_NAME,
error,
}
})?;
let mut response =
authenticated_llm_request(http_client, token_provider, auth_context, |token| {
let mut request = http_client::Request::builder()
.method(Method::POST)
.uri(url.as_ref())
.when_some(app_version.as_ref(), |builder, app_version| {
builder.header(ZED_VERSION_HEADER_NAME, app_version.to_string())
})
.header("Content-Type", "application/json")
.header("Authorization", format!("Bearer {token}"));
if request_status_messages {
request = request
.header(CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, "true")
.header(CLIENT_SUPPORTS_STATUS_STREAM_ENDED_HEADER_NAME, "true");
}
Ok(request.body(body.clone().into())?)
})
.await
.map_err(|error| LanguageModelCompletionError::HttpSend {
provider: PROVIDER_NAME,
host,
error,
})?;
let status = response.status();
if status.is_success() {
let includes_status_messages = request_status_messages
&& response
.headers()
.get(SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME)
.is_some();
return Ok(PerformLlmCompletionResponse {
response,
includes_status_messages,
});
}
if status == StatusCode::PAYMENT_REQUIRED {
return Err(LanguageModelCompletionError::from_provider_response(
PROVIDER_NAME,
Some(status),
None,
"payment required to use this language model; please upgrade your account"
.to_string(),
None,
ProviderErrorCategory::PaymentRequired,
));
}
let mut body = String::new();
let headers = response.headers().clone();
response
.body_mut()
.read_to_string(&mut body)
.await
.map_err(|error| LanguageModelCompletionError::ApiReadResponseError {
provider: PROVIDER_NAME,
error,
})?;
Err(ApiError {
status,
body,
headers,
}
.into())
}
}
impl<TP: CloudLlmTokenProvider + 'static> CloudLanguageModel<TP> {
fn compact_anthropic(
&self,
request: LanguageModelRequest,
cx: &AsyncApp,
) -> BoxFuture<'static, Result<CompactionResult, LanguageModelCompletionError>> {
let thread_id = request.thread_id.clone();
let prompt_id = request.prompt_id.clone();
let app_version = self.app_version.clone();
let mut request = match into_anthropic(
request,
self.model.id.to_string(),
1.0,
self.model.max_output_tokens as u64,
AnthropicModelMode::Default,
AnthropicPromptCacheMode::Automatic,
&ANTHROPIC_PROVIDER_ID,
) {
Ok(request) => request.into_compact_request(),
Err(error) => return async move { Err(error.into()) }.boxed(),
};
if !self.model.supports_fast_mode {
request.speed = None;
}
let http_client = self.http_client.clone();
let token_provider = self.token_provider.clone();
let auth_context = token_provider.auth_context(cx);
let future = self.request_limiter.run(async move {
let PerformLlmCompletionResponse {
response,
includes_status_messages,
} = Self::perform_llm_completion(
&http_client,
&*token_provider,
auth_context,
app_version,
CompletionBody {
thread_id,
prompt_id,
provider: cloud_llm_client::LanguageModelProvider::Anthropic,
model: request.model.clone(),
provider_request: serde_json::to_value(&request).map_err(|error| {
LanguageModelCompletionError::SerializeRequest {
provider: ANTHROPIC_PROVIDER_NAME,
error,
}
})?,
},
)
.await?;
let mut mapper =
AnthropicEventMapper::new(ANTHROPIC_PROVIDER_NAME, ANTHROPIC_PROVIDER_ID);
let stream = map_cloud_completion_events(
Box::pin(response_lines(response, includes_status_messages)),
&ANTHROPIC_PROVIDER_NAME,
move |event| mapper.map_event(event),
);
let (context, usage) =
collect_compaction_result(stream, ANTHROPIC_PROVIDER_NAME).await?;
Ok(CompactionResult { context, usage })
});
future.boxed()
}
}
fn needs_llm_token_refresh(response: &Response<AsyncBody>) -> bool {
response
.headers()
.get(EXPIRED_LLM_TOKEN_HEADER_NAME)
.is_some()
|| response
.headers()
.get(OUTDATED_LLM_TOKEN_HEADER_NAME)
.is_some()
}
#[derive(Debug, Error)]
#[error("cloud language model request failed with status {status}: {body}")]
struct ApiError {
status: StatusCode,
body: String,
headers: HeaderMap<HeaderValue>,
}
/// Represents error responses from Zed's cloud API.
///
/// Example JSON for an upstream HTTP error:
/// ```json
/// {
/// "code": "upstream_http_error",
/// "message": "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout",
/// "upstream_status": 503
/// }
/// ```
#[derive(Debug, serde::Deserialize)]
struct CloudApiError {
code: String,
message: String,
#[serde(default)]
#[serde(deserialize_with = "deserialize_optional_status_code")]
upstream_status: Option<StatusCode>,
#[serde(default)]
retry_after: Option<f64>,
}
fn deserialize_optional_status_code<'de, D>(deserializer: D) -> Result<Option<StatusCode>, D::Error>
where
D: serde::Deserializer<'de>,
{
let opt: Option<u16> = Option::deserialize(deserializer)?;
Ok(opt.and_then(|code| StatusCode::from_u16(code).ok()))
}
impl From<ApiError> for LanguageModelCompletionError {
fn from(error: ApiError) -> Self {
if let Ok(cloud_error) = serde_json::from_str::<CloudApiError>(&error.body) {
if cloud_error.code.starts_with("upstream_http_") {
let status = if let Some(status) = cloud_error.upstream_status {
status
} else if cloud_error.code.ends_with("_error") {
error.status
} else {
// If there's a status code in the code string (e.g. "upstream_http_429")
// then use that; otherwise, see if the JSON contains a status code.
cloud_error
.code
.strip_prefix("upstream_http_")
.and_then(|code_str| code_str.parse::<u16>().ok())
.and_then(|code| StatusCode::from_u16(code).ok())
.unwrap_or(error.status)
};
let category =
ProviderErrorCategory::from_http_status(status, &cloud_error.message);
return LanguageModelCompletionError::from_provider_response(
PROVIDER_NAME,
Some(status),
Some(cloud_error.code),
cloud_error.message,
cloud_error.retry_after.map(Duration::from_secs_f64),
category,
);
}
let category =
ProviderErrorCategory::from_http_status(error.status, &cloud_error.message);
return LanguageModelCompletionError::from_provider_response(
PROVIDER_NAME,
Some(error.status),
Some(cloud_error.code),
cloud_error.message,
None,
category,
);
}
let retry_after = None;
LanguageModelCompletionError::from_http_status(
PROVIDER_NAME,
error.status,
error.body,
retry_after,
)
}
}
impl<TP: CloudLlmTokenProvider + 'static> CloudLanguageModel<TP> {
/// Explicit compaction via OpenAI's dedicated compact operation, proxied
/// through the cloud's `/completions/compact` endpoint.
fn compact_open_ai(
&self,
request: LanguageModelRequest,
cx: &AsyncApp,
) -> BoxFuture<'static, Result<CompactionResult, LanguageModelCompletionError>> {
let thread_id = request.thread_id.clone();
let prompt_id = request.prompt_id.clone();
let app_version = self.app_version.clone();
let model_provider = self.model.provider;
let provider_name = provider_name(&self.model.provider);
let supports_none_reasoning_effort =
self.model.supported_effort_levels.iter().any(|effort| {
open_ai::ReasoningEffort::from_str(&effort.value)
.is_ok_and(|effort| effort == open_ai::ReasoningEffort::None)
});
// Cloud proxies to OpenAI's own infrastructure, so the resulting
// compaction state is owned by (and interchangeable with) OpenAI
// proper, not by the cloud transport.
let request = match into_open_ai_response(
request,
&self.model.id.0,
self.model.supports_parallel_tool_calls,
true,
None,
None,
supports_none_reasoning_effort,
&OPEN_AI_PROVIDER_ID,
) {
Ok(request) => request,
Err(error) => return async move { Err(error.into()) }.boxed(),
};
let compact_request = request.into_compact_request();
let http_client = self.http_client.clone();
let token_provider = self.token_provider.clone();
let auth_context = token_provider.auth_context(cx);
let future = self.request_limiter.run(async move {
let PerformLlmCompletionResponse {
response,
includes_status_messages,
} = Self::perform_llm_compaction(
&http_client,
&*token_provider,
auth_context,
app_version,
CompletionBody {
thread_id,
prompt_id,
provider: model_provider,
model: compact_request.model.clone(),
provider_request: serde_json::to_value(compact_request).map_err(|error| {
LanguageModelCompletionError::SerializeRequest {
provider: provider_name.clone(),
error,
}
})?,
},
)
.await?;
let events = response_lines::<open_ai::responses::CompactedResponse>(
response,
includes_status_messages,
);
futures::pin_mut!(events);
while let Some(event) = events.next().await {
match event.map_err(|error| error.into_completion_error(provider_name.clone()))? {
CompletionEvent::Event(response) => {
let usage = token_usage_from_response_usage(&response.usage);
let context = response
.into_compacted_context(OPEN_AI_PROVIDER_ID)
.map_err(LanguageModelCompletionError::Other)?;
return Ok(CompactionResult { context, usage });
}
CompletionEvent::Status(_) => {}
}
}
Err(LanguageModelCompletionError::StreamEndedUnexpectedly {
provider: provider_name,
})
});
future.boxed()
}
}
impl<TP: CloudLlmTokenProvider + 'static> LanguageModel for CloudLanguageModel<TP> {
fn id(&self) -> LanguageModelId {
self.id.clone()
}
fn name(&self) -> LanguageModelName {
LanguageModelName::from(self.model.display_name.clone())
}
fn provider_id(&self) -> LanguageModelProviderId {
PROVIDER_ID
}
fn provider_name(&self) -> LanguageModelProviderName {
PROVIDER_NAME
}
fn upstream_provider_id(&self) -> LanguageModelProviderId {
use cloud_llm_client::LanguageModelProvider::*;
match self.model.provider {
Anthropic => ANTHROPIC_PROVIDER_ID,
Baseten => BASETEN_PROVIDER_ID,
OpenAi => OPEN_AI_PROVIDER_ID,
Google => GOOGLE_PROVIDER_ID,
XAi => X_AI_PROVIDER_ID,
}
}
fn upstream_provider_name(&self) -> LanguageModelProviderName {
use cloud_llm_client::LanguageModelProvider::*;
match self.model.provider {
Anthropic => ANTHROPIC_PROVIDER_NAME,
Baseten => BASETEN_PROVIDER_NAME,
OpenAi => OPEN_AI_PROVIDER_NAME,
Google => GOOGLE_PROVIDER_NAME,
XAi => X_AI_PROVIDER_NAME,
}
}
fn is_latest(&self) -> bool {
self.model.is_latest
}
fn is_disabled(&self) -> Option<DisabledReason> {
if self.model.is_disabled {
self.model.disabled_reason.clone().map(DisabledReason::new)
} else {
None
}
}
fn requires_data_retention(&self) -> bool {
// Anthropic cannot offer Fable models with Zero Data Retention
self.id
.0
.as_ref()
.starts_with(anthropic::FABLE_MODEL_ID_PREFIX)
}
fn refusal_fallback_model_id(&self) -> Option<&'static str> {
if self
.id
.0
.as_ref()
.starts_with(anthropic::FABLE_MODEL_ID_PREFIX)
{
Some(anthropic::FABLE_FALLBACK_MODEL_ID)
} else {
None
}
}
fn supports_tools(&self) -> bool {
self.model.supports_tools
}
fn supports_images(&self) -> bool {
self.model.supports_images
}
fn supports_thinking(&self) -> bool {
self.model.supports_thinking
}
fn supports_disabling_thinking(&self) -> bool {
self.model.supports_disabling_thinking
}
fn supports_fast_mode(&self) -> bool {
self.model.supports_fast_mode
}
fn supports_server_side_compaction(&self) -> bool {
self.model.supports_server_side_compaction
}
fn supports_explicit_compaction(&self) -> bool {
matches!(
self.model.provider,
cloud_llm_client::LanguageModelProvider::OpenAi
| cloud_llm_client::LanguageModelProvider::Anthropic
) && self.model.supports_server_side_compaction
}
fn minimum_explicit_compaction_input_tokens(&self) -> Option<u64> {
(self.model.provider == cloud_llm_client::LanguageModelProvider::Anthropic
&& self.supports_explicit_compaction())
.then_some(anthropic::MIN_COMPACTION_TRIGGER_TOKENS)
}
fn compact(
&self,
request: LanguageModelRequest,
cx: &AsyncApp,
) -> BoxFuture<'static, Result<CompactionResult, LanguageModelCompletionError>> {
if !self.supports_explicit_compaction() {
return async {
Err(LanguageModelCompletionError::Other(anyhow::anyhow!(
"this cloud model does not support explicit compaction"
)))
}
.boxed();
}
match self.model.provider {
cloud_llm_client::LanguageModelProvider::OpenAi => self.compact_open_ai(request, cx),
cloud_llm_client::LanguageModelProvider::Anthropic => {
self.compact_anthropic(request, cx)
}
// Unreachable while the `supports_explicit_compaction` guard
// above holds, but a provider mismatch should degrade to the
// same unsupported error rather than panic.
cloud_llm_client::LanguageModelProvider::Baseten
| cloud_llm_client::LanguageModelProvider::Google
| cloud_llm_client::LanguageModelProvider::XAi => async {
Err(LanguageModelCompletionError::Other(anyhow::anyhow!(
"this cloud model does not support explicit compaction"
)))
}
.boxed(),
}
}
fn supported_effort_levels(&self) -> Vec<LanguageModelEffortLevel> {
self.model
.supported_effort_levels
.iter()
.map(|effort_level| LanguageModelEffortLevel {
name: effort_level.name.clone().into(),
value: effort_level.value.clone().into(),
is_default: effort_level.is_default.unwrap_or(false),
})
.collect()
}
fn supports_streaming_tools(&self) -> bool {
self.model.supports_streaming_tools
}
fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
match choice {
LanguageModelToolChoice::Auto
| LanguageModelToolChoice::Any
| LanguageModelToolChoice::None => true,
}
}
fn supports_split_token_display(&self) -> bool {
use cloud_llm_client::LanguageModelProvider::*;
matches!(self.model.provider, Baseten | OpenAi | XAi)
}
fn telemetry_id(&self) -> String {
format!("zed.dev/{}", self.model.id)
}
fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
match self.model.provider {
cloud_llm_client::LanguageModelProvider::Anthropic
| cloud_llm_client::LanguageModelProvider::Baseten
| cloud_llm_client::LanguageModelProvider::OpenAi => {
LanguageModelToolSchemaFormat::JsonSchema
}
cloud_llm_client::LanguageModelProvider::Google
| cloud_llm_client::LanguageModelProvider::XAi => {
LanguageModelToolSchemaFormat::JsonSchemaSubset
}
}
}
fn max_token_count(&self) -> u64 {
self.model.max_token_count as u64
}
fn max_output_tokens(&self) -> Option<u64> {
Some(self.model.max_output_tokens as u64)
}
fn stream_completion(
&self,
request: LanguageModelRequest,
cx: &AsyncApp,
) -> BoxFuture<
'static,
Result<
BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
LanguageModelCompletionError,
>,
> {
if self.requires_data_retention() && !self.token_provider.has_data_retention_consent(cx) {
let model_name = self.model.display_name.clone();
return async move {
Err(LanguageModelCompletionError::DataRetentionConsentRequired { model_name })
}
.boxed();
}
let thread_id = request.thread_id.clone();
let prompt_id = request.prompt_id.clone();
let app_version = self.app_version.clone();
let thinking_allowed = request.thinking_allowed;
let enable_thinking = thinking_allowed && self.model.supports_thinking;
let provider_name = provider_name(&self.model.provider);
match self.model.provider {
cloud_llm_client::LanguageModelProvider::Anthropic => {
let effort = request
.thinking_effort
.as_ref()
.and_then(|effort| anthropic::Effort::from_str(effort).ok());
let mut request = match into_anthropic(
request,
self.model.id.to_string(),
1.0,
self.model.max_output_tokens as u64,
if enable_thinking {
AnthropicModelMode::Thinking {
budget_tokens: Some(4_096),
}
} else {
AnthropicModelMode::Default
},
AnthropicPromptCacheMode::Automatic,
// Cloud proxies to Anthropic's own infrastructure, so
// compaction state is owned by (and interchangeable with)
// Anthropic proper, not by the cloud transport.
&ANTHROPIC_PROVIDER_ID,
) {
Ok(request) => request,
Err(error) => return async move { Err(error.into()) }.boxed(),
};
if enable_thinking && effort.is_some() {
request.thinking = Some(anthropic::Thinking::Adaptive {
display: Some(anthropic::AdaptiveThinkingDisplay::Summarized),
});
request.output_config = Some(anthropic::OutputConfig { effort });
}
if !self.model.supports_fast_mode {
request.speed = None;
}
let http_client = self.http_client.clone();
let token_provider = self.token_provider.clone();
let auth_context = token_provider.auth_context(cx);
let future = self.request_limiter.stream(async move {
let PerformLlmCompletionResponse {
response,
includes_status_messages,
} = Self::perform_llm_completion(
&http_client,
&*token_provider,
auth_context,
app_version,
CompletionBody {
thread_id,
prompt_id,
provider: cloud_llm_client::LanguageModelProvider::Anthropic,
model: request.model.clone(),
provider_request: serde_json::to_value(&request).map_err(|error| {
LanguageModelCompletionError::SerializeRequest {
provider: provider_name.clone(),
error,
}
})?,
},
)
.await?;
let mut mapper =
AnthropicEventMapper::new(provider_name.clone(), ANTHROPIC_PROVIDER_ID);
Ok(map_cloud_completion_events(
Box::pin(response_lines(response, includes_status_messages)),
&provider_name,
move |event| mapper.map_event(event),
))
});
async move { Ok(future.await?.boxed()) }.boxed()
}
cloud_llm_client::LanguageModelProvider::OpenAi => {
let http_client = self.http_client.clone();
let token_provider = self.token_provider.clone();
let effort = request
.thinking_effort
.as_ref()
.and_then(|effort| open_ai::ReasoningEffort::from_str(effort).ok())
.filter(|effort| *effort != open_ai::ReasoningEffort::None);
let supports_none_reasoning_effort =
self.model.supported_effort_levels.iter().any(|effort| {
open_ai::ReasoningEffort::from_str(&effort.value)
.is_ok_and(|effort| effort == open_ai::ReasoningEffort::None)
});
let mut request = match into_open_ai_response(
request,
&self.model.id.0,
self.model.supports_parallel_tool_calls,
true,
None,
None,
supports_none_reasoning_effort,
&OPEN_AI_PROVIDER_ID,
) {
Ok(request) => request,
Err(error) => return async move { Err(error.into()) }.boxed(),
};
if enable_thinking && let Some(effort) = effort {
request.reasoning = Some(open_ai::responses::ReasoningConfig {
effort,
summary: Some(open_ai::responses::ReasoningSummaryMode::Auto),
});
}
let auth_context = token_provider.auth_context(cx);
let future = self.request_limiter.stream(async move {
let PerformLlmCompletionResponse {
response,
includes_status_messages,
} = Self::perform_llm_completion(
&http_client,
&*token_provider,
auth_context,
app_version,
CompletionBody {
thread_id,
prompt_id,
provider: cloud_llm_client::LanguageModelProvider::OpenAi,
model: request.model.clone(),
provider_request: serde_json::to_value(&request).map_err(|error| {
LanguageModelCompletionError::SerializeRequest {
provider: provider_name.clone(),
error,
}
})?,
},
)
.await?;
let mut mapper = OpenAiResponseEventMapper::new(OPEN_AI_PROVIDER_ID);
Ok(map_cloud_completion_events(
Box::pin(response_lines(response, includes_status_messages)),
&provider_name,
move |event| mapper.map_event(event),
))
});
async move { Ok(future.await?.boxed()) }.boxed()
}
provider @ (cloud_llm_client::LanguageModelProvider::Baseten
| cloud_llm_client::LanguageModelProvider::XAi) => {
let http_client = self.http_client.clone();
let token_provider = self.token_provider.clone();
let request = match into_open_ai(
request,
&self.model.id.0,
self.model.supports_parallel_tool_calls,
false,
None,
ChatCompletionMaxTokensParameter::MaxCompletionTokens,
None,
false,
) {
Ok(request) => request,
Err(error) => return async move { Err(error.into()) }.boxed(),
};
let auth_context = token_provider.auth_context(cx);
let future = self.request_limiter.stream(async move {
let PerformLlmCompletionResponse {
response,
includes_status_messages,
} = Self::perform_llm_completion(
&http_client,
&*token_provider,
auth_context,
app_version,
CompletionBody {
thread_id,
prompt_id,
provider,
model: request.model.clone(),
provider_request: serde_json::to_value(&request).map_err(|error| {
LanguageModelCompletionError::SerializeRequest {
provider: provider_name.clone(),
error,
}
})?,
},
)
.await?;
let mut mapper = OpenAiEventMapper::new();
Ok(map_cloud_completion_events(
Box::pin(response_lines(response, includes_status_messages)),
&provider_name,
move |event| mapper.map_event(event),
))
});
async move { Ok(future.await?.boxed()) }.boxed()
}
cloud_llm_client::LanguageModelProvider::Google => {
let http_client = self.http_client.clone();
let token_provider = self.token_provider.clone();
let request =
match into_google(request, self.model.id.to_string(), GoogleModelMode::Default)
{
Ok(request) => request,
Err(error) => return async move { Err(error.into()) }.boxed(),
};
let auth_context = token_provider.auth_context(cx);
let future = self.request_limiter.stream(async move {
let PerformLlmCompletionResponse {
response,
includes_status_messages,
} = Self::perform_llm_completion(
&http_client,
&*token_provider,
auth_context,
app_version,
CompletionBody {
thread_id,
prompt_id,
provider: cloud_llm_client::LanguageModelProvider::Google,
model: request.model.model_id.clone(),
provider_request: serde_json::to_value(&request).map_err(|error| {
LanguageModelCompletionError::SerializeRequest {
provider: provider_name.clone(),
error,
}
})?,
},
)
.await?;
let mut mapper = GoogleEventMapper::new();
Ok(map_cloud_completion_events(
Box::pin(response_lines(response, includes_status_messages)),
&provider_name,
move |event| mapper.map_event(event),
))
});
async move { Ok(future.await?.boxed()) }.boxed()
}
}
}
}
pub struct CloudModelProvider<TP: CloudLlmTokenProvider> {
token_provider: Arc<TP>,
http_client: Arc<HttpClientWithUrl>,
app_version: Option<Version>,
models: Vec<Arc<cloud_llm_client::LanguageModel>>,
default_model: Option<Arc<cloud_llm_client::LanguageModel>>,
default_fast_model: Option<Arc<cloud_llm_client::LanguageModel>>,
recommended_models: Vec<Arc<cloud_llm_client::LanguageModel>>,
}
impl<TP: CloudLlmTokenProvider + 'static> CloudModelProvider<TP> {
pub fn new(
token_provider: Arc<TP>,
http_client: Arc<HttpClientWithUrl>,
app_version: Option<Version>,
) -> Self {
Self {
token_provider,
http_client,
app_version,
models: Vec::new(),
default_model: None,
default_fast_model: None,
recommended_models: Vec::new(),
}
}
pub fn refresh_models(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
let http_client = self.http_client.clone();
let token_provider = self.token_provider.clone();
cx.spawn(async move |this, cx| {
let auth_context = token_provider.auth_context(cx);
let response =
Self::fetch_models_request(&http_client, &*token_provider, auth_context).await?;
this.update(cx, |this, cx| {
this.update_models(response);
cx.notify();
})
})
}
async fn fetch_models_request(
http_client: &HttpClientWithUrl,
token_provider: &TP,
auth_context: TP::AuthContext,
) -> Result<ListModelsResponse> {
let url = http_client.build_zed_llm_url("/models", &[])?;
let mut response =
authenticated_llm_request(http_client, token_provider, auth_context, |token| {
Ok(http_client::Request::builder()
.method(Method::GET)
.header(CLIENT_SUPPORTS_X_AI_HEADER_NAME, "true")
.uri(url.as_ref())
.header("Authorization", format!("Bearer {token}"))
.body(AsyncBody::empty())?)
})
.await
.context("failed to send list models request")?;
if response.status().is_success() {
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
Ok(serde_json::from_str(&body)?)
} else {
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
anyhow::bail!(
"error listing models.\nStatus: {:?}\nBody: {body}",
response.status(),
);
}
}
pub fn update_models(&mut self, response: ListModelsResponse) {
let models: Vec<_> = response.models.into_iter().map(Arc::new).collect();
self.default_model = models
.iter()
.find(|model| {
response
.default_model
.as_ref()
.is_some_and(|default_model_id| &model.id == default_model_id)
})
.cloned();
self.default_fast_model = models
.iter()
.find(|model| {
response
.default_fast_model
.as_ref()
.is_some_and(|default_fast_model_id| &model.id == default_fast_model_id)
})
.cloned();
self.recommended_models = response
.recommended_models
.iter()
.filter_map(|id| models.iter().find(|model| &model.id == id))
.cloned()
.collect();
self.models = models;
}
pub fn clear_models(&mut self) {
self.models.clear();
self.default_model = None;
self.default_fast_model = None;
self.recommended_models.clear();
}
pub fn create_model(
&self,
model: &Arc<cloud_llm_client::LanguageModel>,
) -> Arc<dyn LanguageModel> {
Arc::new(CloudLanguageModel::<TP> {
id: LanguageModelId::from(model.id.0.to_string()),
model: model.clone(),
token_provider: self.token_provider.clone(),
http_client: self.http_client.clone(),
app_version: self.app_version.clone(),
request_limiter: RateLimiter::new(4),
})
}
pub fn models(&self) -> &[Arc<cloud_llm_client::LanguageModel>] {
&self.models
}
pub fn default_model(&self) -> Option<&Arc<cloud_llm_client::LanguageModel>> {
self.default_model.as_ref()
}
pub fn default_fast_model(&self) -> Option<&Arc<cloud_llm_client::LanguageModel>> {
self.default_fast_model.as_ref()
}
pub fn recommended_models(&self) -> &[Arc<cloud_llm_client::LanguageModel>] {
&self.recommended_models
}
}
pub fn map_cloud_completion_events<T, F>(
stream: Pin<Box<dyn Stream<Item = Result<CompletionEvent<T>, ResponseStreamError>> + Send>>,
provider: &LanguageModelProviderName,
mut map_callback: F,
) -> BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
where
T: DeserializeOwned + 'static,
F: FnMut(T) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
+ Send
+ 'static,
{
let provider = provider.clone();
let mut stream = stream.fuse();
let mut saw_stream_ended = false;
let mut done = false;
let mut pending = VecDeque::new();
stream::poll_fn(move |cx| {
loop {
if let Some(item) = pending.pop_front() {
return Poll::Ready(Some(item));
}
if done {
return Poll::Ready(None);
}
match stream.poll_next_unpin(cx) {
Poll::Ready(Some(event)) => {
let items = match event {
Err(error) => {
vec![Err(error.into_completion_error(provider.clone()))]
}
Ok(CompletionEvent::Status(CompletionRequestStatus::StreamEnded)) => {
saw_stream_ended = true;
vec![]
}
Ok(CompletionEvent::Status(status)) => {
LanguageModelCompletionEvent::from_completion_request_status(
status,
provider.clone(),
)
.transpose()
.map(|event| vec![event])
.unwrap_or_default()
}
Ok(CompletionEvent::Event(event)) => map_callback(event),
};
pending.extend(items);
}
Poll::Ready(None) => {
done = true;
if !saw_stream_ended {
return Poll::Ready(Some(Err(
LanguageModelCompletionError::StreamEndedUnexpectedly {
provider: provider.clone(),
},
)));
}
}
Poll::Pending => return Poll::Pending,
}
}
})
.boxed()
}
pub fn provider_name(
provider: &cloud_llm_client::LanguageModelProvider,
) -> LanguageModelProviderName {
match provider {
cloud_llm_client::LanguageModelProvider::Anthropic => ANTHROPIC_PROVIDER_NAME,
cloud_llm_client::LanguageModelProvider::Baseten => BASETEN_PROVIDER_NAME,
cloud_llm_client::LanguageModelProvider::OpenAi => OPEN_AI_PROVIDER_NAME,
cloud_llm_client::LanguageModelProvider::Google => GOOGLE_PROVIDER_NAME,
cloud_llm_client::LanguageModelProvider::XAi => X_AI_PROVIDER_NAME,
}
}
/// A failure while reading the streamed completion response body.
///
/// Kept as a typed error (rather than `anyhow::Error`) so the consumer can
/// attach the provider name and build a structured
/// [`LanguageModelCompletionError`] without a runtime downcast.
pub enum ResponseStreamError {
Read(std::io::Error),
Deserialize(serde_json::Error),
}
impl ResponseStreamError {
fn into_completion_error(
self,
provider: LanguageModelProviderName,
) -> LanguageModelCompletionError {
match self {
ResponseStreamError::Read(error) => {
LanguageModelCompletionError::ApiReadResponseError { provider, error }
}
ResponseStreamError::Deserialize(error) => {
LanguageModelCompletionError::DeserializeResponse { provider, error }
}
}
}
}
pub fn response_lines<T: DeserializeOwned>(
response: Response<AsyncBody>,
includes_status_messages: bool,
) -> impl Stream<Item = Result<CompletionEvent<T>, ResponseStreamError>> {
futures::stream::try_unfold(
(String::new(), BufReader::new(response.into_body())),
move |(mut line, mut body)| async move {
match body.read_line(&mut line).await {
Ok(0) => Ok(None),
Ok(_) => {
let event = if includes_status_messages {
serde_json::from_str::<CompletionEvent<T>>(&line)
.map_err(ResponseStreamError::Deserialize)?
} else {
CompletionEvent::Event(
serde_json::from_str::<T>(&line)
.map_err(ResponseStreamError::Deserialize)?,
)
};
line.clear();
Ok(Some((event, (line, body))))
}
Err(error) => Err(ResponseStreamError::Read(error)),
}
},
)
}
#[cfg(test)]
mod tests {
use super::*;
use http_client::FakeHttpClient;
use http_client::http::{HeaderMap, StatusCode};
use language_model::{
LanguageModelCompletionError, LanguageModelRequestMessage, MessageContent,
ProviderErrorCategory, Role, Speed,
};
use serde_json::json;
use std::sync::Mutex;
#[gpui::test]
async fn cloud_explicit_compaction_forwards_supported_request_fields(
cx: &mut gpui::TestAppContext,
) {
let captured_request = Arc::new(Mutex::new(None));
let captured_request_for_handler = captured_request.clone();
let http_client = FakeHttpClient::create(move |request| {
let captured_request = captured_request_for_handler.clone();
async move {
let method = request.method().clone();
let uri = request.uri().to_string();
let authorization = request
.headers()
.get("Authorization")
.and_then(|value| value.to_str().ok())
.map(str::to_string);
let requested_status_messages = request
.headers()
.contains_key(CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME);
let requested_stream_end = request
.headers()
.contains_key(CLIENT_SUPPORTS_STATUS_STREAM_ENDED_HEADER_NAME);
let mut body = request.into_body();
let mut body_text = String::new();
body.read_to_string(&mut body_text).await?;
*captured_request.lock().unwrap() = Some((
method,
uri,
authorization,
requested_status_messages,
requested_stream_end,
body_text,
));
Ok(http_client::Response::builder()
.status(200)
.body(AsyncBody::from(format!(
"{}\n",
json!({
"id": "resp_compact",
"created_at": 1_700_000_000,
"object": "response.compaction",
"output": [{
"type": "compaction",
"id": "cmp_manual",
"encrypted_content": "opaque-state"
}],
"usage": {
"input_tokens": 100,
"input_tokens_details": {"cached_tokens": 20},
"output_tokens": 10,
"output_tokens_details": {"reasoning_tokens": 5},
"total_tokens": 110
}
})
)))?)
}
});
let model = cloud_test_model(http_client);
let request = compact_test_request();
let result = model.compact(request, &cx.to_async()).await.unwrap();
assert_eq!(
result.usage,
language_model::TokenUsage {
input_tokens: 80,
output_tokens: 10,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 20,
}
);
let language_model::CompactedContext::ProviderState(state) = result.context else {
panic!("expected provider compaction state");
};
assert_eq!(
open_ai::responses::provider_compaction_items(&state, &OPEN_AI_PROVIDER_ID).unwrap(),
Some(vec![json!({
"type": "compaction",
"id": "cmp_manual",
"encrypted_content": "opaque-state"
})])
);
let (method, uri, authorization, requested_status_messages, requested_stream_end, body) =
captured_request.lock().unwrap().take().unwrap();
assert_eq!(method, Method::POST);
assert_eq!(uri, "http://test.example/completions/compact?");
assert_eq!(authorization.as_deref(), Some("Bearer test-token"));
assert!(!requested_status_messages);
assert!(!requested_stream_end);
let body = serde_json::from_str::<serde_json::Value>(&body).unwrap();
assert_eq!(body["thread_id"], "thread-123");
assert_eq!(body["provider"], "open_ai");
assert_eq!(body["model"], "gpt-5.4");
assert_eq!(
body["provider_request"],
json!({
"model": "gpt-5.4",
"input": [{
"type": "message",
"role": "user",
"content": [{
"type": "input_text",
"text": "Retain this context."
}]
}],
"prompt_cache_key": "thread-123",
"service_tier": "priority"
})
);
}
#[gpui::test]
async fn cloud_anthropic_explicit_compaction_uses_paused_completion(
cx: &mut gpui::TestAppContext,
) {
let captured_request = Arc::new(Mutex::new(None));
let captured_request_for_handler = captured_request.clone();
let http_client = FakeHttpClient::create(move |request| {
let captured_request = captured_request_for_handler.clone();
async move {
let uri = request.uri().to_string();
let mut body = request.into_body();
let mut body_text = String::new();
body.read_to_string(&mut body_text).await?;
*captured_request.lock().unwrap() = Some((uri, body_text));
let response_lines = [
json!({
"event": {
"type": "message_start",
"message": {
"id": "msg_compact",
"type": "message",
"role": "assistant",
"content": [],
"model": "claude-opus-4-6",
"stop_reason": null,
"stop_sequence": null,
"usage": {
"input_tokens": 0,
"output_tokens": 0
}
}
}
}),
json!({
"event": {
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "compaction",
"content": null,
"encrypted_content": null
}
}
}),
json!({
"event": {
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "compaction_delta",
"content": "Summary of the conversation.",
"encrypted_content": "opaque-state"
}
}
}),
json!({
"event": {
"type": "content_block_stop",
"index": 0
}
}),
json!({
"event": {
"type": "message_delta",
"delta": {
"stop_reason": "compaction",
"stop_sequence": null
},
"usage": {
"input_tokens": 0,
"output_tokens": 0,
"iterations": [{
"type": "compaction",
"input_tokens": 60_000,
"output_tokens": 1_000
}]
}
}
}),
json!({"event": {"type": "message_stop"}}),
]
.into_iter()
.map(|line| line["event"].to_string())
.collect::<Vec<_>>()
.join("\n");
Ok(http_client::Response::builder()
.status(200)
.body(AsyncBody::from(format!("{response_lines}\n")))?)
}
});
let model = cloud_anthropic_test_model(http_client);
let result = model
.compact(compact_test_request(), &cx.to_async())
.await
.unwrap();
assert_eq!(
result.usage,
language_model::TokenUsage {
input_tokens: 60_000,
output_tokens: 1_000,
..Default::default()
}
);
let language_model::CompactedContext::Summary {
content,
provider_state,
} = result.context
else {
panic!("expected summary compaction");
};
assert_eq!(content.as_ref(), "Summary of the conversation.");
assert_eq!(
anthropic::completion::provider_compaction_encrypted_content(
&provider_state.expect("expected opaque provider state"),
&ANTHROPIC_PROVIDER_ID,
)
.unwrap()
.as_deref(),
Some("opaque-state")
);
let (uri, body) = captured_request.lock().unwrap().take().unwrap();
assert_eq!(uri, "http://test.example/completions?");
let body = serde_json::from_str::<serde_json::Value>(&body).unwrap();
assert_eq!(body["provider"], "anthropic");
assert_eq!(
body["provider_request"]["context_management"],
json!({
"edits": [{
"type": "compact_20260112",
"trigger": {
"type": "input_tokens",
"value": anthropic::MIN_COMPACTION_TRIGGER_TOKENS
},
"pause_after_compaction": true
}]
})
);
assert!(body["provider_request"]["tools"].is_null());
}
#[gpui::test]
async fn cloud_explicit_compaction_rejects_output_without_compaction_item(
cx: &mut gpui::TestAppContext,
) {
let http_client = FakeHttpClient::create(|_| async move {
Ok(http_client::Response::builder()
.status(200)
.body(AsyncBody::from(format!(
"{}\n",
json!({
"id": "resp_compact",
"created_at": 1_700_000_000,
"object": "response.compaction",
"output": [{
"type": "message",
"role": "assistant",
"content": "This is not an opaque compaction item."
}],
"usage": {
"input_tokens": 100,
"input_tokens_details": {"cached_tokens": 20},
"output_tokens": 10,
"output_tokens_details": {"reasoning_tokens": 5},
"total_tokens": 110
}
})
)))?)
});
let model = cloud_test_model(http_client);
let error = model
.compact(compact_test_request(), &cx.to_async())
.await
.unwrap_err();
assert!(
matches!(&error, LanguageModelCompletionError::Other(_)),
"expected invalid canonical output to be rejected, got {error:?}"
);
assert!(error.to_string().contains("compaction item"));
}
#[test]
fn test_api_error_conversion_with_upstream_http_error() {
// upstream_http_error with 503 status should become an Overloaded rejection
let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout","upstream_status":503}"#;
let api_error = ApiError {
status: StatusCode::INTERNAL_SERVER_ERROR,
body: error_body.to_string(),
headers: HeaderMap::new(),
};
let completion_error: LanguageModelCompletionError = api_error.into();
match completion_error {
LanguageModelCompletionError::ProviderRejection {
provider,
retry_after,
category: ProviderErrorCategory::Overloaded,
..
} => {
assert_eq!(provider, PROVIDER_NAME);
assert_eq!(retry_after, None);
}
_ => panic!(
"Expected Overloaded rejection for upstream 503, got: {:?}",
completion_error
),
}
// upstream_http_error with 500 status should preserve the rejection
let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the OpenAI API: internal server error","upstream_status":500}"#;
let api_error = ApiError {
status: StatusCode::INTERNAL_SERVER_ERROR,
body: error_body.to_string(),
headers: HeaderMap::new(),
};
let completion_error: LanguageModelCompletionError = api_error.into();
match completion_error {
LanguageModelCompletionError::ProviderRejection {
provider,
status,
code,
message,
retry_after,
category,
} => {
assert_eq!(provider, PROVIDER_NAME);
assert_eq!(status, Some(StatusCode::INTERNAL_SERVER_ERROR));
assert_eq!(code.as_deref(), Some("upstream_http_error"));
assert_eq!(
message,
"Received an error from the OpenAI API: internal server error"
);
assert_eq!(retry_after, None);
assert_eq!(category, ProviderErrorCategory::InternalServer);
}
_ => panic!(
"Expected ProviderRejection for upstream 500, got: {:?}",
completion_error
),
}
// upstream_http_error with 429 status should become a RateLimit rejection
let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the Google API: rate limit exceeded","upstream_status":429}"#;
let api_error = ApiError {
status: StatusCode::INTERNAL_SERVER_ERROR,
body: error_body.to_string(),
headers: HeaderMap::new(),
};
let completion_error: LanguageModelCompletionError = api_error.into();
match completion_error {
LanguageModelCompletionError::ProviderRejection {
provider,
retry_after,
category: ProviderErrorCategory::RateLimit,
..
} => {
assert_eq!(provider, PROVIDER_NAME);
assert_eq!(retry_after, None);
}
_ => panic!(
"Expected RateLimit rejection for upstream 429, got: {:?}",
completion_error
),
}
// Regular 500 error without upstream_http_error should preserve the Zed rejection
let error_body = "Regular internal server error";
let api_error = ApiError {
status: StatusCode::INTERNAL_SERVER_ERROR,
body: error_body.to_string(),
headers: HeaderMap::new(),
};
let completion_error: LanguageModelCompletionError = api_error.into();
match completion_error {
LanguageModelCompletionError::ProviderRejection {
provider,
status,
code,
message,
retry_after,
category,
} => {
assert_eq!(provider, PROVIDER_NAME);
assert_eq!(status, Some(StatusCode::INTERNAL_SERVER_ERROR));
assert_eq!(code, None);
assert_eq!(message, "Regular internal server error");
assert_eq!(retry_after, None);
assert_eq!(category, ProviderErrorCategory::InternalServer);
}
_ => panic!(
"Expected ProviderRejection for regular 500, got: {:?}",
completion_error
),
}
// upstream_http_429 format should be converted to a RateLimit rejection
let error_body = r#"{"code":"upstream_http_429","message":"Upstream Anthropic rate limit exceeded.","retry_after":30.5}"#;
let api_error = ApiError {
status: StatusCode::INTERNAL_SERVER_ERROR,
body: error_body.to_string(),
headers: HeaderMap::new(),
};
let completion_error: LanguageModelCompletionError = api_error.into();
match completion_error {
LanguageModelCompletionError::ProviderRejection {
provider,
retry_after,
category: ProviderErrorCategory::RateLimit,
..
} => {
assert_eq!(provider, PROVIDER_NAME);
assert_eq!(retry_after, Some(Duration::from_secs_f64(30.5)));
}
_ => panic!(
"Expected RateLimit rejection for upstream_http_429, got: {:?}",
completion_error
),
}
// Invalid JSON in error body should fall back to regular error handling
let error_body = "Not JSON at all";
let api_error = ApiError {
status: StatusCode::INTERNAL_SERVER_ERROR,
body: error_body.to_string(),
headers: HeaderMap::new(),
};
let completion_error: LanguageModelCompletionError = api_error.into();
match completion_error {
LanguageModelCompletionError::ProviderRejection {
provider,
status,
code,
message,
..
} => {
assert_eq!(provider, PROVIDER_NAME);
assert_eq!(status, Some(StatusCode::INTERNAL_SERVER_ERROR));
assert_eq!(code, None);
assert_eq!(message, "Not JSON at all");
}
_ => panic!(
"Expected ProviderRejection for invalid JSON, got: {:?}",
completion_error
),
}
}
#[gpui::test]
async fn cloud_transport_errors_include_hostname(cx: &mut gpui::TestAppContext) {
let http_client =
FakeHttpClient::create(|_| async move { Err(anyhow::anyhow!("DNS lookup failed")) });
let model = cloud_test_model(http_client);
let error = model
.compact(compact_test_request(), &cx.to_async())
.await
.expect_err("request should fail");
assert!(matches!(
error,
LanguageModelCompletionError::HttpSend {
host,
error,
..
} if host == "test.example" && error.to_string() == "DNS lookup failed"
));
}
#[test]
fn test_response_stream_error_maps_to_structured_variant() {
// Read/deserialize failures mid-stream must keep their structured
// variant rather than collapsing into `Other` (the source of the
// generic "Request failed." message).
let read = ResponseStreamError::Read(std::io::Error::from(std::io::ErrorKind::BrokenPipe))
.into_completion_error(PROVIDER_NAME);
assert!(
matches!(
read,
LanguageModelCompletionError::ApiReadResponseError { .. }
),
"Expected ApiReadResponseError, got: {read:?}"
);
let deserialize = ResponseStreamError::Deserialize(
serde_json::from_str::<serde_json::Value>("not json").unwrap_err(),
)
.into_completion_error(PROVIDER_NAME);
assert!(
matches!(
deserialize,
LanguageModelCompletionError::DeserializeResponse { .. }
),
"Expected DeserializeResponse, got: {deserialize:?}"
);
}
#[test]
fn cloud_anthropic_supports_explicit_compaction_after_minimum_input() {
let model = cloud_anthropic_test_model(FakeHttpClient::with_404_response());
assert!(model.supports_explicit_compaction());
assert_eq!(
model.minimum_explicit_compaction_input_tokens(),
Some(anthropic::MIN_COMPACTION_TRIGGER_TOKENS)
);
}
fn compact_test_request() -> LanguageModelRequest {
LanguageModelRequest {
thread_id: Some("thread-123".to_string()),
messages: vec![LanguageModelRequestMessage {
role: Role::User,
content: vec![MessageContent::Text("Retain this context.".to_string())],
cache: false,
reasoning_details: None,
}],
speed: Some(Speed::Fast),
..Default::default()
}
}
fn cloud_anthropic_test_model(
http_client: Arc<HttpClientWithUrl>,
) -> CloudLanguageModel<TestTokenProvider> {
CloudLanguageModel {
id: LanguageModelId::from("claude-opus-4-6".to_string()),
model: Arc::new(cloud_llm_client::LanguageModel {
provider: cloud_llm_client::LanguageModelProvider::Anthropic,
id: cloud_llm_client::LanguageModelId(Arc::from("claude-opus-4-6")),
display_name: "Claude Opus 4.6".to_string(),
is_latest: true,
max_token_count: 1_000_000,
max_token_count_in_max_mode: None,
max_output_tokens: 128_000,
supports_tools: true,
supports_images: true,
supports_thinking: true,
supports_disabling_thinking: true,
supports_fast_mode: true,
supports_server_side_compaction: true,
supported_effort_levels: Vec::new(),
supports_streaming_tools: true,
supports_parallel_tool_calls: false,
is_disabled: false,
disabled_reason: None,
}),
token_provider: Arc::new(TestTokenProvider),
http_client,
app_version: None,
request_limiter: RateLimiter::new(4),
}
}
fn cloud_test_model(
http_client: Arc<HttpClientWithUrl>,
) -> CloudLanguageModel<TestTokenProvider> {
CloudLanguageModel {
id: LanguageModelId::from("gpt-5.4".to_string()),
model: Arc::new(cloud_llm_client::LanguageModel {
provider: cloud_llm_client::LanguageModelProvider::OpenAi,
id: cloud_llm_client::LanguageModelId(Arc::from("gpt-5.4")),
display_name: "GPT-5.4".to_string(),
is_latest: true,
max_token_count: 1_000_000,
max_token_count_in_max_mode: None,
max_output_tokens: 128_000,
supports_tools: true,
supports_images: true,
supports_thinking: true,
supports_disabling_thinking: true,
supports_fast_mode: true,
supports_server_side_compaction: true,
supported_effort_levels: Vec::new(),
supports_streaming_tools: true,
supports_parallel_tool_calls: true,
is_disabled: false,
disabled_reason: None,
}),
token_provider: Arc::new(TestTokenProvider),
http_client,
app_version: None,
request_limiter: RateLimiter::new(4),
}
}
struct TestTokenProvider;
impl CloudLlmTokenProvider for TestTokenProvider {
type AuthContext = ();
fn auth_context(&self, _cx: &impl AppContext) -> Self::AuthContext {}
fn cached_token(
&self,
_auth_context: Self::AuthContext,
) -> BoxFuture<'static, Result<String>> {
async { Ok("test-token".to_string()) }.boxed()
}
fn refresh_token(
&self,
_auth_context: Self::AuthContext,
) -> BoxFuture<'static, Result<String>> {
async { Ok("refreshed-test-token".to_string()) }.boxed()
}
fn has_data_retention_consent(&self, _cx: &impl AppContext) -> bool {
false
}
}
}
Reference in new issue View Git Blame Copy Permalink