diff --git a/Cargo.lock b/Cargo.lock index c0711dee737..46c2fc46e2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9811,6 +9811,7 @@ dependencies = [ "async-lock", "aws-config", "aws-credential-types", + "aws-sigv4", "aws_http_client", "base64 0.22.1", "bedrock", diff --git a/Cargo.toml b/Cargo.toml index 3611fce865b..7a7003ec5eb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -543,6 +543,7 @@ aws-credential-types = { version = "1.2.8", features = [ aws-sdk-bedrockruntime = { version = "1.112.0", features = [ "behavior-version-latest", ] } +aws-sigv4 = { version = "1.4.0", features = ["http1"] } aws-smithy-runtime-api = { version = "1.9.2", features = ["http-1x", "client"] } aws-smithy-types = { version = "1.3.4", features = ["http-body-1-x"] } backtrace = "0.3" diff --git a/crates/bedrock/src/models.rs b/crates/bedrock/src/models.rs index ce5241a2dbb..cd035be8198 100644 --- a/crates/bedrock/src/models.rs +++ b/crates/bedrock/src/models.rs @@ -46,7 +46,7 @@ pub struct BedrockModelCacheConfiguration { #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, EnumIter)] -pub enum Model { +pub enum ConverseModel { // Anthropic Claude 4+ models #[serde( rename = "claude-fable-5", @@ -227,7 +227,7 @@ pub enum Model { }, } -impl Model { +impl ConverseModel { pub fn default_fast(_region: &str) -> Self { Self::ClaudeHaiku4_5 } @@ -817,34 +817,162 @@ impl Model { } } +/// The wire protocol used to talk to a [`MantleModel`] on the `bedrock-mantle` endpoint. +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize, PartialEq, EnumIter)] +pub enum MantleProtocol { + /// The OpenAI Chat Completions API (`/chat/completions`). + #[default] + ChatCompletions, + /// The OpenAI Responses API (`/responses`). + Responses, +} + +/// Models only reachable through the `bedrock-mantle` endpoint's +/// OpenAI-compatible APIs, i.e. with no `Converse`/`Invoke` support on +/// `bedrock-runtime`. +#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, EnumIter)] +pub enum MantleModel { + #[serde(rename = "gpt-5.5")] + Gpt5_5, + #[serde(rename = "gpt-5.4")] + Gpt5_4, + #[serde(rename = "grok-4.3")] + Grok4_3, + #[serde(rename = "custom")] + Custom { + name: String, + display_name: Option, + max_tokens: u64, + max_output_tokens: Option, + protocol: MantleProtocol, + supports_tools: bool, + supports_images: bool, + supports_thinking: bool, + }, +} + +impl MantleModel { + /// The model id Zed uses internally (also used as the `name` in settings). + pub fn id(&self) -> &str { + match self { + Self::Gpt5_5 => "gpt-5.5", + Self::Gpt5_4 => "gpt-5.4", + Self::Grok4_3 => "grok-4.3", + Self::Custom { name, .. } => name, + } + } + + /// The model id as expected in Bedrock Mantle request bodies, e.g. `openai.gpt-5.5`. + pub fn request_id(&self) -> &str { + match self { + Self::Gpt5_5 => "openai.gpt-5.5", + Self::Gpt5_4 => "openai.gpt-5.4", + Self::Grok4_3 => "xai.grok-4.3", + Self::Custom { name, .. } => name, + } + } + + pub fn display_name(&self) -> &str { + match self { + Self::Gpt5_5 => "GPT-5.5", + Self::Gpt5_4 => "GPT-5.4", + Self::Grok4_3 => "Grok 4.3", + Self::Custom { + display_name, name, .. + } => display_name.as_deref().unwrap_or(name.as_str()), + } + } + + /// Which OpenAI-compatible API this model must be called through. + pub fn protocol(&self) -> MantleProtocol { + match self { + Self::Gpt5_5 | Self::Gpt5_4 | Self::Grok4_3 => MantleProtocol::Responses, + Self::Custom { protocol, .. } => *protocol, + } + } + + pub fn max_token_count(&self) -> u64 { + match self { + Self::Gpt5_5 | Self::Gpt5_4 => 272_000, + Self::Grok4_3 => 1_000_000, + Self::Custom { max_tokens, .. } => *max_tokens, + } + } + + pub fn max_output_tokens(&self) -> u64 { + match self { + // AWS doesn't document a hard cap for GPT-5.5/5.4 on Mantle. + Self::Gpt5_5 | Self::Gpt5_4 => 128_000, + Self::Grok4_3 => 131_072, + Self::Custom { + max_output_tokens, .. + } => max_output_tokens.unwrap_or(4_096), + } + } + + pub fn supports_tools(&self) -> bool { + match self { + Self::Gpt5_5 | Self::Gpt5_4 | Self::Grok4_3 => true, + Self::Custom { supports_tools, .. } => *supports_tools, + } + } + + pub fn supports_images(&self) -> bool { + match self { + Self::Gpt5_5 | Self::Gpt5_4 | Self::Grok4_3 => true, + Self::Custom { + supports_images, .. + } => *supports_images, + } + } + + pub fn supports_thinking(&self) -> bool { + match self { + Self::Gpt5_5 | Self::Gpt5_4 | Self::Grok4_3 => true, + Self::Custom { + supports_thinking, .. + } => *supports_thinking, + } + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn test_builtin_mantle_models_use_responses_protocol() { + assert_eq!(MantleModel::Gpt5_5.protocol(), MantleProtocol::Responses); + assert_eq!(MantleModel::Gpt5_4.protocol(), MantleProtocol::Responses); + assert_eq!(MantleModel::Grok4_3.protocol(), MantleProtocol::Responses); + } + #[test] fn test_us_region_inference_ids() -> anyhow::Result<()> { assert_eq!( - Model::ClaudeSonnet4_5.cross_region_inference_id("us-east-1", false)?, + ConverseModel::ClaudeSonnet4_5.cross_region_inference_id("us-east-1", false)?, "us.anthropic.claude-sonnet-4-5-20250929-v1:0" ); assert_eq!( - Model::ClaudeSonnet4.cross_region_inference_id("us-west-2", false)?, + ConverseModel::ClaudeSonnet4.cross_region_inference_id("us-west-2", false)?, "us.anthropic.claude-sonnet-4-20250514-v1:0" ); assert_eq!( - Model::ClaudeFable5.cross_region_inference_id("us-east-1", false)?, + ConverseModel::ClaudeFable5.cross_region_inference_id("us-east-1", false)?, "us.anthropic.claude-fable-5" ); assert_eq!( - Model::ClaudeSonnet5.cross_region_inference_id("us-east-1", false)?, + ConverseModel::ClaudeSonnet5.cross_region_inference_id("us-east-1", false)?, "us.anthropic.claude-sonnet-5" ); assert_eq!( - Model::NovaPro.cross_region_inference_id("us-east-2", false)?, + ConverseModel::NovaPro.cross_region_inference_id("us-east-2", false)?, "us.amazon.nova-pro-v1:0" ); assert_eq!( - Model::DeepSeekR1.cross_region_inference_id("us-east-1", false)?, + ConverseModel::DeepSeekR1.cross_region_inference_id("us-east-1", false)?, "us.deepseek.r1-v1:0" ); Ok(()) @@ -853,27 +981,27 @@ mod tests { #[test] fn test_eu_region_inference_ids() -> anyhow::Result<()> { assert_eq!( - Model::ClaudeSonnet4.cross_region_inference_id("eu-west-1", false)?, + ConverseModel::ClaudeSonnet4.cross_region_inference_id("eu-west-1", false)?, "eu.anthropic.claude-sonnet-4-20250514-v1:0" ); assert_eq!( - Model::ClaudeSonnet4_5.cross_region_inference_id("eu-west-1", false)?, + ConverseModel::ClaudeSonnet4_5.cross_region_inference_id("eu-west-1", false)?, "eu.anthropic.claude-sonnet-4-5-20250929-v1:0" ); assert_eq!( - Model::NovaLite.cross_region_inference_id("eu-north-1", false)?, + ConverseModel::NovaLite.cross_region_inference_id("eu-north-1", false)?, "eu.amazon.nova-lite-v1:0" ); assert_eq!( - Model::ClaudeOpus4_6.cross_region_inference_id("eu-west-1", false)?, + ConverseModel::ClaudeOpus4_6.cross_region_inference_id("eu-west-1", false)?, "eu.anthropic.claude-opus-4-6-v1" ); assert_eq!( - Model::ClaudeOpus4_7.cross_region_inference_id("eu-west-1", false)?, + ConverseModel::ClaudeOpus4_7.cross_region_inference_id("eu-west-1", false)?, "eu.anthropic.claude-opus-4-7" ); assert_eq!( - Model::ClaudeOpus4_8.cross_region_inference_id("eu-west-1", false)?, + ConverseModel::ClaudeOpus4_8.cross_region_inference_id("eu-west-1", false)?, "eu.anthropic.claude-opus-4-8" ); Ok(()) @@ -882,19 +1010,19 @@ mod tests { #[test] fn test_inference_profile_only_models_fall_back_to_global() -> anyhow::Result<()> { assert_eq!( - Model::ClaudeFable5.cross_region_inference_id("eu-west-1", false)?, + ConverseModel::ClaudeFable5.cross_region_inference_id("eu-west-1", false)?, "global.anthropic.claude-fable-5" ); assert_eq!( - Model::ClaudeSonnet5.cross_region_inference_id("eu-west-1", false)?, + ConverseModel::ClaudeSonnet5.cross_region_inference_id("eu-west-1", false)?, "global.anthropic.claude-sonnet-5" ); assert_eq!( - Model::ClaudeFable5.cross_region_inference_id("ap-southeast-2", false)?, + ConverseModel::ClaudeFable5.cross_region_inference_id("ap-southeast-2", false)?, "global.anthropic.claude-fable-5" ); assert_eq!( - Model::ClaudeSonnet5.cross_region_inference_id("ap-northeast-1", false)?, + ConverseModel::ClaudeSonnet5.cross_region_inference_id("ap-northeast-1", false)?, "global.anthropic.claude-sonnet-5" ); Ok(()) @@ -903,11 +1031,11 @@ mod tests { #[test] fn test_apac_region_inference_ids() -> anyhow::Result<()> { assert_eq!( - Model::ClaudeSonnet4_5.cross_region_inference_id("ap-south-1", false)?, + ConverseModel::ClaudeSonnet4_5.cross_region_inference_id("ap-south-1", false)?, "apac.anthropic.claude-sonnet-4-5-20250929-v1:0" ); assert_eq!( - Model::NovaLite.cross_region_inference_id("ap-south-1", false)?, + ConverseModel::NovaLite.cross_region_inference_id("ap-south-1", false)?, "apac.amazon.nova-lite-v1:0" ); Ok(()) @@ -916,23 +1044,23 @@ mod tests { #[test] fn test_au_region_inference_ids() -> anyhow::Result<()> { assert_eq!( - Model::ClaudeHaiku4_5.cross_region_inference_id("ap-southeast-2", false)?, + ConverseModel::ClaudeHaiku4_5.cross_region_inference_id("ap-southeast-2", false)?, "au.anthropic.claude-haiku-4-5-20251001-v1:0" ); assert_eq!( - Model::ClaudeSonnet4_5.cross_region_inference_id("ap-southeast-4", false)?, + ConverseModel::ClaudeSonnet4_5.cross_region_inference_id("ap-southeast-4", false)?, "au.anthropic.claude-sonnet-4-5-20250929-v1:0" ); assert_eq!( - Model::ClaudeOpus4_6.cross_region_inference_id("ap-southeast-2", false)?, + ConverseModel::ClaudeOpus4_6.cross_region_inference_id("ap-southeast-2", false)?, "au.anthropic.claude-opus-4-6-v1" ); assert_eq!( - Model::ClaudeOpus4_7.cross_region_inference_id("ap-southeast-2", false)?, + ConverseModel::ClaudeOpus4_7.cross_region_inference_id("ap-southeast-2", false)?, "au.anthropic.claude-opus-4-7" ); assert_eq!( - Model::ClaudeOpus4_8.cross_region_inference_id("ap-southeast-2", false)?, + ConverseModel::ClaudeOpus4_8.cross_region_inference_id("ap-southeast-2", false)?, "au.anthropic.claude-opus-4-8" ); Ok(()) @@ -941,15 +1069,15 @@ mod tests { #[test] fn test_jp_region_inference_ids() -> anyhow::Result<()> { assert_eq!( - Model::ClaudeHaiku4_5.cross_region_inference_id("ap-northeast-1", false)?, + ConverseModel::ClaudeHaiku4_5.cross_region_inference_id("ap-northeast-1", false)?, "jp.anthropic.claude-haiku-4-5-20251001-v1:0" ); assert_eq!( - Model::ClaudeSonnet4_5.cross_region_inference_id("ap-northeast-3", false)?, + ConverseModel::ClaudeSonnet4_5.cross_region_inference_id("ap-northeast-3", false)?, "jp.anthropic.claude-sonnet-4-5-20250929-v1:0" ); assert_eq!( - Model::Nova2Lite.cross_region_inference_id("ap-northeast-1", false)?, + ConverseModel::Nova2Lite.cross_region_inference_id("ap-northeast-1", false)?, "jp.amazon.nova-2-lite-v1:0" ); Ok(()) @@ -958,7 +1086,7 @@ mod tests { #[test] fn test_ca_region_inference_ids() -> anyhow::Result<()> { assert_eq!( - Model::NovaLite.cross_region_inference_id("ca-central-1", false)?, + ConverseModel::NovaLite.cross_region_inference_id("ca-central-1", false)?, "ca.amazon.nova-lite-v1:0" ); Ok(()) @@ -967,11 +1095,11 @@ mod tests { #[test] fn test_gov_region_inference_ids() -> anyhow::Result<()> { assert_eq!( - Model::ClaudeSonnet4_5.cross_region_inference_id("us-gov-east-1", false)?, + ConverseModel::ClaudeSonnet4_5.cross_region_inference_id("us-gov-east-1", false)?, "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0" ); assert_eq!( - Model::ClaudeSonnet4_5.cross_region_inference_id("us-gov-west-1", false)?, + ConverseModel::ClaudeSonnet4_5.cross_region_inference_id("us-gov-west-1", false)?, "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0" ); Ok(()) @@ -980,45 +1108,45 @@ mod tests { #[test] fn test_global_inference_ids() -> anyhow::Result<()> { assert_eq!( - Model::ClaudeSonnet4.cross_region_inference_id("us-east-1", true)?, + ConverseModel::ClaudeSonnet4.cross_region_inference_id("us-east-1", true)?, "global.anthropic.claude-sonnet-4-20250514-v1:0" ); assert_eq!( - Model::ClaudeSonnet4_5.cross_region_inference_id("eu-west-1", true)?, + ConverseModel::ClaudeSonnet4_5.cross_region_inference_id("eu-west-1", true)?, "global.anthropic.claude-sonnet-4-5-20250929-v1:0" ); assert_eq!( - Model::ClaudeHaiku4_5.cross_region_inference_id("ap-south-1", true)?, + ConverseModel::ClaudeHaiku4_5.cross_region_inference_id("ap-south-1", true)?, "global.anthropic.claude-haiku-4-5-20251001-v1:0" ); assert_eq!( - Model::ClaudeOpus4_6.cross_region_inference_id("us-east-1", true)?, + ConverseModel::ClaudeOpus4_6.cross_region_inference_id("us-east-1", true)?, "global.anthropic.claude-opus-4-6-v1" ); assert_eq!( - Model::ClaudeOpus4_7.cross_region_inference_id("us-east-1", true)?, + ConverseModel::ClaudeOpus4_7.cross_region_inference_id("us-east-1", true)?, "global.anthropic.claude-opus-4-7" ); assert_eq!( - Model::ClaudeOpus4_8.cross_region_inference_id("us-east-1", true)?, + ConverseModel::ClaudeOpus4_8.cross_region_inference_id("us-east-1", true)?, "global.anthropic.claude-opus-4-8" ); assert_eq!( - Model::ClaudeFable5.cross_region_inference_id("us-east-1", true)?, + ConverseModel::ClaudeFable5.cross_region_inference_id("us-east-1", true)?, "global.anthropic.claude-fable-5" ); assert_eq!( - Model::ClaudeSonnet5.cross_region_inference_id("us-east-1", true)?, + ConverseModel::ClaudeSonnet5.cross_region_inference_id("us-east-1", true)?, "global.anthropic.claude-sonnet-5" ); assert_eq!( - Model::Nova2Lite.cross_region_inference_id("us-east-1", true)?, + ConverseModel::Nova2Lite.cross_region_inference_id("us-east-1", true)?, "global.amazon.nova-2-lite-v1:0" ); // Models without global support fall back to regional assert_eq!( - Model::NovaPro.cross_region_inference_id("us-east-1", true)?, + ConverseModel::NovaPro.cross_region_inference_id("us-east-1", true)?, "us.amazon.nova-pro-v1:0" ); Ok(()) @@ -1028,27 +1156,27 @@ mod tests { fn test_models_without_cross_region() -> anyhow::Result<()> { // Models without cross-region support return their request_id directly assert_eq!( - Model::Gemma3_4B.cross_region_inference_id("us-east-1", false)?, + ConverseModel::Gemma3_4B.cross_region_inference_id("us-east-1", false)?, "google.gemma-3-4b-it" ); assert_eq!( - Model::MistralLarge3.cross_region_inference_id("eu-west-1", false)?, + ConverseModel::MistralLarge3.cross_region_inference_id("eu-west-1", false)?, "mistral.mistral-large-3-675b-instruct" ); assert_eq!( - Model::Qwen3VL235B.cross_region_inference_id("ap-south-1", false)?, + ConverseModel::Qwen3VL235B.cross_region_inference_id("ap-south-1", false)?, "qwen.qwen3-vl-235b-a22b" ); assert_eq!( - Model::GptOss120B.cross_region_inference_id("us-east-1", false)?, + ConverseModel::GptOss120B.cross_region_inference_id("us-east-1", false)?, "openai.gpt-oss-120b-1:0" ); assert_eq!( - Model::MiniMaxM2.cross_region_inference_id("us-east-1", false)?, + ConverseModel::MiniMaxM2.cross_region_inference_id("us-east-1", false)?, "minimax.minimax-m2" ); assert_eq!( - Model::KimiK2Thinking.cross_region_inference_id("us-east-1", false)?, + ConverseModel::KimiK2Thinking.cross_region_inference_id("us-east-1", false)?, "moonshot.kimi-k2-thinking" ); Ok(()) @@ -1056,7 +1184,7 @@ mod tests { #[test] fn test_custom_model_inference_ids() -> anyhow::Result<()> { - let custom_model = Model::Custom { + let custom_model = ConverseModel::Custom { name: "custom.my-model-v1:0".to_string(), max_tokens: 100000, display_name: Some("My Custom Model".to_string()), @@ -1078,78 +1206,90 @@ mod tests { #[test] fn test_friendly_id_vs_request_id() { - assert_eq!(Model::ClaudeSonnet4_5.id(), "claude-sonnet-4-5"); - assert_eq!(Model::NovaLite.id(), "nova-lite"); - assert_eq!(Model::DeepSeekR1.id(), "deepseek-r1"); - assert_eq!(Model::Llama4Scout17B.id(), "llama-4-scout-17b"); - assert_eq!(Model::ClaudeFable5.id(), "claude-fable-5"); - assert_eq!(Model::ClaudeSonnet5.id(), "claude-sonnet-5"); + assert_eq!(ConverseModel::ClaudeSonnet4_5.id(), "claude-sonnet-4-5"); + assert_eq!(ConverseModel::NovaLite.id(), "nova-lite"); + assert_eq!(ConverseModel::DeepSeekR1.id(), "deepseek-r1"); + assert_eq!(ConverseModel::Llama4Scout17B.id(), "llama-4-scout-17b"); + assert_eq!(ConverseModel::ClaudeFable5.id(), "claude-fable-5"); + assert_eq!(ConverseModel::ClaudeSonnet5.id(), "claude-sonnet-5"); assert_eq!( - Model::ClaudeSonnet4_5.request_id(), + ConverseModel::ClaudeSonnet4_5.request_id(), "anthropic.claude-sonnet-4-5-20250929-v1:0" ); - assert_eq!(Model::NovaLite.request_id(), "amazon.nova-lite-v1:0"); - assert_eq!(Model::DeepSeekR1.request_id(), "deepseek.r1-v1:0"); assert_eq!( - Model::Llama4Scout17B.request_id(), + ConverseModel::NovaLite.request_id(), + "amazon.nova-lite-v1:0" + ); + assert_eq!(ConverseModel::DeepSeekR1.request_id(), "deepseek.r1-v1:0"); + assert_eq!( + ConverseModel::Llama4Scout17B.request_id(), "meta.llama4-scout-17b-instruct-v1:0" ); - assert_eq!(Model::ClaudeFable5.request_id(), "anthropic.claude-fable-5"); assert_eq!( - Model::ClaudeSonnet5.request_id(), + ConverseModel::ClaudeFable5.request_id(), + "anthropic.claude-fable-5" + ); + assert_eq!( + ConverseModel::ClaudeSonnet5.request_id(), "anthropic.claude-sonnet-5" ); // Thinking aliases deserialize to the same model - assert_eq!(Model::ClaudeSonnet4.id(), "claude-sonnet-4"); + assert_eq!(ConverseModel::ClaudeSonnet4.id(), "claude-sonnet-4"); assert_eq!( - Model::from_id("claude-sonnet-4-thinking").unwrap().id(), + ConverseModel::from_id("claude-sonnet-4-thinking") + .unwrap() + .id(), "claude-sonnet-4" ); assert_eq!( - Model::from_id("claude-fable-5-thinking").unwrap().id(), + ConverseModel::from_id("claude-fable-5-thinking") + .unwrap() + .id(), "claude-fable-5" ); assert_eq!( - Model::from_id("claude-sonnet-5-thinking").unwrap().id(), + ConverseModel::from_id("claude-sonnet-5-thinking") + .unwrap() + .id(), "claude-sonnet-5" ); } #[test] fn test_thinking_modes() { - assert!(Model::ClaudeHaiku4_5.supports_thinking()); - assert!(Model::ClaudeSonnet4.supports_thinking()); - assert!(Model::ClaudeSonnet4_5.supports_thinking()); - assert!(Model::ClaudeOpus4_6.supports_thinking()); - assert!(Model::ClaudeFable5.supports_thinking()); + assert!(ConverseModel::ClaudeHaiku4_5.supports_thinking()); + assert!(ConverseModel::ClaudeSonnet4.supports_thinking()); + assert!(ConverseModel::ClaudeSonnet4_5.supports_thinking()); + assert!(ConverseModel::ClaudeOpus4_6.supports_thinking()); + assert!(ConverseModel::ClaudeFable5.supports_thinking()); - assert!(!Model::ClaudeSonnet4.supports_adaptive_thinking()); - assert!(Model::ClaudeOpus4_6.supports_adaptive_thinking()); - assert!(Model::ClaudeSonnet4_6.supports_adaptive_thinking()); - assert!(Model::ClaudeFable5.supports_adaptive_thinking()); - assert!(Model::ClaudeSonnet5.supports_adaptive_thinking()); - assert!(!Model::ClaudeOpus4_7.supports_xhigh_adaptive_thinking()); - assert!(Model::ClaudeFable5.supports_xhigh_adaptive_thinking()); - assert!(Model::ClaudeSonnet5.supports_xhigh_adaptive_thinking()); - assert!(Model::ClaudeOpus4_8.supports_xhigh_adaptive_thinking()); + assert!(!ConverseModel::ClaudeSonnet4.supports_adaptive_thinking()); + assert!(ConverseModel::ClaudeOpus4_6.supports_adaptive_thinking()); + assert!(ConverseModel::ClaudeSonnet4_6.supports_adaptive_thinking()); + assert!(ConverseModel::ClaudeFable5.supports_adaptive_thinking()); + assert!(ConverseModel::ClaudeSonnet5.supports_adaptive_thinking()); + assert!(!ConverseModel::ClaudeOpus4_7.supports_xhigh_adaptive_thinking()); + assert!(ConverseModel::ClaudeFable5.supports_xhigh_adaptive_thinking()); + assert!(ConverseModel::ClaudeSonnet5.supports_xhigh_adaptive_thinking()); + assert!(ConverseModel::ClaudeOpus4_8.supports_xhigh_adaptive_thinking()); assert_eq!(BedrockAdaptiveThinkingEffort::XHigh.as_str(), "xhigh"); assert_eq!( - Model::ClaudeSonnet4.thinking_mode(), + ConverseModel::ClaudeSonnet4.thinking_mode(), BedrockModelMode::Thinking { budget_tokens: Some(4096) } ); assert_eq!( - Model::ClaudeOpus4_6.thinking_mode(), + ConverseModel::ClaudeOpus4_6.thinking_mode(), BedrockModelMode::AdaptiveThinking { effort: BedrockAdaptiveThinkingEffort::High } ); assert_eq!( - Model::ClaudeHaiku4_5.thinking_mode(), + ConverseModel::ClaudeHaiku4_5.thinking_mode(), BedrockModelMode::Thinking { budget_tokens: Some(4096) } @@ -1158,44 +1298,44 @@ mod tests { #[test] fn test_max_token_count() { - assert_eq!(Model::ClaudeSonnet4_5.max_token_count(), 1_000_000); - assert_eq!(Model::ClaudeOpus4_6.max_token_count(), 1_000_000); - assert_eq!(Model::ClaudeFable5.max_token_count(), 1_000_000); - assert_eq!(Model::ClaudeSonnet5.max_token_count(), 1_000_000); - assert_eq!(Model::Llama4Scout17B.max_token_count(), 128_000); - assert_eq!(Model::NovaPremier.max_token_count(), 1_000_000); + assert_eq!(ConverseModel::ClaudeSonnet4_5.max_token_count(), 1_000_000); + assert_eq!(ConverseModel::ClaudeOpus4_6.max_token_count(), 1_000_000); + assert_eq!(ConverseModel::ClaudeFable5.max_token_count(), 1_000_000); + assert_eq!(ConverseModel::ClaudeSonnet5.max_token_count(), 1_000_000); + assert_eq!(ConverseModel::Llama4Scout17B.max_token_count(), 128_000); + assert_eq!(ConverseModel::NovaPremier.max_token_count(), 1_000_000); } #[test] fn test_max_output_tokens() { - assert_eq!(Model::ClaudeSonnet4_5.max_output_tokens(), 64_000); - assert_eq!(Model::ClaudeOpus4_6.max_output_tokens(), 128_000); - assert_eq!(Model::ClaudeFable5.max_output_tokens(), 128_000); - assert_eq!(Model::ClaudeSonnet5.max_output_tokens(), 128_000); - assert_eq!(Model::ClaudeOpus4_1.max_output_tokens(), 32_000); - assert_eq!(Model::Gemma3_4B.max_output_tokens(), 8_192); + assert_eq!(ConverseModel::ClaudeSonnet4_5.max_output_tokens(), 64_000); + assert_eq!(ConverseModel::ClaudeOpus4_6.max_output_tokens(), 128_000); + assert_eq!(ConverseModel::ClaudeFable5.max_output_tokens(), 128_000); + assert_eq!(ConverseModel::ClaudeSonnet5.max_output_tokens(), 128_000); + assert_eq!(ConverseModel::ClaudeOpus4_1.max_output_tokens(), 32_000); + assert_eq!(ConverseModel::Gemma3_4B.max_output_tokens(), 8_192); } #[test] fn test_supports_tool_use() { - assert!(Model::ClaudeSonnet4_5.supports_tool_use()); - assert!(Model::ClaudeFable5.supports_tool_use()); - assert!(Model::NovaPro.supports_tool_use()); - assert!(Model::MistralLarge3.supports_tool_use()); - assert!(!Model::Gemma3_4B.supports_tool_use()); - assert!(Model::Qwen3_32B.supports_tool_use()); - assert!(Model::MiniMaxM2.supports_tool_use()); - assert!(Model::KimiK2_5.supports_tool_use()); - assert!(Model::DeepSeekR1.supports_tool_use()); - assert!(!Model::Llama4Scout17B.supports_tool_use()); + assert!(ConverseModel::ClaudeSonnet4_5.supports_tool_use()); + assert!(ConverseModel::ClaudeFable5.supports_tool_use()); + assert!(ConverseModel::NovaPro.supports_tool_use()); + assert!(ConverseModel::MistralLarge3.supports_tool_use()); + assert!(!ConverseModel::Gemma3_4B.supports_tool_use()); + assert!(ConverseModel::Qwen3_32B.supports_tool_use()); + assert!(ConverseModel::MiniMaxM2.supports_tool_use()); + assert!(ConverseModel::KimiK2_5.supports_tool_use()); + assert!(ConverseModel::DeepSeekR1.supports_tool_use()); + assert!(!ConverseModel::Llama4Scout17B.supports_tool_use()); } #[test] fn test_supports_caching() { - assert!(Model::ClaudeSonnet4_5.supports_caching()); - assert!(Model::ClaudeOpus4_6.supports_caching()); - assert!(Model::ClaudeFable5.supports_caching()); - assert!(!Model::Llama4Scout17B.supports_caching()); - assert!(!Model::NovaPro.supports_caching()); + assert!(ConverseModel::ClaudeSonnet4_5.supports_caching()); + assert!(ConverseModel::ClaudeOpus4_6.supports_caching()); + assert!(ConverseModel::ClaudeFable5.supports_caching()); + assert!(!ConverseModel::Llama4Scout17B.supports_caching()); + assert!(!ConverseModel::NovaPro.supports_caching()); } } diff --git a/crates/language_models/Cargo.toml b/crates/language_models/Cargo.toml index 57a0d530e45..8da9eb332ed 100644 --- a/crates/language_models/Cargo.toml +++ b/crates/language_models/Cargo.toml @@ -18,6 +18,7 @@ anthropic = { workspace = true, features = ["schemars"] } anyhow.workspace = true aws-config = { workspace = true, features = ["behavior-version-latest"] } aws-credential-types = { workspace = true, features = ["hardcoded-credentials"] } +aws-sigv4.workspace = true aws_http_client.workspace = true base64.workspace = true bedrock = { workspace = true, features = ["schemars"] } diff --git a/crates/language_models/src/provider/bedrock.rs b/crates/language_models/src/provider/bedrock.rs index ba3b587151a..fe12a5eb3a1 100644 --- a/crates/language_models/src/provider/bedrock.rs +++ b/crates/language_models/src/provider/bedrock.rs @@ -5,8 +5,11 @@ use anyhow::{Context as _, Result, anyhow}; use async_lock::OnceCell; use aws_config::stalled_stream_protection::StalledStreamProtectionConfig; use aws_config::{BehaviorVersion, Region}; +use aws_credential_types::provider::{ProvideCredentials, SharedCredentialsProvider}; use aws_credential_types::{Credentials, Token}; use aws_http_client::AwsHttpClient; +use aws_sigv4::http_request::{SignableBody, SignableRequest, SigningSettings, sign}; +use aws_sigv4::sign::v4; use bedrock::BedrockSystemContentBlock; use bedrock::bedrock_client::Client as BedrockClient; use bedrock::bedrock_client::config::timeout::TimeoutConfig; @@ -20,37 +23,53 @@ use bedrock::{ BedrockStreamingResponse, BedrockThinkingBlock, BedrockThinkingTextBlock, BedrockTool, BedrockToolChoice, BedrockToolConfig, BedrockToolInputSchema, BedrockToolResultBlock, BedrockToolResultContentBlock, BedrockToolResultStatus, BedrockToolSpec, BedrockToolUseBlock, - Model, value_to_aws_document, + ConverseModel, MantleModel, MantleProtocol, value_to_aws_document, }; use collections::{BTreeMap, HashMap}; use credentials_provider::CredentialsProvider; -use futures::{FutureExt, Stream, StreamExt, future::BoxFuture, stream::BoxStream}; +use futures::{ + AsyncBufReadExt, AsyncReadExt, FutureExt, Stream, StreamExt, future::BoxFuture, io::BufReader, + stream::BoxStream, +}; use gpui::{ App, AsyncApp, Context, Entity, FocusHandle, Subscription, Task, TaskExt, Window, actions, }; use gpui_tokio::Tokio; -use http_client::HttpClient; +use http_client::{ + AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt, + http::{HeaderValue, header::AUTHORIZATION}, +}; use language_model::{ AuthenticateError, EnvVar, IconOrSvg, InlineDescription, LanguageModel, - LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, - LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName, - LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice, - LanguageModelToolResultContent, LanguageModelToolUse, MessageContent, ProviderSettingsView, - RateLimiter, Role, SubPageProviderSettings, TokenUsage, env_var, + LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelEffortLevel, + LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, + LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, + LanguageModelToolChoice, LanguageModelToolResultContent, LanguageModelToolSchemaFormat, + LanguageModelToolUse, MessageContent, ProviderSettingsView, RateLimiter, Role, + SubPageProviderSettings, TokenUsage, env_var, }; +use open_ai::responses::Request as OpenAiResponseRequest; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use serde_json::Value; -use settings::{BedrockAvailableModel as AvailableModel, Settings, SettingsStore}; +use settings::{ + BedrockAvailableModel as AvailableModel, BedrockMantleAvailableModel as MantleAvailableModel, + Settings, SettingsStore, +}; use std::sync::LazyLock; +use std::time::SystemTime; use strum::{EnumIter, IntoEnumIterator, IntoStaticStr}; use ui::{ButtonLink, ConfiguredApiCard, Divider, List, ListBulletItem, prelude::*}; use ui_input::InputField; use util::ResultExt; use crate::AllLanguageModelSettings; -use http_client::CustomHeaders; +use crate::provider::open_ai::{ + ChatCompletionMaxTokensParameter, OpenAiEventMapper, OpenAiResponseEventMapper, into_open_ai, + into_open_ai_response, +}; use language_model::util::{fix_streamed_json, parse_tool_arguments}; +use open_ai::{ReasoningEffort, RequestError, ResponseStreamEvent}; actions!(bedrock, [Tab, TabPrev]); @@ -116,6 +135,7 @@ impl BedrockCredentials { #[derive(Default, Clone, Debug, PartialEq)] pub struct AmazonBedrockSettings { pub available_models: Vec, + pub mantle_available_models: Vec, pub custom_headers: CustomHeaders, pub region: Option, pub endpoint: Option, @@ -151,6 +171,13 @@ impl From for BedrockAuthMethod { } } +fn mantle_protocol_from_settings(value: settings::BedrockMantleProtocolContent) -> MantleProtocol { + match value { + settings::BedrockMantleProtocolContent::ChatCompletions => MantleProtocol::ChatCompletions, + settings::BedrockMantleProtocolContent::Responses => MantleProtocol::Responses, + } +} + #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize, JsonSchema)] #[serde(tag = "type", rename_all = "lowercase")] pub enum ModelMode { @@ -200,6 +227,115 @@ static ZED_BEDROCK_REGION_VAR: LazyLock = env_var!("ZED_AWS_REGION"); static ZED_AWS_ENDPOINT_VAR: LazyLock = env_var!("ZED_AWS_ENDPOINT"); static ZED_BEDROCK_BEARER_TOKEN_VAR: LazyLock = env_var!("ZED_BEDROCK_BEARER_TOKEN"); +/// AWS Regions where the `bedrock-mantle` endpoint is available. +/// See . +const MANTLE_SUPPORTED_REGIONS: &[&str] = &[ + "us-east-2", + "us-east-1", + "us-west-2", + "ap-southeast-3", + "ap-south-1", + "ap-southeast-2", + "ap-northeast-1", + "eu-central-1", + "eu-west-1", + "eu-west-2", + "eu-south-1", + "eu-north-1", + "sa-east-1", + "us-gov-west-1", +]; + +fn mantle_endpoint_url(region: &str) -> String { + format!("https://bedrock-mantle.{region}.api.aws/openai/v1") +} + +enum MantleAuth { + ApiKey { api_key: String }, + SigV4 { credentials: Credentials }, +} + +impl MantleAuth { + fn apply(&self, request: &mut HttpRequest, body: &[u8], region: &str) -> Result<()> { + match self { + MantleAuth::ApiKey { api_key } => { + let value = HeaderValue::from_str(&format!("Bearer {}", api_key.trim())) + .context("building Mantle bearer token authorization header")?; + request.headers_mut().insert(AUTHORIZATION, value); + } + MantleAuth::SigV4 { credentials } => { + sign_mantle_request_sigv4(request, body, credentials, region)?; + } + } + + Ok(()) + } +} + +fn sign_mantle_request_sigv4( + request: &mut HttpRequest, + body: &[u8], + credentials: &Credentials, + region: &str, +) -> Result<()> { + sign_mantle_request_sigv4_at(request, body, credentials, region, SystemTime::now()) +} + +fn sign_mantle_request_sigv4_at( + request: &mut HttpRequest, + body: &[u8], + credentials: &Credentials, + region: &str, + time: SystemTime, +) -> Result<()> { + if !request + .headers() + .contains_key(http_client::http::header::HOST) + && let Some(authority) = request.uri().authority() + { + let host = HeaderValue::from_str(authority.as_str()) + .context("invalid host header derived from Mantle request URI")?; + request + .headers_mut() + .insert(http_client::http::header::HOST, host); + } + + let identity = credentials.clone().into(); + let signing_params: aws_sigv4::http_request::SigningParams = v4::SigningParams::builder() + .identity(&identity) + .region(region) + .name("bedrock-mantle") + .time(time) + .settings(SigningSettings::default()) + .build() + .context("building Mantle SigV4 signing params")? + .into(); + + let method = request.method().as_str(); + let uri = request.uri().to_string(); + let headers = request + .headers() + .iter() + .map(|(name, value)| { + value + .to_str() + .map(|value| (name.as_str(), value)) + .with_context(|| format!("header {name} is not valid UTF-8 and cannot be signed")) + }) + .collect::>>()?; + + let signable_request = + SignableRequest::new(method, uri, headers.into_iter(), SignableBody::Bytes(body)) + .context("constructing Mantle SigV4 request")?; + + let (instructions, _signature) = sign(signable_request, &signing_params) + .context("signing Mantle request with SigV4")? + .into_parts(); + instructions.apply_to_request_http1x(request); + + Ok(()) +} + pub struct State { /// The resolved authentication method. Settings take priority over UX credentials. auth: Option, @@ -407,6 +543,7 @@ impl State { pub struct BedrockLanguageModelProvider { http_client: AwsHttpClient, + plain_http_client: Arc, handle: tokio::runtime::Handle, state: Entity, } @@ -428,13 +565,14 @@ impl BedrockLanguageModelProvider { }); Self { - http_client: AwsHttpClient::new(http_client), + http_client: AwsHttpClient::new(http_client.clone()), + plain_http_client: http_client, handle: Tokio::handle(cx), state, } } - fn create_language_model(&self, model: bedrock::Model) -> Arc { + fn create_language_model(&self, model: bedrock::ConverseModel) -> Arc { Arc::new(BedrockModel { id: LanguageModelId::from(model.id().to_string()), model, @@ -445,6 +583,17 @@ impl BedrockLanguageModelProvider { request_limiter: RateLimiter::new(4), }) } + + fn create_mantle_language_model(&self, model: bedrock::MantleModel) -> Arc { + Arc::new(BedrockMantleModel { + id: LanguageModelId::from(model.id().to_string()), + model, + http_client: self.plain_http_client.clone(), + state: self.state.clone(), + credentials_provider: Arc::new(OnceCell::new()), + request_limiter: RateLimiter::new(4), + }) + } } impl LanguageModelProvider for BedrockLanguageModelProvider { @@ -461,32 +610,29 @@ impl LanguageModelProvider for BedrockLanguageModelProvider { } fn default_model(&self, _cx: &App) -> Option> { - Some(self.create_language_model(bedrock::Model::default())) + Some(self.create_language_model(bedrock::ConverseModel::default())) } fn default_fast_model(&self, cx: &App) -> Option> { let region = self.state.read(cx).get_region(); - Some(self.create_language_model(bedrock::Model::default_fast(region.as_str()))) + Some(self.create_language_model(bedrock::ConverseModel::default_fast(region.as_str()))) } fn provided_models(&self, cx: &App) -> Vec> { + let bedrock_settings = &AllLanguageModelSettings::get_global(cx).bedrock; let mut models = BTreeMap::default(); - for model in bedrock::Model::iter() { - if !matches!(model, bedrock::Model::Custom { .. }) { + for model in bedrock::ConverseModel::iter() { + if !matches!(model, bedrock::ConverseModel::Custom { .. }) { models.insert(model.id().to_string(), model); } } // Override with available models from settings - for model in AllLanguageModelSettings::get_global(cx) - .bedrock - .available_models - .iter() - { + for model in bedrock_settings.available_models.iter() { models.insert( model.name.clone(), - bedrock::Model::Custom { + bedrock::ConverseModel::Custom { name: model.name.clone(), display_name: model.display_name.clone(), max_tokens: model.max_tokens, @@ -502,10 +648,43 @@ impl LanguageModelProvider for BedrockLanguageModelProvider { ); } - models + let mut models: Vec> = models .into_values() .map(|model| self.create_language_model(model)) - .collect() + .collect(); + + let mut mantle_models = BTreeMap::default(); + + for model in bedrock::MantleModel::iter() { + if !matches!(model, bedrock::MantleModel::Custom { .. }) { + mantle_models.insert(model.id().to_string(), model); + } + } + + // Override with available Mantle models from settings + for model in bedrock_settings.mantle_available_models.iter() { + mantle_models.insert( + model.name.clone(), + bedrock::MantleModel::Custom { + name: model.name.clone(), + display_name: model.display_name.clone(), + max_tokens: model.max_tokens, + max_output_tokens: model.max_output_tokens, + protocol: mantle_protocol_from_settings(model.protocol), + supports_tools: model.supports_tools.unwrap_or(false), + supports_images: model.supports_images.unwrap_or(false), + supports_thinking: model.supports_thinking.unwrap_or(false), + }, + ); + } + + models.extend( + mantle_models + .into_values() + .map(|model| self.create_mantle_language_model(model)), + ); + + models } fn is_authenticated(&self, cx: &App) -> bool { @@ -524,7 +703,7 @@ impl LanguageModelProvider for BedrockLanguageModelProvider { .into() }) .description(InlineDescription::Text( - "To use Zed's agent with Bedrock, set a custom authentication strategy in your settings or use static credentials.".into(), + "To use Zed's agent with Bedrock, set a custom authentication strategy in your settings or use static credentials. Mantle-only models (e.g. GPT-5.5, GPT-5.4, Grok 4.3) additionally require IAM permissions for the `bedrock-mantle` endpoint.".into(), )), )) } @@ -540,7 +719,7 @@ impl LanguageModelProviderState for BedrockLanguageModelProvider { struct BedrockModel { id: LanguageModelId, - model: Model, + model: ConverseModel, http_client: AwsHttpClient, handle: tokio::runtime::Handle, client: OnceCell, @@ -840,6 +1019,485 @@ impl LanguageModel for BedrockModel { } } +const MANTLE_SELECTABLE_REASONING_EFFORTS: &[ReasoningEffort] = &[ + ReasoningEffort::Low, + ReasoningEffort::Medium, + ReasoningEffort::High, + ReasoningEffort::XHigh, +]; + +fn mantle_default_reasoning_effort(model: &MantleModel) -> Option { + model.supports_thinking().then_some(ReasoningEffort::Medium) +} + +fn mantle_selected_reasoning_effort( + request: &LanguageModelRequest, + model: &MantleModel, +) -> Option { + if !model.supports_thinking() { + return None; + } + + if request.thinking_allowed { + request + .thinking_effort + .as_deref() + .and_then(|effort| effort.parse::().ok()) + .filter(|effort| *effort != ReasoningEffort::None) + .or_else(|| mantle_default_reasoning_effort(model)) + } else { + Some(ReasoningEffort::None) + } +} + +fn mantle_supported_effort_levels(model: &MantleModel) -> Vec { + let Some(default_effort) = mantle_default_reasoning_effort(model) else { + return Vec::new(); + }; + + MANTLE_SELECTABLE_REASONING_EFFORTS + .iter() + .copied() + .map(|effort| LanguageModelEffortLevel { + name: effort.label().into(), + value: effort.value().into(), + is_default: effort == default_effort, + }) + .collect() +} + +/// Special-cases Mantle authorization failures with a message that points at +/// the separate `bedrock-mantle` IAM policy namespace instead of regular +/// `bedrock-runtime` permissions. +fn map_mantle_error(model: &MantleModel, error: RequestError) -> LanguageModelCompletionError { + if let RequestError::HttpResponseError { status_code, .. } = &error + && *status_code == http_client::http::StatusCode::FORBIDDEN + { + return LanguageModelCompletionError::PermissionError { + provider: PROVIDER_NAME, + message: format!( + "Bedrock Mantle denied this request for {}. Mantle-only models require IAM \ + permissions for the `bedrock-mantle` endpoint (for example via the \ + `AmazonBedrockMantleInferenceAccess` managed policy) in addition to whatever \ + permissions your existing Bedrock credentials already have.", + model.display_name() + ), + }; + } + error.into() +} + +/// Resolves an AWS credentials provider for profile/SSO/automatic auth. +/// Cached in `cell` since building it may read config files from disk; +/// credentials themselves are still re-resolved on every call. Async so this +/// never blocks the foreground thread (unlike `BedrockModel::get_or_init_client`). +async fn resolve_mantle_credentials_provider( + cell: &OnceCell, + profile_name: Option, + region: String, +) -> Result { + let provider = cell + .get_or_try_init(move || async move { + let mut config_builder = + aws_config::defaults(BehaviorVersion::latest()).region(Region::new(region)); + + if let Some(profile_name) = profile_name.filter(|name| !name.is_empty()) { + config_builder = config_builder.profile_name(profile_name); + } + + let config = config_builder.load().await; + config + .credentials_provider() + .context("no AWS credentials provider is configured") + }) + .await + .context("resolving AWS credentials for Bedrock Mantle")?; + Ok(provider.clone()) +} + +/// Resolves provider settings into concrete Mantle request auth. A configured +/// Bedrock API key is sent as bearer auth; every AWS-credential-based method +/// signs the Mantle HTTP request directly with SigV4. +async fn resolve_mantle_auth( + credentials_provider: Arc>, + auth: Option, + region: String, +) -> Result { + match auth { + Some(BedrockAuth::ApiKey { api_key }) => Ok(MantleAuth::ApiKey { api_key }), + Some(BedrockAuth::IamCredentials { + access_key_id, + secret_access_key, + session_token, + }) => Ok(MantleAuth::SigV4 { + credentials: Credentials::new( + access_key_id, + secret_access_key, + session_token, + None, + "zed-bedrock-provider", + ), + }), + Some(BedrockAuth::NamedProfile { profile_name }) + | Some(BedrockAuth::SingleSignOn { profile_name }) => { + let provider = resolve_mantle_credentials_provider( + &credentials_provider, + Some(profile_name), + region.clone(), + ) + .await?; + let credentials = provider + .provide_credentials() + .await + .context("failed to resolve AWS credentials")?; + Ok(MantleAuth::SigV4 { credentials }) + } + Some(BedrockAuth::Automatic) | None => { + let provider = + resolve_mantle_credentials_provider(&credentials_provider, None, region.clone()) + .await?; + let credentials = provider + .provide_credentials() + .await + .context("failed to resolve AWS credentials")?; + Ok(MantleAuth::SigV4 { credentials }) + } + } +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum MantleChatStreamResult { + Ok(ResponseStreamEvent), + Err { error: MantleChatStreamError }, +} + +#[derive(Deserialize)] +struct MantleChatStreamError { + message: String, +} + +fn parse_mantle_chat_stream_line(line: &str) -> Result { + match serde_json::from_str(line) { + Ok(MantleChatStreamResult::Ok(response)) => Ok(response), + Ok(MantleChatStreamResult::Err { error }) => Err(anyhow!(error.message)), + Err(error) => { + log::error!( + "Failed to parse Mantle chat completion stream event: `{}`\nResponse: `{}`", + error, + line, + ); + Err(anyhow!(error)) + } + } +} + +fn parse_mantle_response_stream_line(line: &str) -> Result { + serde_json::from_str(line).map_err(|error| { + log::error!( + "Failed to parse Mantle responses stream event: `{}`\nResponse: `{}`", + error, + line, + ); + anyhow!(error) + }) +} + +async fn stream_mantle_sse( + client: &dyn HttpClient, + provider_name: &str, + url: &str, + region: &str, + auth: &MantleAuth, + request: Request, + extra_headers: &CustomHeaders, + parse_stream_line: fn(&str) -> Result, +) -> std::result::Result>, RequestError> +where + Request: Serialize, + Event: Send + 'static, +{ + let body = serde_json::to_vec(&request).map_err(|error| RequestError::Other(error.into()))?; + let mut request = HttpRequest::builder() + .method(Method::POST) + .uri(url) + .header("Content-Type", "application/json") + .extra_headers(extra_headers) + .body(AsyncBody::from(body.clone())) + .map_err(|error| RequestError::Other(error.into()))?; + + auth.apply(&mut request, &body, region) + .map_err(RequestError::Other)?; + + let mut response = client.send(request).await?; + if response.status().is_success() { + let reader = BufReader::new(response.into_body()); + Ok(reader + .lines() + .filter_map(move |line| async move { + match line { + Ok(line) => { + let line = line + .strip_prefix("data: ") + .or_else(|| line.strip_prefix("data:"))?; + if line == "[DONE]" || line.is_empty() { + None + } else { + Some(parse_stream_line(line)) + } + } + Err(error) => Some(Err(anyhow!(error))), + } + }) + .boxed()) + } else { + let mut body = String::new(); + response + .body_mut() + .read_to_string(&mut body) + .await + .map_err(|error| RequestError::Other(error.into()))?; + + Err(RequestError::HttpResponseError { + provider: provider_name.to_owned(), + status_code: response.status(), + body, + headers: response.headers().clone(), + }) + } +} + +fn strip_unsupported_mantle_response_fields(request: &mut OpenAiResponseRequest) { + request.context_management = None; +} + +struct BedrockMantleModel { + id: LanguageModelId, + model: MantleModel, + http_client: Arc, + state: Entity, + credentials_provider: Arc>, + request_limiter: RateLimiter, +} + +impl BedrockMantleModel { + fn stream_mantle_request( + &self, + request: Request, + cx: &AsyncApp, + endpoint: &'static str, + parse_stream_line: fn(&str) -> Result, + ) -> BoxFuture<'static, Result>, LanguageModelCompletionError>> + where + Request: Serialize + Send + 'static, + Event: Send + 'static, + { + let http_client = self.http_client.clone(); + let model = self.model.clone(); + let credentials_provider = self.credentials_provider.clone(); + let (auth, region) = cx.read_entity(&self.state, |state, _cx| { + (state.auth.clone(), state.get_region()) + }); + let url = format!("{}/{}", mantle_endpoint_url(®ion), endpoint); + let extra_headers = cx.read_entity(&self.state, |_, cx| { + AllLanguageModelSettings::get_global(cx) + .bedrock + .custom_headers + .clone() + }); + let provider_name = PROVIDER_NAME.0.to_string(); + let auth_task = Tokio::spawn_result( + cx, + resolve_mantle_auth(credentials_provider, auth, region.clone()), + ); + + let future = self.request_limiter.stream(async move { + let auth = auth_task + .await + .map_err(LanguageModelCompletionError::Other)?; + stream_mantle_sse( + http_client.as_ref(), + &provider_name, + &url, + ®ion, + &auth, + request, + &extra_headers, + parse_stream_line, + ) + .await + .map_err(|err| map_mantle_error(&model, err)) + }); + + async move { Ok(future.await?.boxed()) }.boxed() + } + + fn stream_completion( + &self, + request: open_ai::Request, + cx: &AsyncApp, + ) -> BoxFuture< + 'static, + Result>, LanguageModelCompletionError>, + > { + self.stream_mantle_request( + request, + cx, + "chat/completions", + parse_mantle_chat_stream_line, + ) + } + + fn stream_response( + &self, + request: OpenAiResponseRequest, + cx: &AsyncApp, + ) -> BoxFuture< + 'static, + Result< + BoxStream<'static, Result>, + LanguageModelCompletionError, + >, + > { + let mut request = request; + strip_unsupported_mantle_response_fields(&mut request); + self.stream_mantle_request(request, cx, "responses", parse_mantle_response_stream_line) + } +} + +impl LanguageModel for BedrockMantleModel { + fn id(&self) -> LanguageModelId { + self.id.clone() + } + + fn name(&self) -> LanguageModelName { + LanguageModelName::from(self.model.display_name().to_string()) + } + + fn provider_id(&self) -> LanguageModelProviderId { + PROVIDER_ID + } + + fn provider_name(&self) -> LanguageModelProviderName { + PROVIDER_NAME + } + + fn supports_tools(&self) -> bool { + self.model.supports_tools() + } + + fn tool_input_format(&self) -> LanguageModelToolSchemaFormat { + LanguageModelToolSchemaFormat::JsonSchemaSubset + } + + fn supports_images(&self) -> bool { + self.model.supports_images() + } + + fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool { + match choice { + LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any => { + self.model.supports_tools() + } + LanguageModelToolChoice::None => true, + } + } + + fn supports_streaming_tools(&self) -> bool { + true + } + + fn supports_thinking(&self) -> bool { + self.model.supports_thinking() + } + + fn supported_effort_levels(&self) -> Vec { + mantle_supported_effort_levels(&self.model) + } + + fn supports_split_token_display(&self) -> bool { + true + } + + fn telemetry_id(&self) -> String { + format!("bedrock-mantle/{}", self.model.id()) + } + + fn max_token_count(&self) -> u64 { + self.model.max_token_count() + } + + fn max_output_tokens(&self) -> Option { + Some(self.model.max_output_tokens()) + } + + fn stream_completion( + &self, + request: LanguageModelRequest, + cx: &AsyncApp, + ) -> BoxFuture< + 'static, + Result< + BoxStream<'static, Result>, + LanguageModelCompletionError, + >, + > { + let region = cx.read_entity(&self.state, |state, _cx| state.get_region()); + + if !MANTLE_SUPPORTED_REGIONS.contains(®ion.as_str()) { + let display_name = self.model.display_name().to_string(); + let supported = MANTLE_SUPPORTED_REGIONS.join(", "); + return futures::future::ready(Err(LanguageModelCompletionError::Other(anyhow!( + "{display_name} is not available in {region} because Bedrock Mantle isn't offered \ + there. Try switching to one of the following regions: {supported}." + )))) + .boxed(); + } + + let model_id = self.model.request_id().to_string(); + let max_output_tokens = Some(self.model.max_output_tokens()); + + match self.model.protocol() { + MantleProtocol::Responses => { + let request = into_open_ai_response( + request, + &model_id, + self.model.supports_tools(), + false, + max_output_tokens, + mantle_default_reasoning_effort(&self.model), + self.model.supports_thinking(), + ); + let completions = self.stream_response(request, cx); + async move { + let mapper = OpenAiResponseEventMapper::new(); + Ok(mapper.map_stream(completions.await?).boxed()) + } + .boxed() + } + MantleProtocol::ChatCompletions => { + let reasoning_effort = mantle_selected_reasoning_effort(&request, &self.model); + let request = into_open_ai( + request, + &model_id, + self.model.supports_tools(), + false, + max_output_tokens, + ChatCompletionMaxTokensParameter::MaxCompletionTokens, + reasoning_effort, + false, + ); + let completions = self.stream_completion(request, cx); + async move { + let mapper = OpenAiEventMapper::new(); + Ok(mapper.map_stream(completions.await?).boxed()) + } + .boxed() + } + } + } +} + fn deny_tool_use_events( events: impl Stream>, ) -> impl Stream> { @@ -895,7 +1553,7 @@ pub fn into_bedrock( } MessageContent::Compaction(_) => None, MessageContent::Thinking { text, signature } => { - if model.contains(Model::DeepSeekR1.request_id()) { + if model.contains(ConverseModel::DeepSeekR1.request_id()) { // DeepSeekR1 doesn't support thinking blocks // And the AWS API demands that you strip them return None; @@ -918,7 +1576,7 @@ pub fn into_bedrock( )) } MessageContent::RedactedThinking(blob) => { - if model.contains(Model::DeepSeekR1.request_id()) { + if model.contains(ConverseModel::DeepSeekR1.request_id()) { // DeepSeekR1 doesn't support thinking blocks // And the AWS API demands that you strip them return None; @@ -1834,4 +2492,218 @@ mod tests { "a cache-marked message with content should end with a cache point" ); } + + #[test] + fn test_sign_mantle_request_sigv4_uses_mantle_service() { + let credentials = Credentials::new( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + None, + None, + "test", + ); + let body = br#"{"model":"openai.gpt-5.5"}"#; + let mut request = HttpRequest::builder() + .method(Method::POST) + .uri("https://bedrock-mantle.us-east-1.api.aws/openai/v1/responses") + .header("Content-Type", "application/json") + .body(AsyncBody::from(body.to_vec())) + .unwrap(); + let time = std::time::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + + sign_mantle_request_sigv4_at(&mut request, body, &credentials, "us-east-1", time).unwrap(); + + assert_eq!( + request + .headers() + .get(http_client::http::header::HOST) + .and_then(|value| value.to_str().ok()), + Some("bedrock-mantle.us-east-1.api.aws") + ); + assert_eq!( + request + .headers() + .get("x-amz-date") + .and_then(|value| value.to_str().ok()), + Some("20231114T221320Z") + ); + let authorization = request + .headers() + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .unwrap(); + assert!(authorization.starts_with("AWS4-HMAC-SHA256 ")); + assert!( + authorization + .contains("Credential=AKIDEXAMPLE/20231114/us-east-1/bedrock-mantle/aws4_request") + ); + assert!(authorization.contains("SignedHeaders=content-type;host")); + assert!(authorization.contains("Signature=")); + } + + #[test] + fn test_mantle_endpoint_url_uses_openai_path_prefix() { + assert_eq!( + mantle_endpoint_url("us-east-1"), + "https://bedrock-mantle.us-east-1.api.aws/openai/v1" + ); + assert_eq!( + mantle_endpoint_url("us-west-2"), + "https://bedrock-mantle.us-west-2.api.aws/openai/v1" + ); + } + + #[test] + fn test_mantle_protocol_from_settings() { + assert_eq!( + mantle_protocol_from_settings(settings::BedrockMantleProtocolContent::ChatCompletions), + MantleProtocol::ChatCompletions + ); + assert_eq!( + mantle_protocol_from_settings(settings::BedrockMantleProtocolContent::Responses), + MantleProtocol::Responses + ); + } + + #[test] + fn test_mantle_supported_regions_matches_docs() { + assert!(MANTLE_SUPPORTED_REGIONS.contains(&"us-east-1")); + assert!(MANTLE_SUPPORTED_REGIONS.contains(&"eu-west-1")); + assert!(!MANTLE_SUPPORTED_REGIONS.contains(&"ap-southeast-1")); + } + + #[test] + fn test_builtin_mantle_models_support_thinking() { + assert!(MantleModel::Gpt5_5.supports_thinking()); + assert!(MantleModel::Gpt5_4.supports_thinking()); + assert!(MantleModel::Grok4_3.supports_thinking()); + assert_eq!( + mantle_default_reasoning_effort(&MantleModel::Gpt5_5), + Some(ReasoningEffort::Medium) + ); + assert_eq!( + mantle_default_reasoning_effort(&MantleModel::Grok4_3), + Some(ReasoningEffort::Medium) + ); + } + + #[test] + fn test_mantle_supported_effort_levels_hide_none() { + let effort_levels = mantle_supported_effort_levels(&MantleModel::Gpt5_5); + let values = effort_levels + .iter() + .map(|level| level.value.as_ref()) + .collect::>(); + + assert_eq!(values, ["low", "medium", "high", "xhigh"]); + assert_eq!( + effort_levels + .iter() + .find(|level| level.is_default) + .map(|level| level.value.as_ref()), + Some("medium") + ); + } + + #[test] + fn test_custom_mantle_model_can_disable_thinking() { + let model = MantleModel::Custom { + name: "custom-mantle-model".to_string(), + display_name: None, + max_tokens: 128_000, + max_output_tokens: None, + protocol: MantleProtocol::Responses, + supports_tools: true, + supports_images: false, + supports_thinking: false, + }; + + assert!(!model.supports_thinking()); + assert_eq!(mantle_default_reasoning_effort(&model), None); + assert!(mantle_supported_effort_levels(&model).is_empty()); + assert_eq!( + mantle_selected_reasoning_effort( + &LanguageModelRequest { + thinking_effort: Some("high".to_string()), + ..Default::default() + }, + &model, + ), + None + ); + } + + #[test] + fn test_disabled_mantle_thinking_serializes_none() { + let request = into_open_ai_response( + LanguageModelRequest { + thinking_allowed: false, + ..Default::default() + }, + MantleModel::Grok4_3.request_id(), + true, + false, + Some(MantleModel::Grok4_3.max_output_tokens()), + mantle_default_reasoning_effort(&MantleModel::Grok4_3), + MantleModel::Grok4_3.supports_thinking(), + ); + + assert_eq!( + serde_json::to_value(&request).unwrap()["reasoning"], + serde_json::json!({ "effort": "none" }) + ); + } + + #[test] + fn test_mantle_reasoning_passes_known_efforts_through() { + for effort in ["low", "medium", "high", "xhigh", "minimal", "max"] { + assert_eq!( + mantle_selected_reasoning_effort( + &LanguageModelRequest { + thinking_allowed: true, + thinking_effort: Some(effort.to_string()), + ..Default::default() + }, + &MantleModel::Gpt5_5, + ) + .map(|effort| effort.value()), + Some(effort) + ); + } + + assert_eq!( + mantle_selected_reasoning_effort( + &LanguageModelRequest { + thinking_allowed: true, + thinking_effort: Some("none".to_string()), + ..Default::default() + }, + &MantleModel::Gpt5_5, + ), + Some(ReasoningEffort::Medium) + ); + } + + #[test] + fn test_strip_unsupported_mantle_response_fields_removes_context_management() { + let mut request = into_open_ai_response( + LanguageModelRequest { + compact_at_tokens: Some(10_000), + ..Default::default() + }, + "openai.gpt-5.5", + true, + false, + Some(128_000), + Some(ReasoningEffort::Medium), + false, + ); + + assert!(request.context_management.is_some()); + strip_unsupported_mantle_response_fields(&mut request); + assert!(request.context_management.is_none()); + + let request = serde_json::to_value(&request).unwrap(); + assert!(request.get("context_management").is_none()); + } } diff --git a/crates/language_models/src/settings.rs b/crates/language_models/src/settings.rs index 1ce7caa3a67..4f360eea95c 100644 --- a/crates/language_models/src/settings.rs +++ b/crates/language_models/src/settings.rs @@ -95,6 +95,7 @@ impl settings::Settings for AllLanguageModelSettings { .collect(), bedrock: AmazonBedrockSettings { available_models: bedrock.available_models.unwrap_or_default(), + mantle_available_models: bedrock.mantle_available_models.unwrap_or_default(), custom_headers: custom_headers_from( "Amazon Bedrock", bedrock.custom_headers, diff --git a/crates/open_ai/src/completion.rs b/crates/open_ai/src/completion.rs index f39a9af218f..0638203d769 100644 --- a/crates/open_ai/src/completion.rs +++ b/crates/open_ai/src/completion.rs @@ -814,7 +814,8 @@ impl OpenAiResponseEventMapper { } events } - ResponsesStreamEvent::ReasoningSummaryTextDelta { delta, .. } => { + ResponsesStreamEvent::ReasoningSummaryTextDelta { delta, .. } + | ResponsesStreamEvent::ReasoningDelta { delta, .. } => { if delta.is_empty() { Vec::new() } else { @@ -973,6 +974,7 @@ impl OpenAiResponseEventMapper { | ResponsesStreamEvent::ContentPartDone { .. } | ResponsesStreamEvent::ReasoningSummaryTextDone { .. } | ResponsesStreamEvent::ReasoningSummaryPartDone { .. } + | ResponsesStreamEvent::ReasoningDone { .. } | ResponsesStreamEvent::Created { .. } | ResponsesStreamEvent::InProgress { .. } | ResponsesStreamEvent::Unknown => Vec::new(), @@ -1371,6 +1373,22 @@ mod tests { )); } + #[test] + fn responses_stream_maps_mantle_reasoning_delta() { + let event = serde_json::from_value::(json!({ + "type": "response.reasoning.delta", + "delta": "checking the contract terms" + })) + .unwrap(); + + let mapped = map_response_events(vec![event]); + assert!(matches!( + &mapped[0], + LanguageModelCompletionEvent::Thinking { text, signature: None } + if text == "checking the contract terms" + )); + } + #[test] fn response_usage_deserializes_cached_tokens() -> Result<()> { let usage: ResponseUsage = serde_json::from_value(json!({ diff --git a/crates/open_ai/src/responses.rs b/crates/open_ai/src/responses.rs index 647f9fcbace..33ec1dfba46 100644 --- a/crates/open_ai/src/responses.rs +++ b/crates/open_ai/src/responses.rs @@ -310,6 +310,23 @@ pub enum StreamEvent { output_index: usize, summary_index: usize, }, + #[serde(rename = "response.reasoning.delta")] + ReasoningDelta { + #[serde(default)] + item_id: Option, + #[serde(default)] + output_index: Option, + delta: String, + }, + #[serde(rename = "response.reasoning.done")] + ReasoningDone { + #[serde(default)] + item_id: Option, + #[serde(default)] + output_index: Option, + #[serde(default)] + text: Option, + }, #[serde(rename = "response.function_call_arguments.delta")] FunctionCallArgumentsDelta { item_id: String, diff --git a/crates/settings_content/src/language_model.rs b/crates/settings_content/src/language_model.rs index f61d80ead75..5abc65dc805 100644 --- a/crates/settings_content/src/language_model.rs +++ b/crates/settings_content/src/language_model.rs @@ -114,6 +114,10 @@ pub struct AnthropicAvailableModel { #[derive(Default, Clone, Debug, Serialize, Deserialize, PartialEq, JsonSchema, MergeFrom)] pub struct AmazonBedrockSettingsContent { pub available_models: Option>, + /// Custom models served through the `bedrock-mantle` endpoint's + /// OpenAI-compatible APIs, in addition to the built-in Mantle models + /// (GPT-5.5, GPT-5.4, Grok 4.3). + pub mantle_available_models: Option>, pub custom_headers: Option>, pub endpoint_url: Option, pub region: Option, @@ -139,6 +143,32 @@ pub struct BedrockAvailableModel { pub mode: Option, } +#[with_fallible_options] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema, MergeFrom)] +pub struct BedrockMantleAvailableModel { + /// The model id as expected in Bedrock Mantle request bodies, e.g. `openai.gpt-5.5`. + pub name: String, + pub display_name: Option, + pub max_tokens: u64, + pub max_output_tokens: Option, + /// The OpenAI-compatible API this model must be called through. + pub protocol: BedrockMantleProtocolContent, + pub supports_tools: Option, + pub supports_images: Option, + /// Whether this custom Mantle model supports OpenAI reasoning effort parameters. + pub supports_thinking: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize, JsonSchema, MergeFrom)] +pub enum BedrockMantleProtocolContent { + /// The OpenAI Chat Completions API (`/chat/completions`). + #[serde(rename = "chat_completions")] + ChatCompletions, + /// The OpenAI Responses API (`/responses`). + #[serde(rename = "responses")] + Responses, +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema, MergeFrom)] pub enum BedrockAuthMethodContent { #[serde(rename = "named_profile")] diff --git a/docs/src/ai/use-a-gateway.md b/docs/src/ai/use-a-gateway.md index e010cfe8159..6cd1a7a614d 100644 --- a/docs/src/ai/use-a-gateway.md +++ b/docs/src/ai/use-a-gateway.md @@ -213,6 +213,38 @@ Some AWS environments require a guardrail on every Bedrock API call. Add `guardr } ``` +### Bedrock Mantle Models {#bedrock-mantle-models} + +Some models, such as GPT-5.5, GPT-5.4, and Grok 4.3, aren't available through Bedrock's Converse API and are only reachable through `bedrock-mantle`, AWS's OpenAI-compatible inference endpoint. Zed routes these models through `bedrock-mantle` automatically; they appear alongside the rest of the Bedrock models in the model picker once you're authenticated, with no extra configuration required. + +Mantle models require IAM permissions for the `bedrock-mantle` endpoint (for example via the `AmazonBedrockMantleInferenceAccess` managed policy) in addition to whatever permissions your existing Bedrock credentials already have, and `bedrock-mantle` is only available in [some AWS Regions](https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.html#regions). Zed surfaces an error naming the current Region and the supported ones if you try to use a Mantle model outside of them. + +#### Custom Bedrock Mantle Models {#bedrock-mantle-custom-models} + +You can add custom models served through `bedrock-mantle` with `mantle_available_models`: + +```json [settings] +{ + "language_models": { + "bedrock": { + "mantle_available_models": [ + { + "name": "openai.gpt-oss-120b", + "display_name": "GPT-OSS 120B", + "max_tokens": 128000, + "protocol": "chat_completions", + "supports_tools": true, + "supports_images": false, + "supports_thinking": true + } + ] + } + } +} +``` + +`protocol` selects which OpenAI-compatible API the model is called through, and must be either `chat_completions` or `responses`. Set `supports_thinking` to `true` for custom Mantle models that accept OpenAI reasoning effort parameters; Zed will then expose `low`, `medium`, `high`, and `xhigh` in the thinking effort picker, while disabling thinking sends `none`. + ## OpenAI-Compatible Gateways {#openai-compatible} If your gateway exposes an OpenAI-compatible API, configure it with [Use API Access](./use-api-access.md#openai-compatible).