From 502628013124201081245aaab68fb3be4c0c9408 Mon Sep 17 00:00:00 2001 From: Shardul Vaidya <31039336+5herlocked@users.noreply.github.com> Date: Thu, 12 Feb 2026 09:32:24 -0500 Subject: [PATCH] bedrock: Enable 1M context window (#48542) Release Notes: - Added `allow_extended_context` to the Bedrock settings which enables 1M context windows on models that support it --------- Co-authored-by: Ona Co-authored-by: Bennet Bo Fenner --- crates/bedrock/src/bedrock.rs | 28 ++++++++++++------- crates/bedrock/src/models.rs | 14 ++++++++++ .../language_models/src/provider/bedrock.rs | 24 ++++++++++++++-- crates/language_models/src/settings.rs | 1 + crates/settings_content/src/language_model.rs | 2 ++ docs/src/ai/llm-providers.md | 21 ++++++++++++++ 6 files changed, 77 insertions(+), 13 deletions(-) diff --git a/crates/bedrock/src/bedrock.rs b/crates/bedrock/src/bedrock.rs index 2ca39326d30f88f59922aeeb34bca40d3ea79f46..92ab097c4925326813a539d986f4fa1d89ca096a 100644 --- a/crates/bedrock/src/bedrock.rs +++ b/crates/bedrock/src/bedrock.rs @@ -31,6 +31,8 @@ use thiserror::Error; pub use crate::models::*; +pub const CONTEXT_1M_BETA_HEADER: &str = "context-1m-2025-08-07"; + pub async fn stream_completion( client: bedrock::Client, request: Request, @@ -39,6 +41,8 @@ pub async fn stream_completion( .model_id(request.model.clone()) .set_messages(request.messages.into()); + let mut additional_fields: HashMap = HashMap::new(); + match request.thinking { Some(Thinking::Enabled { budget_tokens: Some(budget_tokens), @@ -50,24 +54,27 @@ pub async fn stream_completion( Document::Number(AwsNumber::PosInt(budget_tokens)), ), ]); - response = - response.additional_model_request_fields(Document::Object(HashMap::from([( - "thinking".to_string(), - Document::from(thinking_config), - )]))); + additional_fields.insert("thinking".to_string(), Document::from(thinking_config)); } Some(Thinking::Adaptive { effort: _ }) => { let thinking_config = HashMap::from([("type".to_string(), Document::String("adaptive".to_string()))]); - response = - response.additional_model_request_fields(Document::Object(HashMap::from([( - "thinking".to_string(), - Document::from(thinking_config), - )]))); + additional_fields.insert("thinking".to_string(), Document::from(thinking_config)); } _ => {} } + if request.allow_extended_context { + additional_fields.insert( + "anthropic_beta".to_string(), + Document::Array(vec![Document::String(CONTEXT_1M_BETA_HEADER.to_string())]), + ); + } + + if !additional_fields.is_empty() { + response = response.additional_model_request_fields(Document::Object(additional_fields)); + } + if request.tools.as_ref().is_some_and(|t| !t.tools.is_empty()) { response = response.set_tool_config(request.tools); } @@ -178,6 +185,7 @@ pub struct Request { pub temperature: Option, pub top_k: Option, pub top_p: Option, + pub allow_extended_context: bool, } #[derive(Debug, Serialize, Deserialize)] diff --git a/crates/bedrock/src/models.rs b/crates/bedrock/src/models.rs index bd3c1e7337743a334312ba0c5cd314a1f7410583..e073f5f17aa09505d45fd5bffc99c442e22f54bf 100644 --- a/crates/bedrock/src/models.rs +++ b/crates/bedrock/src/models.rs @@ -638,6 +638,20 @@ impl Model { } } + pub fn supports_extended_context(&self) -> bool { + matches!( + self, + Model::ClaudeSonnet4 + | Model::ClaudeSonnet4Thinking + | Model::ClaudeSonnet4_5 + | Model::ClaudeSonnet4_5Thinking + | Model::ClaudeOpus4_5 + | Model::ClaudeOpus4_5Thinking + | Model::ClaudeOpus4_6 + | Model::ClaudeOpus4_6Thinking + ) + } + pub fn cross_region_inference_id( &self, region: &str, diff --git a/crates/language_models/src/provider/bedrock.rs b/crates/language_models/src/provider/bedrock.rs index b45b34e0b75f148b15439704e6cdeef75314aa6a..f16af99b50c578b45cb012d0334267fd5b91fe5c 100644 --- a/crates/language_models/src/provider/bedrock.rs +++ b/crates/language_models/src/provider/bedrock.rs @@ -111,6 +111,7 @@ pub struct AmazonBedrockSettings { pub role_arn: Option, pub authentication_method: Option, pub allow_global: Option, + pub allow_extended_context: Option, } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, EnumIter, IntoStaticStr, JsonSchema)] @@ -382,6 +383,13 @@ impl State { .and_then(|s| s.allow_global) .unwrap_or(false) } + + fn get_allow_extended_context(&self) -> bool { + self.settings + .as_ref() + .and_then(|s| s.allow_extended_context) + .unwrap_or(false) + } } pub struct BedrockLanguageModelProvider { @@ -672,9 +680,14 @@ impl LanguageModel for BedrockModel { LanguageModelCompletionError, >, > { - let (region, allow_global) = cx.read_entity(&self.state, |state, _cx| { - (state.get_region(), state.get_allow_global()) - }); + let (region, allow_global, allow_extended_context) = + cx.read_entity(&self.state, |state, _cx| { + ( + state.get_region(), + state.get_allow_global(), + state.get_allow_extended_context(), + ) + }); let model_id = match self.model.cross_region_inference_id(®ion, allow_global) { Ok(s) => s, @@ -685,6 +698,8 @@ impl LanguageModel for BedrockModel { let deny_tool_calls = request.tool_choice == Some(LanguageModelToolChoice::None); + let use_extended_context = allow_extended_context && self.model.supports_extended_context(); + let request = match into_bedrock( request, model_id, @@ -692,6 +707,7 @@ impl LanguageModel for BedrockModel { self.model.max_output_tokens(), self.model.mode(), self.model.supports_caching(), + use_extended_context, ) { Ok(request) => request, Err(err) => return futures::future::ready(Err(err.into())).boxed(), @@ -747,6 +763,7 @@ pub fn into_bedrock( max_output_tokens: u64, mode: BedrockModelMode, supports_caching: bool, + allow_extended_context: bool, ) -> Result { let mut new_messages: Vec = Vec::new(); let mut system_message = String::new(); @@ -955,6 +972,7 @@ pub fn into_bedrock( temperature: request.temperature.or(Some(default_temperature)), top_k: None, top_p: None, + allow_extended_context, }) } diff --git a/crates/language_models/src/settings.rs b/crates/language_models/src/settings.rs index 62f0025c755e10ea1bdae605d9dcc752298bb5f1..b8f548acbeeac20b4c9af2a8e64de2ed2d805093 100644 --- a/crates/language_models/src/settings.rs +++ b/crates/language_models/src/settings.rs @@ -59,6 +59,7 @@ impl settings::Settings for AllLanguageModelSettings { role_arn: None, // todo(was never a setting for this...) authentication_method: bedrock.authentication_method.map(Into::into), allow_global: bedrock.allow_global, + allow_extended_context: bedrock.allow_extended_context, }, deepseek: DeepSeekSettings { api_url: deepseek.api_url.unwrap(), diff --git a/crates/settings_content/src/language_model.rs b/crates/settings_content/src/language_model.rs index 88f5b1b985756f9c25074591b0146f2ccc715f3c..1f0f338f6d7ac35b9d6862f961bb45f7d2abfb33 100644 --- a/crates/settings_content/src/language_model.rs +++ b/crates/settings_content/src/language_model.rs @@ -63,6 +63,8 @@ pub struct AmazonBedrockSettingsContent { pub profile: Option, pub authentication_method: Option, pub allow_global: Option, + /// Enable the 1M token extended context window beta for supported Anthropic models. + pub allow_extended_context: Option, } #[with_fallible_options] diff --git a/docs/src/ai/llm-providers.md b/docs/src/ai/llm-providers.md index bd353defa2120ff7aba4fe5ed0ade88223f64ea0..6a9b82d7e1dc752e373af576635be22cd44b08ee 100644 --- a/docs/src/ai/llm-providers.md +++ b/docs/src/ai/llm-providers.md @@ -149,6 +149,27 @@ We will support Cross-Region inference for each of the models on a best-effort b For the most up-to-date supported regions and models, refer to the [Supported Models and Regions for Cross Region inference](https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html). +#### Extended Context Window {#bedrock-extended-context} + +Anthropic models on Bedrock support a [1M token extended context window](https://docs.anthropic.com/en/docs/build-with-claude/extended-context) beta. To enable this feature, add `"allow_extended_context": true` to your Bedrock configuration: + +```json [settings] +{ + "language_models": { + "bedrock": { + "authentication_method": "named_profile", + "region": "your-aws-region", + "profile": "your-profile-name", + "allow_extended_context": true + } + } +} +``` + +When enabled, Zed will include the `anthropic_beta` field in requests to Bedrock, enabling the 1M token context window for supported Anthropic models such as Claude Sonnet 4.5 and Claude Opus 4.6. + +> **Note**: Extended context usage may incur additional API costs. Refer to your AWS Bedrock pricing for details. + ### Anthropic {#anthropic} You can use Anthropic models by choosing them via the model dropdown in the Agent Panel.