Improve agent provider resilience

2026-09-02 17:09:08 -05:00
parent b115946534
commit 7a33cc7e56
72 changed files with 997 additions and 320 deletions
@@ -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,
@@ -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 =
@@ -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.
@@ -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),
@@ -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(),
        }
@@ -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(),
@@ -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;