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

Improve agent provider resilience

Browse Source
This commit is contained in:
Ryan Ward
2026-09-02 17:09:08 -05:00
parent b115946534
commit 7a33cc7e56
72 changed files with 997 additions and 320 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/acp/transport.rs
+1
View File
@@ -228,6 +228,7 @@ fn response_translator(
max_context_tokens: None,
capabilities: RuntimeCapabilities::session_runtime(),
empty_output_message: Some("> ACP agent completed without a text response.".to_owned()),
todo_items: None,
})
}
app/src/ai/agent/api/convert_conversation_tests.rs
+1 -1
View File
@@ -2145,7 +2145,7 @@ fn test_internal_command_completion_assessment_restores_output_without_visible_i
};
mark_internal_command_completion_assessment(&mut hidden_assessment);
let provider_history =
crate::ai::bedrock::request_translator::convert_proto_message(&hidden_assessment)
crate::ai::provider::request_translator::convert_proto_message(&hidden_assessment)
.expect("hidden assessment should remain in provider history");
assert_eq!(
provider_history.role,
app/src/ai/agent/conversation.rs
+11 -11
View File
@@ -334,8 +334,8 @@ pub struct AIConversation {
/// Whether the user has pinned this child agent in the orchestration
/// pill bar. Persisted via `AgentConversationData.pinned`.
pinned: bool,
bedrock_message_history: Vec<crate::ai::bedrock::convert::ConversationMessage>,
tool_result_archive: Vec<crate::ai::bedrock::convert::ConversationMessage>,
bedrock_message_history: Vec<crate::ai::provider::convert::ConversationMessage>,
tool_result_archive: Vec<crate::ai::provider::convert::ConversationMessage>,
progressive_summary: Option<String>,
messages_summarized_up_to: usize,
current_context_tokens: u32,
@@ -448,10 +448,10 @@ impl AIConversation {
tasks: Vec<api::Task>,
conversation_data: Option<AgentConversationData>,
) -> Result<Self, RestoreConversationError> {
let bedrock_message_history: Vec<crate::ai::bedrock::convert::ConversationMessage> = tasks
let bedrock_message_history: Vec<crate::ai::provider::convert::ConversationMessage> = tasks
.iter()
.flat_map(|task| task.messages.iter())
.filter_map(crate::ai::bedrock::request_translator::convert_proto_message)
.filter_map(crate::ai::provider::request_translator::convert_proto_message)
.collect();
let (task_store, todo_lists, status) = if tasks.is_empty() {
@@ -800,32 +800,32 @@ impl AIConversation {
self.has_pending_progressive_summary = val;
}
pub fn bedrock_message_history(&self) -> &[crate::ai::bedrock::convert::ConversationMessage] {
pub fn bedrock_message_history(&self) -> &[crate::ai::provider::convert::ConversationMessage] {
&self.bedrock_message_history
}
pub fn bedrock_message_history_mut(
&mut self,
) -> &mut Vec<crate::ai::bedrock::convert::ConversationMessage> {
) -> &mut Vec<crate::ai::provider::convert::ConversationMessage> {
&mut self.bedrock_message_history
}
pub fn tool_result_archive(&self) -> &[crate::ai::bedrock::convert::ConversationMessage] {
pub fn tool_result_archive(&self) -> &[crate::ai::provider::convert::ConversationMessage] {
&self.tool_result_archive
}
pub fn archive_tool_results(
&mut self,
messages: Vec<crate::ai::bedrock::convert::ConversationMessage>,
messages: Vec<crate::ai::provider::convert::ConversationMessage>,
) {
use crate::ai::bedrock::convert::{ContentPart, MessageContent};
use crate::ai::provider::convert::{ContentPart, MessageContent};
// Cap the archive to prevent unbounded growth. The archive is only used by
// `recall_tool_history` which already truncates individual results to 50K chars,
// so retaining the most recent entries is sufficient for lookup.
const MAX_TOOL_RESULT_ARCHIVE_ENTRIES: usize = 400;
let mut pending_tool_uses: Vec<crate::ai::bedrock::convert::ConversationMessage> =
let mut pending_tool_uses: Vec<crate::ai::provider::convert::ConversationMessage> =
Vec::new();
for msg in messages {
@@ -873,7 +873,7 @@ impl AIConversation {
pub fn append_to_bedrock_history(
&mut self,
messages: Vec<crate::ai::bedrock::convert::ConversationMessage>,
messages: Vec<crate::ai::provider::convert::ConversationMessage>,
) {
self.bedrock_message_history.extend(messages);
}
app/src/ai/bedrock/mod.rs
-23
View File
@@ -1,23 +0,0 @@
pub mod client;
pub mod convert;
pub mod crash_log;
pub mod diagnostic;
pub mod discovery;
pub mod external_config;
pub mod models;
pub mod request_translator;
pub mod response_translator;
pub mod settings_view;
#[cfg(test)]
mod convert_tests;
#[cfg(test)]
#[allow(dead_code)]
mod e2e_tests;
#[cfg(test)]
#[allow(dead_code)]
mod integration_tests;
#[cfg(test)]
mod models_tests;
#[cfg(test)]
mod response_translator_tests;
app/src/ai/bedrock/models.rs
-50
View File
@@ -1,50 +0,0 @@
#![allow(dead_code)]
use crate::settings::ai::BedrockModelConfig;
pub fn configured_model_uses_rig(
selected_model_id: &str,
configured_models: &[BedrockModelConfig],
region: &str,
cross_region_inference: bool,
) -> bool {
let selected_model_id = strip_context_marker(selected_model_id);
configured_models.iter().any(|model| {
if !model.use_rig {
return false;
}
let configured_model_id = strip_context_marker(&model.model_id);
if configured_model_id == selected_model_id {
return true;
}
galaxy_agent_rig::resolve_bedrock_model_id(&model.model_id, region, cross_region_inference)
.is_ok_and(|resolved| strip_context_marker(&resolved) == selected_model_id)
})
}
fn strip_context_marker(model_id: &str) -> &str {
model_id
.strip_suffix("[1m]")
.or_else(|| model_id.strip_suffix("[1M]"))
.unwrap_or(model_id)
}
pub fn apply_cross_region_prefix(model_id: &str, region: &str) -> String {
if model_id.starts_with("arn:") {
return model_id.to_string();
}
if model_id.contains('.') && model_id.split('.').next().unwrap_or("").len() <= 6 {
return model_id.to_string();
}
let prefix = match region {
r if r.starts_with("us-") || r.starts_with("ca-") => "us",
r if r.starts_with("eu-") || r == "il-central-1" => "eu",
r if r == "ap-northeast-1" || r == "ap-northeast-3" => "jp",
r if r == "ap-southeast-2" || r == "ap-southeast-4" || r == "ap-southeast-6" => "au",
r if r.starts_with("ap-") => "apac",
_ => return model_id.to_string(),
};
format!("{}.{}", prefix, model_id)
}
app/src/ai/blocklist/action_model.rs
+65 -11
View File
@@ -1754,23 +1754,77 @@ impl BlocklistAIActionModel {
batch: &PendingToolBatch,
ctx: &mut ModelContext<Self>,
) -> Result<(), ProviderActionQueueError> {
let refs = provider_action_correlations(&actions, conversation_id, batch)?;
for ((_, action_id), _) in &refs {
if self
.provider_tool_executions
.contains_key(&(conversation_id, action_id.clone()))
let action_ids = self.provider_action_ids_to_enqueue(&actions, conversation_id, batch)?;
let refs = provider_action_correlations(&actions, conversation_id, batch)?
.into_iter()
.filter(|((_, action_id), _)| action_ids.contains(action_id));
self.provider_tool_executions.extend(refs);
self.executor.update(ctx, |executor, ctx| {
executor.mark_restored_actions(conversation_id, &recovery_action_ids, ctx);
});
let actions: Vec<AIAgentAction> = actions
.into_iter()
.filter(|action| action_ids.contains(&action.id))
.collect();
if !actions.is_empty() {
self.queue_actions(actions, conversation_id, ctx);
}
Ok(())
}
pub(super) fn provider_action_ids_to_enqueue(
&self,
actions: &[AIAgentAction],
conversation_id: AIConversationId,
batch: &PendingToolBatch,
) -> Result<HashSet<AIAgentActionId>, ProviderActionQueueError> {
Self::provider_action_ids_to_enqueue_from(
&self.provider_tool_executions,
actions,
conversation_id,
batch,
)
}
fn provider_action_ids_to_enqueue_from(
existing_correlations: &HashMap<
(AIConversationId, AIAgentActionId),
ProviderToolExecutionRef,
>,
actions: &[AIAgentAction],
conversation_id: AIConversationId,
batch: &PendingToolBatch,
) -> Result<HashSet<AIAgentActionId>, ProviderActionQueueError> {
let refs = provider_action_correlations(actions, conversation_id, batch)?;
let mut action_ids = HashSet::with_capacity(refs.len());
let mut seen_refs = HashSet::with_capacity(refs.len());
for ((_, action_id), execution_ref) in &refs {
if !seen_refs.insert(execution_ref.clone()) {
return Err(ProviderActionQueueError::ExistingCorrelation {
call_id: action_id.to_string(),
});
}
if let Some(existing_ref) =
existing_correlations.get(&(conversation_id, action_id.clone()))
{
if existing_ref == execution_ref {
continue;
}
return Err(ProviderActionQueueError::ExistingCorrelation {
call_id: action_id.to_string(),
});
}
if existing_correlations
.values()
.any(|existing_ref| existing_ref == execution_ref)
{
return Err(ProviderActionQueueError::ExistingCorrelation {
call_id: action_id.to_string(),
});
}
action_ids.insert(action_id.clone());
}
self.provider_tool_executions.extend(refs);
self.executor.update(ctx, |executor, ctx| {
executor.mark_restored_actions(conversation_id, &recovery_action_ids, ctx);
});
self.queue_actions(actions, conversation_id, ctx);
Ok(())
Ok(action_ids)
}
/// Queues the `actions` in the given iterator for the given conversation,
app/src/ai/blocklist/action_model_tests.rs
+20
View File
@@ -104,6 +104,26 @@ fn provider_action_correlations_require_the_exact_unresolved_batch_order() {
);
}
#[test]
fn exact_provider_action_correlation_is_idempotent() {
let conversation_id = AIConversationId::new();
let batch = pending_tool_batch(&["first"]);
let actions = vec![action("first")];
let execution_ref = ProviderToolExecutionRef::new(conversation_id, &batch.work_id, "first");
let existing = HashMap::from([((conversation_id, actions[0].id.clone()), execution_ref)]);
assert!(
BlocklistAIActionModel::provider_action_ids_to_enqueue_from(
&existing,
&actions,
conversation_id,
&batch,
)
.unwrap()
.is_empty()
);
}
#[test]
fn parallel_phase_only_admits_matching_autoexecutable_actions() {
let phase =
app/src/ai/blocklist/block.rs
+5 -1
View File
@@ -4390,7 +4390,9 @@ impl AIBlock {
// If auto-login is enabled, run the login command automatically
if auto_login_enabled {
ctx.emit(AIBlockEvent::RunAwsLoginCommand);
// Try the SDK chain first. This handles credentials that were refreshed by
// another AWS process without unnecessarily launching an interactive SSO flow.
ctx.emit(AIBlockEvent::RefreshAwsCredentials);
}
let model_name = model_name.clone();
@@ -6498,6 +6500,8 @@ pub enum AIBlockEvent {
OpenActiveAgentProfileEditor,
/// Run the configured AWS auth refresh command to fix expired Bedrock credentials
RunAwsLoginCommand,
/// Reload credentials and fall back to the configured auth command if needed.
RefreshAwsCredentials,
/// Emitted when a passive code diff has loaded its diffs and is ready to display.
/// This is used to trigger height recalculation since the diffs are loaded asynchronously
/// after the initial output completes.
app/src/ai/blocklist/controller.rs
+57 -6
View File
@@ -5895,6 +5895,7 @@ impl BlocklistAIController {
return;
}
};
coordinator.set_max_context_tokens(response_config.max_context_tokens);
if let Some(reason) = cancellation_reason {
if !coordinator.run().is_terminal() {
if let Err(error) = coordinator.run_mut().cancel(reason.to_string()) {
@@ -6368,6 +6369,7 @@ impl BlocklistAIController {
return;
}
};
coordinator.set_max_context_tokens(response_config.max_context_tokens);
if let Some(profile) = cli_monitor_profile {
if let Err(error) = coordinator.insert_profile(
CLI_MONITOR_PROVIDER_PROFILE,
@@ -6961,7 +6963,7 @@ impl BlocklistAIController {
ctx: &mut ModelContext<Self>,
) {
let usage_u32 = |tokens: u64| u32::try_from(tokens).unwrap_or(u32::MAX);
let cost_cents = crate::ai::bedrock::response_translator::estimate_cost_cents(
let cost_cents = crate::ai::provider::response_translator::estimate_cost_cents(
usage_u32(usage.input_tokens),
usage_u32(usage.output_tokens),
usage_u32(usage.cached_input_tokens),
@@ -6997,6 +6999,13 @@ impl BlocklistAIController {
return;
}
self.begin_active_provider_progressive_summary(conversation_id, ctx);
if matches!(
self.active_provider_progressive_summaries
.get(&conversation_id),
Some(ActiveProviderProgressiveSummaryState::InFlight { .. })
) {
return;
}
let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else {
return;
};
@@ -7452,6 +7461,24 @@ impl BlocklistAIController {
}
};
match block {
ProviderRunBlock::ContextWindowExceeded => {
slot.run = Some(run);
if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) {
self.fail_active_provider_run(
conversation_id,
format!("failed to persist context compaction boundary: {error}"),
ctx,
);
return;
}
if !self.begin_active_provider_progressive_summary(conversation_id, ctx) {
self.fail_active_provider_run(
conversation_id,
"provider context overflow could not start compaction".to_string(),
ctx,
);
}
}
ProviderRunBlock::ReadyToCallModel => {
slot.run = Some(run);
if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) {
@@ -7539,7 +7566,7 @@ impl BlocklistAIController {
else {
return;
};
let (converted_actions, invalid_results) =
let (mut converted_actions, invalid_results) =
convert_provider_tool_batch(&run.action_context, &batch);
for result in &invalid_results {
if let Err(error) = run
@@ -7586,6 +7613,30 @@ impl BlocklistAIController {
};
}
}
let action_ids_to_enqueue = match self.action_model.update(ctx, |action_model, _| {
action_model.provider_action_ids_to_enqueue(
&converted_actions
.iter()
.map(|(action, _)| action.clone())
.collect::<Vec<_>>(),
conversation_id,
&executable_batch,
)
}) {
Ok(action_ids) => action_ids,
Err(error) => {
self.fail_active_provider_run(conversation_id, error.to_string(), ctx);
return;
}
};
converted_actions.retain(|(action, _)| action_ids_to_enqueue.contains(&action.id));
let queued_call_ids = converted_actions
.iter()
.map(|(action, _)| action.id.to_string())
.collect::<HashSet<_>>();
executable_batch
.calls
.retain(|call| queued_call_ids.contains(&call.call.id));
let stream_id = self.active_provider_runs[&conversation_id]
.stream_id
.clone();
@@ -9522,7 +9573,7 @@ impl BlocklistAIController {
}
Some(warp_multi_agent_api::response_event::stream_finished::Reason::ContextWindowExceeded(_)) => {
let error_message = "Input exceeded context window limit.";
crate::ai::bedrock::crash_log::log_crash(
crate::ai::provider::crash_log::log_crash(
"ContextWindowExceeded",
error_message,
"unknown",
@@ -9615,7 +9666,7 @@ impl BlocklistAIController {
let error_message = format!(
"Response stream finished unexpectedly with internal error: {message}",
);
crate::ai::bedrock::crash_log::log_crash(
crate::ai::provider::crash_log::log_crash(
"InternalError",
&error_message,
"unknown",
@@ -9640,7 +9691,7 @@ impl BlocklistAIController {
}
Some(warp_multi_agent_api::response_event::stream_finished::Reason::MaxTokenLimit(_)) => {
let error_message = "Input exceeded context window limit.";
crate::ai::bedrock::crash_log::log_crash(
crate::ai::provider::crash_log::log_crash(
"MaxTokenLimit",
error_message,
"unknown",
@@ -9940,7 +9991,7 @@ impl BlocklistAIController {
u32::try_from(tokens).unwrap_or(u32::MAX)
};
let cost_cents =
crate::ai::bedrock::response_translator::estimate_cost_cents(
crate::ai::provider::response_translator::estimate_cost_cents(
usage_u32(usage.input_tokens),
usage_u32(usage.output_tokens),
usage_u32(usage.cached_input_tokens),
app/src/ai/blocklist/controller/response_stream.rs
+2 -10
View File
@@ -32,11 +32,11 @@ use crate::ai::agent::conversation::AIConversationId;
#[cfg(not(target_family = "wasm"))]
use crate::ai::agent::AIAgentInput;
use crate::ai::agent::{AIIdentifiers, CancellationReason};
use crate::ai::bedrock::client::BedrockClientConfig;
#[cfg(not(target_family = "wasm"))]
use crate::ai::blocklist::BlocklistAIPermissions;
use crate::ai::llms::{LLMId, LLMPreferences};
use crate::ai::openai::client::OpenAIClientConfig;
use crate::ai::provider::client::BedrockClientConfig;
use crate::ai::provider::ProviderConfig;
#[cfg(not(target_family = "wasm"))]
use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord};
@@ -181,7 +181,6 @@ impl ResponseStream {
reasoning_effort: client_config.reasoning_effort.clone(),
max_input_tokens: client_config.max_input_tokens,
max_output_tokens: client_config.max_output_tokens,
use_rig: client_config.use_rig,
supports_system_messages: client_config.supports_system_messages,
});
}
@@ -191,12 +190,6 @@ impl ResponseStream {
let auth_method = *settings.bedrock_auth_method.value();
let region = settings.bedrock_region.value().clone();
let cross_region_inference = *settings.bedrock_cross_region_inference.value();
let use_rig = crate::ai::bedrock::models::configured_model_uses_rig(
model_id,
settings.bedrock_models.value(),
&region,
cross_region_inference,
);
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
let mut config = BedrockClientConfig {
auth_method,
@@ -206,7 +199,6 @@ impl ResponseStream {
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
cross_region_inference,
use_rig,
};
if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } =
@@ -236,7 +228,7 @@ impl ResponseStream {
format!("bedrock:{:?}:region={region}", config.auth_method)
}
ProviderConfig::OpenAI(config) => {
format!("openai:{:?}:rig={}", config.kind, config.use_rig)
format!("openai:{:?}", config.kind)
}
ProviderConfig::None => "none".to_string(),
}
app/src/ai/blocklist/controller_tests.rs
+1
View File
@@ -267,6 +267,7 @@ fn provider_snapshot(conversation_id: AIConversationId) -> super::ActiveProvider
max_context_tokens: Some(128_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
},
action_context: crate::ai::runtime::ProviderActionContext::new_for_test(
task_id.to_string(),
app/src/ai/blocklist/usage/context_window_view.rs
+1 -1
View File
@@ -3,7 +3,7 @@ use galaxyui::elements::{
};
use galaxyui::{AppContext, Element, Entity, SingletonEntity, TypedActionView, View, ViewContext};
use crate::ai::bedrock::convert::{ContentPart, ConversationMessage, MessageContent, MessageRole};
use crate::ai::provider::convert::{ContentPart, ConversationMessage, MessageContent, MessageRole};
use crate::appearance::Appearance;
use crate::ui_components::blended_colors;
app/src/ai/crosscheck/reviewer.rs
+3 -5
View File
@@ -175,7 +175,6 @@ impl CrosscheckReviewer {
reasoning_effort: client_config.reasoning_effort.clone(),
max_input_tokens: client_config.max_input_tokens,
max_output_tokens: Some(REVIEWER_MAX_OUTPUT_TOKENS),
use_rig: client_config.use_rig,
supports_system_messages: client_config.supports_system_messages,
});
}
@@ -184,7 +183,7 @@ impl CrosscheckReviewer {
if *settings.bedrock_enabled.value() {
let auth_method = *settings.bedrock_auth_method.value();
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
let mut config = crate::ai::bedrock::client::BedrockClientConfig {
let mut config = crate::ai::provider::client::BedrockClientConfig {
auth_method,
profile: settings.bedrock_profile.value().clone(),
region: settings.bedrock_region.value().clone(),
@@ -192,7 +191,6 @@ impl CrosscheckReviewer {
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
use_rig: false,
};
if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } =
@@ -339,9 +337,9 @@ impl CrosscheckReviewer {
async fn invoke_via_bedrock(
agent_output: String,
model_id: String,
config: crate::ai::bedrock::client::BedrockClientConfig,
config: crate::ai::provider::client::BedrockClientConfig,
) -> Result<String, String> {
use crate::ai::bedrock::client::BedrockClient;
use crate::ai::provider::client::BedrockClient;
use crate::ai::provider::types::{ConversationMessage, MessageContent, MessageRole};
let cross_region_inference = config.cross_region_inference;
app/src/ai/llms.rs
+4 -18
View File
@@ -779,7 +779,7 @@ impl LLMPreferences {
if !*settings.bedrock_enabled.value() {
return;
}
let config = crate::ai::bedrock::client::BedrockClientConfig {
let config = crate::ai::provider::client::BedrockClientConfig {
auth_method: *settings.bedrock_auth_method.value(),
profile: settings.bedrock_profile.value().clone(),
region: settings.bedrock_region.value().clone(),
@@ -787,11 +787,10 @@ impl LLMPreferences {
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
use_rig: false,
};
let _ = ctx.spawn(
async move { crate::ai::bedrock::discovery::discover_available_models(config).await },
async move { crate::ai::provider::discovery::discover_available_models(config).await },
|me, result, ctx| match result {
Ok(models) => {
AISettings::handle(ctx).update(ctx, |settings, ctx| {
@@ -839,7 +838,7 @@ impl LLMPreferences {
let cross_region = *settings.bedrock_cross_region_inference.value();
// Check if user wants only 1-hour cache models
use crate::ai::bedrock::external_config::ExternalBedrockConfig;
use crate::ai::provider::external_config::ExternalBedrockConfig;
let external_config = ExternalBedrockConfig::load();
let require_1h_cache = external_config.enable_prompt_caching_1h;
@@ -885,7 +884,7 @@ impl LLMPreferences {
}
let model_id = if cross_region && !region.is_empty() {
super::bedrock::models::apply_cross_region_prefix(&model.model_id, &region)
super::provider::models::apply_cross_region_prefix(&model.model_id, &region)
} else {
model.model_id.clone()
};
@@ -1173,11 +1172,6 @@ impl LLMPreferences {
reasoning_effort: reasoning_effort.clone(),
max_input_tokens: Some(openai_model_context_size(model)),
max_output_tokens: model.max_output_tokens,
use_rig: model.use_rig
|| !matches!(
provider_kind,
OpenAIProviderKind::OpenAI | OpenAIProviderKind::LiteLLM
),
supports_system_messages: model.supports_system_messages(),
};
self.openai_provider_routing
@@ -1841,7 +1835,6 @@ impl LLMPreferences {
max_input_tokens: model.context_size,
max_output_tokens: None,
provider: None,
use_rig: true,
supports_system_messages: Some(true),
capability_overrides: std::collections::HashMap::new(),
reasoning_efforts: Vec::new(),
@@ -2658,15 +2651,12 @@ pub(crate) fn merge_discovered_provider_models(
{
discovered.display_name = existing.display_name.clone();
discovered.enabled = existing.enabled;
discovered.use_rig = existing.use_rig;
if existing.supports_system_messages.is_some() {
discovered.supports_system_messages = existing.supports_system_messages;
}
if discovered.provider.is_none() {
discovered.provider = existing.provider.clone();
}
} else {
discovered.use_rig = true;
}
if discovered.model_id.starts_with("codex-gpt-") {
discovered.supports_system_messages = Some(false);
@@ -2707,7 +2697,6 @@ pub(crate) fn merge_discovered_chatgpt_subscription_models(
.find(|model| model.model_id == discovered.model_id)
{
discovered.enabled = existing.enabled;
discovered.use_rig = existing.use_rig;
if existing.supports_system_messages.is_some() {
discovered.supports_system_messages = existing.supports_system_messages;
}
@@ -2934,7 +2923,6 @@ fn chatgpt_models_from_codex_response(body: &serde_json::Value) -> Vec<OpenAIMod
max_input_tokens,
max_output_tokens: None,
provider: Some("openai".to_string()),
use_rig: true,
supports_system_messages: Some(true),
capability_overrides: std::collections::HashMap::new(),
reasoning_efforts,
@@ -3056,7 +3044,6 @@ async fn fetch_from_litellm_model_info(
max_input_tokens,
max_output_tokens,
provider,
use_rig: false,
supports_system_messages: if model_name.starts_with("codex-gpt-") {
Some(false)
} else {
@@ -3186,7 +3173,6 @@ async fn fetch_from_openai_models(
max_input_tokens,
max_output_tokens,
provider,
use_rig: false,
supports_system_messages: if id.starts_with("codex-gpt-") {
Some(false)
} else {
app/src/ai/llms_tests.rs
+1 -10
View File
@@ -151,7 +151,6 @@ fn openai_model(model_id: &str) -> OpenAIModelConfig {
max_input_tokens: None,
max_output_tokens: None,
provider: None,
use_rig: false,
supports_system_messages: None,
capability_overrides: std::collections::HashMap::new(),
reasoning_efforts: Vec::new(),
@@ -164,7 +163,6 @@ fn bedrock_model(model_id: &str) -> BedrockModelConfig {
model_id: model_id.to_string(),
display_name: model_id.to_string(),
vision_supported: false,
use_rig: true,
}
}
@@ -308,7 +306,6 @@ fn provider_discovery_preserves_local_model_overrides() {
existing.display_name = "My Codex".to_string();
existing.context_size = 100_000;
existing.provider = Some("openai".to_string());
existing.use_rig = true;
// Even stale or incorrect endpoint metadata must not opt ChatGPT-backed
// Codex models back into the system role.
existing.supports_system_messages = Some(true);
@@ -325,7 +322,6 @@ fn provider_discovery_preserves_local_model_overrides() {
assert_eq!(merged[0].display_name, "My Codex");
assert_eq!(merged[0].context_size, 400_000);
assert_eq!(merged[0].max_output_tokens, Some(32_000));
assert!(merged[0].use_rig);
assert_eq!(merged[0].supports_system_messages, Some(false));
assert_eq!(merged[0].provider.as_deref(), Some("openai"));
}
@@ -725,7 +721,7 @@ fn disabled_providers_do_not_leave_models_in_the_runtime_inventory() {
}
#[test]
fn provider_discovery_enables_rig_for_new_models_and_keeps_manual_models() {
fn provider_discovery_keeps_new_and_manual_models() {
let manual = openai_model("manual-model");
let discovered = openai_model("codex-gpt-new");
@@ -733,7 +729,6 @@ fn provider_discovery_enables_rig_for_new_models_and_keeps_manual_models() {
assert_eq!(merged.len(), 2);
assert_eq!(merged[0].model_id, "codex-gpt-new");
assert!(merged[0].use_rig);
assert_eq!(merged[0].supports_system_messages, Some(false));
assert_eq!(merged[1].model_id, "manual-model");
}
@@ -798,7 +793,6 @@ fn chatgpt_codex_models_parse_visible_catalog_entries() {
assert_eq!(models[0].context_size, 828_400);
assert_eq!(models[0].max_input_tokens, Some(258_400));
assert_eq!(models[0].reasoning_efforts, ["low", "xhigh", "ultra"]);
assert!(models[0].use_rig);
assert_eq!(models[0].provider.as_deref(), Some("openai"));
assert_eq!(models[0].supports_system_messages, Some(true));
@@ -828,7 +822,6 @@ fn chatgpt_codex_models_parse_visible_catalog_entries() {
fn chatgpt_catalog_merge_drops_stale_models_but_preserves_overrides() {
let mut existing = openai_model("gpt-5.6-sol");
existing.enabled = false;
existing.use_rig = false;
existing.context_size = 272_000;
existing.max_input_tokens = Some(272_000);
existing.capability_overrides.insert(
@@ -840,7 +833,6 @@ fn chatgpt_catalog_merge_drops_stale_models_but_preserves_overrides() {
let mut discovered = openai_model("gpt-5.6-sol");
discovered.display_name = "GPT-5.6-Sol".to_string();
discovered.vision_supported = true;
discovered.use_rig = true;
discovered.context_size = 828_400;
discovered.max_input_tokens = Some(258_400);
@@ -851,7 +843,6 @@ fn chatgpt_catalog_merge_drops_stale_models_but_preserves_overrides() {
assert_eq!(merged[0].context_size, 828_400);
assert_eq!(merged[0].max_input_tokens, Some(258_400));
assert!(!merged[0].enabled);
assert!(!merged[0].use_rig);
assert_eq!(
merged[0].capability_override("vision"),
crate::settings::ModelCapabilityOverride::Unsupported
app/src/ai/mod.rs
-2
View File
@@ -19,8 +19,6 @@ pub mod auth_secret_types;
#[cfg(not(target_family = "wasm"))]
pub mod aws_credentials;
#[cfg(not(target_family = "wasm"))]
pub mod bedrock;
#[cfg(not(target_family = "wasm"))]
pub(crate) mod bedrock_credentials;
pub(crate) mod block_context;
pub(crate) mod blocklist;
app/src/ai/openai/client.rs
-1
View File
@@ -17,7 +17,6 @@ pub struct OpenAIClientConfig {
pub reasoning_effort: Option<String>,
pub max_input_tokens: Option<u32>,
pub max_output_tokens: Option<u32>,
pub use_rig: bool,
pub supports_system_messages: bool,
}
app/src/ai/openai/response_translator.rs
+2 -2
View File
@@ -10,7 +10,7 @@ use warp_multi_agent_api::response_event::stream_finished;
use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent};
use crate::ai::agent::api::LegacyEvent;
use crate::ai::bedrock::response_translator::{
use crate::ai::provider::response_translator::{
build_create_task, build_stream_init, context_window_for_model,
};
use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageContent, MessageRole};
@@ -531,7 +531,7 @@ fn build_tool_call_message(
tool_input_json: &str,
) -> ResponseEvent {
// Reuse the Bedrock tool call message builder since the proto output is identical
crate::ai::bedrock::response_translator::build_tool_call_message(
crate::ai::provider::response_translator::build_tool_call_message(
task_id,
tool_use_id,
tool_name,
app/src/ai/prompt_builder/mod.rs
+1 -1
View File
@@ -26,7 +26,7 @@ pub use prompts::provider::Provider;
#[allow(unused_imports)]
pub use tools::tools_for_mode;
use crate::ai::bedrock::convert::ToolDefinition;
use crate::ai::provider::convert::ToolDefinition;
/// The fully-resolved prompt configuration ready to send to a model.
#[derive(Debug, Clone)]
app/src/ai/prompt_builder/tests.rs
+1 -1
View File
@@ -103,7 +103,7 @@ mod prompt_builder_tests {
#[test]
fn test_mcp_tools_appended() {
use crate::ai::bedrock::convert::ToolDefinition;
use crate::ai::provider::convert::ToolDefinition;
let mcp_tool = ToolDefinition {
name: "mcp__github__create_pr".to_string(),
app/src/ai/prompt_builder/tools.rs
+1 -1
View File
@@ -3,7 +3,7 @@
//! Each mode has a different set of tools available. Code mode gets everything,
//! Plan mode gets read-only tools, Review mode gets read + search, etc.
use crate::ai::bedrock::convert::ToolDefinition;
use crate::ai::provider::convert::ToolDefinition;
use crate::ai::prompt_builder::mode::Mode;
/// Returns the tool definitions available for the given mode.
app/src/ai/bedrock/client.rs → app/src/ai/provider/client.rs
-1
View File
@@ -38,7 +38,6 @@ pub struct BedrockClientConfig {
pub secret_access_key: String,
pub session_token: Option<String>,
pub cross_region_inference: bool,
pub use_rig: bool,
}
impl BedrockClientConfig {
app/src/ai/bedrock/convert.rs → app/src/ai/provider/convert.rs
View File
app/src/ai/bedrock/convert_request.rs → app/src/ai/provider/convert_request.rs
View File
app/src/ai/bedrock/convert_tests.rs → app/src/ai/provider/convert_tests.rs
View File
app/src/ai/bedrock/crash_log.rs → app/src/ai/provider/crash_log.rs
View File
app/src/ai/bedrock/diagnostic.rs → app/src/ai/provider/diagnostic.rs
View File
app/src/ai/bedrock/discovery.rs → app/src/ai/provider/discovery.rs
-1
View File
@@ -88,7 +88,6 @@ pub async fn discover_available_models(
model_id: model_id.to_owned(),
display_name,
vision_supported,
use_rig: true,
})
}
});
app/src/ai/bedrock/e2e_tests.rs → app/src/ai/provider/e2e_tests.rs
+1 -2
View File
@@ -269,7 +269,6 @@ fn get_test_config() -> Option<BedrockClientConfig> {
secret_access_key: String::new(),
session_token: None,
cross_region_inference: false,
use_rig: false,
})
}
@@ -280,7 +279,7 @@ fn get_test_model() -> String {
}
fn sample_project_path() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("src/ai/bedrock/test_fixtures/sample_project")
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("src/ai/provider/test_fixtures/sample_project")
}
fn agent_tools() -> Vec<ToolDefinition> {
app/src/ai/bedrock/external_config.rs → app/src/ai/provider/external_config.rs
-1
View File
@@ -144,7 +144,6 @@ fn parse_claude_code_model_map(
model_id: arn,
display_name,
vision_supported: true,
use_rig: false,
}
})
.collect()
app/src/ai/bedrock/external_config_tests.rs → app/src/ai/provider/external_config_tests.rs
View File
app/src/ai/bedrock/integration_tests.rs → app/src/ai/provider/integration_tests.rs
-1
View File
@@ -24,7 +24,6 @@ fn get_test_config() -> Option<BedrockClientConfig> {
secret_access_key: String::new(),
session_token: None,
cross_region_inference: false,
use_rig: false,
})
}
app/src/ai/provider/mod.rs
+45 -1
View File
@@ -1,12 +1,56 @@
pub mod types;
use crate::ai::bedrock::client::BedrockClientConfig;
use crate::ai::openai::client::OpenAIClientConfig;
#[cfg(not(target_family = "wasm"))]
use crate::ai::provider::client::BedrockClientConfig;
#[allow(dead_code)]
#[derive(Clone)]
pub enum ProviderConfig {
#[cfg(not(target_family = "wasm"))]
Bedrock(BedrockClientConfig),
OpenAI(OpenAIClientConfig),
None,
}
// Provider-neutral module boundary; Bedrock-specific transport helpers live here without a
// provider-named directory so the provider layer can be reorganized independently of callers.
#[cfg(not(target_family = "wasm"))]
pub mod client;
#[cfg(not(target_family = "wasm"))]
pub mod convert;
#[cfg(not(target_family = "wasm"))]
pub mod crash_log;
#[cfg(not(target_family = "wasm"))]
pub mod diagnostic;
#[cfg(not(target_family = "wasm"))]
pub mod discovery;
#[cfg(not(target_family = "wasm"))]
pub mod external_config;
#[cfg(not(target_family = "wasm"))]
pub mod models;
#[cfg(not(target_family = "wasm"))]
pub mod request_translator;
#[cfg(not(target_family = "wasm"))]
pub mod response_translator;
#[cfg(not(target_family = "wasm"))]
pub mod settings_view;
#[cfg(not(target_family = "wasm"))]
#[cfg(test)]
mod convert_tests;
#[cfg(not(target_family = "wasm"))]
#[cfg(test)]
#[allow(dead_code)]
mod e2e_tests;
#[cfg(not(target_family = "wasm"))]
#[cfg(test)]
#[allow(dead_code)]
mod integration_tests;
#[cfg(not(target_family = "wasm"))]
#[cfg(test)]
mod models_tests;
#[cfg(not(target_family = "wasm"))]
#[cfg(test)]
mod response_translator_tests;
app/src/ai/provider/models.rs
+19
View File
@@ -0,0 +1,19 @@
pub fn apply_cross_region_prefix(model_id: &str, region: &str) -> String {
if model_id.starts_with("arn:") {
return model_id.to_string();
}
if model_id.contains('.') && model_id.split('.').next().unwrap_or("").len() <= 6 {
return model_id.to_string();
}
let prefix = match region {
r if r.starts_with("us-") || r.starts_with("ca-") => "us",
r if r.starts_with("eu-") || r == "il-central-1" => "eu",
r if r == "ap-northeast-1" || r == "ap-northeast-3" => "jp",
r if r == "ap-southeast-2" || r == "ap-southeast-4" || r == "ap-southeast-6" => "au",
r if r.starts_with("ap-") => "apac",
_ => return model_id.to_string(),
};
format!("{}.{}", prefix, model_id)
}
app/src/ai/bedrock/models_tests.rs → app/src/ai/provider/models_tests.rs
-31
View File
@@ -1,6 +1,4 @@
use super::models::*;
use crate::settings::ai::BedrockModelConfig;
#[test]
fn test_cross_region_prefix_us_east() {
assert_eq!(
@@ -83,32 +81,3 @@ fn test_cross_region_prefix_skips_arn() {
let arn = "arn:aws:bedrock:us-east-1:156729649053:application-inference-profile/fma5bsw4oxhy";
assert_eq!(apply_cross_region_prefix(arn, "us-east-1"), arn);
}
#[test]
fn rig_opt_in_matches_context_markers_and_resolved_inference_profiles() {
let configured = vec![BedrockModelConfig {
model_id: "anthropic.claude-test[1m]".to_string(),
display_name: "Claude Test".to_string(),
vision_supported: false,
use_rig: true,
}];
assert!(configured_model_uses_rig(
"us.anthropic.claude-test",
&configured,
"us-east-1",
true,
));
assert!(configured_model_uses_rig(
"anthropic.claude-test[1M]",
&configured,
"us-east-1",
false,
));
assert!(!configured_model_uses_rig(
"anthropic.other-model",
&configured,
"us-east-1",
false,
));
}
app/src/ai/bedrock/request_translator.rs → app/src/ai/provider/request_translator.rs
View File
app/src/ai/bedrock/request_translator_tests.rs → app/src/ai/provider/request_translator_tests.rs
+1 -1
View File
@@ -6,7 +6,7 @@ use super::{
extract_tools, inject_input_messages_into_task, sanitize_messages_for_bedrock,
};
use crate::ai::agent::api::is_internal_command_completion_assessment;
use crate::ai::bedrock::convert::{ContentPart, ConversationMessage, MessageContent, MessageRole};
use crate::ai::provider::convert::{ContentPart, ConversationMessage, MessageContent, MessageRole};
#[test]
fn test_sanitize_messages_prepends_synthetic_tool_result_before_existing_user_text() {
app/src/ai/bedrock/response_translator.rs → app/src/ai/provider/response_translator.rs
View File
app/src/ai/bedrock/response_translator_tests.rs → app/src/ai/provider/response_translator_tests.rs
View File
app/src/ai/bedrock/settings_view.rs → app/src/ai/provider/settings_view.rs
+2 -2
View File
@@ -3,8 +3,8 @@ use galaxyui::elements::{
};
use galaxyui::{AppContext, Element, Entity, SingletonEntity, TypedActionView, View, ViewContext};
use crate::ai::bedrock::convert::CachingConfig;
use crate::ai::bedrock::external_config::ExternalBedrockConfig;
use crate::ai::provider::convert::CachingConfig;
use crate::ai::provider::external_config::ExternalBedrockConfig;
use crate::appearance::Appearance;
use crate::ui_components::blended_colors;
app/src/ai/bedrock/test_fixtures/sample_project/Cargo.toml → app/src/ai/provider/test_fixtures/sample_project/Cargo.toml
View File
app/src/ai/bedrock/test_fixtures/sample_project/README.md → app/src/ai/provider/test_fixtures/sample_project/README.md
View File
app/src/ai/bedrock/test_fixtures/sample_project/src/lib.rs → app/src/ai/provider/test_fixtures/sample_project/src/lib.rs
View File
app/src/ai/bedrock/test_fixtures/sample_project/src/lib_tests.rs → app/src/ai/provider/test_fixtures/sample_project/src/lib_tests.rs
View File
app/src/ai/bedrock/test_fixtures/sample_project/src/main.rs → app/src/ai/provider/test_fixtures/sample_project/src/main.rs
View File
app/src/ai/bedrock/test_fixtures/sample_project/src/utils.rs → app/src/ai/provider/test_fixtures/sample_project/src/utils.rs
View File
app/src/ai/runtime/event_translator.rs
+202 -17
View File
@@ -10,11 +10,11 @@ use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent};
use super::provider_run_coordinator::ProviderRunProjection;
use crate::ai::agent::runtime_activity;
use crate::ai::bedrock::response_translator::{
use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage};
use crate::ai::provider::response_translator::{
build_add_agent_output_message, build_append_text, build_create_task, build_stream_init,
build_user_query_message,
};
use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage};
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub(crate) struct RuntimeResponseConfig {
@@ -26,6 +26,9 @@ pub(crate) struct RuntimeResponseConfig {
pub(crate) max_context_tokens: Option<u32>,
pub(crate) capabilities: RuntimeCapabilities,
pub(crate) empty_output_message: Option<String>,
/// Todo items supplied by the existing task transcript, when one is available.
#[serde(skip)]
pub(crate) todo_items: Option<Vec<api::TodoItem>>,
}
/// Converts the provider-neutral runtime lifecycle into Galaxy's existing
@@ -51,6 +54,7 @@ pub(crate) struct RuntimeResponseTranslator {
pub(crate) struct ProviderRunResponseProjector {
translator: RuntimeResponseTranslator,
has_started_model_turn: bool,
todo_phase: usize,
finished: bool,
}
@@ -59,6 +63,7 @@ impl ProviderRunResponseProjector {
Self {
translator: RuntimeResponseTranslator::new(config),
has_started_model_turn: false,
todo_phase: 0,
finished: false,
}
}
@@ -70,6 +75,8 @@ impl ProviderRunResponseProjector {
Self {
translator: RuntimeResponseTranslator::restored(config, projection_was_initialized),
has_started_model_turn: false,
// Task-list events are part of the already persisted projection.
todo_phase: usize::MAX,
finished: false,
}
}
@@ -87,17 +94,46 @@ impl ProviderRunResponseProjector {
self.translator.begin_followup_turn();
}
self.has_started_model_turn = true;
self.translator.translate(AgentEvent::TurnStarted {
let mut events = self.translator.translate(AgentEvent::TurnStarted {
runtime_request_id: String::new(),
})
})?;
events.extend(self.todo_phase_events());
Ok(events)
}
ProviderRunProjection::ModelEvent { event, .. } => self.translator.translate(event),
ProviderRunProjection::ModelRetry { .. } => {
Ok(self.translator.discard_failed_turn_output())
}
ProviderRunProjection::ModelTurnRequested { .. }
| ProviderRunProjection::ModelTurnFinished { .. }
| ProviderRunProjection::ToolBatchReady { .. } => Ok(Vec::new()),
| ProviderRunProjection::ModelTurnFinished { .. } => Ok(Vec::new()),
ProviderRunProjection::ToolBatchReady { .. } => {
let todo_index = self.todo_phase.saturating_sub(1);
let Some(todo) = self.todo_items().get(todo_index).cloned() else {
return Ok(Vec::new());
};
self.todo_phase += 1;
let mut events = vec![build_todo_update(
&self.translator.config.task_id,
api::message::update_todos::Operation::MarkTodosCompleted(
api::MarkTodosCompleted {
todo_ids: vec![todo.id],
},
),
)];
events.push(build_todo_update(
&self.translator.config.task_id,
api::message::update_todos::Operation::UpdatePendingTodos(
api::UpdatePendingTodos {
updated_pending_todos: self
.todo_items()
.into_iter()
.skip(self.todo_phase)
.collect(),
},
),
));
Ok(events)
}
}
}
@@ -115,17 +151,154 @@ impl ProviderRunResponseProjector {
}
self.finished = true;
match outcome {
ProviderRunOutcome::Completed(completion) => Ok(self
.translator
.finish_provider_run(completion.stop_reason.clone(), aggregate_usage)),
ProviderRunOutcome::Failed(failure) => Ok(self
.translator
.provider_failure(&failure.message, aggregate_usage)),
ProviderRunOutcome::Completed(completion) => {
let mut events = self.todo_completion_events();
events.extend(
self.translator
.finish_provider_run(completion.stop_reason.clone(), aggregate_usage),
);
Ok(events)
}
ProviderRunOutcome::Failed(failure) => Ok(self.translator.provider_failure(
&failure.message,
failure.source.as_ref(),
aggregate_usage,
)),
ProviderRunOutcome::Cancelled { .. } => Ok(self
.translator
.finish_provider_run(StopReason::Cancelled, aggregate_usage)),
}
}
// Keep the direct-provider workflow visible in the existing task list protocol. These are
// response events, so the normal history model remains the sole owner of task-list state.
fn todo_phase_events(&mut self) -> Vec<ResponseEvent> {
let events = match self.todo_phase {
0 => vec![build_todo_update(
&self.translator.config.task_id,
api::message::update_todos::Operation::CreateTodoList(api::CreateTodoList {
initial_todos: self.todo_items(),
}),
)],
_ => Vec::new(),
};
// The first model turn owns the first phase. Tool batches advance it; this keeps
// arbitrary plan lengths aligned with the UpdateTodos protocol.
if self.todo_phase == 0 {
self.todo_phase = 1;
}
events
}
fn todo_completion_events(&self) -> Vec<ResponseEvent> {
if self.todo_phase == 0 || self.todo_phase == usize::MAX {
return Vec::new();
}
let todos = self.todo_items();
if todos.is_empty() {
return Vec::new();
}
vec![
build_todo_update(
&self.translator.config.task_id,
api::message::update_todos::Operation::MarkTodosCompleted(
api::MarkTodosCompleted {
todo_ids: todos.iter().map(|todo| todo.id.clone()).collect(),
},
),
),
build_todo_update(
&self.translator.config.task_id,
api::message::update_todos::Operation::UpdatePendingTodos(
api::UpdatePendingTodos {
updated_pending_todos: Vec::new(),
},
),
),
]
}
fn todo_items(&self) -> Vec<api::TodoItem> {
self.translator
.config
.todo_items
.clone()
.unwrap_or_else(default_workflow_todos)
}
}
fn default_workflow_todos() -> Vec<api::TodoItem> {
[
api::TodoItem {
id: "direct-provider-research".to_owned(),
title: "Research the request".to_owned(),
description: "Inspect the repository and gather relevant evidence".to_owned(),
},
api::TodoItem {
id: "direct-provider-plan".to_owned(),
title: "Create an implementation plan".to_owned(),
description: "Choose an approach grounded in the repository".to_owned(),
},
api::TodoItem {
id: "direct-provider-critique".to_owned(),
title: "Critique the approach".to_owned(),
description: "Check assumptions, risks, and missing cases".to_owned(),
},
api::TodoItem {
id: "direct-provider-revise".to_owned(),
title: "Revise the plan".to_owned(),
description: "Incorporate findings before editing".to_owned(),
},
api::TodoItem {
id: "direct-provider-implement".to_owned(),
title: "Implement the change".to_owned(),
description: "Make the requested edits".to_owned(),
},
api::TodoItem {
id: "direct-provider-verify".to_owned(),
title: "Verify the result".to_owned(),
description: "Run proportionate checks".to_owned(),
},
api::TodoItem {
id: "direct-provider-repair".to_owned(),
title: "Repair validation issues".to_owned(),
description: "Fix failures found during verification".to_owned(),
},
]
.to_vec()
}
fn build_todo_update(
task_id: &str,
operation: api::message::update_todos::Operation,
) -> ResponseEvent {
let message = api::Message {
id: Uuid::new_v4().to_string(),
task_id: task_id.to_owned(),
request_id: String::new(),
timestamp: None,
server_message_data: String::new(),
citations: Vec::new(),
fetched_memories: Vec::new(),
message: Some(api::message::Message::UpdateTodos(
api::message::UpdateTodos {
operation: Some(operation),
},
)),
};
ResponseEvent {
r#type: Some(api::response_event::Type::ClientActions(
api::response_event::ClientActions {
actions: vec![api::ClientAction {
action: Some(api::client_action::Action::AddMessagesToTask(
api::client_action::AddMessagesToTask {
task_id: task_id.to_owned(),
messages: vec![message],
},
)),
}],
},
)),
}
}
impl RuntimeResponseTranslator {
@@ -421,15 +594,27 @@ impl RuntimeResponseTranslator {
self.finished_with_reason(map_stop_reason(reason))
}
fn provider_failure(&mut self, message: &str, aggregate_usage: &Usage) -> Vec<ResponseEvent> {
fn provider_failure(
&mut self,
message: &str,
source: Option<&galaxy_agent_core::AgentError>,
aggregate_usage: &Usage,
) -> Vec<ResponseEvent> {
let mut events = Vec::new();
self.initialize(&mut events);
events.push(self.finished_with_usage(
let reason = if source
.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(),
})
} else {
stream_finished::Reason::InternalError(stream_finished::InternalError {
message: message.to_owned(),
}),
aggregate_usage,
));
})
};
events.push(self.finished_with_usage(reason, aggregate_usage));
events
}
app/src/ai/runtime/event_translator_tests.rs
+117 -1
View File
@@ -17,6 +17,7 @@ fn provider_translator() -> RuntimeResponseTranslator {
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
})
}
@@ -30,6 +31,7 @@ fn session_translator() -> RuntimeResponseTranslator {
max_context_tokens: None,
capabilities: RuntimeCapabilities::session_runtime(),
empty_output_message: Some("> runtime completed without text".to_owned()),
todo_items: None,
})
}
@@ -44,6 +46,7 @@ fn restored_provider_projection_skips_stream_initialization() {
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
};
let mut projector = ProviderRunResponseProjector::restored(config, true);
let work_id = galaxy_agent_core::ExternalWorkId {
@@ -82,6 +85,109 @@ fn restored_provider_projection_skips_stream_initialization() {
));
}
#[test]
fn provider_projection_uses_supplied_todos_and_advances_each_id() {
let todos = vec![
warp_multi_agent_api::TodoItem {
id: "research".to_owned(),
title: "Research".to_owned(),
description: "Inspect".to_owned(),
},
warp_multi_agent_api::TodoItem {
id: "implement".to_owned(),
title: "Implement".to_owned(),
description: "Edit".to_owned(),
},
warp_multi_agent_api::TodoItem {
id: "verify".to_owned(),
title: "Verify".to_owned(),
description: "Check".to_owned(),
},
];
let mut projector = ProviderRunResponseProjector::new(RuntimeResponseConfig {
task_id: "task".to_owned(),
conversation_id: "conversation".to_owned(),
needs_create_task: false,
user_query: None,
model_id: "model".to_owned(),
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: Some(todos),
});
let work_id = galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(1),
};
let initial = projector
.project(ProviderRunProjection::ModelTurnStarted {
work_id: work_id.clone(),
profile: galaxy_agent_core::ProviderRequestProfile::new("base"),
runtime_id: "runtime".to_owned(),
model_id: "model".to_owned(),
runtime_request_id: "request".to_owned(),
retry_attempt: 0,
elapsed_ms: 1,
})
.unwrap();
assert_eq!(initial.len(), 2);
let first = projector
.project(ProviderRunProjection::ToolBatchReady {
batch: galaxy_agent_core::PendingToolBatch {
work_id: galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(1),
},
calls: Vec::new(),
},
})
.unwrap();
let second = projector
.project(ProviderRunProjection::ToolBatchReady {
batch: galaxy_agent_core::PendingToolBatch {
work_id: galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(1),
},
calls: Vec::new(),
},
})
.unwrap();
let ids = |events: &[warp_multi_agent_api::ResponseEvent]| {
events
.iter()
.flat_map(|event| match &event.r#type {
Some(response_event::Type::ClientActions(actions)) => actions
.actions
.iter()
.filter_map(|action| match &action.action {
Some(client_action::Action::AddMessagesToTask(add)) => add
.messages
.iter()
.filter_map(|message| match &message.message {
Some(message::Message::UpdateTodos(update)) => match update
.operation
.as_ref()
{
Some(message::update_todos::Operation::MarkTodosCompleted(
mark,
)) => Some(mark.todo_ids[0].clone()),
_ => None,
},
_ => None,
})
.next(),
_ => None,
})
.collect::<Vec<_>>(),
_ => Vec::new(),
})
.collect::<Vec<_>>()
};
assert_eq!(ids(&first), vec!["research"]);
assert_eq!(ids(&second), vec!["implement"]);
}
#[test]
fn restored_uninitialized_projection_replays_init_before_live_delta() {
let config = RuntimeResponseConfig {
@@ -93,6 +199,7 @@ fn restored_uninitialized_projection_replays_init_before_live_delta() {
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
};
let mut projector = ProviderRunResponseProjector::restored(config, false);
let work_id = galaxy_agent_core::ExternalWorkId {
@@ -120,11 +227,18 @@ fn restored_uninitialized_projection_replays_init_before_live_delta() {
})
.unwrap();
assert_eq!(started.len(), 1);
assert_eq!(started.len(), 2);
assert!(matches!(
started[0].r#type,
Some(response_event::Type::Init(_))
));
let Some(response_event::Type::ClientActions(actions)) = &started[1].r#type else {
panic!("initial provider turn should publish its task list");
};
assert!(matches!(
actions.actions[0].action,
Some(client_action::Action::AddMessagesToTask(_))
));
assert_eq!(delta.len(), 1);
assert!(matches!(
delta[0].r#type,
@@ -143,6 +257,7 @@ fn provider_followup_turn_starts_a_distinct_text_message() {
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
});
let first_work_id = galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
@@ -366,6 +481,7 @@ fn provider_retry_clears_failed_attempt_output_before_new_messages() {
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
});
let work_id = galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
app/src/ai/runtime/provider_run_coordinator.rs
+76
View File
@@ -82,6 +82,7 @@ pub(crate) enum ProviderRunProjection {
pub(crate) enum ProviderRunBlock {
Tools(PendingToolBatch),
ReadyToCallModel,
ContextWindowExceeded,
AwaitingDriver {
work_id: ExternalWorkId,
stop_reason: StopReason,
@@ -175,6 +176,8 @@ pub(crate) struct ProviderRunCoordinator {
profiles: BTreeMap<String, ProviderRunProfile>,
model_start_timeout: Duration,
model_event_idle_timeout: Duration,
context_compaction_requested: bool,
max_context_tokens: Option<u32>,
}
impl ProviderRunCoordinator {
@@ -218,6 +221,8 @@ impl ProviderRunCoordinator {
profiles,
model_start_timeout: PROVIDER_MODEL_START_TIMEOUT,
model_event_idle_timeout: PROVIDER_MODEL_EVENT_IDLE_TIMEOUT,
context_compaction_requested: false,
max_context_tokens: None,
})
}
@@ -239,6 +244,10 @@ impl ProviderRunCoordinator {
self.profiles.get(profile).map(|profile| &profile.request)
}
pub(crate) fn set_max_context_tokens(&mut self, max_context_tokens: Option<u32>) {
self.max_context_tokens = max_context_tokens;
}
pub(crate) fn insert_profile(
&mut self,
profile: impl Into<String>,
@@ -364,8 +373,16 @@ impl ProviderRunCoordinator {
if !self.checkpoint_or_fail(&mut checkpoint).await? {
return self.terminal_block();
}
if self.request_needs_context_compaction(&call) {
self.run.prepare_context_compaction(&call.work_id)?;
return Ok(ProviderRunBlock::ContextWindowExceeded);
}
self.drive_model_call_acknowledged(call, control.clone(), &mut project)
.await?;
if self.context_compaction_requested {
self.context_compaction_requested = false;
return Ok(ProviderRunBlock::ContextWindowExceeded);
}
}
Some(ProviderRunStep::DispatchTools(batch)) => {
if !self.checkpoint_or_fail(&mut checkpoint).await? {
@@ -428,6 +445,17 @@ impl ProviderRunCoordinator {
}
}
fn request_needs_context_compaction(&self, call: &ProviderModelCall) -> bool {
let Some(max_context_tokens) = self.max_context_tokens else {
return false;
};
let Some(profile) = self.profiles.get(call.profile.as_str()) else {
return false;
};
let request = request_for_model_call(profile.request.clone(), call);
estimate_turn_request_tokens(&request) >= u64::from(max_context_tokens)
}
async fn checkpoint_or_fail<C>(
&mut self,
checkpoint: &mut C,
@@ -830,6 +858,11 @@ impl ProviderRunCoordinator {
self.run.cancel("provider model call was cancelled")?;
return Ok(());
}
if reason == StopReason::ContextWindowExceeded {
self.run.prepare_context_compaction(&call.work_id)?;
self.context_compaction_requested = true;
return Ok(());
}
let turn = buffer.complete(reason, advertised_tools);
if let Err(error) = self.run.accept_model_turn(&call.work_id, turn) {
self.run.fail(
@@ -925,6 +958,11 @@ impl ProviderRunCoordinator {
"payload": &error,
}));
let error_message = error.message.clone();
if error.kind == AgentErrorKind::ContextWindowExceeded {
self.run.prepare_context_compaction(&call.work_id)?;
self.context_compaction_requested = true;
return Ok(());
}
let disposition = self
.run
.register_model_failure(&call.work_id, error.clone())?;
@@ -1082,6 +1120,44 @@ fn request_for_model_call(mut template: TurnRequest, call: &ProviderModelCall) -
template
}
const ESTIMATED_CHARS_PER_TOKEN: u64 = 4;
/// Deliberately overestimates request size without requiring provider-specific tokenizers.
fn estimate_turn_request_tokens(request: &TurnRequest) -> u64 {
let mut chars = request
.system_prompt
.as_deref()
.map_or(0, |text| text.chars().count());
chars += request.prompt.as_ref().map_or(0, |prompt| {
serde_json::to_string(prompt)
.unwrap_or_default()
.chars()
.count()
});
chars += request
.messages
.iter()
.map(|message| {
serde_json::to_string(message)
.unwrap_or_default()
.chars()
.count()
})
.sum::<usize>();
chars += request
.tools
.iter()
.map(|tool| {
serde_json::to_string(tool)
.unwrap_or_default()
.chars()
.count()
})
.sum::<usize>();
let input_tokens = (chars as u64).div_ceil(ESTIMATED_CHARS_PER_TOKEN);
input_tokens.saturating_add(request.max_output_tokens.unwrap_or_default())
}
fn tool_event_call_id(event: &ToolEvent) -> Result<&str, ProviderRunCoordinatorError> {
match event {
ToolEvent::Proposed { call } => Ok(&call.id),
app/src/ai/runtime/provider_run_coordinator_tests.rs
+44
View File
@@ -164,6 +164,48 @@ fn request() -> TurnRequest {
request
}
#[test]
fn estimates_turn_request_with_system_tools_messages_and_output() {
let mut request = TurnRequest::new(
"test-model",
vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("12345678".to_string()),
}],
);
request.system_prompt = Some("1234".to_string());
request.tools = vec![ToolDefinition {
name: "tool".to_string(),
description: "description".to_string(),
input_schema: serde_json::json!({"type": "object"}),
}];
request.max_output_tokens = Some(10);
let expected_input = (request.system_prompt.as_deref().unwrap().len()
+ serde_json::to_string(&request.messages[0]).unwrap().len()
+ serde_json::to_string(&request.tools[0]).unwrap().len()) as u64;
assert_eq!(
estimate_turn_request_tokens(&request),
expected_input.div_ceil(ESTIMATED_CHARS_PER_TOKEN) + 10
);
}
#[tokio::test]
async fn proactively_compacts_before_starting_an_oversized_request() {
let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()]));
let mut coordinator = coordinator(runtime.clone());
coordinator.set_max_context_tokens(Some(1));
let (_sender, control) = turn_control();
let block = coordinator
.drive_until_blocked(control, |_| Ok(()))
.await
.unwrap();
assert_eq!(block, ProviderRunBlock::ContextWindowExceeded);
assert!(runtime.requests().is_empty());
}
fn started(id: &str) -> ScriptedEvent {
Ok(AgentEvent::TurnStarted {
runtime_request_id: id.to_string(),
@@ -1166,6 +1208,7 @@ async fn transcript_projector_emits_one_ui_stream_for_the_whole_run() {
max_context_tokens: Some(100_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
});
let mut ui_events = Vec::new();
let (_sender, control) = turn_control();
@@ -1224,6 +1267,7 @@ fn transcript_projector_preserves_provider_failure_message() {
max_context_tokens: Some(100_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
});
let events = projector
.finish(
app/src/ai/runtime/rig.rs
+5 -3
View File
@@ -21,9 +21,9 @@ use super::rig_tool::action_from_tool_call;
use super::ProviderRunProfile;
use crate::ai::agent::api::RequestParams;
use crate::ai::agent::AIAgentAction;
use crate::ai::bedrock::client::BedrockClient;
use crate::ai::bedrock::convert::CachingConfig;
use crate::ai::bedrock::external_config::ExternalBedrockConfig;
use crate::ai::provider::client::BedrockClient;
use crate::ai::provider::convert::CachingConfig;
use crate::ai::provider::external_config::ExternalBedrockConfig;
use crate::ai::provider::types::ConversationMessage;
use crate::ai::runtime::RuntimeResponseConfig;
use crate::settings::OpenAIProviderKind;
@@ -137,6 +137,7 @@ pub(crate) async fn prepare_provider_run(
task_id,
needs_create_task,
user_query,
todo_items,
request,
persistent_messages,
tool_result_archive,
@@ -159,6 +160,7 @@ pub(crate) async fn prepare_provider_run(
max_context_tokens,
capabilities: base_runtime.descriptor().capabilities.clone(),
empty_output_message: None,
todo_items,
};
Ok(PreparedProviderRun {
base_profile: ProviderRunProfile::new(base_runtime, request),
app/src/ai/runtime/rig_request.rs
+31 -3
View File
@@ -11,21 +11,23 @@ use galaxy_agent_core::{
};
use sha2::{Digest as _, Sha256};
use uuid::Uuid;
use warp_multi_agent_api as api;
use warp_multi_agent_api::ToolType;
use crate::ai::agent::api::RequestParams;
use crate::ai::agent::{AIAgentContext, AIAgentInput, MCPContext, UserQueryMode};
use crate::ai::bedrock::request_translator::{
default_tool_definitions, sanitize_messages_for_bedrock, tool_name_is_supported,
};
use crate::ai::openai::client::OpenAIClientConfig;
use crate::ai::openai::request_translator::sanitize_messages_for_openai;
use crate::ai::provider::request_translator::{
default_tool_definitions, sanitize_messages_for_bedrock, tool_name_is_supported,
};
use crate::ai::provider::types::flatten_tool_history_for_no_tools_turn;
pub(crate) struct PreparedRigTurn {
pub task_id: String,
pub needs_create_task: bool,
pub user_query: Option<String>,
pub todo_items: Option<Vec<api::TodoItem>>,
pub request: TurnRequest,
pub persistent_messages: Vec<ConversationMessage>,
pub tool_result_archive: Vec<ConversationMessage>,
@@ -226,6 +228,8 @@ fn prepare_rig_turn_for_provider(
..
} = params;
let todo_items = todo_items_from_tasks(&tasks);
let task_id = root_task_id
.or_else(|| tasks.first().map(|task| task.id.clone()))
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
@@ -299,6 +303,7 @@ fn prepare_rig_turn_for_provider(
task_id,
needs_create_task,
user_query,
todo_items,
request,
persistent_messages,
tool_result_archive,
@@ -307,6 +312,29 @@ fn prepare_rig_turn_for_provider(
}
}
// Reuse a plan emitted by the model when the existing transcript contains one. This keeps
// direct-provider projection aligned with UpdateTodos instead of inventing a second plan.
fn todo_items_from_tasks(tasks: &[api::Task]) -> Option<Vec<api::TodoItem>> {
let mut items = None;
for task in tasks {
for message in &task.messages {
let Some(api::message::Message::UpdateTodos(update)) = &message.message else {
continue;
};
match update.operation.as_ref()? {
api::message::update_todos::Operation::CreateTodoList(create) => {
items = Some(create.initial_todos.clone());
}
api::message::update_todos::Operation::UpdatePendingTodos(update) => {
items = Some(update.updated_pending_todos.clone());
}
api::message::update_todos::Operation::MarkTodosCompleted(_) => {}
}
}
}
items.filter(|items| !items.is_empty())
}
fn input_messages(
inputs: Vec<AIAgentInput>,
tool_results: Vec<ToolResult>,
app/src/ai/runtime/rig_request_tests.rs
-1
View File
@@ -32,7 +32,6 @@ fn config() -> OpenAIClientConfig {
reasoning_effort: None,
max_input_tokens: Some(128_000),
max_output_tokens: Some(8_192),
use_rig: true,
supports_system_messages: true,
}
}
app/src/ai/runtime/rig_tests.rs
-1
View File
@@ -16,7 +16,6 @@ fn openai_config(model: &str) -> OpenAIClientConfig {
reasoning_effort: None,
max_input_tokens: Some(128_000),
max_output_tokens: Some(8_192),
use_rig: true,
supports_system_messages: true,
}
}
app/src/ai/runtime/rig_tool.rs
+1 -1
View File
@@ -102,7 +102,7 @@ pub(crate) fn action_from_tool_call(
optional_bounded_u64(
input,
"wait_seconds",
crate::ai::bedrock::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS,
crate::ai::provider::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS,
)?
.unwrap_or(2),
))),
app/src/code_review/code_review_header/mod.rs
+2 -1
View File
@@ -297,7 +297,8 @@ impl CodeReviewHeader {
.finish(),
);
}
right_section_compact.add_child(ChildView::new(&code_review_header_fields.diff_selector).finish());
right_section_compact
.add_child(ChildView::new(&code_review_header_fields.diff_selector).finish());
right_section_compact.add_child(Container::new(right_subsection_compact.finish()).finish());
Clipped::new(
app/src/settings/ai.rs
-10
View File
@@ -828,11 +828,6 @@ pub struct BedrockModelConfig {
#[serde(default)]
#[schemars(description = "Whether the model supports image/vision input.")]
pub vision_supported: bool,
#[serde(default)]
#[schemars(
description = "Route this model through Galaxy's Rig Bedrock runtime. Disabled by default while compatibility validation is in progress."
)]
pub use_rig: bool,
}
impl settings_value::SettingsValue for BedrockModelConfig {}
@@ -883,11 +878,6 @@ pub struct OpenAIModelConfig {
description = "Optional provider hint (e.g. anthropic, openai, google) for icon display."
)]
pub provider: Option<String>,
#[serde(default)]
#[schemars(
description = "Route this model through Galaxy's Rig runtime. This is an opt-in migration path."
)]
pub use_rig: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schemars(
description = "Whether this endpoint accepts system-role messages. Set false for ChatGPT-backed LiteLLM models that reject them."
app/src/settings/ai_tests.rs
-1
View File
@@ -444,7 +444,6 @@ fn orchestration_is_enabled_when_ai_is_enabled() {
model_id: "test-model".to_string(),
display_name: "Test model".to_string(),
vision_supported: false,
use_rig: false,
}],
ctx,
)
app/src/settings_view/ai_page.rs
-6
View File
@@ -7251,9 +7251,6 @@ impl ProviderSettingsWidget {
if model.effective_vision_supported() {
details.push("Images".to_string());
}
if model.use_rig {
details.push("Rig".to_string());
}
if !model.reasoning_efforts.is_empty() {
details.push(format!("Reasoning: {}", model.reasoning_efforts.join(", ")));
}
@@ -7361,9 +7358,6 @@ impl ProviderSettingsWidget {
if model.vision_supported {
details.push("Images".to_string());
}
if model.use_rig {
details.push("Rig".to_string());
}
Self::render_model_row(
model.display_name.clone(),
model.model_id.clone(),
app/src/settings_view/provider_setup_view.rs
+2 -3
View File
@@ -949,7 +949,7 @@ impl ProviderSetupView {
return;
}
ProviderSetupProviderType::Bedrock => {
let config = crate::ai::bedrock::client::BedrockClientConfig {
let config = crate::ai::provider::client::BedrockClientConfig {
auth_method: self.draft_bedrock.auth_method,
profile: self.draft_bedrock.profile.clone(),
region: self.draft_bedrock.region.clone(),
@@ -957,11 +957,10 @@ impl ProviderSetupView {
secret_access_key: self.draft_bedrock.secret_access_key.clone(),
session_token: None,
cross_region_inference: self.draft_bedrock.cross_region_inference,
use_rig: false,
};
ctx.spawn(
async move {
crate::ai::bedrock::discovery::discover_available_models(config).await
crate::ai::provider::discovery::discover_available_models(config).await
},
move |me, result, ctx| match result {
Ok(models) => {
app/src/terminal/input/slash_commands/mod.rs
+4 -4
View File
@@ -1189,21 +1189,21 @@ impl Input {
.enumerate()
.map(|(i, msg)| {
let content_str = match &msg.content {
crate::ai::bedrock::convert::MessageContent::Text(t) => {
crate::ai::provider::convert::MessageContent::Text(t) => {
if t.len() > 500 {
format!("{}... ({} chars total)", &t[..500], t.len())
} else {
t.clone()
}
}
crate::ai::bedrock::convert::MessageContent::ToolUse {
crate::ai::provider::convert::MessageContent::ToolUse {
name,
tool_use_id,
..
} => {
format!("ToolUse(name={name}, id={tool_use_id})")
}
crate::ai::bedrock::convert::MessageContent::ToolResult {
crate::ai::provider::convert::MessageContent::ToolResult {
tool_use_id,
content,
is_error,
@@ -1217,7 +1217,7 @@ impl Input {
"ToolResult(id={tool_use_id}, err={is_error}): {truncated}"
)
}
crate::ai::bedrock::convert::MessageContent::MultiPart(parts) => {
crate::ai::provider::convert::MessageContent::MultiPart(parts) => {
format!("MultiPart({} parts)", parts.len())
}
};
app/src/terminal/view.rs
+61 -29
View File
@@ -6994,7 +6994,7 @@ impl TerminalView {
}
fn toggle_settings_view(&mut self, ctx: &mut ViewContext<Self>) {
use crate::ai::bedrock::settings_view::SettingsView;
use crate::ai::provider::settings_view::SettingsView;
// If already showing, remove it
if let Some(view_id) = self.settings_view_id.take() {
@@ -10912,7 +10912,7 @@ impl TerminalView {
// Use the configured command, but if it's the bare default ("aws sso login")
// prefer the external config's auth_refresh_command which includes --profile.
let login_command = if settings_command == "aws sso login" {
use crate::ai::bedrock::external_config::ExternalBedrockConfig;
use crate::ai::provider::external_config::ExternalBedrockConfig;
let external = ExternalBedrockConfig::load();
external.auth_refresh_command.unwrap_or(settings_command)
} else {
@@ -10948,34 +10948,37 @@ impl TerminalView {
|_me, result, ctx| {
match result {
Ok(()) => {
log::info!("[bedrock] AWS login completed successfully, refreshing credentials and resuming");
// Refresh credentials from the updated SSO cache
ApiKeyManager::handle(ctx).update(
ctx,
|manager, ctx| {
drop(crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx));
},
);
// Resume the conversation after a short delay to let credentials load
let _ = ctx.spawn(
async move {
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
},
|me, _, ctx| {
// Find the active conversation and resume it directly
let conversation_id = if FeatureFlag::AgentView.is_enabled() {
me.agent_view_controller
.as_ref(ctx)
.agent_view_state()
.active_conversation_id()
} else {
BlocklistAIHistoryModel::as_ref(ctx).last_conversation_id(me.id())
};
if let Some(conversation_id) = conversation_id {
me.handle_resume_conversation(&conversation_id, ctx);
log::info!("[bedrock] AWS login completed successfully; reloading credentials before resuming");
// Refresh through ApiKeyManager and wait for its completion. A fixed
// delay can resume the request while the SDK still has stale SSO data.
let refresh = ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx)
});
let _ = ctx.spawn(refresh, |me, result, ctx| {
match result {
Ok(()) => {
log::info!("[bedrock] AWS credentials refreshed; resuming conversation");
let conversation_id = if FeatureFlag::AgentView.is_enabled() {
me.agent_view_controller
.as_ref(ctx)
.agent_view_state()
.active_conversation_id()
} else {
BlocklistAIHistoryModel::as_ref(ctx)
.last_conversation_id(me.id())
};
if let Some(conversation_id) = conversation_id {
me.handle_resume_conversation(&conversation_id, ctx);
}
}
},
);
Err(error) => {
log::error!(
"[bedrock] AWS login succeeded but credential reload failed; not resuming: {error}"
);
}
}
});
}
Err(e) => {
log::error!("[bedrock] AWS login failed: {e}");
@@ -20646,6 +20649,35 @@ impl TerminalView {
AIBlockEvent::RunAwsLoginCommand => {
self.run_aws_login_command(ctx);
}
AIBlockEvent::RefreshAwsCredentials => {
let refresh = ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx)
});
let _ = ctx.spawn(refresh, |me, result, ctx| {
match result {
Ok(()) => {
log::info!("[bedrock] AWS credentials reloaded; resuming conversation");
let conversation_id = if FeatureFlag::AgentView.is_enabled() {
me.agent_view_controller
.as_ref(ctx)
.agent_view_state()
.active_conversation_id()
} else {
BlocklistAIHistoryModel::as_ref(ctx).last_conversation_id(me.id())
};
if let Some(conversation_id) = conversation_id {
me.handle_resume_conversation(&conversation_id, ctx);
}
}
Err(error) => {
log::warn!(
"[bedrock] AWS credential reload failed; starting configured login command: {error}"
);
me.run_aws_login_command(ctx);
}
}
});
}
}
ctx.notify();
}