public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Fix provider context limits and AWS authentication routing

Browse Source
This commit is contained in:
rkw6086
2026-09-11 11:15:48 -05:00
parent 73d6839ee9
commit 5d2a0ea2de
8 changed files with 270 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/block.rs
+5 -3
View File
@@ -1720,7 +1720,7 @@ impl AIBlock {
me.finish(FinishReason::Complete, ctx);
}
AIBlockOutputStatus::Failed { error, .. } => {
me.maybe_create_aws_bedrock_credentials_error_view(&error, ctx);
me.maybe_create_aws_bedrock_credentials_error_view(&error, false, ctx);
me.finish(FinishReason::Error, ctx);
}
AIBlockOutputStatus::Cancelled { .. } => {
@@ -2119,7 +2119,7 @@ impl AIBlock {
},
ctx
);
self.maybe_create_aws_bedrock_credentials_error_view(&error, ctx);
self.maybe_create_aws_bedrock_credentials_error_view(&error, true, ctx);
// There are no actions to be taken in this block, it is finished.
self.finish(FinishReason::Error, ctx);
}
@@ -4372,6 +4372,7 @@ impl AIBlock {
fn maybe_create_aws_bedrock_credentials_error_view(
&mut self,
error: &RenderableAIError,
attempt_auto_login: bool,
ctx: &mut ViewContext<Self>,
) {
// Only create the view for AWS Bedrock credentials errors
@@ -4386,7 +4387,8 @@ impl AIBlock {
let ai_settings = AISettings::as_ref(ctx);
let login_command = ai_settings.bedrock_auth_refresh_command.value().clone();
let auto_login_enabled = *ai_settings.bedrock_auto_login.value();
// Recreating a historical error view must not restart authentication.
let auto_login_enabled = attempt_auto_login && *ai_settings.bedrock_auto_login.value();
// If auto-login is enabled, run the login command automatically
if auto_login_enabled {
app/src/ai/llms.rs
+55 -25
View File
@@ -1046,7 +1046,9 @@ impl LLMPreferences {
configured_models
};
if !single_provider_models.is_empty() {
// The registry replaces legacy settings. Otherwise stale legacy credentials
// win duplicate model IDs before the configured provider is considered.
if settings.openai_providers.value().is_empty() && !single_provider_models.is_empty() {
let base_url = settings.openai_base_url.value().clone();
let api_key = {
let key = settings.openai_api_key.value().clone();
@@ -1170,7 +1172,7 @@ impl LLMPreferences {
location: provider_location.clone(),
model: Some(model.model_id.clone()),
reasoning_effort: reasoning_effort.clone(),
max_input_tokens: Some(openai_model_context_size(model)),
max_input_tokens: Some(openai_model_context_size(model, provider_kind)),
max_output_tokens: model.max_output_tokens,
supports_system_messages: model.supports_system_messages(),
};
@@ -1487,7 +1489,7 @@ impl LLMPreferences {
#[cfg(not(target_family = "wasm"))]
pub fn fetch_openai_models_from_endpoint(&mut self, ctx: &mut ModelContext<Self>) {
let settings = AISettings::as_ref(ctx);
if !settings.is_openai_provider_enabled() {
if !settings.is_openai_provider_enabled() || !settings.openai_providers.value().is_empty() {
return;
}
@@ -1833,6 +1835,7 @@ impl LLMPreferences {
vision_supported: true,
context_size: model.context_size.unwrap_or(128_000),
max_input_tokens: model.context_size,
default_context_size: None,
max_output_tokens: None,
provider: None,
supports_system_messages: Some(true),
@@ -2626,8 +2629,18 @@ fn get_new_agent_mode_choices(
}
#[cfg(not(target_family = "wasm"))]
fn openai_model_context_size(model: &OpenAIModelConfig) -> u32 {
model.max_input_tokens.unwrap_or(model.context_size)
fn openai_model_context_size(model: &OpenAIModelConfig, provider_kind: OpenAIProviderKind) -> u32 {
// Older saved ChatGPT catalogs stored the default window in max_input_tokens.
// Preserve their larger usable window until the catalog refresh separates them.
if provider_kind == OpenAIProviderKind::ChatGPTSubscription
&& model.default_context_size.is_none()
{
return model.context_size;
}
model
.max_input_tokens
.unwrap_or(model.context_size)
.min(model.context_size)
}
/// Merges endpoint metadata into a provider's configured models without
@@ -2721,25 +2734,24 @@ fn openai_model_context_window(
model: &OpenAIModelConfig,
provider_kind: OpenAIProviderKind,
) -> LLMContextWindow {
let default_context_size = openai_model_context_size(model);
let max_context_size = if provider_kind == OpenAIProviderKind::ChatGPTSubscription {
model.context_size.max(default_context_size)
} else {
default_context_size
};
let max_context_size = openai_model_context_size(model, provider_kind);
let default_context_size = model
.default_context_size
.or_else(|| {
// Compatibility with saved catalogs from before the fields were separated.
(provider_kind == OpenAIProviderKind::ChatGPTSubscription)
.then_some(model.max_input_tokens)
.flatten()
})
.unwrap_or(max_context_size)
.min(max_context_size);
LLMContextWindow {
is_configurable: max_context_size > default_context_size,
min: default_context_size,
max: max_context_size,
// ChatGPT's catalog exposes `context_window` as the effective window
// and `max_context_window` as the model's actual maximum. Use the
// latter as the default so model refreshes do not reset the profile
// to the smaller effective window (for example, 272k).
default_max: if provider_kind == OpenAIProviderKind::ChatGPTSubscription {
max_context_size
} else {
default_context_size
},
// Use the full usable window; the catalog default remains available as
// the smaller configurable budget, not as a hard input ceiling.
default_max: max_context_size,
}
}
@@ -2888,7 +2900,9 @@ fn chatgpt_models_from_codex_response(body: &serde_json::Value) -> Vec<OpenAIMod
u32::try_from(u64::from(context_size) * u64::from(effective_context_percent) / 100)
.unwrap_or(context_size)
};
let max_input_tokens = Some(effective_context_size(default_context_size));
let default_context_size = Some(effective_context_size(default_context_size));
let max_input_tokens =
u32_from_json_any(model, &["max_input_tokens"]).map(effective_context_size);
let context_size = effective_context_size(max_context_size);
let vision_supported = model["input_modalities"]
@@ -2921,6 +2935,7 @@ fn chatgpt_models_from_codex_response(body: &serde_json::Value) -> Vec<OpenAIMod
vision_supported,
context_size,
max_input_tokens,
default_context_size,
max_output_tokens: None,
provider: Some("openai".to_string()),
supports_system_messages: Some(true),
@@ -2991,7 +3006,12 @@ async fn fetch_from_litellm_model_info(
let max_output_tokens = model_info["max_output_tokens"]
.as_u64()
.and_then(|v| u32::try_from(v).ok());
let context_size = max_input_tokens.unwrap_or(200_000);
let context_size = u32_from_json_any(
model_info,
&["max_context_window", "context_window", "max_model_len"],
)
.or(max_input_tokens)
.unwrap_or(200_000);
let vision_supported = model_info["supports_vision"].as_bool().unwrap_or(true);
@@ -3042,6 +3062,7 @@ async fn fetch_from_litellm_model_info(
vision_supported,
context_size,
max_input_tokens,
default_context_size: None,
max_output_tokens,
provider,
supports_system_messages: if model_name.starts_with("codex-gpt-") {
@@ -3122,9 +3143,17 @@ async fn fetch_from_openai_models(
m,
&["max_input_tokens", "input_token_limit", "max_prompt_tokens"],
);
let context_size = u32_from_any(m, &["max_model_len", "context_window", "token_size"])
.or(max_input_tokens)
.unwrap_or(200_000);
let context_size = u32_from_any(
m,
&[
"max_context_window",
"max_model_len",
"context_window",
"token_size",
],
)
.or(max_input_tokens)
.unwrap_or(200_000);
let max_output_tokens = u32_from_any(
m,
&[
@@ -3171,6 +3200,7 @@ async fn fetch_from_openai_models(
.unwrap_or(true),
context_size,
max_input_tokens,
default_context_size: None,
max_output_tokens,
provider,
supports_system_messages: if id.starts_with("codex-gpt-") {
app/src/ai/llms_tests.rs
+104 -8
View File
@@ -149,6 +149,7 @@ fn openai_model(model_id: &str) -> OpenAIModelConfig {
vision_supported: false,
context_size: 200_000,
max_input_tokens: None,
default_context_size: None,
max_output_tokens: None,
provider: None,
supports_system_messages: None,
@@ -210,8 +211,18 @@ fn acp_agent(
}
fn empty_preferences() -> LLMPreferences {
// Production defaults include a UI placeholder; inventory tests start empty.
let mut models_by_feature = ModelsByFeature::default();
models_by_feature.agent_mode.choices.clear();
models_by_feature.coding.choices.clear();
if let Some(cli) = &mut models_by_feature.cli_agent {
cli.choices.clear();
}
if let Some(computer_use) = &mut models_by_feature.computer_use {
computer_use.choices.clear();
}
LLMPreferences {
models_by_feature: ModelsByFeature::default(),
models_by_feature,
last_update: None,
base_llm_for_terminal_view: HashMap::new(),
custom_llms: Vec::new(),
@@ -791,7 +802,8 @@ fn chatgpt_codex_models_parse_visible_catalog_entries() {
assert_eq!(models[0].display_name, "GPT-5.6-Sol");
assert!(models[0].vision_supported);
assert_eq!(models[0].context_size, 828_400);
assert_eq!(models[0].max_input_tokens, Some(258_400));
assert_eq!(models[0].max_input_tokens, None);
assert_eq!(models[0].default_context_size, Some(258_400));
assert_eq!(models[0].reasoning_efforts, ["low", "xhigh", "ultra"]);
assert_eq!(models[0].provider.as_deref(), Some("openai"));
assert_eq!(models[0].supports_system_messages, Some(true));
@@ -799,7 +811,8 @@ fn chatgpt_codex_models_parse_visible_catalog_entries() {
assert_eq!(models[1].model_id, "gpt-text-only");
assert!(!models[1].vision_supported);
assert_eq!(models[1].context_size, 115_200);
assert_eq!(models[1].max_input_tokens, Some(115_200));
assert_eq!(models[1].max_input_tokens, None);
assert_eq!(models[1].default_context_size, Some(115_200));
let configurable = openai_model_context_window(
&models[0],
@@ -812,10 +825,10 @@ fn chatgpt_codex_models_parse_visible_catalog_entries() {
let fixed =
openai_model_context_window(&models[0], crate::settings::OpenAIProviderKind::LiteLLM);
assert!(!fixed.is_configurable);
assert!(fixed.is_configurable);
assert_eq!(fixed.min, 258_400);
assert_eq!(fixed.default_max, 258_400);
assert_eq!(fixed.max, 258_400);
assert_eq!(fixed.default_max, 828_400);
assert_eq!(fixed.max, 828_400);
}
#[test]
@@ -834,17 +847,100 @@ fn chatgpt_catalog_merge_drops_stale_models_but_preserves_overrides() {
discovered.display_name = "GPT-5.6-Sol".to_string();
discovered.vision_supported = true;
discovered.context_size = 828_400;
discovered.max_input_tokens = Some(258_400);
discovered.default_context_size = Some(258_400);
let merged = merge_discovered_chatgpt_subscription_models(&[existing, stale], vec![discovered]);
assert_eq!(merged.len(), 1);
assert_eq!(merged[0].model_id, "gpt-5.6-sol");
assert_eq!(merged[0].context_size, 828_400);
assert_eq!(merged[0].max_input_tokens, Some(258_400));
assert_eq!(merged[0].max_input_tokens, None);
assert_eq!(merged[0].default_context_size, Some(258_400));
assert!(!merged[0].enabled);
assert_eq!(
merged[0].capability_override("vision"),
crate::settings::ModelCapabilityOverride::Unsupported
);
}
#[test]
fn provider_context_budget_respects_hard_input_and_total_limits() {
let mut model = openai_model("limited-model");
model.context_size = 872_000;
model.max_input_tokens = Some(272_000);
let window = openai_model_context_window(&model, OpenAIProviderKind::LiteLLM);
assert_eq!(window.max, 272_000);
assert_eq!(window.default_max, 272_000);
assert!(!window.is_configurable);
model.max_input_tokens = Some(1_000_000);
assert_eq!(
openai_model_context_size(&model, OpenAIProviderKind::LiteLLM),
872_000
);
model.default_context_size = Some(272_000);
model.max_input_tokens = Some(400_000);
let window = openai_model_context_window(&model, OpenAIProviderKind::ChatGPTSubscription);
assert_eq!(window.min, 272_000);
assert_eq!(window.max, 400_000);
assert_eq!(
openai_model_context_size(&model, OpenAIProviderKind::ChatGPTSubscription),
window.max
);
}
#[test]
fn legacy_chatgpt_default_does_not_limit_routing() {
let mut model = openai_model("legacy-model");
model.context_size = 872_000;
model.max_input_tokens = Some(272_000);
let window = openai_model_context_window(&model, OpenAIProviderKind::ChatGPTSubscription);
assert_eq!(window.min, 272_000);
assert_eq!(window.max, 872_000);
assert_eq!(
openai_model_context_size(&model, OpenAIProviderKind::ChatGPTSubscription),
window.max
);
}
#[test]
fn provider_registry_overrides_legacy_endpoint_for_duplicate_models() {
App::test((), |mut app| async move {
initialize_settings_for_tests(&mut app);
AISettings::handle(&app).update(&mut app, |settings, ctx| {
settings.openai_enabled.set_value(true, ctx).unwrap();
settings
.openai_base_url
.set_value("https://stale.test/v1".to_owned(), ctx)
.unwrap();
settings
.openai_models
.set_value(vec![openai_model("shared-model")], ctx)
.unwrap();
settings
.openai_providers
.set_value(
vec![OpenAIProviderConfig {
kind: OpenAIProviderKind::LiteLLM,
enabled: true,
name: "Current LiteLLM".to_owned(),
base_url: "https://current.test/v1".to_owned(),
api_key: Some("current-key".to_owned()),
project_id: None,
location: None,
models: vec![openai_model("shared-model")],
}],
ctx,
)
.unwrap();
});
let mut preferences = empty_preferences();
app.read(|ctx| preferences.inject_openai_models(ctx));
let config = preferences
.openai_client_config_for_model("shared-model")
.unwrap();
assert_eq!(config.base_url, "https://current.test/v1");
assert_eq!(config.api_key.as_deref(), Some("current-key"));
});
}
app/src/ai/runtime/event_translator.rs
+35 -4
View File
@@ -43,6 +43,8 @@ pub(crate) struct RuntimeResponseTranslator {
/// Usage for the most recent model call.
usage: Usage,
context_usage: Option<(u64, u64)>,
// Set before starting each model call, including failures before TurnStarted.
authentication_target: Option<(api::LlmProvider, String)>,
}
/// Projects a multi-turn provider run into one existing Galaxy response stream.
@@ -98,8 +100,16 @@ impl ProviderRunResponseProjector {
ProviderRunProjection::ModelRetry { .. } => {
Ok(self.translator.discard_failed_turn_output())
}
ProviderRunProjection::ModelTurnRequested { .. }
| ProviderRunProjection::ModelTurnFinished { .. } => Ok(Vec::new()),
ProviderRunProjection::ModelTurnRequested {
runtime_id,
model_id,
..
} => {
self.translator.authentication_target =
Some((authentication_provider_for_runtime(&runtime_id), model_id));
Ok(Vec::new())
}
ProviderRunProjection::ModelTurnFinished { .. } => Ok(Vec::new()),
ProviderRunProjection::ToolBatchReady { batch } => {
if !self.plan_tasks_projected {
if let Some(todos) = todos_from_plan_batch(&batch) {
@@ -149,6 +159,19 @@ impl ProviderRunResponseProjector {
}
}
// Runtime IDs identify the transport, unlike model IDs: a Bedrock-backed model
// behind LiteLLM still authenticates against the OpenAI-compatible endpoint.
fn authentication_provider_for_runtime(runtime_id: &str) -> api::LlmProvider {
match runtime_id.split_once(':').map(|(runtime, _)| runtime) {
Some("rig-bedrock") => api::LlmProvider::AwsBedrock,
Some("rig-openai-compatible" | "rig-chatgpt-subscription") => api::LlmProvider::Openai,
Some("rig-anthropic") => api::LlmProvider::Anthropic,
Some("rig-gemini" | "rig-vertex-ai") => api::LlmProvider::Google,
// Unrecognized runtimes must never trigger provider-specific login commands.
Some(_) | None => api::LlmProvider::Unknown,
}
}
// Only explicit multi-step checklists in a plan become tasks. Other document bullets
// (risks, examples, requirements) are not executable tasks.
fn todos_from_plan_batch(
@@ -251,6 +274,7 @@ impl RuntimeResponseTranslator {
has_visible_output: false,
usage: Usage::default(),
context_usage: None,
authentication_target: None,
}
}
@@ -528,8 +552,15 @@ impl RuntimeResponseTranslator {
.is_some_and(|error| error.kind == galaxy_agent_core::AgentErrorKind::Authentication)
{
stream_finished::Reason::InvalidApiKey(stream_finished::InvalidApiKey {
provider: warp_multi_agent_api::LlmProvider::AwsBedrock as i32,
model_name: self.config.model_id.clone(),
provider: self
.authentication_target
.as_ref()
.map_or(api::LlmProvider::Unknown, |(provider, _)| *provider)
as i32,
model_name: self
.authentication_target
.as_ref()
.map_or_else(|| self.config.model_id.clone(), |(_, model)| model.clone()),
})
} else {
stream_finished::Reason::InternalError(stream_finished::InternalError {
app/src/ai/runtime/event_translator_tests.rs
+60
View File
@@ -640,3 +640,63 @@ fn plan_tasks_require_multiple_explicit_checklist_items_in_a_task_section() {
assert_eq!(todos[1].title, "Test stale callbacks");
assert_ne!(todos[0].id, todos[1].id);
}
#[test]
fn authentication_failures_identify_the_requested_transport_before_stream_start() {
for (runtime, expected) in [
(
"rig-openai-compatible",
warp_multi_agent_api::LlmProvider::Openai,
),
(
"rig-chatgpt-subscription",
warp_multi_agent_api::LlmProvider::Openai,
),
("rig-bedrock", warp_multi_agent_api::LlmProvider::AwsBedrock),
(
"rig-anthropic",
warp_multi_agent_api::LlmProvider::Anthropic,
),
("rig-gemini", warp_multi_agent_api::LlmProvider::Google),
("rig-vertex-ai", warp_multi_agent_api::LlmProvider::Google),
("unrecognized", warp_multi_agent_api::LlmProvider::Unknown),
] {
let mut projector = ProviderRunResponseProjector::new(provider_translator().config);
// A previous Bedrock turn must not leak its identity into the next request.
for runtime_id in [
"rig-bedrock:old-model".to_owned(),
format!("{runtime}:new-model"),
] {
projector
.project(ProviderRunProjection::ModelTurnRequested {
work_id: galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(1),
},
profile: galaxy_agent_core::ProviderRequestProfile::new("cli"),
runtime_id,
model_id: "new-model".to_owned(),
retry_attempt: 0,
})
.expect("requested turn projection");
}
let error = galaxy_agent_core::AgentError::new(
galaxy_agent_core::AgentErrorKind::Authentication,
"Unauthorized",
);
let events =
projector
.translator
.provider_failure("Unauthorized", Some(&error), &Usage::default());
let Some(response_event::Type::Finished(finished)) = &events.last().unwrap().r#type else {
panic!("expected a finished event");
};
let Some(response_event::stream_finished::Reason::InvalidApiKey(details)) =
&finished.reason
else {
panic!("expected authentication failure");
};
assert_eq!(details.provider, expected as i32, "{runtime}");
assert_eq!(details.model_name, "new-model");
}
}
app/src/settings/ai.rs
+7 -1
View File
@@ -857,13 +857,19 @@ pub struct OpenAIModelConfig {
default = "default_context_size",
alias = "token_size",
alias = "max_model_len",
alias = "context_window"
alias = "context_window",
alias = "max_context_window"
)]
#[schemars(description = "Maximum context window size in tokens.")]
pub context_size: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schemars(description = "Optional maximum input tokens supported by this model.")]
pub max_input_tokens: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schemars(
description = "Optional default context budget, distinct from the hard input limit."
)]
pub default_context_size: Option<u32>,
#[serde(
default,
alias = "output_token_limit",
app/src/settings_view/ai_page.rs
+3
View File
@@ -7242,6 +7242,9 @@ impl ProviderSettingsWidget {
format!("Context: {}", model.context_size),
];
if let Some(default_context_size) = model.default_context_size {
details.push(format!("Default context: {default_context_size}"));
}
if let Some(max_input_tokens) = model.max_input_tokens {
details.push(format!("Max input: {max_input_tokens}"));
}
app/src/settings_view/provider_setup_view.rs
+1
View File
@@ -782,6 +782,7 @@ impl ProviderSetupView {
if let Ok(context_size) = editor.as_ref(ctx).buffer_text(ctx).parse() {
model.context_size = context_size;
model.max_input_tokens = Some(context_size);
model.default_context_size = Some(context_size);
}
}
me.update_next_button(ctx);