diff --git a/AGENTS.md b/AGENTS.md index 91bf291c..30dbd559 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -43,36 +43,50 @@ Environment variables: ### AI Provider Architecture -Galaxy supports multiple AI backends via a **provider dispatch pattern**. Provider selection -is controlled by settings (`ai.openai.enabled` takes priority over `ai.bedrock.enabled`). +Galaxy supports session-owned ACP backends and direct model providers. Direct-provider selection is +controlled by settings (`ai.openai.enabled` takes priority over `ai.bedrock.enabled`), but provider +configuration never selects lifecycle ownership. ``` -Provider dispatch: response_stream.rs → resolve_provider_config() → ProviderConfig enum - ↓ Bedrock ↓ OpenAI - bedrock/translator.rs openai/translator.rs +Direct provider: controller.rs → prepare_provider_run() → ProviderRunCoordinator + ↓ one model call per turn + AgentRuntime implementation + ↓ tool batch + correlated action execution + ↓ committed results + next ProviderRun turn + +ACP: controller.rs → ResponseStream → acp_output_stream (session-owned lifecycle) ``` -**Shared types** in `app/src/ai/provider/`: +**Durable direct-provider run**: +- `crates/galaxy_agent_core/src/provider_run.rs` — Serializable `ProviderRun` state machine, run/epoch identity, bounded model retries, ordered tool batches, cancellation, and terminal outcomes +- `app/src/ai/runtime/provider_run_coordinator.rs` — Drives one-turn `AgentRuntime` calls, validates model events, projects output, and commits exact tool lifecycle events +- `app/src/ai/runtime/rig.rs` — Builds base/CLI request profiles and resolves the configured one-turn runtime; it does not own follow-through +- `app/src/ai/runtime/rig_request.rs` — Converts controller request state into provider-neutral `TurnRequest` history, tools, prompts, and MCP aliases +- `app/src/ai/runtime/event_translator.rs` — Projects provider-neutral runtime events into Warp response events for UI/history compatibility +- `app/src/ai/blocklist/controller.rs` — Retains active runs, correlates actions by `(conversation_id, run_id, epoch, call_id)`, monitors commands, persists checkpoints, and restores interrupted runs +- `app/src/ai/blocklist/controller/response_stream.rs` — Owns ACP transport and shared UI/history projection only; it must not drive direct-provider retries or follow-up turns + +**Shared provider types** in `app/src/ai/provider/`: - `types.rs` — `ConversationMessage`, `MessageRole`, `MessageContent`, `ContentPart`, `ToolDefinition` - `mod.rs` — `ProviderConfig` enum (Bedrock | OpenAI | None) **Bedrock provider** in `app/src/ai/bedrock/`: -- `translator.rs` — Orchestrator: takes `api::Request` + config, returns `ResponseStream` -- `request_translator.rs` — Converts Warp proto → Bedrock SDK types (messages, system prompt, tools, sanitization) -- `response_translator.rs` — Converts Bedrock stream events → Warp proto `ResponseEvent`s -- `convert.rs` — Re-exports shared types + Bedrock SDK type builders -- `client.rs` — AWS SDK client construction and `converse_stream` call +- `runtime.rs` — Native one-call `AgentRuntime` over `ConverseStream`, including cancellation, reasoning signatures, token usage, and stop/error classification +- `request_translator.rs` — Shared Bedrock message sanitization and tool definitions +- `response_translator.rs` — Compatibility conversion helpers used by tests and background flows +- `convert.rs` — Bedrock request construction and prompt-caching behavior +- `client.rs` — AWS SDK client construction, runtime creation, and independent background streaming calls - `models.rs` — Model registry and cross-region inference prefix logic - `discovery.rs` — AWS profile listing and model discovery (STS identity check + ListFoundationModels) - `diagnostic.rs` — Debug logging (enabled via `GALAXY_BEDROCK_DIAGNOSTICS=1`) - `external_config.rs` — Fallback config from Claude Code/OpenCode settings -**OpenAI/LiteLLM provider** in `app/src/ai/openai/`: -- `translator.rs` — Orchestrator: same pattern as Bedrock, targets OpenAI chat completions API -- `client.rs` — `reqwest`-based HTTP client for `POST /v1/chat/completions` with streaming -- `convert.rs` — `ConversationMessage` → OpenAI JSON format (system/user/assistant/tool roles, function calling) -- `request_translator.rs` — OpenAI-specific message sanitization (lighter than Bedrock's strict alternation rules) -- `response_translator.rs` — SSE stream parser → Warp proto `ResponseEvent`s +**OpenAI-compatible providers**: +- Direct turns use one-call runtimes from `galaxy_agent_rig` selected in `app/src/ai/runtime/rig.rs` for OpenAI/LiteLLM, ChatGPT subscription, Anthropic, Gemini, and Vertex AI +- `app/src/ai/openai/request_translator.rs` sanitizes provider-neutral history for OpenAI-compatible APIs +- `app/src/ai/openai/client.rs`, `convert.rs`, and `response_translator.rs` remain compatibility/background transport helpers, not lifecycle owners **Provider settings** (in settings TOML): - `ai.bedrock.enabled` — Use AWS Bedrock directly (default: true) @@ -117,22 +131,25 @@ context_size = 128000 - Provider name shown as the description label in the model picker; icon shows OpenAI logo for all OpenAI-compatible providers Key invariants: -- Known tools are in `KNOWN_TOOLS` constant in `response_translator.rs` -- Tool definitions are built via `tool_definition_for_name()` in `convert_request.rs`; includes `recall_tool_history` for retrieving past tool results -- Direct-provider normal and plan turns must advertise `read_plan`, `create_plan`, and `edit_plan` when the matching document capabilities are enabled; plan-creation requests should call `create_plan` after research rather than only returning prose or claiming the tool is unavailable -- Unknown/hallucinated tool calls are caught in the stream, paired with synthetic error results, and now emit a visible `AgentOutput` text message to the UI -- `recall_tool_history` is handled inline by direct-provider adapters using a synthetic result from `messages_sent`; the Rig adapter must automatically start a bounded follow-up provider turn after pairing that result, without continuing turns that proposed client-executed tools -- `recall_tool_history` must exclude prior calls to itself from candidates so inline continuation cannot recursively recall synthetic recall results -- Tool result archive: before progressive summarization drains messages, `ConversationMessage::archive_tool_results()` extracts all tool_use/tool_result pairs into a separate `tool_result_archive` vec. `recall_tool_history` searches both live history + archived results, and supports a `tool_use_id` parameter for exact ID lookup -- Prompt caching uses three cache points: system prompt, conversation history (second-to-last message), tool config +- Every direct-provider `AgentRuntime::start_turn` performs exactly one model call; only `ProviderRun` may schedule another turn or retry +- `use_rig` and provider selection may choose request/transport details but must never choose lifecycle ownership +- Direct-provider output may be projected through `ResponseStream`, but provider progress must not depend on response-stream result draining or `AfterStreamFinished` +- Provider actions and results must correlate by `(conversation_id, run_id, epoch, call_id)`; stale or duplicate callbacks must not advance a run +- Active provider runs must checkpoint before external work, persist without credentials, normalize unsafe restored states, and reconcile command state before continuing +- Known tools are in `KNOWN_TOOLS` in `response_translator.rs`; definitions are built by `tool_definition_for_name()` in `convert_request.rs` +- Direct-provider normal and plan turns must advertise `read_plan`, `create_plan`, and `edit_plan` when the matching document capabilities are enabled; plan-creation requests should call `create_plan` after research rather than only returning prose +- Unknown or invalid tool calls receive one correlated synthetic error result and a visible `AgentOutput` message; the durable run owns any continuation +- `recall_tool_history` is an inline completed tool batch. `ProviderRun` commits its synthetic result and starts a bounded next turn without routing it through client action execution +- `recall_tool_history` excludes earlier calls to itself; archived tool results remain searchable by query or exact `tool_use_id` +- Before progressive summarization drains messages, `ConversationMessage::archive_tool_results()` moves tool-use/result pairs into `tool_result_archive` +- Bedrock prompt caching uses three cache points: system prompt, second-to-last history message, and tool configuration - `ensure_tool_results_paired()` enforces Bedrock's invariant that every `tool_use` has a matching `tool_result` -- `inject_input_messages_into_task()` and `extract_user_query_text()` ensure user queries persist for session restore -- The stream emits a `UserQuery` proto message at the start of each response for conversation title -- Progressive summary (if present) is prepended to the messages array as a user/assistant pair in `translator.rs` -- Loop prevention guardrail in `controller.rs` detects repeated tool failures (3+ identical) and injects corrective instructions -- Direct-provider long-running shell follow-ups create unlinked CLI tasks under the root task with an empty subagent tool-call ID; `TaskStore` linearization must include their exchanges chronologically even though no parent `Subagent` output references them -- Direct-provider completed-command assessments are hidden, tool-free root-task turns; CLI monitor exchanges remain on the retained CLI task, while the root assessment output must survive CLI-task deactivation and restoration and its hidden input must remain available to future provider context -- Orchestrated child conversations are leaf workers by default: nested `RunAgents` and legacy `StartAgent` calls must be rejected before autonomous or permission bypasses, and child requests must not advertise delegation tools +- Progressive summaries are prepended to provider requests as a user/assistant pair; background summarization remains independent of the active provider run +- Loop prevention in `controller.rs` detects repeated tool failures (3+ identical) and injects a corrective instruction +- Direct-provider long-running shell follow-ups retain CLI tasks as UI projections while command monitoring and completion stay owned by the same provider run +- Direct-provider completed-command assessments are hidden, tool-free root-task turns whose output and hidden input survive CLI-task deactivation and restoration +- ACP remains a separate session-owned runtime; direct-provider cleanup must not move ACP lifecycle into `ProviderRun` +- Orchestrated child conversations are leaf workers: nested `RunAgents` and legacy `StartAgent` calls are rejected before autonomy or permission bypasses, and child requests do not advertise delegation tools ### Platform Setup - `./script/bootstrap` - Platform-specific setup plus common agent skill installation from `skills-lock.json`; prompts for project/global when an install or update is needed unless a target flag or environment override is provided. diff --git a/app/src/ai/agent/api.rs b/app/src/ai/agent/api.rs index 7fc0f762..09b9dc43 100644 --- a/app/src/ai/agent/api.rs +++ b/app/src/ai/agent/api.rs @@ -21,10 +21,10 @@ use galaxy_core::features::FeatureFlag; use galaxy_core::user_preferences::GetUserPreferences; use galaxyui::{AppContext, EntityId, SingletonEntity as _}; use mcp::TemplatableMCPServerInfo; -pub use r#impl::generate_multi_agent_output; +pub(crate) use r#impl::prepare_direct_provider_params; use serde::Serialize; -use super::{AIAgentAction, AIAgentInput, MCPContext, MCPServer, RequestMetadata, Suggestions}; +use super::{AIAgentInput, MCPContext, MCPServer, RequestMetadata, Suggestions}; use crate::ai::agent::conversation::AIConversationId; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::blocklist::{BlocklistAIPermissions, RequestInput, SessionContext}; @@ -112,7 +112,7 @@ pub struct RequestParams { /// locally so ACP-provided Galaxy tools can be pinned to the exact pane. pub terminal_view_id: Option, pub input: Vec, - /// Normalized results consumed directly by Rig-selected models. + /// Normalized action results appended to direct-provider run history. pub tool_results: Vec, pub conversation_token: Option, pub forked_from_conversation_token: Option, @@ -151,8 +151,8 @@ pub struct RequestParams { pub research_agent_enabled: bool, pub orchestration_enabled: bool, pub supported_tools_override: Option>, - /// The root task ID for the conversation — needed for direct Bedrock streaming - /// since optimistic tasks don't appear in the proto task_context. + /// The root task ID used to anchor direct-provider projection when optimistic tasks are not + /// present in the proto task context. pub root_task_id: Option, /// The conversation ID of the parent agent that spawned this child agent, if any. pub parent_agent_id: Option, @@ -167,9 +167,8 @@ pub struct RequestParams { /// Kept separately so `recall_tool_history` can search archived results even after /// they've been summarized away from live history. pub tool_result_archive: Vec, - /// Populated by direct-provider paths after building the message list. - /// Contains the full messages sent (old history + new input) so the controller - /// can store them back into the conversation for the next request cycle. + /// Populated while preparing a direct-provider run with the durable transcript that the + /// controller persists for restoration and future turns. pub messages_sent: std::sync::Arc>>, /// Global rules (name, content) from the local CloudModel (AIFact/AIMemory). @@ -177,15 +176,10 @@ pub struct RequestParams { pub global_rules: Vec<(String, String)>, } -/// Provider/runtime events consumed by the local conversation controller. -/// -/// The legacy response envelope remains at the UI boundary while the local Rig runtime emits -/// executable tool proposals directly as Galaxy domain actions. This avoids translating Rig tool -/// calls into protobuf only to immediately translate them back before execution. +/// Response event projected into the local conversation controller. #[derive(Debug)] pub enum StreamEvent { Response(warp_multi_agent_api::ResponseEvent), - ToolProposed(AIAgentAction), } pub type Event = Result>; diff --git a/app/src/ai/agent/api/convert_conversation.rs b/app/src/ai/agent/api/convert_conversation.rs index 4d118147..741bfc94 100644 --- a/app/src/ai/agent/api/convert_conversation.rs +++ b/app/src/ai/agent/api/convert_conversation.rs @@ -73,6 +73,7 @@ pub fn convert_conversation_data_to_ai_conversation( let agent_conversation_data = match restoration_mode { RestorationMode::Fork => AgentConversationData { agent_backend: Default::default(), + active_provider_run_json: None, server_conversation_token: None, conversation_usage_metadata: usage_metadata, reverted_action_ids: None, @@ -96,6 +97,7 @@ pub fn convert_conversation_data_to_ai_conversation( }, RestorationMode::Continue => AgentConversationData { agent_backend: Default::default(), + active_provider_run_json: None, server_conversation_token: Some( metadata.server_conversation_token.as_str().to_string(), ), diff --git a/app/src/ai/agent/api/impl.rs b/app/src/ai/agent/api/impl.rs index 9106c32a..e75f2342 100644 --- a/app/src/ai/agent/api/impl.rs +++ b/app/src/ai/agent/api/impl.rs @@ -1,270 +1,10 @@ -use std::collections::HashMap; -use std::sync::Arc; - -use futures_util::StreamExt; use galaxy_core::features::FeatureFlag; use warp_multi_agent_api as api; -use super::convert_to::convert_input; -use super::{ConvertToAPITypeError, RequestParams, ResponseStream, StreamEvent}; -use crate::ai::agent::{redaction, AIAgentInput}; -use crate::ai::openai::translator as openai_translator; -use crate::ai::provider::ProviderConfig; -use crate::server::server_api::AIApiError; +use super::RequestParams; +use crate::ai::agent::redaction; use crate::terminal::model::session::SessionType; -pub async fn generate_multi_agent_output( - provider_config: ProviderConfig, - mut params: RequestParams, - cancellation_rx: futures::channel::oneshot::Receiver<()>, -) -> Result { - let supported_tools_override = params.supported_tools_override.take(); - let mut supported_tools = supported_tools_override - .clone() - .unwrap_or_else(|| get_supported_tools(¶ms)); - remove_orchestration_tools_if_disabled(&mut supported_tools, params.orchestration_enabled); - let mut supported_cli_agent_tools = - supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(¶ms)); - remove_orchestration_tools_if_disabled( - &mut supported_cli_agent_tools, - params.orchestration_enabled, - ); - if params.should_redact_secrets { - redaction::redact_inputs(&mut params.input); - } - - match &provider_config { - ProviderConfig::OpenAI(config) if config.use_rig => { - return Ok(crate::ai::runtime::rig_openai_response_stream( - config.clone(), - params, - supported_tools, - supported_cli_agent_tools, - cancellation_rx, - )); - } - ProviderConfig::Bedrock(config) if config.use_rig => { - return match crate::ai::runtime::rig_bedrock_response_stream( - config.clone(), - params, - supported_tools, - supported_cli_agent_tools, - cancellation_rx, - ) - .await - { - Ok(stream) => Ok(stream), - Err(error) => { - log::error!("[rig/bedrock] Runtime error: {error}"); - let error = Arc::new(crate::server::server_api::AIApiError::Stream { - stream_type: "rig_bedrock", - source: error, - }); - let (sender, receiver) = async_channel::unbounded(); - let _ = sender.send(Err(error)).await; - Ok(Box::pin(receiver)) - } - }; - } - ProviderConfig::OpenAI(_) | ProviderConfig::Bedrock(_) | ProviderConfig::None => {} - } - - let mut logging_metadata = HashMap::new(); - if let Some(ref metadata) = params.metadata { - logging_metadata.insert( - "is_autodetected_user_query".to_owned(), - prost_types::Value { - kind: Some(prost_types::value::Kind::BoolValue( - metadata.is_autodetected_user_query, - )), - }, - ); - logging_metadata.insert( - "entrypoint".to_owned(), - prost_types::Value { - kind: Some(prost_types::value::Kind::StringValue( - metadata.entrypoint.entrypoint(), - )), - }, - ); - logging_metadata.insert( - "is_auto_resume_after_error".to_owned(), - prost_types::Value { - kind: Some(prost_types::value::Kind::BoolValue( - metadata.is_auto_resume_after_error, - )), - }, - ); - } - - let emit_user_query_message = !params - .input - .iter() - .any(|input| matches!(input, AIAgentInput::CommandCompletionAssessment { .. })); - let mut request = api::Request { - task_context: Some(api::request::TaskContext { - tasks: params.tasks, - }), - input: Some(convert_input(params.input)?), - settings: Some(api::request::Settings { - model_config: Some(api::request::settings::ModelConfig { - base: params.model.clone().into(), - cli_agent: params.cli_agent_model.clone().into(), - computer_use_agent: params.computer_use_model.clone().into(), - base_model_context_window_limit: params.context_window_limit.unwrap_or(0), - ..Default::default() - }), - rules_enabled: params.is_memory_enabled, - warp_drive_context_enabled: params.warp_drive_context_enabled, - web_context_retrieval_enabled: true, - supports_parallel_tool_calls: true, - use_anthropic_text_editor_tools: false, - planning_enabled: params.planning_enabled, - supports_create_files: true, - supported_tools: supported_tools.into_iter().map(Into::into).collect(), - supports_long_running_commands: true, - should_preserve_file_content_in_history: true, - supports_todos_ui: true, - supports_linked_code_blocks: FeatureFlag::LinkedCodeBlocks.is_enabled(), - supports_started_child_task_message: true, - // Galaxy's direct providers only receive tools with local schemas and - // executors. Hosted-only suggestion/orchestration capability bits must - // remain false so models do not plan around unavailable Warp services. - supports_suggest_prompt: false, - supports_read_image_files: FeatureFlag::ReadImageFiles.is_enabled(), - supports_reasoning_message: true, - api_keys: params.api_keys, - autonomy_level: params.autonomy_level.into(), - isolation_level: params.isolation_level.into(), - web_search_enabled: params.web_search_enabled, - supported_cli_agent_tools: supported_cli_agent_tools - .into_iter() - .map(Into::into) - .collect(), - supports_v4a_file_diffs: FeatureFlag::V4AFileDiffs.is_enabled(), - supports_summarization_via_message_replacement: - FeatureFlag::SummarizationViaMessageReplacement.is_enabled(), - supports_bundled_skills: FeatureFlag::BundledSkills.is_enabled(), - supports_research_agent: params.research_agent_enabled, - supports_orchestration_v2: false, - supports_background_computer_use: FeatureFlag::BackgroundComputerUse.is_enabled() - && computer_use::background_supported(), - custom_model_providers: params.custom_model_providers, - custom_model_routers: params.custom_model_routers, - }), - metadata: Some(api::request::Metadata { - logging: logging_metadata, - conversation_id: params - .conversation_token - .as_ref() - .map(|token| token.as_str().to_string()) - .unwrap_or_default(), - ambient_agent_task_id: params - .ambient_agent_task_id - .map(|id| id.to_string()) - .unwrap_or_default(), - forked_from_conversation_id: if params.conversation_token.is_none() { - // We only include this param on our initial request to the server - // (when the forked conversation has not been assigned a new id yet). - params - .forked_from_conversation_token - .map(|token| token.as_str().to_string()) - .unwrap_or_default() - } else { - String::new() - }, - parent_agent_id: params.parent_agent_id.unwrap_or_default(), - agent_name: params.agent_name.unwrap_or_default(), - }), - existing_suggestions: params - .existing_suggestions - .map(|suggestions| suggestions.into()), - mcp_context: params.mcp_context.map(Into::into), - }; - - match provider_config { - ProviderConfig::OpenAI(config) => { - let translator_request = openai_translator::TranslatorRequest { - config, - model_id: params.model.as_str().to_string(), - root_task_id: params.root_task_id.clone(), - message_history: params.message_history.clone(), - tool_result_archive: params.tool_result_archive.clone(), - progressive_summary: params.progressive_summary.clone(), - messages_sent: params.messages_sent.clone(), - global_rules: params.global_rules.clone(), - emit_user_query_message, - }; - - match openai_translator::execute(translator_request, &mut request).await { - Ok(stream) => { - let output_stream = stream - .map(|event| event.map(StreamEvent::Response)) - .take_until(cancellation_rx); - Ok(Box::pin(output_stream)) - } - Err(e) => { - log::error!("[openai] Translator error: {e}"); - let err = Arc::new( - crate::server::server_api::AIApiError::Stream { - stream_type: "openai_chat_completions", - source: anyhow::anyhow!("{e}"), - } - .into_quota_limit_if_provider_budget_exhausted(), - ); - let (tx, rx) = async_channel::unbounded(); - let _ = tx.send(Err(err)).await; - Ok(Box::pin(rx)) - } - } - } - ProviderConfig::Bedrock(config) => { - let translator_request = crate::ai::bedrock::translator::TranslatorRequest { - config, - model_id: params.model.as_str().to_string(), - root_task_id: params.root_task_id.clone(), - bedrock_message_history: params.message_history.clone(), - bedrock_tool_result_archive: params.tool_result_archive.clone(), - bedrock_progressive_summary: params.progressive_summary.clone(), - bedrock_messages_sent: params.messages_sent.clone(), - global_rules: params.global_rules.clone(), - emit_user_query_message, - }; - - match crate::ai::bedrock::translator::execute(translator_request, &mut request).await { - Ok(stream) => { - let output_stream = stream - .map(|event| event.map(StreamEvent::Response)) - .take_until(cancellation_rx); - Ok(Box::pin(output_stream)) - } - Err(e) => { - log::error!("[bedrock] Translator error: {e}"); - let err = Arc::new(crate::server::server_api::AIApiError::Stream { - stream_type: "bedrock", - source: anyhow::anyhow!("{e}"), - }); - let (tx, rx) = async_channel::unbounded(); - let _ = tx.send(Err(err)).await; - Ok(Box::pin(rx)) - } - } - } - ProviderConfig::None => { - // No provider configured — do not fall back to Warp's cloud API. - let err = Arc::new(crate::server::server_api::AIApiError::Stream { - stream_type: "none", - source: anyhow::anyhow!( - "No AI runtime configured. Enable an agent runtime or model provider in settings." - ), - }); - let (tx, rx) = async_channel::unbounded(); - let _ = tx.send(Err(err)).await; - Ok(Box::pin(rx)) - } - } -} - fn remove_orchestration_tools_if_disabled( supported_tools: &mut Vec, orchestration_enabled: bool, @@ -280,6 +20,26 @@ fn remove_orchestration_tools_if_disabled( }); } +pub(crate) fn prepare_direct_provider_params( + params: &mut RequestParams, +) -> (Vec, Vec) { + let supported_tools_override = params.supported_tools_override.take(); + let mut supported_tools = supported_tools_override + .clone() + .unwrap_or_else(|| get_supported_tools(params)); + remove_orchestration_tools_if_disabled(&mut supported_tools, params.orchestration_enabled); + let mut supported_cli_agent_tools = + supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(params)); + remove_orchestration_tools_if_disabled( + &mut supported_cli_agent_tools, + params.orchestration_enabled, + ); + if params.should_redact_secrets { + redaction::redact_inputs(&mut params.input); + } + (supported_tools, supported_cli_agent_tools) +} + fn get_supported_tools(params: &RequestParams) -> Vec { let mut supported_tools = vec![ api::ToolType::Grep, diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index a6a2e1e5..46422064 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -229,6 +229,9 @@ pub struct AIConversation { /// Runtime responsible for executing this conversation. agent_backend: AgentBackend, + /// Opaque, versioned snapshot of the active direct-provider run. + active_provider_run_json: Option, + /// The server-generated unique "token" for this conversation. /// /// This must be roundtripped to the server when sending follow-ups within a given conversation. @@ -380,6 +383,7 @@ impl AIConversation { has_opened_code_review: false, conversation_usage_metadata: ConversationUsageMetadata::default(), agent_backend, + active_provider_run_json: None, server_conversation_token: None, task_id: None, forked_from_server_conversation_token: None, @@ -539,6 +543,7 @@ impl AIConversation { let ( agent_backend, + active_provider_run_json, server_conversation_token, forked_from_server_conversation_token, conversation_usage_metadata, @@ -589,6 +594,7 @@ impl AIConversation { }; ( data.agent_backend, + data.active_provider_run_json, server_conversation_token, forked_from_server_conversation_token, conversation_usage_metadata, @@ -611,6 +617,7 @@ impl AIConversation { AgentBackend::default(), None, None, + None, ConversationUsageMetadata::default(), HashSet::new(), Vec::new(), @@ -663,6 +670,7 @@ impl AIConversation { has_opened_code_review: false, conversation_usage_metadata, agent_backend, + active_provider_run_json, server_conversation_token, task_id: run_id.as_deref().and_then(|id| id.parse().ok()), forked_from_server_conversation_token, @@ -705,6 +713,14 @@ impl AIConversation { &self.agent_backend } + pub(crate) fn active_provider_run_json(&self) -> Option<&str> { + self.active_provider_run_json.as_deref() + } + + pub(crate) fn set_active_provider_run_json(&mut self, snapshot: Option) { + self.active_provider_run_json = snapshot; + } + /// Updates the backend of a conversation that has not produced agent output. /// /// Provider failures without output are safe to retry through a newly enabled runtime. Once @@ -2109,6 +2125,81 @@ impl AIConversation { Ok(()) } + /// Reopens an exact restored exchange for continued provider projection. + /// + /// This only restores the process-local stream association; it never adds input or provider + /// history, so the persisted provider run remains the sole continuation source of truth. + pub(crate) fn provider_projection_target( + &self, + response_stream_id: &ResponseStreamId, + ) -> Option<(TaskId, AIAgentExchangeId)> { + let mut exchanges = self + .added_exchanges_by_response + .get(response_stream_id)? + .iter(); + let target = exchanges.next()?; + exchanges + .next() + .is_none() + .then(|| (target.task_id.clone(), target.exchange_id)) + } + + pub(crate) fn rebind_provider_projection( + &mut self, + task_id: &TaskId, + exchange_id: AIAgentExchangeId, + response_stream_id: ResponseStreamId, + terminal_surface_id: EntityId, + ctx: &mut ModelContext, + ) -> Result<(), UpdateConversationError> { + let Some(task) = self.task_store.get(task_id) else { + return Err(UpdateConversationError::TaskNotFound); + }; + if !task.exchanges().any(|exchange| exchange.id == exchange_id) { + return if self.exchange_with_id(exchange_id).is_some() { + Err(UpdateConversationError::ExchangeTaskMismatch) + } else { + Err(UpdateConversationError::ExchangeNotFound) + }; + } + if self + .added_exchanges_by_response + .contains_key(&response_stream_id) + { + return Err(UpdateConversationError::ResponseStreamAlreadyBound); + } + + let exchange = self.get_exchange_to_update(exchange_id)?; + let previous_status = std::mem::replace( + &mut exchange.output_status, + AIAgentOutputStatus::Streaming { output: None }, + ); + let output = match previous_status { + AIAgentOutputStatus::Streaming { output } => output, + AIAgentOutputStatus::Finished { finished_output } => match finished_output { + FinishedAIAgentOutput::Cancelled { output, .. } + | FinishedAIAgentOutput::Error { output, .. } => output, + FinishedAIAgentOutput::Success { output } => Some(output), + }, + }; + exchange.output_status = AIAgentOutputStatus::Streaming { output }; + exchange.finish_time = None; + self.added_exchanges_by_response.insert( + response_stream_id, + Vec1::new(AddedExchange { + task_id: task_id.clone(), + exchange_id, + }), + ); + ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange { + exchange_id, + terminal_surface_id, + conversation_id: self.id, + is_hidden: self.hidden_exchanges.contains(&exchange_id), + }); + Ok(()) + } + pub fn append_reassigned_exchange( &mut self, response_stream_id: &ResponseStreamId, @@ -2255,6 +2346,9 @@ impl AIConversation { action: AIAgentAction, ctx: &mut ModelContext, ) -> Result<(), UpdateConversationError> { + if self.contains_action(&action.id) { + return Ok(()); + } let added_exchanges = self .added_exchanges_by_response .get(stream_id) @@ -3886,6 +3980,7 @@ impl AIConversation { .collect(), conversation_data: AgentConversationData { agent_backend: self.agent_backend.clone(), + active_provider_run_json: self.active_provider_run_json.clone(), server_conversation_token: self .server_conversation_token .clone() @@ -4763,6 +4858,10 @@ fn cleanup_conversation_search_temp_dir( pub enum UpdateConversationError { #[error("Exchange not found.")] ExchangeNotFound, + #[error("Exchange does not belong to the persisted task.")] + ExchangeTaskMismatch, + #[error("Response stream is already bound to an exchange.")] + ResponseStreamAlreadyBound, #[error("Could not update task: {0:?}")] UpdateTask(#[from] UpdateTaskError), #[error("Could not update upgrade optimistic task for server task: {0:?}")] diff --git a/app/src/ai/agent/conversation_tests.rs b/app/src/ai/agent/conversation_tests.rs index e5bfcd62..04b1151c 100644 --- a/app/src/ai/agent/conversation_tests.rs +++ b/app/src/ai/agent/conversation_tests.rs @@ -193,6 +193,19 @@ fn restored_conversation_defaults_autoexecute_override_when_not_persisted() { ); } +#[test] +fn restored_conversation_retains_active_provider_run_json() { + let snapshot = r#"{"version":1,"run":{"state":"awaiting_model"}}"#; + let conversation_data = AgentConversationData { + active_provider_run_json: Some(snapshot.to_string()), + ..Default::default() + }; + + let conversation = restored_conversation(Some(conversation_data)); + + assert_eq!(conversation.active_provider_run_json(), Some(snapshot)); +} + #[test] fn restored_conversation_uses_persisted_last_event_sequence() { let conversation_data: AgentConversationData = diff --git a/app/src/ai/bedrock/client.rs b/app/src/ai/bedrock/client.rs index 44a4bb9f..f3ed6941 100644 --- a/app/src/ai/bedrock/client.rs +++ b/app/src/ai/bedrock/client.rs @@ -6,13 +6,13 @@ use aws_credential_types::provider::ProvideCredentials; use aws_sdk_bedrockruntime::config::Region; use aws_sdk_bedrockruntime::Client as BedrockRuntimeClient; use galaxy_agent_core::AgentError; -use galaxy_agent_rig::{BedrockRigConfig, BedrockRuntime}; use super::convert::{build_converse_request, CachingConfig, ConversationMessage, ToolDefinition}; use super::diagnostic::BedrockDiagnosticLogger; use super::external_config::ExternalBedrockConfig; use super::models::apply_cross_region_prefix; use super::response_translator::bedrock_stream_to_response_events; +use super::runtime::BedrockAgentRuntime; use crate::ai::agent::api::LegacyResponseStream; use crate::settings::ai::BedrockAuthMethod; @@ -171,25 +171,20 @@ impl BedrockClient { }) } - /// Builds the Phase 4 Rig runtime from the AWS SDK client whose region and - /// credentials Galaxy already resolved. This does not change production - /// routing; callers opt in only after the Bedrock parity suite passes. - pub fn rig_runtime( + pub(crate) fn agent_runtime( &self, model: String, cross_region_inference: bool, - prompt_caching: bool, max_output_tokens: Option, - ) -> Result { - BedrockRuntime::from_aws_client( + caching_config: CachingConfig, + ) -> Result { + BedrockAgentRuntime::new( self.runtime_client.clone(), - BedrockRigConfig { - model, - region: self.region.clone(), - cross_region_inference, - prompt_caching, - max_output_tokens, - }, + model, + self.region.clone(), + cross_region_inference, + max_output_tokens, + caching_config, ) } diff --git a/app/src/ai/bedrock/convert_tests.rs b/app/src/ai/bedrock/convert_tests.rs index dbf847ff..25011584 100644 --- a/app/src/ai/bedrock/convert_tests.rs +++ b/app/src/ai/bedrock/convert_tests.rs @@ -258,7 +258,10 @@ fn test_system_prompt_separated_from_messages() { None, None, None, - CachingConfig::default(), + CachingConfig { + enabled: false, + extended_ttl_requested: false, + }, ); assert_eq!(result.system.len(), 1); @@ -334,7 +337,10 @@ fn test_tool_definitions_produce_tool_config() { None, None, None, - CachingConfig::default(), + CachingConfig { + enabled: false, + extended_ttl_requested: false, + }, ); assert!(result.tool_config.is_some()); diff --git a/app/src/ai/bedrock/mod.rs b/app/src/ai/bedrock/mod.rs index 4ebfe0a3..afaccfc6 100644 --- a/app/src/ai/bedrock/mod.rs +++ b/app/src/ai/bedrock/mod.rs @@ -7,8 +7,8 @@ pub mod external_config; pub mod models; pub mod request_translator; pub mod response_translator; +pub mod runtime; pub mod settings_view; -pub mod translator; #[cfg(test)] mod convert_tests; diff --git a/app/src/ai/bedrock/runtime.rs b/app/src/ai/bedrock/runtime.rs new file mode 100644 index 00000000..fc377a7a --- /dev/null +++ b/app/src/ai/bedrock/runtime.rs @@ -0,0 +1,463 @@ +use std::collections::BTreeMap; + +use async_trait::async_trait; +use aws_sdk_bedrockruntime::operation::converse_stream::ConverseStreamOutput; +use aws_sdk_bedrockruntime::types::{ + ContentBlockDelta, ContentBlockStart, ConverseStreamOutput as AwsStreamEvent, + ReasoningContentBlockDelta, StopReason as AwsStopReason, +}; +use aws_sdk_bedrockruntime::Client as AwsBedrockClient; +use futures::{FutureExt, StreamExt}; +use galaxy_agent_core::{ + AgentError, AgentErrorKind, AgentEvent, AgentEventStream, AgentRuntime, RuntimeCapabilities, + RuntimeDescriptor, RuntimeKind, StopReason, ToolCall, ToolEvent, TurnCommand, TurnControl, + TurnRequest, Usage, +}; +use uuid::Uuid; + +use super::convert::{build_converse_request, CachingConfig, ConvertedRequest}; + +const DEFAULT_MAX_OUTPUT_TOKENS: u64 = 64_000; + +#[derive(Clone)] +pub(crate) struct BedrockAgentRuntime { + client: AwsBedrockClient, + resolved_model: String, + max_output_tokens: Option, + caching_config: CachingConfig, + descriptor: RuntimeDescriptor, +} + +impl BedrockAgentRuntime { + pub(crate) fn new( + client: AwsBedrockClient, + configured_model: String, + region: String, + cross_region_inference: bool, + max_output_tokens: Option, + caching_config: CachingConfig, + ) -> Result { + let resolved_model = galaxy_agent_rig::resolve_bedrock_model_id( + &configured_model, + ®ion, + cross_region_inference, + )?; + let descriptor = RuntimeDescriptor { + id: format!("bedrock:{resolved_model}"), + display_name: format!("Bedrock / {resolved_model}"), + kind: RuntimeKind::Provider, + capabilities: RuntimeCapabilities::provider(), + }; + Ok(Self { + client, + resolved_model, + max_output_tokens, + caching_config, + descriptor, + }) + } +} + +#[async_trait] +impl AgentRuntime for BedrockAgentRuntime { + fn descriptor(&self) -> &RuntimeDescriptor { + &self.descriptor + } + + async fn start_turn( + &self, + request: TurnRequest, + control: TurnControl, + ) -> Result { + let converted = + convert_turn_request(request, self.max_output_tokens, self.caching_config.clone()); + let mut request = self + .client + .converse_stream() + .model_id(&self.resolved_model) + .set_system(Some(converted.system)) + .set_messages(Some(converted.messages)) + .inference_config(converted.inference_config); + if let Some(tool_config) = converted.tool_config { + request = request.tool_config(tool_config); + } + + let runtime_request_id = Uuid::new_v4().to_string(); + let send_future = request.send().fuse(); + let initial_control = control.clone(); + let control_future = initial_control.receive().fuse(); + futures::pin_mut!(send_future, control_future); + let output = futures::select_biased! { + command = control_future => match command { + Ok(TurnCommand::Cancel) => { + return Ok(stopped_before_stream(runtime_request_id)); + } + Ok(TurnCommand::Steer { .. }) | Err(_) => { + send_future.await.map_err(map_bedrock_error)? + } + }, + result = send_future => result.map_err(map_bedrock_error)?, + }; + + Ok(translate_bedrock_stream( + output, + runtime_request_id, + control, + )) + } +} + +fn convert_turn_request( + request: TurnRequest, + configured_max_output_tokens: Option, + caching_config: CachingConfig, +) -> ConvertedRequest { + let max_output_tokens = request + .max_output_tokens + .or(configured_max_output_tokens) + .unwrap_or(DEFAULT_MAX_OUTPUT_TOKENS) + .min(i32::MAX as u64) as i32; + build_converse_request( + request.messages, + request.system_prompt, + None, + request.tools, + max_output_tokens, + None, + None, + None, + caching_config, + ) +} + +fn translate_bedrock_stream( + mut output: ConverseStreamOutput, + runtime_request_id: String, + control: TurnControl, +) -> AgentEventStream { + let events = async_stream::stream! { + yield Ok(AgentEvent::TurnStarted { + runtime_request_id, + }); + + let mut translator = BedrockStreamTranslator::default(); + let mut control_open = true; + loop { + let next_event = output.stream.recv().fuse(); + let next_command = if control_open { + futures::future::Either::Left(control.receive()) + } else { + futures::future::Either::Right(futures::future::pending()) + } + .fuse(); + futures::pin_mut!(next_event, next_command); + + let event = futures::select_biased! { + command = next_command => { + match command { + Ok(TurnCommand::Cancel) => { + yield Ok(AgentEvent::TurnStopped { + reason: StopReason::Cancelled, + }); + return; + } + Ok(TurnCommand::Steer { .. }) => continue, + Err(_) => { + control_open = false; + continue; + } + } + } + event = next_event => event, + }; + + match event { + Ok(Some(event)) => match translator.translate(event) { + Ok(events) => { + for event in events { + yield Ok(event); + } + } + Err(error) => { + yield Err(error); + return; + } + }, + Ok(None) => match translator.finish() { + Ok(events) => { + for event in events { + yield Ok(event); + } + return; + } + Err(error) => { + yield Err(error); + return; + } + }, + Err(error) => { + yield Err(map_bedrock_error(error)); + return; + } + } + } + }; + Box::pin(events) +} + +#[derive(Default)] +struct BedrockStreamTranslator { + content_blocks: BTreeMap, + stop_reason: Option, +} + +impl BedrockStreamTranslator { + fn translate(&mut self, event: AwsStreamEvent) -> Result, AgentError> { + match event { + AwsStreamEvent::MessageStart(_) => Ok(Vec::new()), + AwsStreamEvent::ContentBlockStart(start) => { + let Some(block_start) = start.start() else { + return Ok(Vec::new()); + }; + let ContentBlockStart::ToolUse(tool) = block_start else { + return Err(protocol_error( + "Bedrock started an unsupported output content block", + )); + }; + let index = start.content_block_index(); + if self + .content_blocks + .insert( + index, + PendingContentBlock::Tool { + id: tool.tool_use_id().to_string(), + name: tool.name().to_string(), + input: String::new(), + }, + ) + .is_some() + { + return Err(protocol_error(format!( + "Bedrock started content block {index} more than once" + ))); + } + Ok(Vec::new()) + } + AwsStreamEvent::ContentBlockDelta(delta) => { + let Some(delta_value) = delta.delta() else { + return Err(protocol_error("Bedrock emitted an empty content delta")); + }; + let index = delta.content_block_index(); + match delta_value { + ContentBlockDelta::Text(text) => { + Ok(vec![AgentEvent::TextDelta { text: text.clone() }]) + } + ContentBlockDelta::ReasoningContent(reasoning) => { + let block = self.content_blocks.entry(index).or_insert_with(|| { + PendingContentBlock::Reasoning { + text: String::new(), + signature: None, + } + }); + let PendingContentBlock::Reasoning { text, signature } = block else { + return Err(protocol_error(format!( + "Bedrock mixed reasoning and tool data in content block {index}" + ))); + }; + match reasoning { + ReasoningContentBlockDelta::Text(delta) => { + text.push_str(delta); + Ok(vec![AgentEvent::ReasoningDelta { + text: delta.clone(), + }]) + } + ReasoningContentBlockDelta::Signature(delta) => { + signature.get_or_insert_with(String::new).push_str(delta); + Ok(Vec::new()) + } + ReasoningContentBlockDelta::RedactedContent(_) => Ok(Vec::new()), + _ => Err(protocol_error("Bedrock emitted an unknown reasoning delta")), + } + } + ContentBlockDelta::ToolUse(tool_delta) => { + let Some(PendingContentBlock::Tool { input, .. }) = + self.content_blocks.get_mut(&index) + else { + return Err(protocol_error(format!( + "Bedrock emitted tool input before starting content block {index}" + ))); + }; + input.push_str(tool_delta.input()); + Ok(Vec::new()) + } + ContentBlockDelta::Citation(_) => Ok(Vec::new()), + ContentBlockDelta::Image(_) => { + Err(protocol_error("Bedrock emitted unsupported image output")) + } + ContentBlockDelta::ToolResult(_) => Err(protocol_error( + "Bedrock emitted an unexpected tool-result delta", + )), + _ => Err(protocol_error("Bedrock emitted an unknown content delta")), + } + } + AwsStreamEvent::ContentBlockStop(stop) => { + let index = stop.content_block_index(); + let Some(block) = self.content_blocks.remove(&index) else { + return Ok(Vec::new()); + }; + match block { + PendingContentBlock::Tool { id, name, input } => { + let arguments = serde_json::from_str(&input).map_err(|error| { + protocol_error(format!( + "Bedrock returned invalid JSON for tool '{name}' ({id}): {error}" + )) + })?; + Ok(vec![AgentEvent::Tool { + event: ToolEvent::Proposed { + call: ToolCall { + id, + name, + arguments, + }, + }, + }]) + } + PendingContentBlock::Reasoning { text, signature } => { + Ok(vec![AgentEvent::ReasoningCompleted { text, signature }]) + } + } + } + AwsStreamEvent::MessageStop(stop) => { + if self.stop_reason.is_some() { + return Err(protocol_error( + "Bedrock emitted more than one message-stop event", + )); + } + self.stop_reason = Some(map_stop_reason(stop.stop_reason())); + Ok(Vec::new()) + } + AwsStreamEvent::Metadata(metadata) => { + let Some(usage) = metadata.usage() else { + return Ok(Vec::new()); + }; + Ok(vec![AgentEvent::UsageUpdated { + usage: Usage { + input_tokens: nonnegative_tokens(usage.input_tokens()), + output_tokens: nonnegative_tokens(usage.output_tokens()), + cached_input_tokens: nonnegative_tokens( + usage.cache_read_input_tokens().unwrap_or(0), + ), + cache_creation_input_tokens: nonnegative_tokens( + usage.cache_write_input_tokens().unwrap_or(0), + ), + }, + }]) + } + _ => Err(protocol_error("Bedrock emitted an unknown stream event")), + } + } + + fn finish(self) -> Result, AgentError> { + if !self.content_blocks.is_empty() { + return Err(protocol_error( + "Bedrock stream ended with incomplete content blocks", + )); + } + let reason = self + .stop_reason + .ok_or_else(|| protocol_error("Bedrock stream ended before the message-stop event"))?; + Ok(vec![AgentEvent::TurnStopped { reason }]) + } +} + +#[derive(Debug)] +enum PendingContentBlock { + Tool { + id: String, + name: String, + input: String, + }, + Reasoning { + text: String, + signature: Option, + }, +} + +fn map_stop_reason(reason: &AwsStopReason) -> StopReason { + match reason { + AwsStopReason::EndTurn | AwsStopReason::StopSequence | AwsStopReason::ToolUse => { + StopReason::Completed + } + AwsStopReason::MaxTokens => StopReason::MaxTokens, + AwsStopReason::ModelContextWindowExceeded => StopReason::ContextWindowExceeded, + AwsStopReason::ContentFiltered | AwsStopReason::GuardrailIntervened => StopReason::Refusal, + AwsStopReason::MalformedModelOutput | AwsStopReason::MalformedToolUse => { + StopReason::Other(reason.as_str().to_string()) + } + other => StopReason::Other(other.as_str().to_string()), + } +} + +fn nonnegative_tokens(value: i32) -> u64 { + u64::try_from(value).unwrap_or_default() +} + +fn stopped_before_stream(runtime_request_id: String) -> AgentEventStream { + Box::pin(futures::stream::iter([ + Ok(AgentEvent::TurnStarted { runtime_request_id }), + Ok(AgentEvent::TurnStopped { + reason: StopReason::Cancelled, + }), + ])) +} + +fn map_bedrock_error(error: impl std::fmt::Display + std::fmt::Debug) -> AgentError { + let display = error.to_string(); + let debug = format!("{error:?}"); + let message = if debug.len() > display.len() { + debug + } else { + display + }; + let normalized = message.to_ascii_lowercase(); + let kind = if normalized.contains("accessdenied") + || normalized.contains("access denied") + || normalized.contains("unauthorized") + || normalized.contains("credential") + { + AgentErrorKind::Authentication + } else if normalized.contains("throttl") || normalized.contains("rate limit") { + AgentErrorKind::RateLimited + } else if normalized.contains("context window") + || normalized.contains("too many tokens") + || normalized.contains("modelcontextwindowexceeded") + { + AgentErrorKind::ContextWindowExceeded + } else if normalized.contains("validation") + || normalized.contains("resource not found") + || normalized.contains("resourcenotfound") + { + AgentErrorKind::InvalidRequest + } else if normalized.contains("timeout") + || normalized.contains("dispatchfailure") + || normalized.contains("connection") + { + AgentErrorKind::Transport + } else { + AgentErrorKind::Provider + }; + let mut error = AgentError::new(kind, message); + error.recoverable = matches!( + kind, + AgentErrorKind::RateLimited | AgentErrorKind::Transport + ); + error +} + +fn protocol_error(message: impl Into) -> AgentError { + AgentError::new(AgentErrorKind::Protocol, message) +} + +#[cfg(test)] +#[path = "runtime_tests.rs"] +mod tests; diff --git a/app/src/ai/bedrock/runtime_tests.rs b/app/src/ai/bedrock/runtime_tests.rs new file mode 100644 index 00000000..7310d1fe --- /dev/null +++ b/app/src/ai/bedrock/runtime_tests.rs @@ -0,0 +1,350 @@ +use aws_sdk_bedrockruntime::types::{ + CacheTtl, ContentBlock, ContentBlockDelta, ContentBlockDeltaEvent, ContentBlockStart, + ContentBlockStartEvent, ContentBlockStopEvent, ConverseStreamMetadataEvent, + ConverseStreamOutput as AwsStreamEvent, MessageStopEvent, ReasoningContentBlockDelta, + StopReason as AwsStopReason, SystemContentBlock, TokenUsage, Tool, ToolUseBlockDelta, + ToolUseBlockStart, +}; +use galaxy_agent_core::{ + AgentErrorKind, AgentEvent, ConversationMessage, MessageContent, MessageRole, StopReason, + ToolDefinition, ToolEvent, TurnRequest, Usage, +}; +use serde_json::json; + +use super::*; + +fn turn_request() -> TurnRequest { + let mut request = TurnRequest::new( + "anthropic.claude-test", + vec![ + ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("first".to_string()), + }, + ConversationMessage { + role: MessageRole::Assistant, + content: MessageContent::Text("response".to_string()), + }, + ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("continue".to_string()), + }, + ], + ); + request.system_prompt = Some("system".to_string()); + request.tools = vec![ToolDefinition { + name: "read_files".to_string(), + description: "Read files".to_string(), + input_schema: json!({"type": "object"}), + }]; + request +} + +fn cache_ttls(converted: &ConvertedRequest) -> Vec> { + let mut ttls = Vec::new(); + for message in &converted.messages { + for block in message.content() { + if let ContentBlock::CachePoint(point) = block { + ttls.push(point.ttl().cloned()); + } + } + } + for block in &converted.system { + if let SystemContentBlock::CachePoint(point) = block { + ttls.push(point.ttl().cloned()); + } + } + if let Some(tool_config) = &converted.tool_config { + for tool in tool_config.tools() { + if let Tool::CachePoint(point) = tool { + ttls.push(point.ttl().cloned()); + } + } + } + ttls +} + +#[test] +fn one_turn_transport_preserves_disabled_default_and_one_hour_cache_modes() { + let disabled = convert_turn_request( + turn_request(), + Some(4096), + CachingConfig { + enabled: false, + extended_ttl_requested: false, + }, + ); + assert!(cache_ttls(&disabled).is_empty()); + + let default = convert_turn_request(turn_request(), Some(4096), CachingConfig::default()); + assert_eq!(cache_ttls(&default), vec![None, None, None]); + + let one_hour = convert_turn_request( + turn_request(), + Some(4096), + CachingConfig { + enabled: true, + extended_ttl_requested: true, + }, + ); + assert_eq!( + cache_ttls(&one_hour), + vec![ + Some(CacheTtl::OneHour), + Some(CacheTtl::OneHour), + Some(CacheTtl::OneHour), + ] + ); +} + +#[test] +fn one_turn_transport_prefers_request_output_limit() { + let mut request = turn_request(); + request.max_output_tokens = Some(8192); + let converted = convert_turn_request(request, Some(4096), CachingConfig::default()); + assert_eq!(converted.inference_config.max_tokens(), Some(8192)); +} + +fn tool_start(index: i32, id: &str, name: &str) -> AwsStreamEvent { + AwsStreamEvent::ContentBlockStart( + ContentBlockStartEvent::builder() + .content_block_index(index) + .start(ContentBlockStart::ToolUse( + ToolUseBlockStart::builder() + .tool_use_id(id) + .name(name) + .build() + .unwrap(), + )) + .build() + .unwrap(), + ) +} + +fn ordinary_start(index: i32) -> AwsStreamEvent { + AwsStreamEvent::ContentBlockStart( + ContentBlockStartEvent::builder() + .content_block_index(index) + .build() + .unwrap(), + ) +} + +fn content_delta(index: i32, delta: ContentBlockDelta) -> AwsStreamEvent { + AwsStreamEvent::ContentBlockDelta( + ContentBlockDeltaEvent::builder() + .content_block_index(index) + .delta(delta) + .build() + .unwrap(), + ) +} + +fn content_stop(index: i32) -> AwsStreamEvent { + AwsStreamEvent::ContentBlockStop( + ContentBlockStopEvent::builder() + .content_block_index(index) + .build() + .unwrap(), + ) +} + +fn message_stop(reason: AwsStopReason) -> AwsStreamEvent { + AwsStreamEvent::MessageStop( + MessageStopEvent::builder() + .stop_reason(reason) + .build() + .unwrap(), + ) +} + +fn metadata(usage: Usage) -> AwsStreamEvent { + AwsStreamEvent::Metadata( + ConverseStreamMetadataEvent::builder() + .usage( + TokenUsage::builder() + .input_tokens(usage.input_tokens as i32) + .output_tokens(usage.output_tokens as i32) + .total_tokens((usage.input_tokens + usage.output_tokens) as i32) + .cache_read_input_tokens(usage.cached_input_tokens as i32) + .cache_write_input_tokens(usage.cache_creation_input_tokens as i32) + .build() + .unwrap(), + ) + .build(), + ) +} + +#[test] +fn stream_translator_accepts_ordinary_content_block_starts() { + let mut translator = BedrockStreamTranslator::default(); + assert!(translator.translate(ordinary_start(0)).unwrap().is_empty()); + assert_eq!( + translator + .translate(content_delta( + 0, + ContentBlockDelta::Text("response".to_string()), + )) + .unwrap(), + vec![AgentEvent::TextDelta { + text: "response".to_string(), + }] + ); + assert!(translator.translate(content_stop(0)).unwrap().is_empty()); +} + +#[test] +fn stream_translator_correlates_tools_by_content_index() { + let mut translator = BedrockStreamTranslator::default(); + translator + .translate(tool_start(2, "call-2", "grep")) + .unwrap(); + translator + .translate(tool_start(1, "call-1", "read_files")) + .unwrap(); + translator + .translate(content_delta( + 1, + ContentBlockDelta::ToolUse( + ToolUseBlockDelta::builder() + .input("{\"files\":[\"Cargo.toml\"]}") + .build() + .unwrap(), + ), + )) + .unwrap(); + translator + .translate(content_delta( + 2, + ContentBlockDelta::ToolUse( + ToolUseBlockDelta::builder() + .input("{\"query\":\"ProviderRun\"}") + .build() + .unwrap(), + ), + )) + .unwrap(); + + let first = translator.translate(content_stop(1)).unwrap(); + let second = translator.translate(content_stop(2)).unwrap(); + assert!(matches!( + first.as_slice(), + [AgentEvent::Tool { + event: ToolEvent::Proposed { call } + }] if call.id == "call-1" + && call.name == "read_files" + && call.arguments == json!({"files": ["Cargo.toml"]}) + )); + assert!(matches!( + second.as_slice(), + [AgentEvent::Tool { + event: ToolEvent::Proposed { call } + }] if call.id == "call-2" + && call.name == "grep" + && call.arguments == json!({"query": "ProviderRun"}) + )); +} + +#[test] +fn stream_translator_defers_stop_until_usage_metadata_arrives() { + let mut translator = BedrockStreamTranslator::default(); + assert!(translator + .translate(message_stop(AwsStopReason::EndTurn)) + .unwrap() + .is_empty()); + + let expected_usage = Usage { + input_tokens: 10, + output_tokens: 4, + cached_input_tokens: 7, + cache_creation_input_tokens: 3, + }; + assert_eq!( + translator + .translate(metadata(expected_usage.clone())) + .unwrap(), + vec![AgentEvent::UsageUpdated { + usage: expected_usage, + }] + ); + assert_eq!( + translator.finish().unwrap(), + vec![AgentEvent::TurnStopped { + reason: StopReason::Completed, + }] + ); +} + +#[test] +fn stream_translator_preserves_reasoning_text_and_signature() { + let mut translator = BedrockStreamTranslator::default(); + assert_eq!( + translator + .translate(content_delta( + 0, + ContentBlockDelta::ReasoningContent(ReasoningContentBlockDelta::Text( + "inspect".to_string(), + )), + )) + .unwrap(), + vec![AgentEvent::ReasoningDelta { + text: "inspect".to_string(), + }] + ); + translator + .translate(content_delta( + 0, + ContentBlockDelta::ReasoningContent(ReasoningContentBlockDelta::Signature( + "signature".to_string(), + )), + )) + .unwrap(); + assert_eq!( + translator.translate(content_stop(0)).unwrap(), + vec![AgentEvent::ReasoningCompleted { + text: "inspect".to_string(), + signature: Some("signature".to_string()), + }] + ); +} + +#[test] +fn stream_translator_rejects_invalid_tool_json() { + let mut translator = BedrockStreamTranslator::default(); + translator + .translate(tool_start(0, "call", "read_files")) + .unwrap(); + translator + .translate(content_delta( + 0, + ContentBlockDelta::ToolUse( + ToolUseBlockDelta::builder() + .input("not-json") + .build() + .unwrap(), + ), + )) + .unwrap(); + let error = translator.translate(content_stop(0)).unwrap_err(); + assert_eq!(error.kind, AgentErrorKind::Protocol); +} + +#[test] +fn bedrock_stop_reasons_map_to_domain_reasons() { + assert_eq!( + map_stop_reason(&AwsStopReason::ToolUse), + StopReason::Completed + ); + assert_eq!( + map_stop_reason(&AwsStopReason::MaxTokens), + StopReason::MaxTokens + ); + assert_eq!( + map_stop_reason(&AwsStopReason::ModelContextWindowExceeded), + StopReason::ContextWindowExceeded + ); + assert_eq!( + map_stop_reason(&AwsStopReason::GuardrailIntervened), + StopReason::Refusal + ); +} diff --git a/app/src/ai/bedrock/translator.rs b/app/src/ai/bedrock/translator.rs deleted file mode 100644 index 07dd9c92..00000000 --- a/app/src/ai/bedrock/translator.rs +++ /dev/null @@ -1,217 +0,0 @@ -#![allow(dead_code)] - -use std::sync::{Arc, Mutex}; - -use warp_multi_agent_api as api; - -use crate::ai::agent::api::LegacyResponseStream; -use crate::ai::bedrock::client::{BedrockClient, BedrockClientConfig, BedrockError}; -use crate::ai::bedrock::convert::ConversationMessage; -use crate::ai::bedrock::diagnostic::BedrockDiagnosticLogger; -use crate::ai::bedrock::request_translator; -use crate::ai::provider::types::flatten_tool_history_for_no_tools_turn; - -pub struct TranslatorRequest { - pub config: BedrockClientConfig, - pub model_id: String, - pub root_task_id: Option, - pub bedrock_message_history: Vec, - pub bedrock_tool_result_archive: Vec, - pub bedrock_progressive_summary: Option, - pub bedrock_messages_sent: Arc>>, - /// Global rules (name, content) from the local CloudModel. - pub global_rules: Vec<(String, String)>, - /// Whether the native input should be emitted as a transcript-visible user query. - pub emit_user_query_message: bool, -} - -pub async fn execute( - params: TranslatorRequest, - request: &mut api::Request, -) -> Result { - let config = params.config.with_external_fallbacks(); - let cross_region_inference = config.cross_region_inference; - let bedrock = BedrockClient::from_config(config).await?; - - let task_id = params.root_task_id.unwrap_or_else(|| { - request - .task_context - .as_ref() - .and_then(|tc| tc.tasks.first()) - .map(|t| t.id.clone()) - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()) - }); - - let needs_create_task = request - .task_context - .as_ref() - .map(|tc| tc.tasks.is_empty()) - .unwrap_or(true); - - // Use the model from params (selected in UI or defaulted from ANTHROPIC_MODEL) - let mut model_id = params.model_id; - if model_id.is_empty() || model_id == "auto" { - // Fall back to default if nothing is set - model_id = "us.anthropic.claude-opus-4-6[1m]".to_string(); - } - - log::info!("[bedrock] Translator: model={model_id}, task_id={task_id}, needs_create_task={needs_create_task}"); - - let diagnostic_logger = - BedrockDiagnosticLogger::try_new(&model_id, "", "", &task_id).map(Arc::new); - - if let Some(ref logger) = diagnostic_logger { - logger.log_protobuf_input(request); - } - - request_translator::inject_input_messages_into_task(request); - - let new_input_messages = request_translator::extract_new_input_messages(request); - let new_input_count = new_input_messages.len(); - - let mut messages = Vec::new(); - - // Prepend progressive summary as the first message pair if present - if let Some(ref summary) = params.bedrock_progressive_summary { - use crate::ai::bedrock::convert::{MessageContent, MessageRole}; - messages.push(ConversationMessage { - role: MessageRole::User, - content: MessageContent::Text(format!( - "\n{}\n\n\n\ - The above summarizes earlier conversation history. The detailed messages below are the most recent exchanges.", - summary - )), - }); - messages.push(ConversationMessage { - role: MessageRole::Assistant, - content: MessageContent::Text( - "Understood, I have the prior context. Continuing with the recent conversation." - .to_string(), - ), - }); - } - - let history_len = params.bedrock_message_history.len(); - messages.extend(params.bedrock_message_history); - - if !new_input_messages.is_empty() { - log::info!( - "[bedrock] Appending {} new input messages to history of {}", - new_input_messages.len(), - history_len - ); - messages.extend(new_input_messages); - } - - for message in &mut messages { - message.truncate_tool_results_for_provider_request(); - } - - request_translator::sanitize_messages_for_bedrock(&mut messages); - - let system_prompt = request_translator::extract_system_prompt(request, ¶ms.global_rules); - let tools = request_translator::extract_tools(request); - if tools_are_inline_only(&tools) { - flatten_tool_history_for_no_tools_turn(&mut messages); - } - - log::info!( - "[bedrock] Sending {} messages, system_prompt={}, progressive_summary={}, tools={}", - messages.len(), - system_prompt.is_some(), - params.bedrock_progressive_summary.is_some(), - tools.len() - ); - - for (i, msg) in messages.iter().enumerate() { - let content_desc = describe_message_content(&msg.content); - log::info!( - "[bedrock] msg[{}]: role={:?}, content={}", - i, - msg.role, - content_desc - ); - } - - let user_query_text = params - .emit_user_query_message - .then(|| request_translator::extract_user_query_text(request)) - .flatten(); - - let stream = bedrock - .converse_stream( - &model_id, - &task_id, - needs_create_task, - messages.clone(), - system_prompt, - None, // progressive summary is in messages array, not system prompt - tools, - 64000, - None, - cross_region_inference, - user_query_text, - diagnostic_logger, - params.bedrock_messages_sent.clone(), - params.bedrock_tool_result_archive, - ) - .await?; - - if let Ok(mut sent) = params.bedrock_messages_sent.lock() { - // Only persist the actual conversation history (history + new inputs), not the - // ephemeral prepended summary pair, so we don't duplicate the summary on every - // subsequent write-back. The summary is prepended at request time each turn. - let persistent_count = history_len + new_input_count; - if persistent_count > 0 && messages.len() >= persistent_count { - *sent = messages.split_off(messages.len() - persistent_count); - } else { - *sent = messages; - } - } - - Ok(stream) -} - -fn tools_are_inline_only(tools: &[crate::ai::bedrock::convert::ToolDefinition]) -> bool { - tools.iter().all(|tool| tool.name == "recall_tool_history") -} - -fn describe_message_content(content: &crate::ai::bedrock::convert::MessageContent) -> String { - use crate::ai::bedrock::convert::{ContentPart, MessageContent}; - match content { - MessageContent::Text(t) => format!("Text({}chars)", t.len()), - MessageContent::ToolUse { - tool_use_id, name, .. - } => format!("ToolUse(name={}, id={})", name, tool_use_id), - MessageContent::ToolResult { - tool_use_id, - is_error, - .. - } => format!("ToolResult(id={}, is_error={})", tool_use_id, is_error), - MessageContent::MultiPart(parts) => { - let part_descs: Vec = parts - .iter() - .map(|p| match p { - ContentPart::Text(t) => format!("Text({})", t.len()), - ContentPart::Reasoning { text, signature } => { - format!( - "Reasoning({}chars,signed={})", - text.len(), - signature.is_some() - ) - } - ContentPart::Image { data, mime_type } => { - format!("Image({mime_type},{}bytes)", data.len()) - } - ContentPart::ToolUse { - name, tool_use_id, .. - } => format!("ToolUse({},{})", name, tool_use_id), - ContentPart::ToolResult { tool_use_id, .. } => { - format!("ToolResult({})", tool_use_id) - } - }) - .collect(); - format!("MultiPart[{}]", part_descs.join(", ")) - } - } -} diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index 6a5e70be..5f3bb970 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -35,7 +35,8 @@ pub use execute::{ }; use futures::future::{join_all, BoxFuture}; use galaxy_agent_core::{ - PermissionDecision, PermissionKind, PermissionRequest, ToolEvent, ToolResult, ToolResultStatus, + ExternalWorkId, PendingToolBatch, PermissionDecision, PermissionKind, PermissionRequest, + ToolEvent, ToolResult, ToolResultStatus, }; use galaxyui::{AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; use itertools::Itertools; @@ -71,6 +72,7 @@ use crate::ai::document::ai_document_model::AIDocumentModel; use crate::ai::get_relevant_files::controller::GetRelevantFilesController; #[cfg(not(target_family = "wasm"))] use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord}; +use crate::ai::runtime::ProviderToolExecutionRef; use crate::terminal::model::session::active_session::ActiveSession; use crate::terminal::model_events::ModelEventDispatcher; use crate::terminal::TerminalModel; @@ -172,6 +174,22 @@ struct RunningActions { action_ids: Vec, } +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub(super) enum ProviderActionQueueError { + #[error("provider action set mismatch: expected {expected:?}, received {received:?}")] + ActionSetMismatch { + expected: Vec, + received: Vec, + }, + #[error("provider action '{call_id}' is already correlated to active work")] + ExistingCorrelation { call_id: String }, +} + +type ProviderActionCorrelation = ( + (AIConversationId, AIAgentActionId), + ProviderToolExecutionRef, +); + impl RunningActions { fn new(phase: RunningActionPhase, action_id: AIAgentActionId) -> Self { Self { @@ -268,6 +286,13 @@ fn permission_kind_for_action(action: &AIAgentActionType) -> PermissionKind { } } +fn sort_action_results_by_order( + results: &mut [Arc], + action_order: &HashMap, +) { + results.sort_by_key(|result| action_order.get(&result.id).copied().unwrap_or(usize::MAX)); +} + fn domain_tool_result(action_result: &AIAgentActionResult, permission_denied: bool) -> ToolResult { let status = if permission_denied { ToolResultStatus::Denied @@ -626,6 +651,10 @@ pub struct BlocklistAIActionModel { /// than reconstructing them from the legacy request protobuf. finished_tool_results: HashMap>, + /// Provider-owned action results retained until their exact tool batch is fully committed. + provider_finished_action_results: + HashMap<(AIConversationId, ExternalWorkId), Vec>>, + /// Original order for the current batch of actions. /// /// We maintain this so that even though we might process actions in parallel, @@ -635,6 +664,10 @@ pub struct BlocklistAIActionModel { /// Permission-card rejections that still need a correlated completion event. denied_permissions: HashSet<(AIConversationId, AIAgentActionId)>, + /// Durable provider work identity for actions owned by an active provider run. + provider_tool_executions: + HashMap<(AIConversationId, AIAgentActionId), ProviderToolExecutionRef>, + /// Past actions and their corresponding statuses from previous AI exchanges. past_action_results: HashMap>, @@ -669,10 +702,18 @@ impl BlocklistAIActionModel { ) }); ctx.subscribe_to_model(&executor, move |me, _, event, ctx| match event { - BlocklistAIActionExecutorEvent::ExecutingAction { action_id } => { - ctx.emit(BlocklistAIActionEvent::ExecutingAction(action_id.clone())); + BlocklistAIActionExecutorEvent::ExecutingAction { + action_id, + conversation_id, + } => { + let execution_ref = me.provider_tool_execution_ref(*conversation_id, action_id); + ctx.emit(BlocklistAIActionEvent::ExecutingAction { + action_id: action_id.clone(), + execution_ref: execution_ref.clone(), + }); ctx.emit(BlocklistAIActionEvent::ToolLifecycle { action_id: action_id.clone(), + execution_ref, event: ToolEvent::Started { call_id: action_id.to_string(), }, @@ -710,11 +751,13 @@ impl BlocklistAIActionModel { pending_actions: Default::default(), finished_action_results: Default::default(), finished_tool_results: Default::default(), + provider_finished_action_results: Default::default(), executor, past_action_results: HashMap::new(), running_actions: Default::default(), action_order: Default::default(), denied_permissions: Default::default(), + provider_tool_executions: Default::default(), terminal_view_id, pending_preprocessed_actions: Default::default(), is_view_only: false, @@ -752,7 +795,10 @@ impl BlocklistAIActionModel { action_id.clone(), RunningActionPhase::Serial, ); - ctx.emit(BlocklistAIActionEvent::ExecutingAction(action_id.clone())); + ctx.emit(BlocklistAIActionEvent::ExecutingAction { + action_id: action_id.clone(), + execution_ref: self.provider_tool_execution_ref(conversation_id, action_id), + }); } /// Returns true if the action model is operating in view-only mode (used for shared-session viewers). @@ -942,9 +988,14 @@ impl BlocklistAIActionModel { fn sort_finished_results(&mut self, conversation_id: AIConversationId) { if let Some(action_order) = self.action_order.get(&conversation_id) { if let Some(finished_results) = self.finished_action_results.get_mut(&conversation_id) { - finished_results.sort_by_key(|result| { - action_order.get(&result.id).copied().unwrap_or(usize::MAX) - }); + sort_action_results_by_order(finished_results, action_order); + } + for ((finished_conversation_id, _), finished_results) in + &mut self.provider_finished_action_results + { + if *finished_conversation_id == conversation_id { + sort_action_results_by_order(finished_results, action_order); + } } if let Some(tool_results) = self.finished_tool_results.get_mut(&conversation_id) { let tool_order = action_order @@ -1113,6 +1164,7 @@ impl BlocklistAIActionModel { // Search through all conversations' finished action results self.finished_action_results .values() + .chain(self.provider_finished_action_results.values()) .flat_map(|results| results.iter()) .find(|result| &result.id == id) .or_else(|| self.past_action_results.get(id)) @@ -1288,11 +1340,14 @@ impl BlocklistAIActionModel { "reason": format!("{reason:?}"), }), ); - ctx.emit(BlocklistAIActionEvent::ActionBlockedOnUserConfirmation( - action.id.clone(), - )); + let execution_ref = self.provider_tool_execution_ref(conversation_id, &action.id); + ctx.emit(BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { + action_id: action.id.clone(), + execution_ref: execution_ref.clone(), + }); ctx.emit(BlocklistAIActionEvent::ToolLifecycle { action_id: action.id.clone(), + execution_ref, event: ToolEvent::PermissionRequested { request: PermissionRequest { id: permission_request_id(&action.id), @@ -1389,6 +1444,7 @@ impl BlocklistAIActionModel { ); ctx.emit(BlocklistAIActionEvent::ToolLifecycle { action_id: action_id.clone(), + execution_ref: self.provider_tool_execution_ref(conversation_id, &action_id), event: ToolEvent::PermissionResolved { request_id: permission_request_id(&action_id), call_id: action_id.to_string(), @@ -1461,6 +1517,40 @@ impl BlocklistAIActionModel { }) } + fn provider_tool_execution_ref( + &self, + conversation_id: AIConversationId, + action_id: &AIAgentActionId, + ) -> Option { + self.provider_tool_executions + .get(&(conversation_id, action_id.clone())) + .cloned() + } + + #[allow(dead_code)] + pub(super) fn queue_provider_actions( + &mut self, + actions: Vec, + conversation_id: AIConversationId, + batch: &PendingToolBatch, + ctx: &mut ModelContext, + ) -> 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())) + { + return Err(ProviderActionQueueError::ExistingCorrelation { + call_id: action_id.to_string(), + }); + } + } + self.provider_tool_executions.extend(refs); + self.queue_actions(actions, conversation_id, ctx); + Ok(()) + } + /// Queues the `actions` in the given iterator for the given conversation, /// to be dispatched in the order in which they appear in the iterator. pub(super) fn queue_actions( @@ -1553,11 +1643,18 @@ impl BlocklistAIActionModel { // as otherwise tools get stuck in a pending state on the viewer's side of things. This check // must be scoped to the current conversation as some providers generate tool call IDs that // only unique within a conversation. - if self + let has_finished_result = self .finished_action_results .get(&conversation_id) - .is_some_and(|results| results.iter().any(|r| r.id == action_id)) - { + .is_some_and(|results| results.iter().any(|result| result.id == action_id)) + || self + .provider_finished_action_results + .iter() + .filter(|((finished_conversation_id, _), _)| { + *finished_conversation_id == conversation_id + }) + .any(|(_, results)| results.iter().any(|result| result.id == action_id)); + if has_finished_result { continue; } @@ -1577,7 +1674,10 @@ impl BlocklistAIActionModel { .entry(conversation_id) .or_default() .push_back(action); - ctx.emit(BlocklistAIActionEvent::QueuedAction(action_id)); + ctx.emit(BlocklistAIActionEvent::QueuedAction { + execution_ref: self.provider_tool_execution_ref(conversation_id, &action_id), + action_id, + }); } self.try_to_execute_available_actions(conversation_id, ctx); } @@ -1681,6 +1781,14 @@ impl BlocklistAIActionModel { executor.cancel_all_running_async_actions_for_conversation(conversation_id, reason, ctx) }); + if let Some(preprocessing) = self.pending_preprocessed_actions.remove(&conversation_id) { + self.provider_tool_executions + .retain(|(correlated_conversation_id, action_id), _| { + *correlated_conversation_id != conversation_id + || !preprocessing.contains(action_id) + }); + } + let Some(actions_to_cancel) = self.pending_actions.get_mut(&conversation_id) else { return; }; @@ -1752,10 +1860,14 @@ impl BlocklistAIActionModel { ); ctx.emit(BlocklistAIActionEvent::ToolLifecycle { action_id: pending_action.id.clone(), + execution_ref: self + .provider_tool_execution_ref(conversation_id, &pending_action.id), event: ToolEvent::PermissionResolved { request_id: permission_request_id(&pending_action.id), call_id: pending_action.id.to_string(), - decision: PermissionDecision::Denied { reason: None }, + decision: PermissionDecision::Denied { + reason: Some("Permission denied by the user.".to_string()), + }, }, }); } @@ -1817,11 +1929,40 @@ impl BlocklistAIActionModel { .unwrap_or_default() } + pub(super) fn provider_finished_action_results( + &self, + conversation_id: AIConversationId, + work_id: &ExternalWorkId, + ) -> Vec> { + self.provider_finished_action_results + .get(&(conversation_id, work_id.clone())) + .cloned() + .unwrap_or_default() + } + + pub(super) fn archive_provider_finished_action_results( + &mut self, + conversation_id: AIConversationId, + work_id: &ExternalWorkId, + ) { + let results = self + .provider_finished_action_results + .remove(&(conversation_id, work_id.clone())) + .unwrap_or_default(); + for result in results { + self.past_action_results.insert(result.id.clone(), result); + } + } + /// Clears finished action results for a conversation. Used when reverting. pub(super) fn clear_finished_action_results(&mut self, conversation_id: AIConversationId) { self.action_order.remove(&conversation_id); self.finished_action_results.remove(&conversation_id); self.finished_tool_results.remove(&conversation_id); + self.provider_finished_action_results + .retain(|(finished_conversation_id, _), _| { + *finished_conversation_id != conversation_id + }); } #[cfg(test)] @@ -1921,14 +2062,19 @@ impl BlocklistAIActionModel { } } + let execution_ref = self + .provider_tool_executions + .remove(&(conversation_id, action_result.id.clone())); let permission_denied = self .denied_permissions .remove(&(conversation_id, action_result.id.clone())); let tool_result = domain_tool_result(&action_result, permission_denied); - self.finished_tool_results - .entry(conversation_id) - .or_default() - .push(tool_result.clone()); + if execution_ref.is_none() { + self.finished_tool_results + .entry(conversation_id) + .or_default() + .push(tool_result.clone()); + } #[cfg(not(target_family = "wasm"))] log_tool_event( ctx, @@ -1955,17 +2101,29 @@ impl BlocklistAIActionModel { "error": action_result_error_summary(&action_result.result), }), ); - ctx.emit(BlocklistAIActionEvent::ToolLifecycle { - action_id: action_result.id.clone(), - event: ToolEvent::Completed { - result: tool_result, - }, - }); + // Permission denial completes provider-owned calls when the permission decision is + // applied, so emitting a second correlated completion would violate exactly-once delivery. + if execution_ref.is_none() || !permission_denied { + ctx.emit(BlocklistAIActionEvent::ToolLifecycle { + action_id: action_result.id.clone(), + execution_ref: execution_ref.clone(), + event: ToolEvent::Completed { + result: tool_result, + }, + }); + } - self.finished_action_results - .entry(conversation_id) - .or_default() - .push(action_result); + if let Some(execution_ref) = &execution_ref { + self.provider_finished_action_results + .entry((conversation_id, execution_ref.work_id())) + .or_default() + .push(action_result); + } else { + self.finished_action_results + .entry(conversation_id) + .or_default() + .push(action_result); + } if self .running_actions @@ -1986,6 +2144,7 @@ impl BlocklistAIActionModel { action_id, conversation_id, cancellation_reason, + execution_ref: execution_ref.clone(), }); if self @@ -1999,8 +2158,10 @@ impl BlocklistAIActionModel { // completion (no cancellation reason) is resolved by the controller's // follow-up handling. Stamping here for any of those would clobber the real // status and message. - if cancellation_reason - .is_some_and(|r| matches!(r.conversation_outcome(), CancellationOutcome::Cancelled)) + if execution_ref.is_none() + && cancellation_reason.is_some_and(|r| { + matches!(r.conversation_outcome(), CancellationOutcome::Cancelled) + }) { BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { // Treat action result as authoritative for determining status. @@ -2110,23 +2271,63 @@ impl BlocklistAIActionModel { } } +fn provider_action_correlations( + actions: &[AIAgentAction], + conversation_id: AIConversationId, + batch: &PendingToolBatch, +) -> Result, ProviderActionQueueError> { + let expected = batch.unresolved_call_ids(); + let received = actions + .iter() + .map(|action| action.id.to_string()) + .collect::>(); + if expected != received { + return Err(ProviderActionQueueError::ActionSetMismatch { expected, received }); + } + + Ok(actions + .iter() + .map(|action| { + ( + (conversation_id, action.id.clone()), + ProviderToolExecutionRef::new( + conversation_id, + &batch.work_id, + action.id.to_string(), + ), + ) + }) + .collect()) +} + #[derive(Debug, Clone)] pub enum BlocklistAIActionEvent { /// Emitted when the action with the given ID is enqueued for execution. - QueuedAction(AIAgentActionId), + QueuedAction { + action_id: AIAgentActionId, + execution_ref: Option, + }, /// Emitted when the action with the given ID requires user confirmation to execute. - ActionBlockedOnUserConfirmation(AIAgentActionId), + ActionBlockedOnUserConfirmation { + action_id: AIAgentActionId, + execution_ref: Option, + }, /// Emitted when the action with the given ID begins execution. - ExecutingAction(AIAgentActionId), + ExecutingAction { + action_id: AIAgentActionId, + execution_ref: Option, + }, /// Emitted when the action with the given ID has finished. FinishedAction { action_id: AIAgentActionId, conversation_id: AIConversationId, cancellation_reason: Option, + execution_ref: Option, }, /// Provider-neutral permission and execution lifecycle event for runtime consumers. ToolLifecycle { action_id: AIAgentActionId, + execution_ref: Option, event: ToolEvent, }, InitProject(AIAgentActionId), @@ -2142,10 +2343,10 @@ pub enum BlocklistAIActionEvent { impl BlocklistAIActionEvent { pub fn action_id(&self) -> &AIAgentActionId { match self { - BlocklistAIActionEvent::QueuedAction(action_id) => action_id, - BlocklistAIActionEvent::ActionBlockedOnUserConfirmation(action_id) => action_id, - BlocklistAIActionEvent::ExecutingAction(action_id) => action_id, - BlocklistAIActionEvent::FinishedAction { action_id, .. } => action_id, + BlocklistAIActionEvent::QueuedAction { action_id, .. } + | BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { action_id, .. } + | BlocklistAIActionEvent::ExecutingAction { action_id, .. } + | BlocklistAIActionEvent::FinishedAction { action_id, .. } => action_id, BlocklistAIActionEvent::ToolLifecycle { action_id, .. } => action_id, BlocklistAIActionEvent::InitProject(action_id) => action_id, BlocklistAIActionEvent::ToggleCodeReview(action_id) => action_id, diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index 7f90c954..ac34e26b 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -688,6 +688,7 @@ impl BlocklistAIActionExecutor { ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction { action_id: action_id.clone(), + conversation_id, }); ctx.emit(BlocklistAIActionExecutorEvent::FinishedAction { result: Arc::new(AIAgentActionResult { @@ -904,6 +905,7 @@ impl BlocklistAIActionExecutor { ); ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction { action_id: action_id.clone(), + conversation_id, }); log::info!("[tool-debug] try_to_execute_action: spawning ASYNC execution for action_id={:?}", action_id); ctx.spawn(execute_future, move |me, result, ctx| { @@ -932,6 +934,7 @@ impl BlocklistAIActionExecutor { AnyActionExecution::Sync(action_result) => { ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction { action_id: action_id.clone(), + conversation_id, }); ctx.emit(BlocklistAIActionExecutorEvent::FinishedAction { result: Arc::new(AIAgentActionResult { @@ -1140,9 +1143,10 @@ impl Entity for BlocklistAIActionExecutor { } pub enum BlocklistAIActionExecutorEvent { - /// Emitted when an action is execution starts. + /// Emitted when an action begins execution. ExecutingAction { action_id: AIAgentActionId, + conversation_id: AIConversationId, }, /// Emitted when an action has finished. diff --git a/app/src/ai/blocklist/action_model_tests.rs b/app/src/ai/blocklist/action_model_tests.rs index e8a8346e..469f8ddb 100644 --- a/app/src/ai/blocklist/action_model_tests.rs +++ b/app/src/ai/blocklist/action_model_tests.rs @@ -4,7 +4,8 @@ use std::sync::Arc; use super::*; use crate::ai::agent::task::TaskId; use crate::ai::agent::{ - AIAgentActionResultType, AnyFileContent, FileContext, GrepResult, ReadFilesResult, + AIAgentAction, AIAgentActionResultType, AIAgentActionType, AnyFileContent, FileContext, + GrepResult, ReadFilesResult, }; fn make_action_result(id: &str) -> Arc { @@ -23,6 +24,36 @@ fn action_result(id: &str, result: AIAgentActionResultType) -> AIAgentActionResu } } +fn action(id: &str) -> AIAgentAction { + AIAgentAction { + id: AIAgentActionId::from(id.to_string()), + action: AIAgentActionType::InitProject, + task_id: TaskId::new("task".to_string()), + requires_result: true, + tool_name: Some("init_project".to_string()), + } +} + +fn pending_tool_batch(call_ids: &[&str]) -> PendingToolBatch { + PendingToolBatch { + work_id: galaxy_agent_core::ExternalWorkId { + run_id: galaxy_agent_core::ProviderRunId::new("run"), + epoch: galaxy_agent_core::RunEpoch::new(7), + }, + calls: call_ids + .iter() + .map(|call_id| galaxy_agent_core::PendingToolCall { + call: galaxy_agent_core::ToolCall { + id: (*call_id).to_string(), + name: "init_project".to_string(), + arguments: serde_json::json!({}), + }, + state: galaxy_agent_core::PendingToolCallState::Proposed, + }) + .collect(), + } +} + fn count_startable_actions_for_pass(phases: &[(RunningActionPhase, bool)]) -> usize { let mut current_phase = None; let mut count = 0; @@ -45,6 +76,34 @@ fn count_startable_actions_for_pass(phases: &[(RunningActionPhase, bool)]) -> us count } +#[test] +fn provider_action_correlations_require_the_exact_unresolved_batch_order() { + let conversation_id = AIConversationId::new(); + let batch = pending_tool_batch(&["first", "second"]); + let actions = vec![action("first"), action("second")]; + + let correlations = provider_action_correlations(&actions, conversation_id, &batch).unwrap(); + assert_eq!(correlations.len(), 2); + assert_eq!(correlations[0].0, (conversation_id, actions[0].id.clone())); + assert_eq!(correlations[0].1.run_id, batch.work_id.run_id); + assert_eq!(correlations[0].1.epoch, batch.work_id.epoch); + assert_eq!(correlations[0].1.call_id, "first"); + + let error = provider_action_correlations( + &[action("second"), action("first")], + conversation_id, + &batch, + ) + .unwrap_err(); + assert_eq!( + error, + ProviderActionQueueError::ActionSetMismatch { + expected: vec!["first".to_string(), "second".to_string()], + received: vec!["second".to_string(), "first".to_string()], + } + ); +} + #[test] fn parallel_phase_only_admits_matching_autoexecutable_actions() { let phase = @@ -94,8 +153,7 @@ fn finished_results_stay_in_original_action_order() { make_action_result("second"), ]; - finished_results - .sort_by_key(|result| action_order.get(&result.id).copied().unwrap_or(usize::MAX)); + sort_action_results_by_order(&mut finished_results, &action_order); assert_eq!( finished_results[0].id, diff --git a/app/src/ai/blocklist/block.rs b/app/src/ai/blocklist/block.rs index 2b49f68f..aacfdaef 100644 --- a/app/src/ai/blocklist/block.rs +++ b/app/src/ai/blocklist/block.rs @@ -4723,7 +4723,7 @@ impl AIBlock { } match event { - BlocklistAIActionEvent::ExecutingAction(..) => { + BlocklistAIActionEvent::ExecutingAction { .. } => { match &me.autonomy_setting_speedbump { AutonomySettingSpeedbump::ShouldShowForAutoexecutingReadonlyCommands { action_id: speedbump_action_id, @@ -4793,7 +4793,7 @@ impl AIBlock { _ => {} } } - BlocklistAIActionEvent::ActionBlockedOnUserConfirmation(..) => { + BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { .. } => { ctx.emit(AIBlockEvent::ActionBlockedOnUserConfirmation); } BlocklistAIActionEvent::FinishedAction { action_id, .. } => { @@ -4950,7 +4950,7 @@ impl AIBlock { } ctx.notify(); } - BlocklistAIActionEvent::QueuedAction(action_id) => { + BlocklistAIActionEvent::QueuedAction { action_id, .. } => { // Update search codebase view status when action is queued if let Some(view) = me.search_codebase_view.get(action_id) { view.update(ctx, |view, ctx| { diff --git a/app/src/ai/blocklist/block/cli_controller.rs b/app/src/ai/blocklist/block/cli_controller.rs index 67dc39f1..d42579ef 100644 --- a/app/src/ai/blocklist/block/cli_controller.rs +++ b/app/src/ai/blocklist/block/cli_controller.rs @@ -15,6 +15,7 @@ use crate::ai::agent::{ }; use crate::ai::blocklist::agent_view::{AgentViewController, AgentViewEntryOrigin}; use crate::ai::blocklist::context_model::block_context_from_terminal_model; +use crate::ai::blocklist::controller::PendingProviderCommandCompletion; use crate::ai::blocklist::{ BlocklistAIActionEvent, BlocklistAIActionModel, BlocklistAIController, BlocklistAIControllerEvent, BlocklistAIHistoryEvent, @@ -56,6 +57,7 @@ struct PendingCommandCompletion { initial_requested_command_action_id: Option, prompt: String, completed_command: RunningCommand, + exit_code: i32, final_turn_started: bool, } @@ -169,7 +171,7 @@ impl CLISubagentController { }); ctx.subscribe_to_model(action_model, |me, _, event, ctx| match event { - BlocklistAIActionEvent::ActionBlockedOnUserConfirmation(_) => { + BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { .. } => { let mut terminal_model = me.terminal_model.lock(); let active_block = terminal_model.block_list_mut().active_block_mut(); active_block.update_is_agent_blocked(true); @@ -181,7 +183,7 @@ impl CLISubagentController { agent_has_control: active_block.is_agent_in_control(), }); } - BlocklistAIActionEvent::ExecutingAction(..) => { + BlocklistAIActionEvent::ExecutingAction { .. } => { let mut terminal_model = me.terminal_model.lock(); let active_block = terminal_model.block_list_mut().active_block_mut(); active_block.update_is_agent_blocked(false); @@ -303,6 +305,7 @@ impl CLISubagentController { requested_command_id: requested_command_action_id.clone(), is_alt_screen_active: false, }, + exit_code, final_turn_started: false, }) } @@ -319,16 +322,40 @@ impl CLISubagentController { }; drop(terminal_model); - let Some(has_last_snapshot) = me + let provider_consumed_completion = completion.as_ref().is_some_and(|completion| { + me.controller.update(ctx, |controller, ctx| { + controller.accept_provider_command_completion( + completion.conversation_id, + PendingProviderCommandCompletion::new( + completion.completed_command.block_id.clone(), + completion.initial_requested_command_action_id.clone(), + completion.completed_command.command.clone(), + completion.completed_command.grid_contents.clone(), + completion.exit_code, + ), + ctx, + ) + }) + }); + let has_last_snapshot = me .active_subagents_by_block .get(&block_id) - .map(|state| state.last_snapshot_at.is_some()) - else { - return; - }; + .is_some_and(|state| state.last_snapshot_at.is_some()); if has_last_snapshot { ctx.emit(CLISubagentEvent::UpdatedLastSnapshot); } + if provider_consumed_completion { + me.finish_subagent( + &block_id, + conversation_id, + requested_command_action_id, + ctx, + ); + return; + } + if !me.active_subagents_by_block.contains_key(&block_id) { + return; + } // A Stop takeover intentionally cancels the subagent. The command may still // finish later, but that completion must not start a new assessment turn. Also @@ -483,7 +510,11 @@ impl CLISubagentController { if self .controller .as_ref(ctx) - .has_active_stream_for_conversation(conversation_id, ctx) + .has_active_provider_run(conversation_id) + || self + .controller + .as_ref(ctx) + .has_active_stream_for_conversation(conversation_id, ctx) || self .action_model .as_ref(ctx) @@ -737,13 +768,7 @@ impl CLISubagentController { .collect() }; self.controller.update(ctx, |controller, ctx| { - controller.resume_conversation( - conversation_id, - /*can_attempt_resume_on_error*/ true, - /*is_auto_resume_after_error*/ false, - resume_context, - ctx, - ); + controller.resume_conversation(conversation_id, resume_context, ctx); }); } } diff --git a/app/src/ai/blocklist/block/status_bar.rs b/app/src/ai/blocklist/block/status_bar.rs index 3227609d..a157f276 100644 --- a/app/src/ai/blocklist/block/status_bar.rs +++ b/app/src/ai/blocklist/block/status_bar.rs @@ -329,7 +329,7 @@ impl BlocklistAIStatusBar { ); ctx.subscribe_to_model(&action_model, |_, _, event, ctx| match event { - BlocklistAIActionEvent::ExecutingAction(..) + BlocklistAIActionEvent::ExecutingAction { .. } | BlocklistAIActionEvent::FinishedAction { .. } => ctx.notify(), _ => (), }); diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 0df05e0e..85edc94d 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -10,7 +10,7 @@ mod pending_response_streams; pub mod response_stream; pub(super) mod shared_session; mod slash_command; -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; #[cfg(not(target_family = "wasm"))] use std::path::PathBuf; use std::sync::Arc; @@ -19,7 +19,12 @@ use std::time::Duration; use ai::skills::SkillPathOrigin; use anyhow::anyhow; use chrono::{DateTime, Local}; -use galaxy_agent_core::ToolLoopGuard; +use futures::channel::oneshot; +use galaxy_agent_core::{ + turn_control, ExternalWorkId, PendingToolBatch, ProviderRun, ProviderRunFailureKind, + ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunState, ToolLoopGuard, + TurnCommand, TurnCommandSender, TurnRequest, +}; use galaxy_core::assertions::safe_assert; use input_context::{input_context_for_request, parse_context_attachments}; use itertools::Itertools; @@ -50,13 +55,15 @@ use crate::ai::agent::task::TaskId; #[cfg(not(target_family = "wasm"))] use crate::ai::agent::AIAgentActionTypeDiscriminants; use crate::ai::agent::{ - extract_user_query_mode, AIAgentAction, AIAgentActionResult, AIAgentActionResultType, - AIAgentAttachment, AIAgentContext, AIAgentExchangeId, AIAgentInput, AIAgentOutputStatus, - AIIdentifiers, CancellationOutcome, CancellationReason, DocumentContentAttachmentSource, - EntrypointType, FileContext, FinishedAIAgentOutput, PassiveSuggestionResultType, - PassiveSuggestionTrigger, PassiveSuggestionTriggerType, RenderableAIError, + extract_user_query_mode, AIAgentAction, AIAgentActionId, AIAgentActionResult, + AIAgentActionResultType, AIAgentActionType, AIAgentAttachment, AIAgentContext, + AIAgentExchangeId, AIAgentInput, AIAgentOutputStatus, AIIdentifiers, CancellationOutcome, + CancellationReason, DocumentContentAttachmentSource, EntrypointType, FileContext, + FinishedAIAgentOutput, PassiveSuggestionResultType, PassiveSuggestionTrigger, + PassiveSuggestionTriggerType, ReadShellCommandOutputResult, RenderableAIError, RequestCommandOutputResult, RequestCost, RequestMetadata, RunningCommand, StaticQueryType, - TransientNetworkErrorKind, UserQueryMode, + TransferShellCommandControlToUserResult, TransientNetworkErrorKind, UserQueryMode, + WriteToLongRunningShellCommandResult, }; use crate::ai::agent_events::AgentMessageEventMetadata; #[cfg(not(target_family = "wasm"))] @@ -66,14 +73,19 @@ use crate::ai::document::ai_document_model::{ AIDocumentId, AIDocumentModel, AIDocumentUserEditStatus, }; use crate::ai::llms::{LLMId, LLMPreferences}; -use crate::ai::provider::types::ContentPart; +use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageContent}; #[cfg(not(target_family = "wasm"))] use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord}; +use crate::ai::runtime::{ + prepare_provider_run, provider_runtime_for_request, PreparedProviderRun, ProviderActionContext, + ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderRunResponseProjector, + ProviderToolExecutionRef, ProviderToolLifecycleOutcome, RuntimeResponseConfig, + BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE, +}; use crate::ai::AIRequestUsageModel; use crate::cloud_object::model::persistence::CloudModel; use crate::features::FeatureFlag; use crate::global_resource_handles::GlobalResourceHandlesProvider; -use crate::network::NetworkStatus; use crate::notebooks::editor::model::FileLinkResolutionContext; use crate::persistence::model::AgentBackend; use crate::persistence::ModelEvent; @@ -83,14 +95,13 @@ use crate::server::server_api::AIApiError; use crate::server::server_api::ServerApiProvider; use crate::server::telemetry::TelemetryEvent; use crate::terminal::model::block::{ - formatted_terminal_contents_for_input, BlockId, CURSOR_MARKER, + formatted_terminal_contents_for_input, BlockId, BlockState, CURSOR_MARKER, }; use crate::terminal::model::session::active_session::ActiveSession; use crate::terminal::model::session::SessionType; use crate::terminal::model::terminal_model::TerminalModel; use crate::terminal::view::inline_banner::ZeroStatePromptSuggestionType; use crate::terminal::ShellLaunchData; -use crate::workspace::OneTimeModalModel; use crate::workspaces::update_manager::TeamUpdateManager; use crate::workspaces::user_workspaces::UserWorkspaces; @@ -191,60 +202,12 @@ fn remote_action_summaries(actions: &[AIAgentAction]) -> Vec .collect() } -#[derive(Debug, Clone)] -struct FailedToolProposal { - action_id: String, - task_id: String, - tool_name: String, - requires_result: bool, - error: String, -} - -impl FailedToolProposal { - fn new(action: &AIAgentAction, error: impl Into) -> Self { - Self { - action_id: action.id.to_string(), - task_id: action.task_id.to_string(), - tool_name: failed_proposal_tool_name(action), - requires_result: action.requires_result, - error: error.into(), - } - } - - #[cfg(not(target_family = "wasm"))] - fn to_remote_log_value(&self) -> serde_json::Value { - serde_json::json!({ - "action_id": self.action_id, - "task_id": self.task_id, - "tool_name": self.tool_name, - "requires_result": self.requires_result, - "error": remote_logging::sanitize_error(&self.error), - }) - } -} - -fn failed_proposal_tool_name(action: &AIAgentAction) -> String { - if let Some(tool_name) = action.tool_name.clone() { - return tool_name; - } - #[cfg(not(target_family = "wasm"))] - { - format!("{:?}", AIAgentActionTypeDiscriminants::from(&action.action)) - } - #[cfg(target_family = "wasm")] - { - "unknown".to_string() - } -} - #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ToolQueueDecision { Cancelled, UnfinishedExchange, - BlockedFailedToolProposal, BlockedActiveChildAgents, NoActions, - QueueActionsWithStreamSnapshotFallback, QueueActions, } @@ -253,29 +216,20 @@ impl ToolQueueDecision { match self { Self::Cancelled => "cancelled", Self::UnfinishedExchange => "unfinished_exchange", - Self::BlockedFailedToolProposal => "blocked_failed_tool_proposal", Self::BlockedActiveChildAgents => "blocked_active_child_agents", Self::NoActions => "no_actions", - Self::QueueActionsWithStreamSnapshotFallback => { - "queue_actions_with_stream_snapshot_fallback" - } Self::QueueActions => "queue_actions", } } fn will_queue_actions(self) -> bool { - matches!( - self, - Self::QueueActions | Self::QueueActionsWithStreamSnapshotFallback - ) + matches!(self, Self::QueueActions) } #[cfg(not(target_family = "wasm"))] fn remote_log_level(self) -> RemoteLogLevel { match self { - Self::BlockedFailedToolProposal - | Self::BlockedActiveChildAgents - | Self::QueueActionsWithStreamSnapshotFallback => RemoteLogLevel::Warn, + Self::BlockedActiveChildAgents => RemoteLogLevel::Warn, Self::Cancelled | Self::UnfinishedExchange | Self::NoActions | Self::QueueActions => { RemoteLogLevel::Info } @@ -286,23 +240,17 @@ impl ToolQueueDecision { fn tool_queue_decision( has_cancellation: bool, has_unfinished_exchange: bool, - has_failed_tool_proposal: bool, has_active_child_agents: bool, candidate_action_count: usize, - queued_from_stream_snapshot_count: usize, ) -> ToolQueueDecision { if has_cancellation { ToolQueueDecision::Cancelled } else if has_unfinished_exchange { ToolQueueDecision::UnfinishedExchange - } else if has_failed_tool_proposal { - ToolQueueDecision::BlockedFailedToolProposal } else if has_active_child_agents { ToolQueueDecision::BlockedActiveChildAgents } else if candidate_action_count == 0 { ToolQueueDecision::NoActions - } else if queued_from_stream_snapshot_count > 0 { - ToolQueueDecision::QueueActionsWithStreamSnapshotFallback } else { ToolQueueDecision::QueueActions } @@ -575,6 +523,715 @@ impl RequestInput { } } +struct ActiveProviderRun { + coordinator: ProviderRunCoordinator, + projector: ProviderRunResponseProjector, + response_config: RuntimeResponseConfig, + action_context: ProviderActionContext, + messages_sent: Arc>>, + persistence_offset: usize, +} + +impl ActiveProviderRun { + fn set_task_id(&mut self, task_id: &TaskId) { + let task_id = task_id.to_string(); + self.projector.set_task_id(task_id.clone()); + self.response_config.task_id.clone_from(&task_id); + self.action_context.set_task_id(task_id); + } +} + +const ACTIVE_PROVIDER_RUN_SNAPSHOT_VERSION: u32 = 1; + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +struct ProviderProjectionTarget { + task_id: TaskId, + exchange_id: AIAgentExchangeId, +} + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +struct ProviderCommandMonitorState { + run_id: ProviderRunId, + originating_work_id: ExternalWorkId, + originating_call_id: String, + initial_requested_command_action_id: AIAgentActionId, + block_id: BlockId, + command: String, + cli_task_id: TaskId, +} + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +struct PendingProviderMonitorObservation { + block_id: BlockId, + cli_task_id: TaskId, +} + +struct RestoredProviderCommandEvidence { + conversation_id: Option, + requested_command_action_id: Option, + cli_task_id: Option, + command: String, + state: BlockState, + output: String, + exit_code: i32, +} + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub(super) struct PendingProviderCommandCompletion { + block_id: BlockId, + initial_requested_command_action_id: Option, + command: String, + output: String, + exit_code: i32, +} + +impl PendingProviderCommandCompletion { + pub(super) fn new( + block_id: BlockId, + initial_requested_command_action_id: Option, + command: String, + output: String, + exit_code: i32, + ) -> Self { + Self { + block_id, + initial_requested_command_action_id, + command, + output, + exit_code, + } + } + + fn observation(&self) -> MessageContent { + let output = if self.output.is_empty() { + "(no output)" + } else { + self.output.as_str() + }; + MessageContent::Text(format!( + "The monitored command has finished with exit code {}. Continue the original objective \ + using this as evidence; a nonzero exit is not automatic run completion.\n\nCommand:\n{}\n\nFinal output:\n{}", + self.exit_code, self.command, output + )) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +enum ProviderCommandResult { + Snapshot { + block_id: BlockId, + command: Option, + }, + Finished { + block_id: BlockId, + command: Option, + output: String, + exit_code: i32, + }, +} + +struct ActiveProviderRunSlot { + stream_id: ResponseStreamId, + response_stream: ModelHandle, + did_input_contain_user_query: bool, + run_id: ProviderRunId, + root_task_id: TaskId, + projection_target: ProviderProjectionTarget, + run: Option, + checkpoint: Option, + turn_control: Option, + cancellation_reason: Option, + committed_provider_batch: Option, + command_action_refs: HashMap, + command_monitor: Option, + pending_monitor_observation: Option, + pending_command_completion: Option, + monitor_prose_continuations: usize, +} + +#[derive(Clone)] +struct ActiveProviderRunCheckpoint { + run: ProviderRun, + base_request: TurnRequest, + cli_monitor_request: Option, + response_config: RuntimeResponseConfig, + action_context: ProviderActionContext, + persistence_offset: usize, +} + +struct PreparedRestoredProviderRun { + snapshot: ActiveProviderRunSnapshot, + profiles: BTreeMap, +} + +impl ActiveProviderRunCheckpoint { + fn from_active_run(run: &ActiveProviderRun) -> Result { + let base_request = run + .coordinator + .profile_request(BASE_PROVIDER_PROFILE) + .cloned() + .ok_or_else(|| "provider run is missing its base request profile".to_string())?; + let cli_monitor_request = run + .coordinator + .profile_request(CLI_MONITOR_PROVIDER_PROFILE) + .cloned(); + Ok(Self { + run: run.coordinator.run().clone(), + base_request, + cli_monitor_request, + response_config: run.response_config.clone(), + action_context: run.action_context.clone(), + persistence_offset: run.persistence_offset, + }) + } + + fn with_run(&self, run: ProviderRun) -> Self { + Self { + run, + ..self.clone() + } + } +} + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +struct ActiveProviderRunSnapshot { + version: u32, + run: ProviderRun, + base_request: TurnRequest, + cli_monitor_request: Option, + response_config: RuntimeResponseConfig, + action_context: ProviderActionContext, + projection_target: ProviderProjectionTarget, + root_task_id: TaskId, + did_input_contain_user_query: bool, + persistence_offset: usize, + committed_provider_batch: Option, + command_action_refs: HashMap, + command_monitor: Option, + pending_monitor_observation: Option, + pending_command_completion: Option, + monitor_prose_continuations: usize, +} + +impl ActiveProviderRunSnapshot { + fn from_slot(slot: &ActiveProviderRunSlot) -> Result { + let checkpoint = match slot.run.as_ref() { + Some(run) => ActiveProviderRunCheckpoint::from_active_run(run)?, + None => slot + .checkpoint + .clone() + .ok_or_else(|| "provider run is not prepared".to_string())?, + }; + Self::from_slot_and_checkpoint(slot, checkpoint) + } + + fn from_slot_and_checkpoint( + slot: &ActiveProviderRunSlot, + checkpoint: ActiveProviderRunCheckpoint, + ) -> Result { + if checkpoint.run.id() != &slot.run_id { + return Err("provider run snapshot identity mismatch".to_string()); + } + Ok(Self { + version: ACTIVE_PROVIDER_RUN_SNAPSHOT_VERSION, + run: checkpoint.run, + base_request: checkpoint.base_request, + cli_monitor_request: checkpoint.cli_monitor_request, + response_config: checkpoint.response_config, + action_context: checkpoint.action_context, + projection_target: slot.projection_target.clone(), + root_task_id: slot.root_task_id.clone(), + did_input_contain_user_query: slot.did_input_contain_user_query, + persistence_offset: checkpoint.persistence_offset, + committed_provider_batch: slot.committed_provider_batch.clone(), + command_action_refs: slot.command_action_refs.clone(), + command_monitor: slot.command_monitor.clone(), + pending_monitor_observation: slot.pending_monitor_observation.clone(), + pending_command_completion: slot.pending_command_completion.clone(), + monitor_prose_continuations: slot.monitor_prose_continuations, + }) + } + + fn parse(json: &str) -> Result { + let snapshot: Self = serde_json::from_str(json) + .map_err(|error| format!("invalid active provider run snapshot: {error}"))?; + if snapshot.version != ACTIVE_PROVIDER_RUN_SNAPSHOT_VERSION { + return Err(format!( + "unsupported active provider run snapshot version {}", + snapshot.version + )); + } + Ok(snapshot) + } + + fn validate(&self, conversation_id: AIConversationId) -> Result<(), String> { + let run_id = self.run.id(); + if self.persistence_offset > self.run.transcript().len() { + return Err("provider run persistence offset exceeds transcript length".to_string()); + } + if self.base_request.model.as_str() != self.response_config.model_id { + return Err("provider run base model does not match response projection".to_string()); + } + if self.action_context.task_id() != self.response_config.task_id { + return Err("provider run action and response task IDs do not match".to_string()); + } + let current_task_id = self.action_context.task_id(); + let task_id_is_valid = current_task_id == &*self.root_task_id + || current_task_id == &*self.projection_target.task_id + || self + .command_monitor + .as_ref() + .is_some_and(|monitor| current_task_id == &*monitor.cli_task_id); + if !task_id_is_valid { + return Err( + "provider run current task is not owned by its projection or monitor".to_string(), + ); + } + + match self.run.profile().as_str() { + BASE_PROVIDER_PROFILE => {} + CLI_MONITOR_PROVIDER_PROFILE if self.cli_monitor_request.is_some() => {} + CLI_MONITOR_PROVIDER_PROFILE => { + return Err( + "provider run uses the CLI profile without a persisted request".to_string(), + ); + } + profile => { + return Err(format!( + "provider run uses unknown request profile '{profile}'" + )); + } + } + + if self.command_monitor.is_some() && self.cli_monitor_request.is_none() { + return Err("provider command monitor is missing its CLI request profile".to_string()); + } + if self + .committed_provider_batch + .as_ref() + .is_some_and(|work_id| &work_id.run_id != run_id) + { + return Err("committed provider batch belongs to a different run".to_string()); + } + for (action_id, execution_ref) in &self.command_action_refs { + if execution_ref.conversation_id != conversation_id + || &execution_ref.run_id != run_id + || execution_ref.call_id != action_id.to_string() + { + return Err(format!( + "provider command correlation for action {action_id} has invalid identity" + )); + } + } + + match &self.command_monitor { + Some(monitor) => { + if monitor.run_id != *run_id + || monitor.originating_work_id.run_id != *run_id + || monitor.originating_call_id + != monitor.initial_requested_command_action_id.to_string() + { + return Err("provider command monitor has invalid run identity".to_string()); + } + let Some(execution_ref) = self + .command_action_refs + .get(&monitor.initial_requested_command_action_id) + else { + return Err( + "provider command monitor is missing its action correlation".to_string() + ); + }; + if execution_ref.work_id() != monitor.originating_work_id + || execution_ref.call_id != monitor.originating_call_id + { + return Err( + "provider command monitor action correlation does not match".to_string() + ); + } + if self + .pending_monitor_observation + .as_ref() + .is_some_and(|observation| { + observation.block_id != monitor.block_id + || observation.cli_task_id != monitor.cli_task_id + }) + { + return Err("provider monitor observation does not match its owner".to_string()); + } + if self + .pending_command_completion + .as_ref() + .is_some_and(|completion| { + completion.block_id != monitor.block_id + || completion.initial_requested_command_action_id.as_ref() + != Some(&monitor.initial_requested_command_action_id) + }) + { + return Err("provider command completion does not match its owner".to_string()); + } + } + None => { + if self.pending_monitor_observation.is_some() + || self.pending_command_completion.is_some() + { + return Err( + "provider command evidence is missing its monitor owner".to_string() + ); + } + } + } + Ok(()) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ProviderFinishedActionDisposition { + Ignore, + AwaitBatchCommit, + Resume, +} + +fn normalize_restored_provider_snapshot( + snapshot: &mut ActiveProviderRunSnapshot, +) -> Result<(), String> { + let normalization = snapshot + .run + .normalize_after_restore() + .map_err(|error| error.to_string())?; + let interrupted_call_ids = normalization + .interrupted_call_ids + .iter() + .map(String::as_str) + .collect::>(); + snapshot + .command_action_refs + .retain(|_, execution_ref| !interrupted_call_ids.contains(execution_ref.call_id.as_str())); + if normalization.committed_tool_batch { + snapshot.committed_provider_batch = None; + } + if let Some(committed_work_id) = snapshot.committed_provider_batch.as_ref() { + let has_unreconciled_command = snapshot.command_monitor.is_none() + && snapshot + .command_action_refs + .values() + .any(|execution_ref| execution_ref.work_id() == *committed_work_id); + if has_unreconciled_command { + return Err( + "restored provider command batch completed without durable terminal evidence" + .to_string(), + ); + } + snapshot.committed_provider_batch = None; + } + Ok(()) +} + +fn apply_restored_provider_command_evidence( + conversation_id: AIConversationId, + snapshot: &mut ActiveProviderRunSnapshot, + evidence: RestoredProviderCommandEvidence, +) -> Result<(), String> { + let Some(monitor) = snapshot.command_monitor.as_ref() else { + return Ok(()); + }; + if evidence.conversation_id != Some(conversation_id) + || evidence.requested_command_action_id.as_ref() + != Some(&monitor.initial_requested_command_action_id) + || evidence.cli_task_id.as_ref() != Some(&monitor.cli_task_id) + { + return Err("restored provider command block identity does not match".to_string()); + } + if monitor.command.trim().is_empty() || evidence.command != monitor.command { + return Err("restored provider command text does not match".to_string()); + } + match evidence.state { + BlockState::BeforeExecution | BlockState::Executing => { + snapshot.pending_command_completion = None; + snapshot.pending_monitor_observation = Some(PendingProviderMonitorObservation { + block_id: monitor.block_id.clone(), + cli_task_id: monitor.cli_task_id.clone(), + }); + } + BlockState::DoneWithExecution | BlockState::DoneWithNoExecution => { + snapshot.pending_monitor_observation = None; + snapshot.pending_command_completion = Some(PendingProviderCommandCompletion { + block_id: monitor.block_id.clone(), + initial_requested_command_action_id: Some( + monitor.initial_requested_command_action_id.clone(), + ), + command: evidence.command, + output: evidence.output, + exit_code: evidence.exit_code, + }); + } + BlockState::Background | BlockState::Static => { + return Err("restored provider command block has an invalid state".to_string()); + } + } + Ok(()) +} + +fn provider_execution_matches_active_work( + run_id: &ProviderRunId, + active_work_id: Option<&ExternalWorkId>, + execution_ref: &ProviderToolExecutionRef, +) -> bool { + &execution_ref.run_id == run_id + && active_work_id.is_some_and(|work_id| { + work_id.run_id == execution_ref.run_id && work_id.epoch == execution_ref.epoch + }) +} + +fn provider_finished_action_disposition( + run_id: &ProviderRunId, + active_work_id: Option<&ExternalWorkId>, + committed_work_id: Option<&ExternalWorkId>, + execution_ref: &ProviderToolExecutionRef, +) -> ProviderFinishedActionDisposition { + let execution_work_id = execution_ref.work_id(); + if &execution_ref.run_id != run_id { + ProviderFinishedActionDisposition::Ignore + } else if committed_work_id == Some(&execution_work_id) { + ProviderFinishedActionDisposition::Resume + } else if active_work_id == Some(&execution_work_id) { + ProviderFinishedActionDisposition::AwaitBatchCommit + } else { + ProviderFinishedActionDisposition::Ignore + } +} + +fn is_provider_command_action(action: &AIAgentActionType) -> bool { + matches!( + action, + AIAgentActionType::RequestCommandOutput { .. } + | AIAgentActionType::WriteToLongRunningShellCommand { .. } + | AIAgentActionType::ReadShellCommandOutput { .. } + | AIAgentActionType::TransferShellCommandControlToUser { .. } + ) +} + +fn provider_command_completion_matches( + slot_run_id: &ProviderRunId, + command_action_refs: &HashMap, + command_monitor: Option<&ProviderCommandMonitorState>, + block_id: &BlockId, + initial_requested_command_action_id: Option<&AIAgentActionId>, +) -> bool { + if let Some(monitor) = command_monitor { + return monitor.run_id == *slot_run_id + && monitor.block_id == *block_id + && initial_requested_command_action_id + .is_none_or(|action_id| monitor.initial_requested_command_action_id == *action_id); + } + + initial_requested_command_action_id + .and_then(|action_id| command_action_refs.get(action_id)) + .is_some_and(|execution_ref| execution_ref.run_id == *slot_run_id) +} + +fn reconcile_provider_completion_with_snapshot( + completion: Option<&mut PendingProviderCommandCompletion>, + block_id: &BlockId, + expected_initial_action_id: &AIAgentActionId, + snapshot_command: Option<&str>, + fallback_command: Option<&str>, +) -> Result { + let Some(completion) = completion else { + return Ok(false); + }; + if completion.block_id != *block_id + || completion + .initial_requested_command_action_id + .as_ref() + .is_some_and(|action_id| action_id != expected_initial_action_id) + { + return Err("provider command completion did not match committed snapshot".to_owned()); + } + if completion.command.is_empty() { + completion.command = snapshot_command + .or(fallback_command) + .unwrap_or_default() + .to_owned(); + } + Ok(true) +} + +fn classify_provider_command_result( + result: &AIAgentActionResultType, +) -> Option { + match result { + AIAgentActionResultType::RequestCommandOutput(result) => match result { + RequestCommandOutputResult::Completed { + block_id, + command, + output, + exit_code, + .. + } => Some(ProviderCommandResult::Finished { + block_id: block_id.clone(), + command: Some(command.clone()), + output: output.clone(), + exit_code: exit_code.value(), + }), + RequestCommandOutputResult::LongRunningCommandSnapshot { + block_id, command, .. + } => Some(ProviderCommandResult::Snapshot { + block_id: block_id.clone(), + command: Some(command.clone()), + }), + RequestCommandOutputResult::CancelledBeforeExecution + | RequestCommandOutputResult::Denylisted { .. } => None, + }, + AIAgentActionResultType::WriteToLongRunningShellCommand(result) => match result { + WriteToLongRunningShellCommandResult::Snapshot { block_id, .. } => { + Some(ProviderCommandResult::Snapshot { + block_id: block_id.clone(), + command: None, + }) + } + WriteToLongRunningShellCommandResult::CommandFinished { + block_id, + output, + exit_code, + .. + } => Some(ProviderCommandResult::Finished { + block_id: block_id.clone(), + command: None, + output: output.clone(), + exit_code: exit_code.value(), + }), + WriteToLongRunningShellCommandResult::Cancelled + | WriteToLongRunningShellCommandResult::Error(_) => None, + }, + AIAgentActionResultType::ReadShellCommandOutput(result) => match result { + ReadShellCommandOutputResult::LongRunningCommandSnapshot { + block_id, command, .. + } => Some(ProviderCommandResult::Snapshot { + block_id: block_id.clone(), + command: Some(command.clone()), + }), + ReadShellCommandOutputResult::CommandFinished { + block_id, + command, + output, + exit_code, + .. + } => Some(ProviderCommandResult::Finished { + block_id: block_id.clone(), + command: Some(command.clone()), + output: output.clone(), + exit_code: exit_code.value(), + }), + ReadShellCommandOutputResult::Cancelled | ReadShellCommandOutputResult::Error(_) => { + None + } + }, + AIAgentActionResultType::TransferShellCommandControlToUser(result) => match result { + TransferShellCommandControlToUserResult::Snapshot { block_id, .. } => { + Some(ProviderCommandResult::Snapshot { + block_id: block_id.clone(), + command: None, + }) + } + TransferShellCommandControlToUserResult::CommandFinished { + block_id, + output, + exit_code, + .. + } => Some(ProviderCommandResult::Finished { + block_id: block_id.clone(), + command: None, + output: output.clone(), + exit_code: exit_code.value(), + }), + TransferShellCommandControlToUserResult::Cancelled + | TransferShellCommandControlToUserResult::Error(_) => None, + }, + _ => None, + } +} + +const MAX_PROVIDER_MONITOR_PROSE_CONTINUATIONS: usize = 1; + +enum ProviderBoundaryDisposition { + Advance { completed_block_id: Option }, + Park, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ProviderBoundaryPhase { + Ready, + AwaitingDriver, + Unsafe, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ProviderBoundaryIntent { + ApplyCompletion, + ApplyMonitorObservation, + RetryMonitor, + CompleteRun, + Advance, + Park, +} + +fn provider_boundary_phase(state: &ProviderRunState) -> ProviderBoundaryPhase { + match state { + ProviderRunState::ReadyToCallModel => ProviderBoundaryPhase::Ready, + ProviderRunState::AwaitingDriver { .. } => ProviderBoundaryPhase::AwaitingDriver, + ProviderRunState::AwaitingModel { .. } + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingTools { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => ProviderBoundaryPhase::Unsafe, + } +} + +fn provider_boundary_intent( + phase: ProviderBoundaryPhase, + has_committed_batch: bool, + has_completion: bool, + has_monitor_observation: bool, + is_cli_profile: bool, + has_monitor: bool, + monitor_prose_continuations: usize, +) -> ProviderBoundaryIntent { + if has_committed_batch || phase == ProviderBoundaryPhase::Unsafe { + return ProviderBoundaryIntent::Park; + } + if has_completion { + return ProviderBoundaryIntent::ApplyCompletion; + } + if has_monitor_observation { + return ProviderBoundaryIntent::ApplyMonitorObservation; + } + match phase { + ProviderBoundaryPhase::Ready => ProviderBoundaryIntent::Advance, + ProviderBoundaryPhase::AwaitingDriver if is_cli_profile && has_monitor => { + if monitor_prose_continuations < MAX_PROVIDER_MONITOR_PROSE_CONTINUATIONS { + ProviderBoundaryIntent::RetryMonitor + } else { + ProviderBoundaryIntent::Park + } + } + ProviderBoundaryPhase::AwaitingDriver => ProviderBoundaryIntent::CompleteRun, + ProviderBoundaryPhase::Unsafe => ProviderBoundaryIntent::Park, + } +} + +enum ProviderDriveMessage { + Response(warp_multi_agent_api::ResponseEvent), + Checkpoint { + checkpoint: ActiveProviderRunCheckpoint, + acknowledgement: oneshot::Sender>, + }, + Blocked { + run: ActiveProviderRun, + result: Result, + }, +} + /// Controller for Blocklist AI. /// /// This is responsible for managing and updating blocklist AI state for a single terminal surface. @@ -586,6 +1243,8 @@ pub struct BlocklistAIController { terminal_model: Arc>, in_flight_response_streams: PendingResponseStreams, + active_provider_runs: HashMap, + restoring_provider_runs: HashSet, /// The ID of the terminal surface this controller is associated with. terminal_surface_id: EntityId, @@ -603,9 +1262,6 @@ pub struct BlocklistAIController { /// Set by the agent driver based on the workspace directory (e.g. `{working_dir}/.warp-core/attachments`). attachments_download_dir: Option, - /// Pending auto-resume tasks that are waiting for network connectivity. - /// These should be cancelled when a new request is sent for the same conversation. - pending_auto_resume_handles: HashMap, /// Pending dormant Claude wake preparations for success-idle child conversations. #[cfg_attr(target_family = "wasm", allow(dead_code))] pending_local_claude_wakes: HashMap, @@ -614,15 +1270,8 @@ pub struct BlocklistAIController { /// Conversations with finished action results that should not be drained /// until active child agents in their orchestration subtree finish. pending_child_blocked_follow_ups: HashSet, - /// Tool proposals that arrived in a provider stream but failed to attach to - /// conversation history. If a proposal cannot be attached, executing it via - /// the stream snapshot fallback would create orphaned tool history. - failed_tool_proposals_by_stream: HashMap>, - /// Per-conversation loop detection state for preventing recursive tool failures. loop_detection: HashMap, - /// Per-conversation error retry count for injecting corrective messages on failure. - error_retry_counts: HashMap, /// Passive suggestion results that should be included with the next request /// for a given conversation (e.g. accepted/iterated code diffs that weren't /// auto-resumed). @@ -835,14 +1484,28 @@ impl BlocklistAIController { ctx: &mut ModelContext, ) -> Self { ctx.subscribe_to_model(&action_model, move |me, _, event, ctx| { + if let BlocklistAIActionEvent::ToolLifecycle { + execution_ref: Some(execution_ref), + event, + .. + } = event + { + me.handle_provider_tool_lifecycle(execution_ref, event, ctx); + return; + } let BlocklistAIActionEvent::FinishedAction { conversation_id, cancellation_reason, + execution_ref, .. } = event else { return; }; + if let Some(execution_ref) = execution_ref { + me.handle_provider_actions_finished(*conversation_id, execution_ref, ctx); + return; + } // `FinalizedExternally` (e.g. shell exit) means the conversation status and message // is set elsewhere through a dedicated path, so we must not trigger a follow-up or update conversation status here. let cancellation_outcome = @@ -962,27 +1625,52 @@ impl BlocklistAIController { }); let history_model = BlocklistAIHistoryModel::handle(ctx); - ctx.subscribe_to_model(&history_model, |me, _, event, ctx| { - let BlocklistAIHistoryEvent::UpdatedConversationStatus { + ctx.subscribe_to_model(&history_model, |me, _, event, ctx| match event { + BlocklistAIHistoryEvent::RestoredConversations { + terminal_surface_id, + conversation_ids, + } if *terminal_surface_id == me.terminal_surface_id => { + me.schedule_restored_provider_runs(conversation_ids, ctx); + } + BlocklistAIHistoryEvent::UpdatedConversationStatus { terminal_surface_id, new_status, .. - } = event - else { - return; - }; - if *terminal_surface_id != me.terminal_surface_id || !new_status.is_done() { - return; - } - - let pending_parents = me - .pending_child_blocked_follow_ups - .iter() - .copied() - .collect::>(); - for parent_id in pending_parents { - me.maybe_resume_child_blocked_follow_up(parent_id, ctx); + } if *terminal_surface_id == me.terminal_surface_id && new_status.is_done() => { + let pending_parents = me + .pending_child_blocked_follow_ups + .iter() + .copied() + .collect::>(); + for parent_id in pending_parents { + me.maybe_resume_child_blocked_follow_up(parent_id, ctx); + } } + BlocklistAIHistoryEvent::StartedNewConversation { .. } + | BlocklistAIHistoryEvent::CreatedSubtask { .. } + | BlocklistAIHistoryEvent::UpgradedTask { .. } + | BlocklistAIHistoryEvent::AppendedExchange { .. } + | BlocklistAIHistoryEvent::ReassignedExchange { .. } + | BlocklistAIHistoryEvent::UpdatedStreamingExchange { .. } + | BlocklistAIHistoryEvent::UpdatedConversationStatus { .. } + | BlocklistAIHistoryEvent::SetActiveConversation { .. } + | BlocklistAIHistoryEvent::ClearedActiveConversation { .. } + | BlocklistAIHistoryEvent::ClearedConversationsForTerminalSurface { .. } + | BlocklistAIHistoryEvent::UpdatedTodoList { .. } + | BlocklistAIHistoryEvent::UpdatedAutoexecuteOverride { .. } + | BlocklistAIHistoryEvent::SplitConversation { .. } + | BlocklistAIHistoryEvent::RemoveConversation { .. } + | BlocklistAIHistoryEvent::DeletedConversation { .. } + | BlocklistAIHistoryEvent::RestoredConversations { .. } + | BlocklistAIHistoryEvent::UpdatedConversationMetadata { .. } + | BlocklistAIHistoryEvent::UpdatedConversationTitle { .. } + | BlocklistAIHistoryEvent::UpdatedConversationArtifacts { .. } + | BlocklistAIHistoryEvent::ConversationServerTokenAssigned { .. } + | BlocklistAIHistoryEvent::ConversationTransferredBetweenTerminalSurfaces { .. } + | BlocklistAIHistoryEvent::NewConversationRequestComplete { .. } + | BlocklistAIHistoryEvent::OrchestrationConfigUpdated { .. } + | BlocklistAIHistoryEvent::ConversationUsageMetadataUpdated { .. } + | BlocklistAIHistoryEvent::LocalSharedSessionEstablished { .. } => {} }); ctx.subscribe_to_model(&conversation_selection, |me, _, event, ctx| { @@ -1051,19 +1739,18 @@ impl BlocklistAIController { active_session, terminal_model, in_flight_response_streams: PendingResponseStreams::new(), + active_provider_runs: HashMap::new(), + restoring_provider_runs: HashSet::new(), terminal_surface_id, should_refresh_available_llms_on_stream_finish: false, shared_session_state: shared_session::SharedSessionState::default(), ambient_agent_task_id: None, attachments_download_dir: None, - pending_auto_resume_handles: HashMap::new(), pending_local_claude_wakes: HashMap::new(), pending_passive_follow_ups: HashSet::new(), pending_child_blocked_follow_ups: HashSet::new(), - failed_tool_proposals_by_stream: HashMap::new(), pending_passive_suggestion_results: HashMap::new(), loop_detection: HashMap::new(), - error_retry_counts: HashMap::new(), crosscheck_reviewer, } } @@ -1369,7 +2056,6 @@ impl BlocklistAIController { entrypoint: entrypoint_type, is_auto_resume_after_error: false, }), - /*can_attempt_resume_on_error*/ true, is_queued_prompt, ctx, ); @@ -1659,7 +2345,6 @@ impl BlocklistAIController { entrypoint: EntrypointType::AgentInitiated, is_auto_resume_after_error: false, }), - /*can_attempt_resume_on_error*/ false, /*is_queued_prompt*/ false, ctx, ) @@ -1759,8 +2444,8 @@ impl BlocklistAIController { } /// Nudges a CLI monitor that ended a turn without proposing a polling action. The running - /// command is attached through normal long-running-command detection so Rig and the legacy - /// provider path both receive the monitor-specific prompt and tool set. + /// command is attached through normal long-running-command detection so the provider run + /// receives the monitor-specific prompt and tool set. pub fn send_cli_monitor_nudge( &mut self, conversation_id: AIConversationId, @@ -1795,9 +2480,8 @@ impl BlocklistAIController { queued_query_id: Option, ctx: &mut ModelContext, ) { - // User sending a new query resets loop detection and error retry state — fresh context. + // User sending a new query resets loop detection for the fresh context. self.loop_detection.remove(&conversation_id); - self.error_retry_counts.remove(&conversation_id); // Reset any in-flight crosscheck review for this conversation. self.crosscheck_reviewer.update(ctx, |reviewer, _| { reviewer.reset_review(conversation_id); @@ -2328,13 +3012,8 @@ impl BlocklistAIController { .extend(event_inputs); } - let result = self.send_request_input( - request_input, - None, - /*can_attempt_resume_on_error*/ true, - /*is_queued_prompt*/ false, - ctx, - ); + let result = + self.send_request_input(request_input, None, /*is_queued_prompt*/ false, ctx); if has_piggybacked_events && result.is_err() { OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { @@ -2452,8 +3131,7 @@ impl BlocklistAIController { Failing action: {}\n\n\ Take a completely different approach to accomplish the goal. \ If you cannot find an alternative, explain to the user what is failing and why.", - looping_entry.threshold, - looping_entry.description + looping_entry.threshold, looping_entry.description ); Some(warning) } else { @@ -2717,7 +3395,6 @@ impl BlocklistAIController { ctx, ), None, - /*can_attempt_resume_on_error*/ true, /*is_queued_prompt*/ false, ctx, ) @@ -2779,8 +3456,6 @@ impl BlocklistAIController { pub fn resume_conversation( &mut self, conversation_id: AIConversationId, - can_attempt_resume_on_error: bool, - is_auto_resume_after_error: bool, additional_context: Vec, ctx: &mut ModelContext, ) { @@ -2817,15 +3492,6 @@ impl BlocklistAIController { ); let inputs = vec![AIAgentInput::ResumeConversation { context }]; - let metadata = if is_auto_resume_after_error { - Some(RequestMetadata { - is_autodetected_user_query: false, - entrypoint: EntrypointType::ResumeConversation, - is_auto_resume_after_error: true, - }) - } else { - None - }; let _ = self.send_request_input( RequestInput::for_task( inputs, @@ -2836,16 +3502,12 @@ impl BlocklistAIController { self.terminal_surface_id, ctx, ), - metadata, - can_attempt_resume_on_error, + None, /*is_queued_prompt*/ false, ctx, ); } - /// Schedules an auto-resume-after-error for the conversation once the network is online - /// and the auto-handoff sleep modal is closed, so the resume doesn't race the user's - /// enable/dismiss decision on wake. /// Handles the completion of a crosscheck review cycle. /// /// If the reviewer provided feedback, it is injected as a synthetic user @@ -2932,7 +3594,6 @@ impl BlocklistAIController { ctx, ), None, - /*can_attempt_resume_on_error*/ false, /*is_queued_prompt*/ false, ctx, ); @@ -2983,7 +3644,6 @@ impl BlocklistAIController { entrypoint: EntrypointType::AgentInitiated, is_auto_resume_after_error: false, }), - /*can_attempt_resume_on_error*/ false, /*is_queued_prompt*/ false, ctx, ) { @@ -3109,38 +3769,6 @@ impl BlocklistAIController { .join("\n\n") } - fn schedule_auto_resume_after_error( - &mut self, - conversation_id: AIConversationId, - ctx: &mut ModelContext, - ) { - let wait_for_online = NetworkStatus::as_ref(ctx).wait_until_online(); - let wait_for_modal_closed = - OneTimeModalModel::as_ref(ctx).wait_until_auto_handoff_sleep_modal_closed(); - let wait = async move { - wait_for_online.await; - // Await the modal second: the future reads live modal state at - // poll time, so a modal surfaced on wake (after connectivity - // returns) is still observed. - wait_for_modal_closed.await; - }; - let handle = ctx.spawn(wait, move |me, _, ctx| { - // Clean up the pending handle now that the resume is executing. - me.pending_auto_resume_handles.remove(&conversation_id); - me.resume_conversation( - conversation_id, - // Don't allow a second resume-on-error to prevent a persistent loop. - /*can_attempt_resume_on_error*/ - false, - /*is_auto_resume_after_error*/ true, - vec![], - ctx, - ); - }); - self.pending_auto_resume_handles - .insert(conversation_id, handle); - } - pub fn send_passive_code_diff_request( &mut self, query: String, @@ -3182,7 +3810,6 @@ impl BlocklistAIController { }, is_auto_resume_after_error: false, }), - /*can_attempt_resume_on_error*/ true, /*is_queued_prompt*/ false, ctx, ) @@ -3223,8 +3850,8 @@ impl BlocklistAIController { { let Some(conversation) = history_model.conversation(&conversation_id) else { return Err(anyhow!( - "Tried to build passive suggestions request params for non-existent conversation with ID {conversation_id:?}" - )); + "Tried to build passive suggestions request params for non-existent conversation with ID {conversation_id:?}" + )); }; let task_id = conversation.get_root_task_id().clone(); let conversation_data = api::ConversationData { @@ -3260,8 +3887,8 @@ impl BlocklistAIController { (conversation_id, task_id, conversation_data) } else { return Err(anyhow!( - "Tried to use agent response completed trigger to generate passive suggestions without a conversation ID" - )); + "Tried to use agent response completed trigger to generate passive suggestions without a conversation ID" + )); }; let inputs = vec![AIAgentInput::TriggerPassiveSuggestion { @@ -3345,7 +3972,6 @@ impl BlocklistAIController { }, is_auto_resume_after_error: false, }), - /*can_attempt_resume_on_error*/ true, /*is_queued_prompt*/ false, ctx, ) @@ -3414,7 +4040,6 @@ impl BlocklistAIController { &mut self, mut request_input: RequestInput, query_metadata: Option, - can_attempt_resume_on_error: bool, is_queued_prompt: bool, ctx: &mut ModelContext, ) -> anyhow::Result<(AIConversationId, ResponseStreamId)> { @@ -3469,21 +4094,9 @@ impl BlocklistAIController { request_input.computer_use_model_id = acp_model_id; } - // Cancel any pending auto-resume for this conversation, since the user is sending a new - // request. - if let Some(handle) = self - .pending_auto_resume_handles - .remove(&request_input.conversation_id) - { - handle.abort(); - } - - // Passive background requests never auto-resume: a resume would issue a fresh - // turn on a conversation the user never sees. let is_passive_request = request_input .all_inputs() .any(|input| input.is_passive_request()); - let can_attempt_resume_on_error = can_attempt_resume_on_error && !is_passive_request; // Make sure there's no existing response stream for the conversation. If // there is, something has gone wrong. @@ -3595,6 +4208,15 @@ impl BlocklistAIController { } let server_conversation_token_for_identifiers = conversation_data.server_conversation_token.clone(); + let provider_configs = matches!(&agent_backend, AgentBackend::Provider).then(|| { + ( + ResponseStream::resolve_provider_config(request_params.model.as_str(), ctx), + ResponseStream::resolve_provider_config( + request_params.cli_agent_model.as_str(), + ctx, + ), + ) + }); let response_stream = ctx.add_model(|ctx| { // Create AIIdentifiers for the response stream @@ -3605,13 +4227,16 @@ impl BlocklistAIController { client_exchange_id: None, model_id: Some(request_params.model.clone()), }; - ResponseStream::new( - request_params.clone(), - ai_identifiers, - agent_backend.clone(), - can_attempt_resume_on_error, - ctx, - ) + if provider_configs.is_some() { + ResponseStream::new_provider_projection(request_params.clone(), ai_identifiers, ctx) + } else { + ResponseStream::new( + request_params.clone(), + ai_identifiers, + agent_backend.clone(), + ctx, + ) + } }); let response_stream_id = response_stream.as_ref(ctx).id().clone(); let response_stream_clone = response_stream.clone(); @@ -3662,15 +4287,77 @@ impl BlocklistAIController { } }); + let provider_projection_target = if provider_configs.is_some() { + let (task_id, exchange_id) = history_model + .as_ref(ctx) + .conversation(&conversation_data.id) + .and_then(|conversation| { + conversation.provider_projection_target(&response_stream_id) + }) + .ok_or_else(|| { + anyhow!( + "direct-provider response stream does not have exactly one projection target" + ) + })?; + Some(ProviderProjectionTarget { + task_id, + exchange_id, + }) + } else { + None + }; self.in_flight_response_streams.register_new_stream( response_stream_id.clone(), conversation_data.id, - response_stream, + response_stream.clone(), CancellationReason::FollowUpSubmitted { is_for_same_conversation: true, }, ctx, ); + if let Some((base_provider_config, cli_provider_config)) = provider_configs { + let provider_run_id = ProviderRunId::new(format!( + "{}:{}", + conversation_data.id, + response_stream_id.as_str() + )); + let root_task_id = history_model + .as_ref(ctx) + .conversation(&conversation_data.id) + .expect("conversation exists while starting provider run") + .get_root_task_id() + .clone(); + self.active_provider_runs.insert( + conversation_data.id, + ActiveProviderRunSlot { + stream_id: response_stream_id.clone(), + response_stream, + did_input_contain_user_query: input_contains_user_query, + run_id: provider_run_id, + root_task_id, + projection_target: provider_projection_target + .expect("provider projection target was validated"), + run: None, + checkpoint: None, + turn_control: None, + cancellation_reason: None, + committed_provider_batch: None, + command_action_refs: HashMap::new(), + command_monitor: None, + pending_monitor_observation: None, + pending_command_completion: None, + monitor_prose_continuations: 0, + }, + ); + self.prepare_active_provider_run( + conversation_data.id, + response_stream_id.clone(), + base_provider_config, + cli_provider_config, + request_params.clone(), + ctx, + ); + } // Skip the context reset for a fired queued-prompt row (`is_queued_prompt`): its // attachments came from the row, not the live staging, so the live `pending_attachments` @@ -3717,6 +4404,1549 @@ impl BlocklistAIController { Ok((conversation_data.id, response_stream_id)) } + fn schedule_restored_provider_runs( + &mut self, + conversation_ids: &[AIConversationId], + ctx: &mut ModelContext, + ) { + let history_model = BlocklistAIHistoryModel::handle(ctx); + let conversation_ids = conversation_ids + .iter() + .copied() + .filter(|conversation_id| { + !self.active_provider_runs.contains_key(conversation_id) + && !self.restoring_provider_runs.contains(conversation_id) + && history_model + .as_ref(ctx) + .conversation(conversation_id) + .is_some_and(|conversation| { + conversation.active_provider_run_json().is_some() + }) + }) + .collect::>(); + if conversation_ids.is_empty() { + return; + } + self.restoring_provider_runs + .extend(conversation_ids.iter().copied()); + + // RestoredConversations is emitted before terminal views finish rebuilding their blocks. + let _ = ctx.spawn(async {}, move |me, _, ctx| { + for conversation_id in conversation_ids { + me.restore_active_provider_run(conversation_id, ctx); + } + }); + } + + fn restore_active_provider_run( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + if self.active_provider_runs.contains_key(&conversation_id) { + self.restoring_provider_runs.remove(&conversation_id); + return; + } + let history_model = BlocklistAIHistoryModel::handle(ctx); + let Some(snapshot_json) = history_model + .as_ref(ctx) + .conversation(&conversation_id) + .and_then(AIConversation::active_provider_run_json) + .map(str::to_owned) + else { + self.restoring_provider_runs.remove(&conversation_id); + return; + }; + let mut snapshot = match ActiveProviderRunSnapshot::parse(&snapshot_json) { + Ok(snapshot) => snapshot, + Err(error) => { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + }; + if let Err(error) = snapshot.validate(conversation_id) { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + let history_validation = history_model + .as_ref(ctx) + .conversation(&conversation_id) + .ok_or_else(|| "restored provider conversation is missing".to_string()) + .and_then(|conversation| { + if conversation.agent_backend() != &AgentBackend::Provider { + return Err( + "restored provider run belongs to a non-provider conversation".to_string(), + ); + } + if conversation.get_root_task_id() != &snapshot.root_task_id { + return Err( + "restored provider run root task does not match history".to_string() + ); + } + let Some(task) = conversation.get_task(&snapshot.projection_target.task_id) else { + return Err("restored provider projection task is missing".to_string()); + }; + if !task + .exchanges() + .any(|exchange| exchange.id == snapshot.projection_target.exchange_id) + { + return Err( + "restored provider projection exchange is missing from its task" + .to_string(), + ); + } + Ok(()) + }); + if let Err(error) = history_validation { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + + if let Err(error) = normalize_restored_provider_snapshot(&mut snapshot) { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + if let Err(error) = self.reconcile_restored_provider_command(conversation_id, &mut snapshot) + { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + if let Err(error) = snapshot.validate(conversation_id) { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + if let Err(error) = self.persist_provider_run_snapshot(conversation_id, &snapshot, ctx) { + self.fail_restored_provider_run(conversation_id, error, ctx); + return; + } + + let base_provider_config = + ResponseStream::resolve_provider_config(snapshot.base_request.model.as_str(), ctx); + let cli_provider_config = snapshot + .cli_monitor_request + .as_ref() + .map(|request| ResponseStream::resolve_provider_config(request.model.as_str(), ctx)); + let _ = ctx.spawn( + async move { + let base_runtime = + provider_runtime_for_request(base_provider_config, &snapshot.base_request) + .await?; + let mut profiles = BTreeMap::new(); + profiles.insert( + BASE_PROVIDER_PROFILE.to_string(), + ProviderRunProfile::new(base_runtime, snapshot.base_request.clone()), + ); + if let Some(cli_monitor_request) = snapshot.cli_monitor_request.as_ref() { + let cli_provider_config = cli_provider_config.ok_or_else(|| { + anyhow!("restored CLI provider request is missing its provider config") + })?; + let cli_runtime = + provider_runtime_for_request(cli_provider_config, cli_monitor_request) + .await?; + profiles.insert( + CLI_MONITOR_PROVIDER_PROFILE.to_string(), + ProviderRunProfile::new(cli_runtime, cli_monitor_request.clone()), + ); + } + Ok::<_, anyhow::Error>(PreparedRestoredProviderRun { snapshot, profiles }) + }, + move |me, result, ctx| { + me.handle_prepared_restored_provider_run(conversation_id, result, ctx); + }, + ); + } + + fn reconcile_restored_provider_command( + &self, + conversation_id: AIConversationId, + snapshot: &mut ActiveProviderRunSnapshot, + ) -> Result<(), String> { + let Some(monitor) = snapshot.command_monitor.as_ref() else { + return Ok(()); + }; + let evidence = { + let terminal_model = self.terminal_model.lock(); + let block = terminal_model + .block_list() + .block_with_id(&monitor.block_id) + .ok_or_else(|| "restored provider command block is missing".to_string())?; + RestoredProviderCommandEvidence { + conversation_id: block.ai_conversation_id(), + requested_command_action_id: block.requested_command_action_id().cloned(), + cli_task_id: block.cli_subagent_task_id().cloned(), + command: block.command_to_string(), + state: block.state(), + output: block.output_to_string(), + exit_code: block.exit_code().value(), + } + }; + apply_restored_provider_command_evidence(conversation_id, snapshot, evidence) + } + + fn persist_provider_run_snapshot( + &self, + conversation_id: AIConversationId, + snapshot: &ActiveProviderRunSnapshot, + ctx: &mut ModelContext, + ) -> Result<(), String> { + let json = serde_json::to_string(snapshot) + .map_err(|error| format!("failed to serialize restored provider run: {error}"))?; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model + .persist_active_provider_run_json(conversation_id, Some(json), ctx) + .map_err(|error| format!("failed to persist restored provider run: {error:?}")) + }) + } + + fn handle_prepared_restored_provider_run( + &mut self, + conversation_id: AIConversationId, + result: anyhow::Result, + ctx: &mut ModelContext, + ) { + if !self.restoring_provider_runs.contains(&conversation_id) + || self.active_provider_runs.contains_key(&conversation_id) + { + self.restoring_provider_runs.remove(&conversation_id); + return; + } + let PreparedRestoredProviderRun { snapshot, profiles } = match result { + Ok(prepared) => prepared, + Err(error) => { + self.fail_restored_provider_run(conversation_id, error.to_string(), ctx); + return; + } + }; + let ActiveProviderRunSnapshot { + version: _, + run: provider_run, + base_request: _, + cli_monitor_request: _, + response_config, + action_context, + projection_target, + root_task_id, + did_input_contain_user_query, + persistence_offset, + committed_provider_batch, + command_action_refs, + command_monitor, + pending_monitor_observation, + pending_command_completion, + monitor_prose_continuations, + } = snapshot; + let run_id = provider_run.id().clone(); + let transcript = provider_run.transcript(); + let offset = persistence_offset.min(transcript.len()); + let messages_sent = Arc::new(std::sync::Mutex::new(transcript[offset..].to_vec())); + let coordinator = match ProviderRunCoordinator::new(provider_run, profiles) { + Ok(coordinator) => coordinator, + Err(error) => { + self.fail_restored_provider_run(conversation_id, error.to_string(), ctx); + return; + } + }; + let model = LLMId::from(response_config.model_id.as_str()); + let ai_identifiers = AIIdentifiers { + client_conversation_id: Some(conversation_id), + model_id: Some(model.clone()), + ..AIIdentifiers::default() + }; + let response_stream = ctx.add_model(|ctx| { + ResponseStream::new_restored_provider_projection( + model, + messages_sent.clone(), + ai_identifiers, + ctx, + ) + }); + let stream_id = response_stream.as_ref(ctx).id().clone(); + let response_stream_clone = response_stream.clone(); + ctx.subscribe_to_model(&response_stream, move |me, _, event, ctx| { + me.handle_response_stream_event( + did_input_contain_user_query, + event, + &response_stream_clone, + ctx, + ); + }); + let rebind_result = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.rebind_provider_projection( + conversation_id, + &projection_target.task_id, + projection_target.exchange_id, + stream_id.clone(), + self.terminal_surface_id, + ctx, + ) + }); + if let Err(error) = rebind_result { + ctx.unsubscribe_from_model(&response_stream); + self.fail_restored_provider_run( + conversation_id, + format!("failed to rebind restored provider projection: {error:?}"), + ctx, + ); + return; + } + + self.in_flight_response_streams.register_new_stream( + stream_id.clone(), + conversation_id, + response_stream.clone(), + CancellationReason::FollowUpSubmitted { + is_for_same_conversation: true, + }, + ctx, + ); + self.active_provider_runs.insert( + conversation_id, + ActiveProviderRunSlot { + stream_id, + response_stream, + did_input_contain_user_query, + run_id, + root_task_id, + projection_target, + run: Some(ActiveProviderRun { + coordinator, + projector: ProviderRunResponseProjector::restored(response_config.clone()), + response_config, + action_context, + messages_sent, + persistence_offset, + }), + checkpoint: None, + turn_control: None, + cancellation_reason: None, + committed_provider_batch, + command_action_refs, + command_monitor, + pending_monitor_observation, + pending_command_completion, + monitor_prose_continuations, + }, + ); + self.restoring_provider_runs.remove(&conversation_id); + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.update_conversation_status( + self.terminal_surface_id, + conversation_id, + ConversationStatus::InProgress, + ctx, + ); + }); + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run(conversation_id, error, ctx); + return; + } + self.resume_restored_provider_run(conversation_id, ctx); + } + + fn resume_restored_provider_run( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + let should_advance_boundary = self + .active_provider_runs + .get(&conversation_id) + .and_then(|slot| slot.run.as_ref()) + .is_some_and(|run| { + matches!( + run.coordinator.run().state(), + ProviderRunState::ReadyToCallModel | ProviderRunState::AwaitingDriver { .. } + ) + }); + if !should_advance_boundary { + self.drive_active_provider_run(conversation_id, ctx); + return; + } + + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return; + }; + let Some(mut run) = slot.run.take() else { + return; + }; + let boundary = match Self::advance_provider_at_safe_boundary(slot, &mut run) { + Ok(boundary) => boundary, + Err(error) => { + let message = format!("failed to resume restored provider run: {error}"); + let _ = run + .coordinator + .run_mut() + .fail(ProviderRunFailureKind::Restore, message); + ProviderBoundaryDisposition::Advance { + completed_block_id: None, + } + } + }; + slot.run = Some(run); + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run(conversation_id, error, ctx); + return; + } + match boundary { + ProviderBoundaryDisposition::Advance { completed_block_id } => { + if let Some(block_id) = completed_block_id { + self.deactivate_provider_cli_task(conversation_id, &block_id, ctx); + } + self.drive_active_provider_run(conversation_id, ctx); + } + ProviderBoundaryDisposition::Park => {} + } + } + + fn fail_restored_provider_run( + &mut self, + conversation_id: AIConversationId, + message: String, + ctx: &mut ModelContext, + ) { + self.restoring_provider_runs.remove(&conversation_id); + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.update_conversation_status_with_error( + self.terminal_surface_id, + conversation_id, + ConversationStatus::Error, + Some(RenderableAIError::Other { + error_message: format!("Failed to restore active provider run: {message}"), + will_attempt_resume: false, + waiting_for_network: false, + is_user_error: false, + }), + ctx, + ); + }); + // Clearing the snapshot writes the conversation after its error status has been updated. + if let Err(error) = self.clear_persisted_active_provider_run(conversation_id, ctx) { + log::error!("Failed to clear unrestorable provider run: {error}"); + } + } + + fn prepare_active_provider_run( + &mut self, + conversation_id: AIConversationId, + stream_id: ResponseStreamId, + base_provider_config: crate::ai::provider::ProviderConfig, + cli_provider_config: crate::ai::provider::ProviderConfig, + request_params: api::RequestParams, + ctx: &mut ModelContext, + ) { + let _ = ctx.spawn( + async move { + prepare_provider_run(base_provider_config, cli_provider_config, request_params) + .await + }, + move |me, result, ctx| { + me.handle_prepared_provider_run(conversation_id, stream_id, result, ctx); + }, + ); + } + + fn handle_prepared_provider_run( + &mut self, + conversation_id: AIConversationId, + stream_id: ResponseStreamId, + result: anyhow::Result, + ctx: &mut ModelContext, + ) { + let Some(slot) = self.active_provider_runs.get(&conversation_id) else { + return; + }; + if slot.stream_id != stream_id { + return; + } + let provider_run_id = slot.run_id.clone(); + let prepared = match result { + Ok(prepared) => prepared, + Err(error) => { + self.fail_provider_startup(conversation_id, stream_id, error.to_string(), ctx); + return; + } + }; + let PreparedProviderRun { + base_profile, + cli_monitor_profile, + tool_result_archive, + messages_sent, + persistence_offset, + response_config, + action_context, + } = prepared; + let mut coordinator = match ProviderRunCoordinator::from_request( + provider_run_id, + base_profile.runtime, + base_profile.request, + tool_result_archive, + ProviderRunLimits::default(), + ) { + Ok(coordinator) => coordinator, + Err(error) => { + self.fail_provider_startup(conversation_id, stream_id, error.to_string(), ctx); + return; + } + }; + if let Some(profile) = cli_monitor_profile { + if let Err(error) = coordinator.insert_profile( + CLI_MONITOR_PROVIDER_PROFILE, + profile.runtime, + profile.request, + ) { + self.fail_provider_startup(conversation_id, stream_id, error.to_string(), ctx); + return; + } + } + let cancellation_reason = self + .active_provider_runs + .get(&conversation_id) + .and_then(|slot| slot.cancellation_reason); + let mut run = ActiveProviderRun { + coordinator, + projector: ProviderRunResponseProjector::new(response_config.clone()), + response_config, + action_context, + messages_sent, + persistence_offset, + }; + if let Some(reason) = cancellation_reason { + if let Err(error) = run.coordinator.run_mut().cancel(reason.to_string()) { + log::error!("Failed to cancel provider run during startup: {error}"); + } + } + if let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) { + slot.run = Some(run); + } + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run(conversation_id, error, ctx); + return; + } + self.drive_active_provider_run(conversation_id, ctx); + } + + fn persist_active_provider_run( + &self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) -> Result<(), String> { + let slot = self + .active_provider_runs + .get(&conversation_id) + .ok_or_else(|| "active provider run disappeared before persistence".to_string())?; + let snapshot = ActiveProviderRunSnapshot::from_slot(slot)?; + let json = serde_json::to_string(&snapshot) + .map_err(|error| format!("failed to serialize active provider run: {error}"))?; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model + .persist_active_provider_run_json(conversation_id, Some(json), ctx) + .map_err(|error| format!("failed to persist active provider run: {error:?}")) + }) + } + + fn clear_persisted_active_provider_run( + &self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) -> Result<(), String> { + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model + .persist_active_provider_run_json(conversation_id, None, ctx) + .map_err(|error| format!("failed to clear active provider run: {error:?}")) + }) + } + + fn drive_active_provider_run( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return; + }; + let Some(mut run) = slot.run.take() else { + return; + }; + let checkpoint_template = match ActiveProviderRunCheckpoint::from_active_run(&run) { + Ok(checkpoint) => checkpoint, + Err(error) => { + slot.run = Some(run); + self.fail_active_provider_run(conversation_id, error, ctx); + return; + } + }; + slot.checkpoint = Some(checkpoint_template.clone()); + let stream_id = slot.stream_id.clone(); + let (turn_control_sender, turn_control) = turn_control(); + if slot.cancellation_reason.is_some() { + let _ = turn_control_sender.try_send(TurnCommand::Cancel); + } + slot.turn_control = Some(turn_control_sender); + + let (sender, receiver) = async_channel::unbounded(); + ctx.spawn_stream_local( + receiver, + move |me, message, ctx| { + me.handle_provider_drive_message(conversation_id, &stream_id, message, ctx); + }, + |_, _| {}, + ); + let _ = ctx.spawn( + async move { + let projection_sender = sender.clone(); + let checkpoint_sender = sender.clone(); + let result = run + .coordinator + .drive_until_blocked_with_checkpoint( + turn_control, + |projection| { + for event in run.projector.project(projection)? { + projection_sender + .try_send(ProviderDriveMessage::Response(event)) + .map_err(|_| { + "provider response projection receiver was closed" + .to_string() + })?; + } + Ok(()) + }, + move |provider_run| { + let checkpoint_sender = checkpoint_sender.clone(); + let checkpoint = checkpoint_template.with_run(provider_run); + Box::pin(async move { + let (acknowledgement, receiver) = oneshot::channel(); + checkpoint_sender + .send(ProviderDriveMessage::Checkpoint { + checkpoint, + acknowledgement, + }) + .await + .map_err(|_| { + "provider checkpoint receiver was closed".to_string() + })?; + receiver.await.map_err(|_| { + "provider checkpoint acknowledgement was dropped".to_string() + })? + }) + }, + ) + .await + .map_err(|error| error.to_string()); + let _ = sender + .send(ProviderDriveMessage::Blocked { run, result }) + .await; + }, + |_, _, _| {}, + ); + } + + fn handle_provider_drive_message( + &mut self, + conversation_id: AIConversationId, + stream_id: &ResponseStreamId, + message: ProviderDriveMessage, + ctx: &mut ModelContext, + ) { + let Some(slot) = self.active_provider_runs.get(&conversation_id) else { + return; + }; + if &slot.stream_id != stream_id { + return; + } + match message { + ProviderDriveMessage::Response(event) => { + let response_stream = slot.response_stream.clone(); + let did_input_contain_user_query = slot.did_input_contain_user_query; + let event = ResponseStream::projected_event(event); + self.handle_response_stream_event( + did_input_contain_user_query, + &event, + &response_stream, + ctx, + ); + } + ProviderDriveMessage::Checkpoint { + checkpoint, + acknowledgement, + } => { + let result = match self.active_provider_runs.get_mut(&conversation_id) { + Some(slot) if checkpoint.run.id() == &slot.run_id => { + slot.checkpoint = Some(checkpoint); + self.persist_active_provider_run(conversation_id, ctx) + } + Some(_) => Err("provider checkpoint run identity did not match".to_string()), + None => Err("provider run disappeared before checkpoint".to_string()), + }; + let _ = acknowledgement.send(result); + } + ProviderDriveMessage::Blocked { run, result } => { + self.handle_provider_run_blocked(conversation_id, run, result, ctx); + } + } + } + + fn advance_provider_at_safe_boundary( + slot: &mut ActiveProviderRunSlot, + run: &mut ActiveProviderRun, + ) -> Result { + let state = run.coordinator.run().state(); + let phase = provider_boundary_phase(state); + let intent = provider_boundary_intent( + phase, + slot.committed_provider_batch.is_some(), + slot.pending_command_completion.is_some(), + slot.pending_monitor_observation.is_some(), + run.coordinator.run().profile().as_str() == CLI_MONITOR_PROVIDER_PROFILE, + slot.command_monitor.is_some(), + slot.monitor_prose_continuations, + ); + if intent == ProviderBoundaryIntent::Park { + return Ok(ProviderBoundaryDisposition::Park); + } + if intent == ProviderBoundaryIntent::Advance { + return Ok(ProviderBoundaryDisposition::Advance { + completed_block_id: None, + }); + } + + let ready_work_id = run.coordinator.run().ready_work_id(); + let awaiting_driver_work_id = match state { + ProviderRunState::AwaitingDriver { work_id, .. } => Some(work_id.clone()), + ProviderRunState::ReadyToCallModel + | ProviderRunState::AwaitingModel { .. } + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingTools { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => None, + }; + + if intent == ProviderBoundaryIntent::ApplyCompletion { + let completion = slot + .pending_command_completion + .take() + .expect("boundary intent checked completion mailbox"); + let Some(work_id) = ready_work_id.as_ref().or(awaiting_driver_work_id.as_ref()) else { + slot.pending_command_completion = Some(completion); + return Ok(ProviderBoundaryDisposition::Park); + }; + let completed_block_id = completion.block_id.clone(); + if let Some(monitor) = slot.command_monitor.as_ref() { + log::debug!( + "Completing provider command monitor run={:?} work={:?} call={} block={:?}", + monitor.run_id, + monitor.originating_work_id, + monitor.originating_call_id, + monitor.block_id + ); + } + run.set_task_id(&slot.root_task_id); + let observation = completion.observation(); + if ready_work_id.is_some() { + run.coordinator + .run_mut() + .continue_ready_with_observation(work_id, observation, BASE_PROVIDER_PROFILE) + .map_err(|error| error.to_string())?; + } else { + run.coordinator + .run_mut() + .continue_with_observation(work_id, observation, BASE_PROVIDER_PROFILE) + .map_err(|error| error.to_string())?; + } + if let Some(action_id) = completion.initial_requested_command_action_id.as_ref() { + slot.command_action_refs.remove(action_id); + } + slot.command_monitor = None; + slot.pending_monitor_observation = None; + slot.monitor_prose_continuations = 0; + return Ok(ProviderBoundaryDisposition::Advance { + completed_block_id: Some(completed_block_id), + }); + } + + if intent == ProviderBoundaryIntent::ApplyMonitorObservation { + let observation = slot + .pending_monitor_observation + .take() + .expect("boundary intent checked monitor mailbox"); + let work_id = ready_work_id + .as_ref() + .or(awaiting_driver_work_id.as_ref()) + .expect("ready boundary must have work identity"); + let Some(monitor) = slot + .command_monitor + .as_ref() + .filter(|monitor| monitor.block_id == observation.block_id) + else { + return Err("provider command monitor observation lost its owner".to_string()); + }; + run.set_task_id(&observation.cli_task_id); + let message = MessageContent::Text(format!( + "The command is still running. Continue monitoring block {:?} with the CLI tools \ + and do not claim completion until final command evidence is available.\n\nCommand:\n{}", + monitor.block_id, monitor.command + )); + if ready_work_id.is_some() { + run.coordinator + .run_mut() + .continue_ready_with_observation(work_id, message, CLI_MONITOR_PROVIDER_PROFILE) + .map_err(|error| error.to_string())?; + } else { + run.coordinator + .run_mut() + .continue_with_observation(work_id, message, CLI_MONITOR_PROVIDER_PROFILE) + .map_err(|error| error.to_string())?; + } + slot.monitor_prose_continuations = 0; + return Ok(ProviderBoundaryDisposition::Advance { + completed_block_id: None, + }); + } + + let work_id = awaiting_driver_work_id.expect("driver boundary must have work identity"); + if intent == ProviderBoundaryIntent::RetryMonitor { + let monitor = slot + .command_monitor + .as_ref() + .expect("boundary intent checked monitor"); + run.set_task_id(&monitor.cli_task_id); + run.coordinator + .run_mut() + .continue_with_observation( + &work_id, + MessageContent::Text(format!( + "The command in block {:?} is still active. Poll it now with a CLI tool; \ + do not respond with only an acknowledgement.", + monitor.block_id + )), + CLI_MONITOR_PROVIDER_PROFILE, + ) + .map_err(|error| error.to_string())?; + slot.monitor_prose_continuations += 1; + return Ok(ProviderBoundaryDisposition::Advance { + completed_block_id: None, + }); + } + + run.coordinator + .run_mut() + .complete(&work_id) + .map_err(|error| error.to_string())?; + Ok(ProviderBoundaryDisposition::Advance { + completed_block_id: None, + }) + } + + fn handle_provider_run_blocked( + &mut self, + conversation_id: AIConversationId, + mut run: ActiveProviderRun, + result: Result, + ctx: &mut ModelContext, + ) { + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return; + }; + slot.turn_control = None; + let cancellation_reason = slot.cancellation_reason; + if let Some(reason) = cancellation_reason { + if !run.coordinator.run().is_terminal() { + let _ = run.coordinator.run_mut().cancel(reason.to_string()); + } + slot.run = Some(run); + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + log::error!("Failed to persist cancelled provider run: {error}"); + } + self.drive_active_provider_run(conversation_id, ctx); + return; + } + let block = match result { + Ok(block) => block, + Err(message) => { + if let Err(error) = run + .coordinator + .run_mut() + .fail(ProviderRunFailureKind::ExternalWork, message) + { + log::error!("Failed to record provider driver failure: {error}"); + } + slot.run = Some(run); + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + log::error!("Failed to persist failed provider run: {error}"); + } + self.drive_active_provider_run(conversation_id, ctx); + return; + } + }; + match block { + ProviderRunBlock::Tools(batch) => { + 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 provider tool batch: {error}"), + ctx, + ); + return; + } + self.queue_provider_tool_batch(conversation_id, batch, ctx); + } + ProviderRunBlock::AwaitingDriver { .. } => { + let disposition = match Self::advance_provider_at_safe_boundary(slot, &mut run) { + Ok(disposition) => disposition, + Err(error) => { + let message = format!("failed to advance provider run: {error}"); + let _ = run + .coordinator + .run_mut() + .fail(ProviderRunFailureKind::Protocol, message); + ProviderBoundaryDisposition::Advance { + completed_block_id: None, + } + } + }; + 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 provider driver transition: {error}"), + ctx, + ); + return; + } + match disposition { + ProviderBoundaryDisposition::Advance { completed_block_id } => { + if let Some(block_id) = completed_block_id { + self.deactivate_provider_cli_task(conversation_id, &block_id, ctx); + } + self.drive_active_provider_run(conversation_id, ctx); + } + ProviderBoundaryDisposition::Park => {} + } + } + ProviderRunBlock::Done(outcome) => { + slot.run = Some(run); + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + log::error!("Failed to persist terminal provider run: {error}"); + } + let run = self + .active_provider_runs + .get_mut(&conversation_id) + .and_then(|slot| slot.run.take()) + .expect("terminal provider run was just restored to its slot"); + self.finish_active_provider_run(conversation_id, run, outcome, ctx); + } + } + } + + fn queue_provider_tool_batch( + &mut self, + conversation_id: AIConversationId, + batch: PendingToolBatch, + ctx: &mut ModelContext, + ) { + let conversion = self + .active_provider_runs + .get(&conversation_id) + .and_then(|slot| slot.run.as_ref()) + .map(|run| { + batch + .calls + .iter() + .filter(|pending| pending.state.result().is_none()) + .map(|pending| run.action_context.action_from_tool_call(&pending.call)) + .collect::, _>>() + }); + let actions = match conversion { + Some(Ok(actions)) => actions, + Some(Err(message)) => { + self.fail_active_provider_run(conversation_id, message, ctx); + return; + } + None => return, + }; + let stream_id = self.active_provider_runs[&conversation_id] + .stream_id + .clone(); + for action in &actions { + let apply_result = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.apply_domain_tool_proposal( + &stream_id, + conversation_id, + self.terminal_surface_id, + action.clone(), + ctx, + ) + }); + if let Err(error) = apply_result { + self.fail_active_provider_run( + conversation_id, + format!("failed to attach provider tool proposal: {error:?}"), + ctx, + ); + return; + } + } + if let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) { + slot.command_action_refs.extend( + actions + .iter() + .filter(|action| is_provider_command_action(&action.action)) + .map(|action| { + ( + action.id.clone(), + ProviderToolExecutionRef::new( + conversation_id, + &batch.work_id, + action.id.to_string(), + ), + ) + }), + ); + } + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run( + conversation_id, + format!("failed to persist queued provider actions: {error}"), + ctx, + ); + return; + } + let queue_result = self.action_model.update(ctx, |action_model, ctx| { + action_model.queue_provider_actions(actions, conversation_id, &batch, ctx) + }); + if let Err(error) = queue_result { + self.fail_active_provider_run(conversation_id, error.to_string(), ctx); + } + } + + fn handle_provider_tool_lifecycle( + &mut self, + execution_ref: &ProviderToolExecutionRef, + event: &galaxy_agent_core::ToolEvent, + ctx: &mut ModelContext, + ) { + let conversation_id = execution_ref.conversation_id; + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return; + }; + if slot.cancellation_reason.is_some() { + return; + } + let Some(run) = slot.run.as_mut() else { + return; + }; + if !provider_execution_matches_active_work( + run.coordinator.run().id(), + run.coordinator.run().active_work_id(), + execution_ref, + ) { + return; + } + let should_drive = match run.coordinator.apply_tool_lifecycle(execution_ref, event) { + Ok(ProviderToolLifecycleOutcome::Pending) => false, + Ok(ProviderToolLifecycleOutcome::BatchCommitted) => { + slot.committed_provider_batch = Some(execution_ref.work_id()); + false + } + Err(error) => { + let message = format!("invalid provider tool lifecycle: {error}"); + if let Err(fail_error) = run + .coordinator + .run_mut() + .fail(ProviderRunFailureKind::Protocol, message) + { + log::error!("Failed to record provider tool lifecycle failure: {fail_error}"); + } + true + } + }; + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run( + conversation_id, + format!("failed to persist provider tool lifecycle: {error}"), + ctx, + ); + return; + } + if should_drive { + self.drive_active_provider_run(conversation_id, ctx); + } + } + + fn handle_provider_command_action_results( + &mut self, + conversation_id: AIConversationId, + work_id: &ExternalWorkId, + results: &[Arc], + ctx: &mut ModelContext, + ) -> Result<(), String> { + for result in results { + let Some(command_result) = classify_provider_command_result(&result.result) else { + continue; + }; + let Some(action_ref) = self + .active_provider_runs + .get(&conversation_id) + .and_then(|slot| slot.command_action_refs.get(&result.id)) + .cloned() + else { + continue; + }; + if action_ref.work_id() != *work_id { + continue; + } + + match command_result { + ProviderCommandResult::Snapshot { block_id, command } => { + let completion_already_pending = { + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return Ok(()); + }; + let existing_monitor = slot + .command_monitor + .as_ref() + .filter(|monitor| monitor.block_id == block_id); + let expected_initial_action_id = existing_monitor + .map(|monitor| monitor.initial_requested_command_action_id.clone()) + .unwrap_or_else(|| result.id.clone()); + let fallback_command = + existing_monitor.map(|monitor| monitor.command.clone()); + reconcile_provider_completion_with_snapshot( + slot.pending_command_completion.as_mut(), + &block_id, + &expected_initial_action_id, + command.as_deref(), + fallback_command.as_deref(), + )? + }; + if completion_already_pending { + continue; + } + let cli_task_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.create_cli_subagent_task_for_conversation( + block_id.clone(), + conversation_id, + self.terminal_surface_id, + ctx, + ) + }); + let cli_task_id = cli_task_id.map_err(|error| { + format!( + "failed to create provider CLI task for block {block_id:?}: {error:?}" + ) + })?; + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return Ok(()); + }; + if slot.run_id != action_ref.run_id + || slot.committed_provider_batch.as_ref() != Some(work_id) + { + continue; + } + let existing_monitor = slot + .command_monitor + .as_ref() + .filter(|monitor| monitor.block_id == block_id); + let initial_requested_command_action_id = existing_monitor + .map(|monitor| monitor.initial_requested_command_action_id.clone()) + .unwrap_or_else(|| result.id.clone()); + let command = command + .or_else(|| existing_monitor.map(|monitor| monitor.command.clone())) + .unwrap_or_default(); + slot.command_monitor = Some(ProviderCommandMonitorState { + run_id: slot.run_id.clone(), + originating_work_id: work_id.clone(), + originating_call_id: result.id.to_string(), + initial_requested_command_action_id, + block_id: block_id.clone(), + command, + cli_task_id: cli_task_id.clone(), + }); + slot.pending_monitor_observation = Some(PendingProviderMonitorObservation { + block_id, + cli_task_id, + }); + } + ProviderCommandResult::Finished { + block_id, + command, + output, + exit_code, + } => { + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return Ok(()); + }; + let Some(monitor) = slot + .command_monitor + .as_ref() + .filter(|monitor| monitor.block_id == block_id) + else { + continue; + }; + if slot.pending_command_completion.is_none() { + slot.pending_command_completion = Some(PendingProviderCommandCompletion { + block_id, + initial_requested_command_action_id: Some( + monitor.initial_requested_command_action_id.clone(), + ), + command: command.unwrap_or_else(|| monitor.command.clone()), + output, + exit_code, + }); + } + } + } + } + if let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) { + let retained_action_ids = [ + slot.command_monitor + .as_ref() + .map(|monitor| monitor.initial_requested_command_action_id.clone()), + slot.pending_command_completion + .as_ref() + .and_then(|completion| completion.initial_requested_command_action_id.clone()), + ]; + slot.command_action_refs.retain(|action_id, execution_ref| { + execution_ref.work_id() != *work_id + || retained_action_ids + .iter() + .flatten() + .any(|retained| retained == action_id) + }); + } + self.persist_active_provider_run(conversation_id, ctx)?; + Ok(()) + } + + fn deactivate_provider_cli_task( + &mut self, + conversation_id: AIConversationId, + block_id: &BlockId, + ctx: &mut ModelContext, + ) { + let result = BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, _| { + history_model.deactivate_cli_subagent_task_for_conversation(block_id, conversation_id) + }); + if let Err(error) = result { + log::error!("Failed to deactivate provider CLI task for block {block_id:?}: {error:?}"); + } + } + + fn handle_provider_actions_finished( + &mut self, + conversation_id: AIConversationId, + execution_ref: &ProviderToolExecutionRef, + ctx: &mut ModelContext, + ) { + let disposition = self + .active_provider_runs + .get(&conversation_id) + .and_then(|slot| { + let run = slot.run.as_ref()?; + Some(provider_finished_action_disposition( + run.coordinator.run().id(), + run.coordinator.run().active_work_id(), + slot.committed_provider_batch.as_ref(), + execution_ref, + )) + }) + .unwrap_or(ProviderFinishedActionDisposition::Ignore); + if disposition == ProviderFinishedActionDisposition::AwaitBatchCommit { + return; + } + let work_id = execution_ref.work_id(); + let results = (disposition == ProviderFinishedActionDisposition::Resume).then(|| { + self.action_model + .as_ref(ctx) + .provider_finished_action_results(conversation_id, &work_id) + }); + let command_result = results.as_deref().map(|results| { + self.handle_provider_command_action_results(conversation_id, &work_id, results, ctx) + }); + self.action_model.update(ctx, |action_model, _| { + action_model.archive_provider_finished_action_results(conversation_id, &work_id); + }); + if let Some(Err(error)) = command_result { + self.fail_active_provider_run(conversation_id, error, ctx); + return; + } + if disposition != ProviderFinishedActionDisposition::Resume { + return; + } + + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return; + }; + slot.committed_provider_batch = None; + let Some(mut run) = slot.run.take() else { + return; + }; + let boundary = match Self::advance_provider_at_safe_boundary(slot, &mut run) { + Ok(boundary) => boundary, + Err(error) => { + let message = format!("failed to resume provider tool batch: {error}"); + let _ = run + .coordinator + .run_mut() + .fail(ProviderRunFailureKind::Protocol, message); + ProviderBoundaryDisposition::Advance { + completed_block_id: None, + } + } + }; + 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 provider batch continuation: {error}"), + ctx, + ); + return; + } + match boundary { + ProviderBoundaryDisposition::Advance { completed_block_id } => { + if let Some(block_id) = completed_block_id { + self.deactivate_provider_cli_task(conversation_id, &block_id, ctx); + } + self.drive_active_provider_run(conversation_id, ctx); + } + ProviderBoundaryDisposition::Park => {} + } + } + + fn fail_active_provider_run( + &mut self, + conversation_id: AIConversationId, + message: String, + ctx: &mut ModelContext, + ) { + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return; + }; + let Some(run) = slot.run.as_mut() else { + return; + }; + if let Err(error) = run + .coordinator + .run_mut() + .fail(ProviderRunFailureKind::Protocol, message) + { + log::error!("Failed to terminate provider run: {error}"); + } + self.drive_active_provider_run(conversation_id, ctx); + } + + fn finish_active_provider_run( + &mut self, + conversation_id: AIConversationId, + mut run: ActiveProviderRun, + outcome: ProviderRunOutcome, + ctx: &mut ModelContext, + ) { + if let Ok(mut messages_sent) = run.messages_sent.lock() { + let transcript = run.coordinator.run().transcript(); + let offset = run.persistence_offset.min(transcript.len()); + *messages_sent = transcript[offset..].to_vec(); + } + let events = match run.projector.finish(&outcome) { + Ok(events) => events, + Err(message) => { + self.fail_provider_startup( + conversation_id, + self.active_provider_runs[&conversation_id] + .stream_id + .clone(), + format!("failed to finish provider response projection: {message}"), + ctx, + ); + return; + } + }; + let Some(slot) = self.active_provider_runs.get(&conversation_id) else { + return; + }; + let stream_id = slot.stream_id.clone(); + let response_stream = slot.response_stream.clone(); + let did_input_contain_user_query = slot.did_input_contain_user_query; + for event in events { + let event = ResponseStream::projected_event(event); + self.handle_response_stream_event( + did_input_contain_user_query, + &event, + &response_stream, + ctx, + ); + } + if matches!(outcome, ProviderRunOutcome::Cancelled { .. }) { + let cancellation_reason = + self.active_provider_runs[&conversation_id].cancellation_reason; + if let Some(reason) = cancellation_reason { + let status = match reason.conversation_outcome() { + CancellationOutcome::KeepInProgress => ConversationStatus::InProgress, + CancellationOutcome::Succeeded => ConversationStatus::Success, + CancellationOutcome::Cancelled => ConversationStatus::Cancelled, + CancellationOutcome::FinalizedExternally => { + self.cleanup_active_provider_run( + conversation_id, + &stream_id, + &response_stream, + ctx, + ); + return; + } + }; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.update_conversation_status( + self.terminal_surface_id, + conversation_id, + status, + ctx, + ); + }); + } + } + self.cleanup_active_provider_run(conversation_id, &stream_id, &response_stream, ctx); + } + + fn fail_provider_startup( + &mut self, + conversation_id: AIConversationId, + stream_id: ResponseStreamId, + message: String, + ctx: &mut ModelContext, + ) { + let response_stream = self + .active_provider_runs + .get(&conversation_id) + .filter(|slot| slot.stream_id == stream_id) + .map(|slot| slot.response_stream.clone()); + let Some(response_stream) = response_stream else { + return; + }; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.mark_response_stream_completed_with_error( + RenderableAIError::Other { + error_message: message, + will_attempt_resume: false, + waiting_for_network: false, + is_user_error: false, + }, + false, + &stream_id, + conversation_id, + self.terminal_surface_id, + ctx, + ); + }); + self.cleanup_active_provider_run(conversation_id, &stream_id, &response_stream, ctx); + } + + fn cleanup_active_provider_run( + &mut self, + conversation_id: AIConversationId, + stream_id: &ResponseStreamId, + response_stream: &ModelHandle, + ctx: &mut ModelContext, + ) { + if let Err(error) = self.clear_persisted_active_provider_run(conversation_id, ctx) { + log::error!("Failed to clear persisted provider run during cleanup: {error}"); + } + self.active_provider_runs.remove(&conversation_id); + self.restoring_provider_runs.remove(&conversation_id); + self.in_flight_response_streams.cleanup_stream(stream_id); + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, _| { + if let Some(conversation) = history_model.conversation_mut(&conversation_id) { + conversation.cleanup_completed_response_stream(stream_id); + } + }); + ctx.unsubscribe_from_model(response_stream); + ctx.emit(BlocklistAIControllerEvent::FinishedReceivingOutput { + stream_id: stream_id.clone(), + conversation_id, + }); + AIRequestUsageModel::handle(ctx).update(ctx, |request_usage_model, ctx| { + request_usage_model.refresh_request_usage_async(ctx); + }); + self.maybe_refresh_ai_overages(ctx); + } + + fn cancel_active_provider_run( + &mut self, + conversation_id: AIConversationId, + reason: CancellationReason, + ctx: &mut ModelContext, + ) -> bool { + let Some(mut slot) = self.active_provider_runs.remove(&conversation_id) else { + return false; + }; + slot.cancellation_reason = Some(reason); + if let Some(turn_control) = &slot.turn_control { + let _ = turn_control.try_send(TurnCommand::Cancel); + } + if let Some(mut run) = slot.run.take() { + if !run.coordinator.run().is_terminal() { + let _ = run.coordinator.run_mut().cancel(reason.to_string()); + } + if let Ok(mut messages_sent) = run.messages_sent.lock() { + let transcript = run.coordinator.run().transcript(); + let offset = run.persistence_offset.min(transcript.len()); + *messages_sent = transcript[offset..].to_vec(); + } + } + + self.action_model.update(ctx, |action_model, ctx| { + action_model.cancel_all_pending_actions(conversation_id, Some(reason), ctx); + }); + + let cancellation_outcome = reason.conversation_outcome(); + if FeatureFlag::AgentSharedSessions.is_enabled() + && !matches!(cancellation_outcome, CancellationOutcome::KeepInProgress) + { + self.send_cancellation_to_viewers(ctx); + } + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.mark_response_stream_cancelled( + &slot.stream_id, + conversation_id, + self.terminal_surface_id, + reason, + ctx, + ); + }); + if matches!(cancellation_outcome, CancellationOutcome::Cancelled) { + self.set_input_mode_for_cancellation(ctx); + } + + self.cleanup_active_provider_run( + conversation_id, + &slot.stream_id, + &slot.response_stream, + ctx, + ); + true + } + + fn cancel_active_provider_run_for_stream( + &mut self, + stream_id: &ResponseStreamId, + reason: CancellationReason, + ctx: &mut ModelContext, + ) -> bool { + let conversation_id = + self.active_provider_runs + .iter() + .find_map(|(conversation_id, slot)| { + (&slot.stream_id == stream_id).then_some(*conversation_id) + }); + conversation_id.is_some_and(|conversation_id| { + self.cancel_active_provider_run(conversation_id, reason, ctx) + }) + } + /// Cancels a pending AI request response stream, given the exchange ID, if it exists. /// Returns true if a pending stream was found and canceled, false otherwise. pub fn try_cancel_pending_response_stream( @@ -3725,8 +5955,103 @@ impl BlocklistAIController { reason: CancellationReason, ctx: &mut ModelContext, ) -> bool { - self.in_flight_response_streams - .try_cancel_stream(stream_id, reason, ctx) + self.cancel_active_provider_run_for_stream(stream_id, reason, ctx) + || self + .in_flight_response_streams + .try_cancel_stream(stream_id, reason, ctx) + } + + pub(super) fn has_active_provider_run(&self, conversation_id: AIConversationId) -> bool { + self.active_provider_runs.contains_key(&conversation_id) + } + + pub(super) fn accept_provider_command_completion( + &mut self, + conversation_id: AIConversationId, + mut completion: PendingProviderCommandCompletion, + ctx: &mut ModelContext, + ) -> bool { + let should_wake = { + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return false; + }; + if !provider_command_completion_matches( + &slot.run_id, + &slot.command_action_refs, + slot.command_monitor.as_ref(), + &completion.block_id, + completion.initial_requested_command_action_id.as_ref(), + ) { + return false; + } + + if completion.command.is_empty() { + completion.command = slot + .command_monitor + .as_ref() + .map(|monitor| monitor.command.clone()) + .unwrap_or_default(); + } + slot.pending_command_completion = Some(completion); + slot.committed_provider_batch.is_none() + && slot.run.as_ref().is_some_and(|run| { + matches!( + run.coordinator.run().state(), + ProviderRunState::ReadyToCallModel + | ProviderRunState::AwaitingDriver { .. } + ) + }) + }; + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run( + conversation_id, + format!("failed to persist provider command completion: {error}"), + ctx, + ); + return true; + } + if !should_wake { + return true; + } + + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return true; + }; + let Some(mut run) = slot.run.take() else { + return true; + }; + let boundary = match Self::advance_provider_at_safe_boundary(slot, &mut run) { + Ok(boundary) => boundary, + Err(error) => { + let message = format!("failed to apply provider command completion: {error}"); + let _ = run + .coordinator + .run_mut() + .fail(ProviderRunFailureKind::Protocol, message); + ProviderBoundaryDisposition::Advance { + completed_block_id: None, + } + } + }; + 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 provider command boundary: {error}"), + ctx, + ); + return true; + } + match boundary { + ProviderBoundaryDisposition::Advance { completed_block_id } => { + if let Some(block_id) = completed_block_id { + self.deactivate_provider_cli_task(conversation_id, &block_id, ctx); + } + self.drive_active_provider_run(conversation_id, ctx); + } + ProviderBoundaryDisposition::Park => {} + } + true } pub fn has_active_stream_for_conversation( @@ -3768,25 +6093,18 @@ impl BlocklistAIController { reason: CancellationReason, ctx: &mut ModelContext, ) { - // Cancel any pending auto-resume for this conversation. - if let Some(handle) = self.pending_auto_resume_handles.remove(&conversation_id) { - handle.abort(); - } - // Discard any queued passive suggestion results for this conversation. self.pending_passive_suggestion_results .remove(&conversation_id); - if !self - .in_flight_response_streams - .try_cancel_streams_for_conversation(conversation_id, reason, ctx) + let cancelled_provider = self.cancel_active_provider_run(conversation_id, reason, ctx); + if !cancelled_provider + && !self + .in_flight_response_streams + .try_cancel_streams_for_conversation(conversation_id, reason, ctx) { // No active stream whose cancellation would mark the conversation `Cancelled`. - // A parked auto-resume was aborted above; nothing else will move the - // conversation out of TransientError, so surface the cancellation directly. - // - // TODO(REMOTE-1950): Track the parked auto-resume as a first-class cancelable so its - // cancellation flows through the same `AfterStreamFinished` path, dropping this special case. + // Surface cancellation directly when the conversation is parked in a transient error. if matches!( reason.conversation_outcome(), CancellationOutcome::Cancelled @@ -3947,8 +6265,10 @@ impl BlocklistAIController { reason: CancellationReason, ctx: &mut ModelContext, ) -> bool { - self.in_flight_response_streams - .try_cancel_stream(response_stream_id, reason, ctx) + self.cancel_active_provider_run_for_stream(response_stream_id, reason, ctx) + || self + .in_flight_response_streams + .try_cancel_stream(response_stream_id, reason, ctx) } fn handle_response_stream_event( @@ -3978,58 +6298,6 @@ impl BlocklistAIController { }; let history_model = BlocklistAIHistoryModel::handle(ctx); match event { - Ok(api::StreamEvent::ToolProposed(action)) => { - #[cfg(not(target_family = "wasm"))] - let action_log_context = ( - action.id.to_string(), - action.task_id.to_string(), - remote_action_tool_name(&action), - action.requires_result, - ); - let failed_proposal = FailedToolProposal::new(&action, String::new()); - let apply_result = history_model.update(ctx, |history_model, ctx| { - history_model.apply_domain_tool_proposal( - &stream_id, - conversation_id, - self.terminal_surface_id, - action, - ctx, - ) - }); - if let Err(error) = apply_result { - log::error!( - "Failed to apply Rig tool proposal to conversation: {error:?}" - ); - let mut failed_proposal = failed_proposal; - failed_proposal.error = format!("{error:?}"); - self.failed_tool_proposals_by_stream - .entry(stream_id.clone()) - .or_default() - .push(failed_proposal); - #[cfg(not(target_family = "wasm"))] - { - let (action_id, task_id, tool_name, requires_result) = - action_log_context; - remote_logging::log_model_event( - ctx, - RemoteLogRecord { - level: RemoteLogLevel::Error, - message: "Tool proposal apply failed".to_string(), - context: serde_json::json!({ - "event": "tool_proposal_apply_failed", - "stream_id": stream_id.as_str(), - "conversation_id": conversation_id.to_string(), - "action_id": action_id, - "task_id": task_id, - "tool_name": tool_name, - "requires_result": requires_result, - "error": remote_logging::sanitize_error(format!("{error:?}")), - }), - }, - ); - } - } - } Ok(api::StreamEvent::Response(event)) => { // If this controller is part of a shared session, forward the entire response event to viewers first. if FeatureFlag::AgentSharedSessions.is_enabled() @@ -4200,179 +6468,20 @@ impl BlocklistAIController { }); } - // Check if this error is eligible for corrective retry. - // Similar to loop detection, inject a message telling the LLM - // to try a different approach rather than just failing. - // Exclude errors that are proxy/config issues (cache_control, - // BadRequestError from LiteLLM) since the LLM can't fix those. - let error_str = format!("{e}"); - let is_proxy_config_error = error_str.contains("cache_control") - || error_str.contains("tool_use` ids were found without") - || error_str.contains("BadRequestError"); - let is_corrective_retry_candidate = !is_proxy_config_error - && !matches!(e.as_ref(), AIApiError::QuotaLimit { .. }) - && (error_str.contains("ValidationException") - || error_str.contains("context window") - || error_str.contains("too many tokens") - || error_str.contains("input is too long") - || error_str.contains("throttl") - || error_str.contains("ThrottlingException")); - - const MAX_ERROR_RETRIES: usize = 2; - let retry_count = - self.error_retry_counts.entry(conversation_id).or_insert(0); - let should_corrective_retry = - response_stream.as_ref(ctx).allows_corrective_retries() - && is_corrective_retry_candidate - && *retry_count < MAX_ERROR_RETRIES; - - if should_corrective_retry { - *retry_count += 1; - let retry_num = *retry_count; - log::warn!( - "[error-retry] Attempting corrective retry {}/{} for conversation {:?}: {}", - retry_num, - MAX_ERROR_RETRIES, + history_model.update(ctx, |history_model, ctx| { + history_model.mark_response_stream_completed_with_error( + (&e).into(), + /*recovery_pending*/ false, + &stream_id, conversation_id, - error_str + self.terminal_surface_id, + ctx, ); - - // Mark the error on the conversation but with recovery pending - let renderable_error = RenderableAIError::Other { - error_message: format!( - "Error encountered, retrying with different approach (attempt {}/{})", - retry_num, MAX_ERROR_RETRIES - ), - will_attempt_resume: true, - waiting_for_network: false, - is_user_error: false, - }; - history_model.update(ctx, |history_model, ctx| { - history_model.mark_response_stream_completed_with_error( - renderable_error, - /*recovery_pending*/ true, - &stream_id, - conversation_id, - self.terminal_surface_id, - ctx, - ); - }); - - // Inject a corrective message and resume - if let Some(conversation) = - BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) - { - let root_task_id = conversation.get_root_task_id().clone(); - let corrective_msg = format!( - "[SYSTEM] The previous request resulted in an error: {}\n\n\ - Please try a completely different approach to accomplish the goal. \ - If the error is related to context size, reduce the amount of content \ - you are working with (read fewer files, use smaller commands, break \ - the task into smaller steps). If you cannot find an alternative, \ - explain to the user what is failing and why.", - error_str - ); - - let inputs = vec![AIAgentInput::UserQuery { - query: corrective_msg, - context: Arc::from([]), - static_query_type: None, - referenced_attachments: HashMap::new(), - user_query_mode: UserQueryMode::Normal, - running_command: None, - intended_agent: None, - }]; - - let _ = self.send_request_input( - RequestInput::for_task( - inputs, - root_task_id, - &self.active_session, - self.get_current_response_initiator(), - conversation_id, - self.terminal_surface_id, - ctx, - ), - None, - /*can_attempt_resume_on_error*/ false, - /*is_queued_prompt*/ false, - ctx, - ); - } - } else { - // Clear retry count on non-retryable errors or exhausted retries - self.error_retry_counts.remove(&conversation_id); - - // A resume scheduled for this failure keeps the conversation in - // the non-terminal TransientError status instead of Error. - let recovery_pending = response_stream - .as_ref(ctx) - .should_resume_conversation_after_stream_finished(); - let mut renderable_error: RenderableAIError = (&e).into(); - if let RenderableAIError::Other { - will_attempt_resume, - waiting_for_network, - .. - } - | RenderableAIError::TransientNetworkError { - will_attempt_resume, - waiting_for_network, - .. - } = &mut renderable_error - { - // Rendering-only hints; state machine consumers key off the - // TransientError conversation status instead. - *will_attempt_resume |= recovery_pending; - if recovery_pending { - let network_status = NetworkStatus::as_ref(ctx); - *waiting_for_network = !network_status.is_online(); - } - } - - history_model.update(ctx, |history_model, ctx| { - history_model.mark_response_stream_completed_with_error( - renderable_error, - recovery_pending, - &stream_id, - conversation_id, - self.terminal_surface_id, - ctx, - ); - }); - } + }); } } } - ResponseStreamEvent::WaitingForNetwork { waiting } => { - let Some(conversation_id) = BlocklistAIHistoryModel::as_ref(ctx) - .conversation_for_response_stream(&stream_id) - else { - log::warn!("Could not find conversation for response stream: {stream_id:?}"); - return; - }; - // Mirror the parked-retry state on the conversation: TransientError while - // waiting for connectivity, back to InProgress when the retry fires. - // This event is only emitted after a recoverable request failure parks a - // retry while offline (see `defer_retry_until_online`), so treating - // `waiting` as a transient-error state is always correct here. - let status = if *waiting { - ConversationStatus::TransientError - } else { - ConversationStatus::InProgress - }; - BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { - history_model.update_conversation_status( - self.terminal_surface_id, - conversation_id, - status, - ctx, - ); - }); - } - ResponseStreamEvent::AfterStreamFinished { - cancellation, - proposed_actions, - } => { + ResponseStreamEvent::AfterStreamFinished { cancellation } => { // Cancellations provide conversation_id (survives truncation); otherwise use dynamic lookup. let conversation_id = match &cancellation { Some(stream_cancellation) => stream_cancellation.conversation_id, @@ -4454,7 +6563,9 @@ impl BlocklistAIController { } log::info!( "[bedrock-debug] AfterStreamFinished: stream_id={:?}, conversation_id={:?}, new_exchange_ids count={}", - stream_id, conversation_id, new_exchange_ids.len() + stream_id, + conversation_id, + new_exchange_ids.len() ); let mut was_passive_request = false; let mut is_any_exchange_unfinished = false; @@ -4483,7 +6594,8 @@ impl BlocklistAIController { let msg_count = output.get().messages.len(); log::info!( "[bedrock-debug] AfterStreamFinished: output has {} messages, {} actions", - msg_count, action_count + msg_count, + action_count ); for msg in output.get().messages.iter() { log::info!( @@ -4496,31 +6608,13 @@ impl BlocklistAIController { } let history_action_count = actions_to_queue.len(); - let proposed_action_count = proposed_actions.len(); - let failed_tool_proposals = self - .failed_tool_proposals_by_stream - .remove(&stream_id) - .unwrap_or_default(); - let mut queued_action_ids = actions_to_queue - .iter() - .map(|action| action.id.clone()) - .collect::>(); - let mut queued_from_stream_snapshot_count = 0; - for action in proposed_actions { - if queued_action_ids.insert(action.id.clone()) { - queued_from_stream_snapshot_count += 1; - actions_to_queue.push(action.clone()); - } - } let active_child_conversation_ids = active_descendant_conversation_ids(history_model.as_ref(ctx), conversation_id); let queue_decision = tool_queue_decision( cancellation.is_some(), is_any_exchange_unfinished, - !failed_tool_proposals.is_empty(), !active_child_conversation_ids.is_empty(), actions_to_queue.len(), - queued_from_stream_snapshot_count, ); #[cfg(not(target_family = "wasm"))] { @@ -4535,19 +6629,12 @@ impl BlocklistAIController { "conversation_id": conversation_id.to_string(), "decision": queue_decision.label(), "history_action_count": history_action_count, - "proposed_action_count": proposed_action_count, "candidate_action_count": actions_to_queue.len(), "will_queue_action_count": if queue_decision.will_queue_actions() { actions_to_queue.len() } else { 0 }, - "queued_from_stream_snapshot_count": queued_from_stream_snapshot_count, - "failed_tool_proposal_count": failed_tool_proposals.len(), - "failed_tool_proposals": failed_tool_proposals - .iter() - .map(FailedToolProposal::to_remote_log_value) - .collect::>(), "active_descendant_conversation_ids": active_child_conversation_ids .iter() .map(ToString::to_string) @@ -4558,7 +6645,6 @@ impl BlocklistAIController { .as_ref() .map(|stream_cancellation| format!("{:?}", stream_cancellation.reason)), "queued_tools": remote_action_summaries(&actions_to_queue), - "proposed_tools": remote_action_summaries(proposed_actions), }), }, ); @@ -4618,11 +6704,6 @@ impl BlocklistAIController { ctx, ); }); - } else if !failed_tool_proposals.is_empty() { - log::warn!( - "Skipping tool queue for stream {stream_id:?}: failed tool proposal attach count={}", - failed_tool_proposals.len() - ); } else if !active_child_conversation_ids.is_empty() { log::info!( "Skipping tool queue for conversation {conversation_id:?}: active child conversations remain: {:?}", @@ -4639,7 +6720,8 @@ impl BlocklistAIController { } else { log::warn!( "[bedrock-debug] AfterStreamFinished: NO actions to queue, was_passive={}, is_any_unfinished={}", - was_passive_request, is_any_exchange_unfinished + was_passive_request, + is_any_exchange_unfinished ); // If this is a child conversation (has a parent) and the // stream ended with EndTurn and no actions, the child agent @@ -4759,14 +6841,6 @@ impl BlocklistAIController { self.handle_pending_events_ready(conversation_id, ctx); } - // Before cleaning up the response stream, check if we should attempt to resume. - if response_stream - .as_ref(ctx) - .should_resume_conversation_after_stream_finished() - { - self.schedule_auto_resume_after_error(conversation_id, ctx); - } - // Clean up the response stream tracking entry now that the stream is complete. history_model.update(ctx, |history_model, _| { if let Some(conversation) = history_model.conversation_mut(&conversation_id) { diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index a84ccc38..f6263fa6 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -16,7 +16,7 @@ use galaxy_agent_core::RuntimeCapabilities; use galaxy_agent_core::TurnCommand; #[cfg(not(target_family = "wasm"))] use galaxy_core::features::FeatureFlag; -use galaxyui::{Entity, ModelContext, SingletonEntity}; +use galaxyui::{AppContext, Entity, ModelContext, SingletonEntity}; use settings::Setting; use uuid::Uuid; use warp_multi_agent_api::response_event; @@ -31,7 +31,7 @@ use crate::ai::agent::api::{self, ConvertToAPITypeError}; use crate::ai::agent::conversation::AIConversationId; #[cfg(not(target_family = "wasm"))] use crate::ai::agent::AIAgentInput; -use crate::ai::agent::{AIAgentAction, AIIdentifiers, CancellationReason}; +use crate::ai::agent::{AIIdentifiers, CancellationReason}; use crate::ai::bedrock::client::BedrockClientConfig; #[cfg(not(target_family = "wasm"))] use crate::ai::blocklist::BlocklistAIPermissions; @@ -40,57 +40,13 @@ use crate::ai::openai::client::OpenAIClientConfig; use crate::ai::provider::ProviderConfig; #[cfg(not(target_family = "wasm"))] use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord}; -use crate::ai::runtime::ProviderRuntime; -use crate::network::NetworkStatus; #[cfg(not(target_family = "wasm"))] use crate::pane_group::PaneGroup; use crate::persistence::model::AgentBackend; use crate::server::server_api::AIApiError; #[cfg(not(target_family = "wasm"))] use crate::settings::LocalControlSettings; -use crate::{report_error, send_telemetry_from_ctx, AISettings, BlocklistAIHistoryModel}; - -/// Maximum number of times a single MAA request is re-sent before the failure is -/// surfaced. -const MAX_RETRIES: usize = 3; - -/// What to do about a failed or truncated MAA response attempt. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum RecoveryAction { - /// Re-send the same request immediately. - RetryNow, - /// Re-send the same request once connectivity returns. - RetryWhenOnline, - /// Resume the conversation with a fresh request after the stream completes. - Resume, - /// Surface the error; the conversation ends in error. - Fail, -} - -/// Decides how to recover from a failed response-stream attempt. -/// -/// Before any client actions have been received, the request can be re-sent verbatim -/// (immediately, or once connectivity returns). After actions have streamed, -/// re-sending is unsafe, so recovery uses a fresh `ResumeConversation` request. -fn recovery_action( - has_received_client_actions: bool, - is_recoverable: bool, - has_retry_budget: bool, - can_attempt_resume_on_error: bool, - is_online: bool, -) -> RecoveryAction { - if !has_received_client_actions && is_recoverable && has_retry_budget { - if is_online { - RecoveryAction::RetryNow - } else { - RecoveryAction::RetryWhenOnline - } - } else if has_received_client_actions && is_recoverable && can_attempt_resume_on_error { - RecoveryAction::Resume - } else { - RecoveryAction::Fail - } -} +use crate::{report_error, AISettings, BlocklistAIHistoryModel}; #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ResponseStreamId(String); @@ -120,13 +76,10 @@ struct AcpRequestControl { turn_control: AcpTurnControlSlot, } -/// Model wrapping an agent API response stream. +/// Projects an ACP response stream or durable provider-run output into conversation history. /// -/// Emits events when the output corresponding to the stream is updated, typically after receiving -/// each response chunk. -/// -/// Handles retries internally - retries are only attempted if no ClientActions events have been -/// received yet, ensuring we don't retry after the AI has started executing actions. +/// Direct-provider lifecycle and retry ownership lives in `ProviderRun`; this model owns only ACP +/// transport plus the common UI/history projection boundary. pub struct ResponseStream { id: ResponseStreamId, runtime_capabilities: RuntimeCapabilities, @@ -134,24 +87,15 @@ pub struct ResponseStream { acp_session_metadata: Arc>, #[cfg(not(target_family = "wasm"))] acp_turn_control: AcpTurnControlSlot, - params: api::RequestParams, - retry_count: usize, - /// One-time fallback from the profile's thinking model to its coding model. - coding_model_fallback_attempted: bool, + params: Option, + projection_model: LLMId, + projection_messages_sent: + Arc>>, start_time: DateTime, time_to_latest_event: TimeDelta, cancellation_tx: Option>, - /// Store the original error for telemetry when retries succeed - original_error: Option, - /// Track whether we've received any client actions - /// If true, we cannot retry on subsequent errors since actions may have been executed + /// Whether the ACP stream emitted client actions, retained for failure diagnostics. has_received_client_actions: bool, - /// Domain tool proposals observed directly from the response stream for the current request. - /// - /// The controller normally queues actions by reading them back from history after the stream - /// finishes. Keeping this snapshot prevents a final tool proposal from being lost if stream - /// completion is handled before that proposal has been applied to history. - proposed_actions: Vec, /// AI identifiers for telemetry emission ai_identifiers: AIIdentifiers, #[cfg(not(target_family = "wasm"))] @@ -159,18 +103,6 @@ pub struct ResponseStream { #[cfg(not(target_family = "wasm"))] remote_log_provider: String, - /// Whether this request can attempt to resume the conversation on error. - /// This is true for all requests except those that are themselves the result of a resume - /// triggered by a previous error. - can_attempt_resume_on_error: bool, - - /// Whether we should attempt to resume the conversation after the stream finishes. - /// - /// This is set when a transient network/server failure occurs after client actions - /// have been received (so an in-request retry is unsafe) and - /// `can_attempt_resume_on_error` is true. - should_resume_conversation_after_stream_finished: bool, - /// Whether a `StreamFinished` event was received for the current request. A /// stream that completes without one was truncated in transit. stream_finished_received: bool, @@ -179,17 +111,8 @@ pub struct ResponseStream { /// request, so stream completion doesn't synthesize a second failure for it. error_event_emitted: bool, - /// Whether a retry is parked waiting for connectivity. While set, completion of - /// the failed attempt's underlying stream is ignored. - deferred_retry_pending: bool, - - /// Unique, internal id for the current request. - /// - /// This ensures that the model never emits events for a request that was already cancelled (or - /// retried) and is still receiving lagging events. - /// - /// Note this is unique compared to `id`; this is unique across retry requests while the response - /// stream id remains stable. + /// Unique internal ID for the active ACP request. Clearing it on cancellation causes late + /// transport events to be discarded while the stable response-stream ID remains a projection ID. current_request_id: Option, } @@ -215,30 +138,25 @@ impl ResponseStream { acp_session_metadata: Arc::new(Mutex::new(AcpSessionMetadata::default())), #[cfg(not(target_family = "wasm"))] acp_turn_control: Arc::new(Mutex::new(None)), - params: api::RequestParams::new_for_test(), - retry_count: 0, - coding_model_fallback_attempted: false, + params: Some(api::RequestParams::new_for_test()), + projection_model: LLMId::from("test-model"), + projection_messages_sent: Arc::new(std::sync::Mutex::new(Vec::new())), start_time: Local::now(), time_to_latest_event: TimeDelta::seconds(0), cancellation_tx: Some(cancellation_tx), - original_error: None, has_received_client_actions: false, - proposed_actions: Vec::new(), ai_identifiers: AIIdentifiers::default(), #[cfg(not(target_family = "wasm"))] remote_log_backend: "provider".to_string(), #[cfg(not(target_family = "wasm"))] remote_log_provider: "test".to_string(), - can_attempt_resume_on_error: false, - should_resume_conversation_after_stream_finished: false, stream_finished_received: false, error_event_emitted: false, - deferred_retry_pending: false, current_request_id: Some(Uuid::new_v4()), } } - fn resolve_provider_config(model_id: &str, ctx: &ModelContext) -> ProviderConfig { + pub(super) fn resolve_provider_config(model_id: &str, ctx: &AppContext) -> ProviderConfig { let settings = AISettings::as_ref(ctx); // Check if this specific model has an OpenAI-compatible routing entry. @@ -271,21 +189,12 @@ 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 mut use_rig = crate::ai::bedrock::models::configured_model_uses_rig( + let use_rig = crate::ai::bedrock::models::configured_model_uses_rig( model_id, settings.bedrock_models.value(), ®ion, cross_region_inference, ); - if use_rig - && crate::ai::bedrock::external_config::ExternalBedrockConfig::load() - .enable_prompt_caching_1h - { - log::warn!( - "[rig/bedrock] Using the compatibility runtime because Rig does not yet expose Bedrock's one-hour cache TTL" - ); - use_rig = false; - } let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx); let mut config = BedrockClientConfig { auth_method, @@ -346,7 +255,7 @@ impl ResponseStream { ); context.insert( "model_id".to_string(), - serde_json::json!(self.params.model.as_str()), + serde_json::json!(self.projection_model.as_str()), ); context.insert( "backend".to_string(), @@ -356,18 +265,10 @@ impl ResponseStream { "provider".to_string(), serde_json::json!(self.remote_log_provider), ); - context.insert( - "retry_count".to_string(), - serde_json::json!(self.retry_count), - ); context.insert( "has_received_client_actions".to_string(), serde_json::json!(self.has_received_client_actions), ); - context.insert( - "can_attempt_resume_on_error".to_string(), - serde_json::json!(self.can_attempt_resume_on_error), - ); context.insert( "identifiers".to_string(), serde_json::to_value(&self.ai_identifiers).unwrap_or_else(|_| serde_json::json!({})), @@ -385,7 +286,6 @@ impl ResponseStream { ai_identifiers: &AIIdentifiers, backend: &str, provider: &str, - can_attempt_resume_on_error: bool, ) { remote_logging::log_model_event( ctx, @@ -413,7 +313,6 @@ impl ResponseStream { "ask_user_question_enabled": params.ask_user_question_enabled, "orchestration_enabled": params.orchestration_enabled, "is_remote_session": params.session_context.is_remote(), - "can_attempt_resume_on_error": can_attempt_resume_on_error, "identifiers": serde_json::to_value(ai_identifiers).unwrap_or_else(|_| serde_json::json!({})), }), }, @@ -683,127 +582,166 @@ impl ResponseStream { ); } - fn spawn_provider_request( + pub(super) fn new_provider_projection( params: api::RequestParams, - provider_config: ProviderConfig, - request_id: Uuid, - cancellation_rx: oneshot::Receiver<()>, + ai_identifiers: AIIdentifiers, ctx: &mut ModelContext, - ) { - let _ = ctx.spawn( - async move { - ProviderRuntime::new(provider_config) - .start_turn(params, cancellation_rx) - .await - }, - move |me, stream, ctx| { - me.handle_response_stream_result(request_id, stream, ctx); - }, + ) -> Self { + let start_time = Local::now(); + let request_id = Uuid::new_v4(); + let response_stream_id = ResponseStreamId(Uuid::new_v4().to_string()); + #[cfg(not(target_family = "wasm"))] + let provider_config = Self::resolve_provider_config(params.model.as_str(), ctx); + #[cfg(not(target_family = "wasm"))] + let remote_log_backend = "provider".to_string(); + #[cfg(not(target_family = "wasm"))] + let remote_log_provider = Self::remote_log_provider_for_config(&provider_config); + #[cfg(not(target_family = "wasm"))] + Self::log_llm_request_started( + ctx, + &response_stream_id, + request_id, + ¶ms, + &ai_identifiers, + &remote_log_backend, + &remote_log_provider, ); + + let projection_model = params.model.clone(); + let projection_messages_sent = params.messages_sent.clone(); + Self { + id: response_stream_id, + runtime_capabilities: RuntimeCapabilities::provider(), + #[cfg(not(target_family = "wasm"))] + acp_session_metadata: Arc::new(Mutex::new(AcpSessionMetadata::default())), + #[cfg(not(target_family = "wasm"))] + acp_turn_control: Arc::new(Mutex::new(None)), + params: Some(params), + projection_model, + projection_messages_sent, + start_time, + time_to_latest_event: TimeDelta::seconds(0), + cancellation_tx: None, + has_received_client_actions: false, + ai_identifiers, + #[cfg(not(target_family = "wasm"))] + remote_log_backend, + #[cfg(not(target_family = "wasm"))] + remote_log_provider, + stream_finished_received: false, + error_event_emitted: false, + current_request_id: None, + } + } + + pub(super) fn new_restored_provider_projection( + model: LLMId, + messages_sent: Arc>>, + ai_identifiers: AIIdentifiers, + ctx: &mut ModelContext, + ) -> Self { + #[cfg(not(target_family = "wasm"))] + let provider_config = Self::resolve_provider_config(model.as_str(), ctx); + Self { + id: ResponseStreamId(Uuid::new_v4().to_string()), + runtime_capabilities: RuntimeCapabilities::provider(), + #[cfg(not(target_family = "wasm"))] + acp_session_metadata: Arc::new(Mutex::new(AcpSessionMetadata::default())), + #[cfg(not(target_family = "wasm"))] + acp_turn_control: Arc::new(Mutex::new(None)), + params: None, + projection_model: model, + projection_messages_sent: messages_sent, + start_time: Local::now(), + time_to_latest_event: TimeDelta::seconds(0), + cancellation_tx: None, + has_received_client_actions: false, + ai_identifiers, + #[cfg(not(target_family = "wasm"))] + remote_log_backend: "provider".to_string(), + #[cfg(not(target_family = "wasm"))] + remote_log_provider: Self::remote_log_provider_for_config(&provider_config), + stream_finished_received: false, + error_event_emitted: false, + current_request_id: None, + } + } + + pub(super) fn projected_event( + event: warp_multi_agent_api::ResponseEvent, + ) -> ResponseStreamEvent { + ResponseStreamEvent::ReceivedEvent(Consumable::new(Ok(api::StreamEvent::Response(event)))) } pub fn new( params: api::RequestParams, ai_identifiers: AIIdentifiers, agent_backend: AgentBackend, - can_attempt_resume_on_error: bool, ctx: &mut ModelContext, ) -> Self { + let AgentBackend::Acp(backend) = agent_backend else { + unreachable!("direct providers must use the durable provider-run projection"); + }; let (cancellation_tx, cancellation_rx) = oneshot::channel(); let start_time = Local::now(); - let request_id = Uuid::new_v4(); let response_stream_id = ResponseStreamId(Uuid::new_v4().to_string()); - let runtime_capabilities = match &agent_backend { - AgentBackend::Provider => RuntimeCapabilities::provider(), - AgentBackend::Acp(_) => RuntimeCapabilities::session_runtime(), - }; + let runtime_capabilities = RuntimeCapabilities::session_runtime(); #[cfg(not(target_family = "wasm"))] let acp_session_metadata = Arc::new(Mutex::new(AcpSessionMetadata::default())); #[cfg(not(target_family = "wasm"))] let acp_turn_control = Arc::new(Mutex::new(None)); #[cfg(not(target_family = "wasm"))] - let remote_log_backend; + let remote_log_backend = "acp".to_string(); #[cfg(not(target_family = "wasm"))] - let remote_log_provider; - match &agent_backend { - AgentBackend::Provider => { - let provider_config = Self::resolve_provider_config(params.model.as_str(), ctx); - #[cfg(not(target_family = "wasm"))] - { - remote_log_backend = "provider".to_string(); - remote_log_provider = Self::remote_log_provider_for_config(&provider_config); - Self::log_llm_request_started( - ctx, - &response_stream_id, - request_id, - ¶ms, - &ai_identifiers, - &remote_log_backend, - &remote_log_provider, - can_attempt_resume_on_error, - ); - } - Self::spawn_provider_request( - params.clone(), - provider_config, - request_id, + let remote_log_provider = if backend.agent_id.is_empty() { + "acp".to_string() + } else { + format!("acp:{}", backend.agent_id) + }; + #[cfg(not(target_family = "wasm"))] + { + Self::log_llm_request_started( + ctx, + &response_stream_id, + request_id, + ¶ms, + &ai_identifiers, + &remote_log_backend, + &remote_log_provider, + ); + Self::spawn_acp_request( + backend.clone(), + params.clone(), + ai_identifiers + .client_conversation_id + .map(|id| format!("{id:?}")) + .unwrap_or_else(|| Uuid::new_v4().to_string()), + request_id, + AcpRequestControl { cancellation_rx, - ctx, - ); - } - AgentBackend::Acp(backend) => { - #[cfg(not(target_family = "wasm"))] - { - remote_log_backend = "acp".to_string(); - remote_log_provider = if backend.agent_id.is_empty() { - "acp".to_string() - } else { - format!("acp:{}", backend.agent_id) - }; - Self::log_llm_request_started( - ctx, - &response_stream_id, - request_id, - ¶ms, - &ai_identifiers, - &remote_log_backend, - &remote_log_provider, - can_attempt_resume_on_error, - ); - } - #[cfg(not(target_family = "wasm"))] - Self::spawn_acp_request( - backend.clone(), - params.clone(), - ai_identifiers - .client_conversation_id - .map(|id| format!("{id:?}")) - .unwrap_or_else(|| Uuid::new_v4().to_string()), - request_id, - AcpRequestControl { - cancellation_rx, - session_metadata: acp_session_metadata.clone(), - turn_control: acp_turn_control.clone(), - }, - ctx, - ); - #[cfg(target_family = "wasm")] - { - let error = Arc::new(AIApiError::Stream { - stream_type: "acp", - source: anyhow!("ACP is unavailable in the web client"), - }); - let stream = Box::pin(futures::stream::once(async move { Err(error) })); - let _ = ctx.spawn( - async move { Ok::<_, ConvertToAPITypeError>(stream) }, - move |me, stream, ctx| { - me.handle_response_stream_result(request_id, stream, ctx); - }, - ); - } - } + session_metadata: acp_session_metadata.clone(), + turn_control: acp_turn_control.clone(), + }, + ctx, + ); } + #[cfg(target_family = "wasm")] + { + let error = Arc::new(AIApiError::Stream { + stream_type: "acp", + source: anyhow!("ACP is unavailable in the web client"), + }); + let stream = Box::pin(futures::stream::once(async move { Err(error) })); + let _ = ctx.spawn( + async move { Ok::<_, ConvertToAPITypeError>(stream) }, + move |me, stream, ctx| { + me.handle_response_stream_result(request_id, stream, ctx); + }, + ); + } + let projection_model = params.model.clone(); + let projection_messages_sent = params.messages_sent.clone(); Self { id: response_stream_id, runtime_capabilities, @@ -811,25 +749,20 @@ impl ResponseStream { acp_session_metadata, #[cfg(not(target_family = "wasm"))] acp_turn_control, - params: params.clone(), + params: Some(params), + projection_model, + projection_messages_sent, start_time, time_to_latest_event: TimeDelta::seconds(0), cancellation_tx: Some(cancellation_tx), - retry_count: 0, - coding_model_fallback_attempted: false, - original_error: None, has_received_client_actions: false, - proposed_actions: Vec::new(), ai_identifiers, #[cfg(not(target_family = "wasm"))] remote_log_backend, #[cfg(not(target_family = "wasm"))] remote_log_provider, - can_attempt_resume_on_error, - should_resume_conversation_after_stream_finished: false, stream_finished_received: false, error_event_emitted: false, - deferred_retry_pending: false, current_request_id: Some(request_id), } } @@ -847,18 +780,18 @@ impl ResponseStream { } pub(super) fn has_error_tool_results(&self) -> bool { - self.params - .tool_results - .iter() - .any(galaxy_agent_core::ToolResult::is_error) + self.params.as_ref().is_some_and(|params| { + params + .tool_results + .iter() + .any(galaxy_agent_core::ToolResult::is_error) + }) } pub(super) fn tool_result_count(&self) -> usize { - self.params.tool_results.len() - } - - pub fn allows_corrective_retries(&self) -> bool { - self.runtime_capabilities.corrective_retries + self.params + .as_ref() + .map_or(0, |params| params.tool_results.len()) } #[cfg(not(target_family = "wasm"))] @@ -885,8 +818,11 @@ impl ResponseStream { { return false; } + let Some(params) = self.params.as_ref() else { + return false; + }; let mut model_text = display_text.clone(); - self.params.redact_text_for_model(&mut model_text); + params.redact_text_for_model(&mut model_text); self.acp_turn_control .lock() .ok() @@ -911,117 +847,16 @@ impl ResponseStream { &self, ) -> &std::sync::Arc>> { - &self.params.messages_sent + &self.projection_messages_sent } /// Returns the model ID associated with this response stream's request. pub fn model_id(&self) -> &str { - self.params.model.as_str() + self.projection_model.as_str() } pub(super) fn llm_id(&self) -> &LLMId { - &self.params.model - } - - /// Returns true if we should attempt to resume the conversation after the stream finishes. - pub fn should_resume_conversation_after_stream_finished(&self) -> bool { - self.should_resume_conversation_after_stream_finished - } - - /// Helper function to emit AgentModeError telemetry for error that is retryable (not user visible). - fn emit_retryable_agent_mode_error_telemetry( - &self, - error: String, - ctx: &mut ModelContext, - ) { - send_telemetry_from_ctx!( - crate::TelemetryEvent::AgentModeError { - identifiers: self.ai_identifiers.clone(), - error, - is_user_visible: false, - will_attempt_to_resume: false, - }, - ctx - ); - } - - fn retry(&mut self, ctx: &mut ModelContext) { - self.retry_count += 1; - // Reset per-attempt state for the new attempt. - self.has_received_client_actions = false; - self.proposed_actions.clear(); - self.stream_finished_received = false; - self.error_event_emitted = false; - self.deferred_retry_pending = false; - - let (cancellation_tx, cancellation_rx) = oneshot::channel(); - if let Some(old_cancellation_tx) = self.cancellation_tx.take() { - let _ = old_cancellation_tx.send(()); - } - self.cancellation_tx = Some(cancellation_tx); - - let request_id = Uuid::new_v4(); - self.current_request_id = Some(request_id); - #[cfg(not(target_family = "wasm"))] - self.log_galaxy_decision( - request_id, - "retry_request", - serde_json::json!({ - "retry_count": self.retry_count, - "model_id": self.params.model.as_str(), - }), - ctx, - ); - let params = self.params.clone(); - let provider_config = Self::resolve_provider_config(params.model.as_str(), ctx); - let _ = ctx.spawn( - async move { - ProviderRuntime::new(provider_config) - .start_turn(params, cancellation_rx) - .await - }, - move |me, stream, ctx| { - me.handle_response_stream_result(request_id, stream, ctx); - }, - ); - } - - fn should_fallback_to_coding_model( - &self, - error: &Arc, - ) -> bool { - if !self.runtime_capabilities.model_selection - || !self.runtime_capabilities.request_retries - || self.coding_model_fallback_attempted - || self.has_received_client_actions - { - return false; - } - let coding_model = self.params.coding_model.as_str(); - !coding_model.is_empty() - && coding_model != self.params.model.as_str() - && matches!( - error.as_ref(), - crate::server::server_api::AIApiError::QuotaLimit { .. } - ) - } - - fn retry_with_coding_model(&mut self, ctx: &mut ModelContext) { - #[cfg(not(target_family = "wasm"))] - if let Some(request_id) = self.current_request_id { - self.log_galaxy_decision( - request_id, - "fallback_to_coding_model", - serde_json::json!({ - "from_model_id": self.params.model.as_str(), - "to_model_id": self.params.coding_model.as_str(), - }), - ctx, - ); - } - self.coding_model_fallback_attempted = true; - self.params.model = self.params.coding_model.clone(); - self.retry(ctx); + &self.projection_model } /// Cancels the stream. The conversation_id is preserved in the emitted event for async handling. @@ -1041,7 +876,6 @@ impl ResponseStream { reason, conversation_id, }), - proposed_actions: self.proposed_actions.clone(), }); } @@ -1083,7 +917,7 @@ impl ResponseStream { ctx, ); self.error_event_emitted = true; - self.report_request_failure(&error, NetworkStatus::as_ref(ctx).is_online()); + self.report_request_failure(&error); ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(Err( error, )))); @@ -1099,83 +933,21 @@ impl ResponseStream { ctx: &mut ModelContext, ) { if self.current_request_id.is_none_or(|id| id != request_id) { - log::info!( - "[bedrock-debug] handle_response_stream_event: stale request_id, dropping event" - ); + log::debug!("Dropping event for stale ACP request {request_id}"); return; } self.time_to_latest_event = Local::now().signed_duration_since(self.start_time); match &event { - Ok(api::StreamEvent::ToolProposed(action)) => { - self.has_received_client_actions = true; - self.proposed_actions.push(action.clone()); - log::debug!( - "Rig proposed domain tool action {} for task {}", - action.id, - action.task_id - ); - #[cfg(not(target_family = "wasm"))] - { - let mut context = - self.common_remote_log_context("llm_tool_proposed", request_id); - context.insert( - "action_id".to_string(), - serde_json::json!(action.id.to_string()), - ); - context.insert( - "task_id".to_string(), - serde_json::json!(action.task_id.to_string()), - ); - context.insert( - "tool_name".to_string(), - serde_json::json!(action_tool_name(action)), - ); - context.insert( - "requires_result".to_string(), - serde_json::json!(action.requires_result), - ); - remote_logging::log_model_event( - ctx, - RemoteLogRecord { - level: RemoteLogLevel::Info, - message: "LLM proposed tool".to_string(), - context: serde_json::Value::Object(context), - }, - ); - } - #[cfg(not(target_family = "wasm"))] - self.log_raw_model_response( - request_id, - "tool_proposed", - format!("{action:#?}"), - ctx, - ); - ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(event))); - } Ok(api::StreamEvent::Response(response_event)) => { let event_type_name = match &response_event.r#type { Some(warp_multi_agent_api::response_event::Type::Init(_)) => "Init", - Some(warp_multi_agent_api::response_event::Type::ClientActions(a)) => { - log::debug!( - "[bedrock] ResponseStream received ClientActions with {} actions", - a.actions.len() - ); + Some(warp_multi_agent_api::response_event::Type::ClientActions(_)) => { "ClientActions" } - Some(warp_multi_agent_api::response_event::Type::Finished(f)) => { - log::info!( - "[bedrock-debug] ResponseStream received Finished (reason={:?})", - f.reason - .as_ref() - .map(|r| format!("{r:?}")) - .unwrap_or("None".into()) - ); - "Finished" - } + Some(warp_multi_agent_api::response_event::Type::Finished(_)) => "Finished", None => "None", }; - log::info!("[bedrock-debug] ResponseStream emitting event type={event_type_name}"); #[cfg(not(target_family = "wasm"))] self.log_raw_model_response( request_id, @@ -1221,237 +993,47 @@ impl ResponseStream { self.stream_finished_received = true; #[cfg(not(target_family = "wasm"))] self.log_llm_response_finished(request_id, finished_event, ctx); - // Emit retry success telemetry on successful completion - if matches!( - finished_event.reason, - Some(warp_multi_agent_api::response_event::stream_finished::Reason::Done(_)) | None - ) { - // Emit retry success telemetry if this was a successful completion after retries - if self.retry_count > 0 { - if let Some(original_error) = &self.original_error { - send_telemetry_from_ctx!( - crate::TelemetryEvent::AgentModeRequestRetrySucceeded { - identifiers: self.ai_identifiers.clone(), - retry_count: self.retry_count, - original_error: original_error.clone(), - }, - ctx - ); - } - } - } } } } ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(event))); } Err(e) => { - log::info!("[bedrock-debug] ResponseStream received ERROR: {e:?}"); - // Store original error if this is the first error - if self.retry_count == 0 { - self.original_error = Some(format!("{e:?}")); - } - - if self.should_fallback_to_coding_model(e) { - log::warn!( - "Thinking model rate-limited; retrying with the profile coding model" - ); - #[cfg(not(target_family = "wasm"))] - self.log_llm_request_error( - request_id, - e.as_ref(), - "fallback_to_coding_model", - ctx, - ); - self.retry_with_coding_model(ctx); - return; - } - - let is_online = NetworkStatus::as_ref(ctx).is_online(); - match recovery_action( - self.has_received_client_actions, - e.is_recoverable() && self.runtime_capabilities.request_retries, - self.retry_count < MAX_RETRIES, - self.can_attempt_resume_on_error, - is_online, - ) { - RecoveryAction::RetryNow => { - log::warn!( - "MultiAgent request failed, retrying (attempt {}/{}) - Error: {e:?}", - self.retry_count + 1, - MAX_RETRIES - ); - #[cfg(not(target_family = "wasm"))] - self.log_llm_request_error(request_id, e.as_ref(), "retry_now", ctx); - // Only emit error telemetry here if we're retrying. - // Final errors that aren't being retried are emitted elsewhere. - self.emit_retryable_agent_mode_error_telemetry(format!("{e:?}"), ctx); - self.retry(ctx); - // Don't emit the error event, we're retrying - return; - } - RecoveryAction::RetryWhenOnline => { - log::warn!( - "MultiAgent request failed while offline; retrying (attempt {}/{}) once connectivity returns - Error: {e:?}", - self.retry_count + 1, - MAX_RETRIES - ); - #[cfg(not(target_family = "wasm"))] - self.log_llm_request_error( - request_id, - e.as_ref(), - "retry_when_online", - ctx, - ); - self.emit_retryable_agent_mode_error_telemetry(format!("{e:?}"), ctx); - self.defer_retry_until_online(ctx); - return; - } - RecoveryAction::Resume => { - // Recoverable failure after client actions: we'll resume the - // conversation once the stream finishes rather than surface the - // error, so the UI suppresses the banner. Log it so the - // auto-recovery isn't completely silent. - log::warn!( - "MultiAgent request failed after client actions; resuming conversation after stream finishes - Error: {e:?}" - ); - #[cfg(not(target_family = "wasm"))] - self.log_llm_request_error( - request_id, - e.as_ref(), - "resume_after_stream", - ctx, - ); - // The resume spawn itself waits for connectivity. - self.should_resume_conversation_after_stream_finished = true; - } - RecoveryAction::Fail => { - #[cfg(not(target_family = "wasm"))] - self.log_llm_request_error(request_id, e.as_ref(), "fail", ctx); - } - } + log::info!("ResponseStream received error: {e:?}"); + #[cfg(not(target_family = "wasm"))] + self.log_llm_request_error(request_id, e.as_ref(), "fail", ctx); self.error_event_emitted = true; - - self.report_request_failure(e, is_online); - + self.report_request_failure(e); ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(event))); } } } fn on_response_stream_complete(&mut self, request_id: Uuid, ctx: &mut ModelContext) { - log::info!("[bedrock-debug] on_response_stream_complete called (request_id={request_id})"); if self.current_request_id.is_none_or(|id| id != request_id) { - log::info!("[bedrock-debug] on_response_stream_complete: stale request_id, ignoring"); + log::debug!("Ignoring completion for stale ACP request {request_id}"); return; } - // A retry is parked waiting for connectivity; the request is logically still - // active, so don't complete the stream for the failed attempt. - if self.deferred_retry_pending { - return; - } - - // The server always sends a StreamFinished event before ending the response, - // but a transport cut between chunks surfaces as a clean EOF. Synthesize the - // failure and recover like any transient error. + // ACP sends StreamFinished before closing. A clean EOF without it is a + // truncated protocol response and cannot be retried safely by this projection. if !self.stream_finished_received && !self.error_event_emitted { - log::warn!( - "generate_multi_agent_output stream ended without emitting StreamFinished event." - ); + log::warn!("ACP response stream ended without emitting StreamFinished"); let unexpected_eof = Arc::new(AIApiError::UnexpectedEof); - let is_online = NetworkStatus::as_ref(ctx).is_online(); - match recovery_action( - self.has_received_client_actions, - unexpected_eof.is_recoverable() && self.runtime_capabilities.request_retries, - self.retry_count < MAX_RETRIES, - self.can_attempt_resume_on_error, - is_online, - ) { - RecoveryAction::RetryNow => { - log::warn!( - "MultiAgent request failed, retrying (attempt {}/{}) - Error: {unexpected_eof:?}", - self.retry_count + 1, - MAX_RETRIES - ); - #[cfg(not(target_family = "wasm"))] - self.log_llm_request_error( - request_id, - unexpected_eof.as_ref(), - "retry_now", - ctx, - ); - self.emit_retryable_agent_mode_error_telemetry( - format!("{unexpected_eof:?}"), - ctx, - ); - self.retry(ctx); - return; - } - RecoveryAction::RetryWhenOnline => { - log::warn!( - "MultiAgent request failed while offline; retrying (attempt {}/{}) once connectivity returns - Error: {unexpected_eof:?}", - self.retry_count + 1, - MAX_RETRIES - ); - #[cfg(not(target_family = "wasm"))] - self.log_llm_request_error( - request_id, - unexpected_eof.as_ref(), - "retry_when_online", - ctx, - ); - self.emit_retryable_agent_mode_error_telemetry( - format!("{unexpected_eof:?}"), - ctx, - ); - self.defer_retry_until_online(ctx); - return; - } - RecoveryAction::Resume => { - // Recoverable truncation after client actions: we'll resume the - // conversation once the stream finishes rather than surface the - // error, so the UI suppresses the banner. Log it so the - // auto-recovery isn't completely silent. - log::warn!( - "MultiAgent request truncated after client actions; resuming conversation after stream finishes - Error: {unexpected_eof:?}" - ); - #[cfg(not(target_family = "wasm"))] - self.log_llm_request_error( - request_id, - unexpected_eof.as_ref(), - "resume_after_stream", - ctx, - ); - self.should_resume_conversation_after_stream_finished = true; - self.error_event_emitted = true; - self.report_request_failure(&unexpected_eof, is_online); - ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(Err( - unexpected_eof, - )))); - } - RecoveryAction::Fail => { - #[cfg(not(target_family = "wasm"))] - self.log_llm_request_error(request_id, unexpected_eof.as_ref(), "fail", ctx); - self.error_event_emitted = true; - self.report_request_failure(&unexpected_eof, is_online); - ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(Err( - unexpected_eof, - )))); - } - } + #[cfg(not(target_family = "wasm"))] + self.log_llm_request_error(request_id, unexpected_eof.as_ref(), "fail", ctx); + self.error_event_emitted = true; + self.report_request_failure(&unexpected_eof); + ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(Err( + unexpected_eof, + )))); } - ctx.emit(ResponseStreamEvent::AfterStreamFinished { - cancellation: None, - proposed_actions: self.proposed_actions.clone(), - }); + ctx.emit(ResponseStreamEvent::AfterStreamFinished { cancellation: None }); self.cancellation_tx = None; } - /// Reports a non-retried request failure to crash reporting with classification - /// tags. - #[cfg_attr(not(feature = "crash_reporting"), expect(unused_variables))] - fn report_request_failure(&self, error: &Arc, is_online: bool) { + /// Reports a terminal response-stream failure with classification tags. + fn report_request_failure(&self, error: &Arc) { #[cfg(feature = "crash_reporting")] sentry::with_scope( |scope| { @@ -1461,44 +1043,13 @@ impl ResponseStream { ); scope.set_tag("error", format!("{error:?}")); scope.set_tag("is_recoverable", error.is_recoverable()); - scope.set_tag( - "will_attempt_resume", - self.should_resume_conversation_after_stream_finished, - ); - scope.set_tag("is_online", is_online); - scope.set_tag("retry_count", self.retry_count); }, || { - report_error!(anyhow!(error.clone()).context(format!( - "MultiAgent request failed after {} retries", - self.retry_count - ))); + report_error!(anyhow!(error.clone()).context("ACP response stream failed")); }, ); #[cfg(not(feature = "crash_reporting"))] - { - report_error!(anyhow!(error.clone()).context(format!( - "MultiAgent request failed after {} retries", - self.retry_count - ))); - } - } - - /// Parks a retry until connectivity returns; cancellation invalidates the parked - /// retry through `current_request_id`. - fn defer_retry_until_online(&mut self, ctx: &mut ModelContext) { - self.deferred_retry_pending = true; - ctx.emit(ResponseStreamEvent::WaitingForNetwork { waiting: true }); - let request_id_at_defer = self.current_request_id; - let wait_for_online = NetworkStatus::as_ref(ctx).wait_until_online(); - let _ = ctx.spawn(wait_for_online, move |me, _, ctx| { - // Cancelled or superseded while waiting — drop the parked retry. - if request_id_at_defer.is_none() || me.current_request_id != request_id_at_defer { - return; - } - ctx.emit(ResponseStreamEvent::WaitingForNetwork { waiting: false }); - me.retry(ctx); - }); + report_error!(anyhow!(error.clone()).context("ACP response stream failed")); } } @@ -1545,16 +1096,6 @@ fn raw_model_request_payload(params: &api::RequestParams) -> String { serde_json::to_string_pretty(&payload).unwrap_or_else(|_| format!("{payload:#?}")) } -#[cfg(not(target_family = "wasm"))] -fn action_tool_name(action: &crate::ai::agent::AIAgentAction) -> String { - action.tool_name.clone().unwrap_or_else(|| { - format!( - "{:?}", - crate::ai::agent::AIAgentActionTypeDiscriminants::from(&action.action) - ) - }) -} - #[cfg(not(target_family = "wasm"))] fn stream_finished_reason_name( reason: &Option, @@ -1682,21 +1223,9 @@ pub struct StreamCancellation { #[derive(Debug, Clone)] pub enum ResponseStreamEvent { ReceivedEvent(Consumable), - /// A retry is parked until connectivity returns (`waiting: true`) or has just - /// fired (`waiting: false`). The controller mirrors this on the conversation - /// status (`TransientError` ↔ `InProgress`). - /// - /// Only emitted from `defer_retry_until_online`, i.e. always after a recoverable - /// request failure while offline — never speculatively before an attempt. Consumers - /// can therefore treat `waiting: true` as a transient-error (reconnecting) state. - WaitingForNetwork { - waiting: bool, - }, AfterStreamFinished { /// Some for cancellation (with context), None for natural completion (uses dynamic lookup). cancellation: Option, - /// Domain tool proposals observed directly from the stream before it finished. - proposed_actions: Vec, }, } diff --git a/app/src/ai/blocklist/controller/response_stream_tests.rs b/app/src/ai/blocklist/controller/response_stream_tests.rs index c58ec5a0..e5bed3f4 100644 --- a/app/src/ai/blocklist/controller/response_stream_tests.rs +++ b/app/src/ai/blocklist/controller/response_stream_tests.rs @@ -1,88 +1,4 @@ -use super::{is_interactive_remote_command, recovery_action, RecoveryAction}; - -// Argument order: has_received_client_actions, is_recoverable, has_retry_budget, -// can_attempt_resume_on_error, is_online. - -#[test] -fn pre_action_failures_retry() { - assert_eq!( - recovery_action(false, true, true, true, true), - RecoveryAction::RetryNow - ); - // Resume eligibility is irrelevant pre-actions. - assert_eq!( - recovery_action(false, true, true, false, true), - RecoveryAction::RetryNow - ); -} - -#[test] -fn pre_action_failures_wait_for_connectivity_when_offline() { - assert_eq!( - recovery_action(false, true, true, true, false), - RecoveryAction::RetryWhenOnline - ); -} - -#[test] -fn pre_action_budget_exhaustion_is_terminal() { - // The request has already been retried MAX_RETRIES times; stop. - assert_eq!( - recovery_action(false, true, false, true, true), - RecoveryAction::Fail - ); - assert_eq!( - recovery_action(false, true, false, true, false), - RecoveryAction::Fail - ); -} - -#[test] -fn non_recoverable_pre_action_failure_is_terminal() { - assert_eq!( - recovery_action(false, false, true, true, true), - RecoveryAction::Fail - ); -} - -#[test] -fn post_action_recoverable_failures_resume() { - assert_eq!( - recovery_action(true, true, true, true, true), - RecoveryAction::Resume - ); - // Offline doesn't change the decision; the resume spawn waits for connectivity. - assert_eq!( - recovery_action(true, true, true, true, false), - RecoveryAction::Resume - ); - // The in-request retry budget is irrelevant once actions have executed. - assert_eq!( - recovery_action(true, true, false, true, true), - RecoveryAction::Resume - ); -} - -#[test] -fn post_action_failures_without_resume_eligibility_are_terminal() { - // Resume requests themselves run with can_attempt_resume_on_error=false, - // bounding recovery to a single resume. - assert_eq!( - recovery_action(true, true, true, false, true), - RecoveryAction::Fail - ); -} - -#[test] -fn non_recoverable_post_action_failure_is_terminal() { - // A non-recoverable error (e.g. a client error) ends the conversation even - // after actions have executed. - assert_eq!( - recovery_action(true, false, true, true, true), - RecoveryAction::Fail - ); -} - +use super::is_interactive_remote_command; #[test] fn raw_interactive_ssh_is_treated_as_remote_for_acp() { for command in [ diff --git a/app/src/ai/blocklist/controller/slash_command.rs b/app/src/ai/blocklist/controller/slash_command.rs index 1bf12ae3..68e7ebd8 100644 --- a/app/src/ai/blocklist/controller/slash_command.rs +++ b/app/src/ai/blocklist/controller/slash_command.rs @@ -194,7 +194,6 @@ impl SlashCommandRequest { entrypoint, is_auto_resume_after_error: false, }), - /*can_attempt_resume_on_error*/ true, is_queued_prompt, ctx, ) { diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 4d7c83b7..a1b45f2c 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -1,8 +1,15 @@ -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap}; use std::sync::{Arc, Mutex}; use ai::agent::action::{AskUserQuestionItem, AskUserQuestionType}; use chrono::Local; +use galaxy_agent_core::{ + CompletedModelTurn, ContentPart, ConversationMessage, ExternalWorkId, MessageContent, + MessageRole, PermissionKind, PermissionRequest, ProviderRun, ProviderRunId, ProviderRunLimits, + ProviderRunState, ProviderRunStep, RunEpoch, RuntimeCapabilities, StopReason, ToolCall, + TurnRequest, Usage, +}; +use galaxy_core::command::ExitCode; use uuid::Uuid; use warp_multi_agent_api::response_event; use warpui::{App, EntityId, SingletonEntity}; @@ -10,9 +17,11 @@ use warpui::{App, EntityId, SingletonEntity}; use crate::ai::agent::conversation::{AIConversationId, ConversationStatus}; use crate::ai::agent::task::TaskId; use crate::ai::agent::{ - AIAgentAction, AIAgentActionId, AIAgentActionType, AIAgentAttachment, AIAgentContext, - AIAgentInput, CancellationReason, ImageContext, PassiveSuggestionTrigger, RunningCommand, - UserQueryMode, + AIAgentAction, AIAgentActionId, AIAgentActionResultType, AIAgentActionType, AIAgentAttachment, + AIAgentContext, AIAgentExchangeId, AIAgentInput, CancellationReason, ImageContext, + PassiveSuggestionTrigger, ReadShellCommandOutputResult, RequestCommandOutputResult, + RunningCommand, ShellCommandError, TransferShellCommandControlToUserResult, UserQueryMode, + WriteToLongRunningShellCommandResult, }; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::blocklist::{ @@ -21,7 +30,7 @@ use crate::ai::blocklist::{ }; use crate::ai::llms::LLMId; use crate::persistence::model::{AcpConversationData, AgentBackend}; -use crate::terminal::model::block::BlockId; +use crate::terminal::model::block::{BlockId, BlockState}; use crate::test_util::settings::initialize_history_persistence_for_tests; use crate::test_util::terminal::{add_window_with_terminal, initialize_app_for_terminal_view}; @@ -82,6 +91,907 @@ fn live_steering_eligibility() -> super::LiveSteeringEligibility { } } +fn provider_execution_ref( + conversation_id: AIConversationId, + run_id: &str, + epoch: u64, +) -> crate::ai::runtime::ProviderToolExecutionRef { + crate::ai::runtime::ProviderToolExecutionRef { + conversation_id, + run_id: ProviderRunId::new(run_id), + epoch: RunEpoch::new(epoch), + call_id: "call".to_owned(), + } +} + +fn provider_snapshot(conversation_id: AIConversationId) -> super::ActiveProviderRunSnapshot { + let task_id = TaskId::new("root-task".to_owned()); + let messages = vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("Finish the task".to_owned()), + }]; + super::ActiveProviderRunSnapshot { + version: super::ACTIVE_PROVIDER_RUN_SNAPSHOT_VERSION, + run: ProviderRun::new( + "restored-run", + messages.clone(), + crate::ai::runtime::BASE_PROVIDER_PROFILE, + ProviderRunLimits::default(), + ), + base_request: TurnRequest::new("provider-model", messages), + cli_monitor_request: None, + response_config: crate::ai::runtime::RuntimeResponseConfig { + task_id: task_id.to_string(), + conversation_id: conversation_id.to_string(), + needs_create_task: false, + user_query: None, + model_id: "provider-model".to_owned(), + max_context_tokens: Some(128_000), + capabilities: RuntimeCapabilities::provider(), + empty_output_message: None, + }, + action_context: crate::ai::runtime::ProviderActionContext::new_for_test( + task_id.to_string(), + ), + projection_target: super::ProviderProjectionTarget { + task_id: task_id.clone(), + exchange_id: AIAgentExchangeId::new(), + }, + root_task_id: task_id, + did_input_contain_user_query: true, + persistence_offset: 0, + committed_provider_batch: None, + command_action_refs: HashMap::new(), + command_monitor: None, + pending_monitor_observation: None, + pending_command_completion: None, + monitor_prose_continuations: 0, + } +} + +fn start_snapshot_tool( + snapshot: &mut super::ActiveProviderRunSnapshot, + call_id: &str, +) -> ExternalWorkId { + let Some(ProviderRunStep::CallModel(call)) = snapshot.run.next_step().unwrap() else { + panic!("expected provider model call"); + }; + snapshot + .run + .accept_model_turn( + &call.work_id, + CompletedModelTurn { + assistant_content: vec![ContentPart::Text("I will run a command.".to_owned())], + tool_calls: vec![ToolCall { + id: call_id.to_owned(), + name: "run_shell_command".to_owned(), + arguments: serde_json::json!({"command": "sleep 10"}), + }], + usage: Usage::default(), + stop_reason: StopReason::Completed, + advertised_tools: BTreeSet::from(["run_shell_command".to_owned()]), + }, + ) + .unwrap(); + let Some(ProviderRunStep::DispatchTools(batch)) = snapshot.run.next_step().unwrap() else { + panic!("expected provider tool batch"); + }; + batch.work_id +} + +fn attach_snapshot_command_monitor( + snapshot: &mut super::ActiveProviderRunSnapshot, + conversation_id: AIConversationId, +) -> (AIAgentActionId, BlockId, TaskId) { + let action_id = AIAgentActionId::from("command-call".to_owned()); + let block_id = BlockId::new(); + let cli_task_id = TaskId::new("cli-task".to_owned()); + let work_id = snapshot.run.ready_work_id().expect("ready work identity"); + snapshot.cli_monitor_request = Some(TurnRequest::new("provider-model", Vec::new())); + snapshot.command_action_refs.insert( + action_id.clone(), + crate::ai::runtime::ProviderToolExecutionRef::new( + conversation_id, + &work_id, + action_id.to_string(), + ), + ); + snapshot.command_monitor = Some(super::ProviderCommandMonitorState { + run_id: snapshot.run.id().clone(), + originating_work_id: work_id, + originating_call_id: action_id.to_string(), + initial_requested_command_action_id: action_id.clone(), + block_id: block_id.clone(), + command: "sleep 10".to_owned(), + cli_task_id: cli_task_id.clone(), + }); + snapshot.action_context.set_task_id(cli_task_id.to_string()); + snapshot.response_config.task_id = cli_task_id.to_string(); + (action_id, block_id, cli_task_id) +} + +#[test] +fn provider_snapshot_parse_and_validation_reject_corrupt_restore_identity() { + let conversation_id = AIConversationId::new(); + let snapshot = provider_snapshot(conversation_id); + let json = serde_json::to_string(&snapshot).unwrap(); + assert!(super::ActiveProviderRunSnapshot::parse(&json) + .unwrap() + .validate(conversation_id) + .is_ok()); + + let mut unsupported_version = serde_json::to_value(&snapshot).unwrap(); + unsupported_version["version"] = serde_json::json!(99); + assert!( + super::ActiveProviderRunSnapshot::parse(&unsupported_version.to_string()) + .unwrap_err() + .contains("unsupported active provider run snapshot version") + ); + + let mut invalid_offset = snapshot.clone(); + invalid_offset.persistence_offset = invalid_offset.run.transcript().len() + 1; + assert_eq!( + invalid_offset.validate(conversation_id).unwrap_err(), + "provider run persistence offset exceeds transcript length" + ); + + let mut mismatched_model = snapshot.clone(); + mismatched_model.response_config.model_id = "different-model".to_owned(); + assert_eq!( + mismatched_model.validate(conversation_id).unwrap_err(), + "provider run base model does not match response projection" + ); + + let mut orphaned_task = snapshot; + orphaned_task.action_context.set_task_id("orphan-task"); + orphaned_task.response_config.task_id = "orphan-task".to_owned(); + assert_eq!( + orphaned_task.validate(conversation_id).unwrap_err(), + "provider run current task is not owned by its projection or monitor" + ); +} + +#[test] +fn restored_committed_command_requires_durable_terminal_owner() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let action_id = AIAgentActionId::from("command-call".to_owned()); + let work_id = snapshot.run.ready_work_id().expect("ready work identity"); + snapshot.committed_provider_batch = Some(work_id.clone()); + snapshot.command_action_refs.insert( + action_id.clone(), + crate::ai::runtime::ProviderToolExecutionRef::new( + conversation_id, + &work_id, + action_id.to_string(), + ), + ); + + assert_eq!( + super::normalize_restored_provider_snapshot(&mut snapshot).unwrap_err(), + "restored provider command batch completed without durable terminal evidence" + ); +} + +#[test] +fn restore_normalization_removes_interrupted_command_correlation() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let call_id = "command-call"; + let work_id = start_snapshot_tool(&mut snapshot, call_id); + snapshot.run.start_tool(&work_id, call_id).unwrap(); + let action_id = AIAgentActionId::from(call_id.to_owned()); + snapshot.command_action_refs.insert( + action_id.clone(), + crate::ai::runtime::ProviderToolExecutionRef::new( + conversation_id, + &work_id, + action_id.to_string(), + ), + ); + + super::normalize_restored_provider_snapshot(&mut snapshot).unwrap(); + + assert!(!snapshot.command_action_refs.contains_key(&action_id)); + assert!(matches!( + snapshot.run.state(), + ProviderRunState::ReadyToCallModel + )); + let MessageContent::MultiPart(parts) = &snapshot.run.transcript().last().unwrap().content + else { + panic!("interrupted command result should be committed"); + }; + assert!(matches!( + &parts[0], + ContentPart::ToolResult { + tool_use_id, + is_error: true, + .. + } if tool_use_id == call_id + )); +} + +#[test] +fn restore_normalization_reproposes_permission_without_losing_correlation() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let call_id = "command-call"; + let work_id = start_snapshot_tool(&mut snapshot, call_id); + snapshot + .run + .request_tool_permission( + &work_id, + PermissionRequest { + id: "permission-1".to_owned(), + call_id: call_id.to_owned(), + kind: PermissionKind::Execute, + reason: None, + }, + ) + .unwrap(); + let action_id = AIAgentActionId::from(call_id.to_owned()); + snapshot.command_action_refs.insert( + action_id.clone(), + crate::ai::runtime::ProviderToolExecutionRef::new( + conversation_id, + &work_id, + action_id.to_string(), + ), + ); + + super::normalize_restored_provider_snapshot(&mut snapshot).unwrap(); + + assert!(snapshot.command_action_refs.contains_key(&action_id)); + let ProviderRunState::AwaitingTools { batch } = snapshot.run.state() else { + panic!("permission reset should keep the tool batch pending"); + }; + assert!(matches!( + batch.calls[0].state, + galaxy_agent_core::PendingToolCallState::Proposed + )); +} + +#[test] +fn restored_active_command_rebuilds_monitor_observation() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let (action_id, block_id, cli_task_id) = + attach_snapshot_command_monitor(&mut snapshot, conversation_id); + snapshot.pending_command_completion = Some(super::PendingProviderCommandCompletion { + block_id: block_id.clone(), + initial_requested_command_action_id: Some(action_id.clone()), + command: "stale".to_owned(), + output: "stale".to_owned(), + exit_code: 1, + }); + + super::apply_restored_provider_command_evidence( + conversation_id, + &mut snapshot, + super::RestoredProviderCommandEvidence { + conversation_id: Some(conversation_id), + requested_command_action_id: Some(action_id), + cli_task_id: Some(cli_task_id.clone()), + command: "sleep 10".to_owned(), + state: BlockState::Executing, + output: "running".to_owned(), + exit_code: 0, + }, + ) + .unwrap(); + + assert!(snapshot.pending_command_completion.is_none()); + let observation = snapshot + .pending_monitor_observation + .expect("active command should restore monitoring"); + assert_eq!(observation.block_id, block_id); + assert_eq!(observation.cli_task_id, cli_task_id); +} + +#[test] +fn restored_completed_command_rebuilds_exact_completion_evidence() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let (action_id, block_id, cli_task_id) = + attach_snapshot_command_monitor(&mut snapshot, conversation_id); + snapshot.pending_monitor_observation = Some(super::PendingProviderMonitorObservation { + block_id: block_id.clone(), + cli_task_id: cli_task_id.clone(), + }); + + super::apply_restored_provider_command_evidence( + conversation_id, + &mut snapshot, + super::RestoredProviderCommandEvidence { + conversation_id: Some(conversation_id), + requested_command_action_id: Some(action_id.clone()), + cli_task_id: Some(cli_task_id), + command: "sleep 10".to_owned(), + state: BlockState::DoneWithExecution, + output: "done".to_owned(), + exit_code: 17, + }, + ) + .unwrap(); + + assert!(snapshot.pending_monitor_observation.is_none()); + let completion = snapshot + .pending_command_completion + .expect("completed command should restore final evidence"); + assert_eq!(completion.block_id, block_id); + assert_eq!( + completion.initial_requested_command_action_id, + Some(action_id) + ); + assert_eq!(completion.command, "sleep 10"); + assert_eq!(completion.output, "done"); + assert_eq!(completion.exit_code, 17); +} + +#[test] +fn restored_command_rejects_mismatched_or_invalid_terminal_evidence() { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let (action_id, _block_id, cli_task_id) = + attach_snapshot_command_monitor(&mut snapshot, conversation_id); + + assert_eq!( + super::apply_restored_provider_command_evidence( + conversation_id, + &mut snapshot, + super::RestoredProviderCommandEvidence { + conversation_id: Some(AIConversationId::new()), + requested_command_action_id: Some(action_id.clone()), + cli_task_id: Some(cli_task_id.clone()), + command: "sleep 10".to_owned(), + state: BlockState::Executing, + output: String::new(), + exit_code: 0, + }, + ) + .unwrap_err(), + "restored provider command block identity does not match" + ); + assert_eq!( + super::apply_restored_provider_command_evidence( + conversation_id, + &mut snapshot, + super::RestoredProviderCommandEvidence { + conversation_id: Some(conversation_id), + requested_command_action_id: Some(action_id), + cli_task_id: Some(cli_task_id), + command: "sleep 10".to_owned(), + state: BlockState::Background, + output: String::new(), + exit_code: 0, + }, + ) + .unwrap_err(), + "restored provider command block has an invalid state" + ); +} + +#[test] +fn provider_restore_failure_is_visible_and_clears_persisted_run() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let terminal_surface_id = terminal.id(); + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + let conversation_id = history_model.start_new_conversation( + terminal_surface_id, + false, + false, + false, + ctx, + ); + history_model + .persist_active_provider_run_json( + conversation_id, + Some("corrupt snapshot".to_owned()), + ctx, + ) + .unwrap(); + conversation_id + }); + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.restoring_provider_runs.insert(conversation_id); + controller.fail_restored_provider_run( + conversation_id, + "snapshot identity mismatch".to_owned(), + ctx, + ); + assert!(!controller + .restoring_provider_runs + .contains(&conversation_id)); + }); + + let history_model = BlocklistAIHistoryModel::handle(ctx); + let conversation = history_model + .as_ref(ctx) + .conversation(&conversation_id) + .expect("failed restored conversation should remain visible"); + assert_eq!(conversation.status(), &ConversationStatus::Error); + assert!(conversation.active_provider_run_json().is_none()); + let error = conversation + .status_error() + .expect("restore failure should retain a structured error"); + assert!(error + .to_string() + .contains("Failed to restore active provider run: snapshot identity mismatch")); + assert!(matches!( + error, + crate::ai::agent::RenderableAIError::Other { + will_attempt_resume: false, + waiting_for_network: false, + is_user_error: false, + .. + } + )); + }); + }); +} + +#[test] +fn provider_lifecycle_requires_exact_active_work_identity() { + let conversation_id = AIConversationId::new(); + let active_work = ExternalWorkId { + run_id: ProviderRunId::new("current"), + epoch: RunEpoch::new(3), + }; + + assert!(super::provider_execution_matches_active_work( + &active_work.run_id, + Some(&active_work), + &provider_execution_ref(conversation_id, "current", 3), + )); + assert!(!super::provider_execution_matches_active_work( + &active_work.run_id, + Some(&active_work), + &provider_execution_ref(conversation_id, "current", 2), + )); + assert!(!super::provider_execution_matches_active_work( + &active_work.run_id, + Some(&active_work), + &provider_execution_ref(conversation_id, "old", 3), + )); +} + +#[test] +fn provider_finished_action_only_resumes_its_committed_batch() { + let conversation_id = AIConversationId::new(); + let run_id = ProviderRunId::new("current"); + let work_id = ExternalWorkId { + run_id: run_id.clone(), + epoch: RunEpoch::new(3), + }; + let current = provider_execution_ref(conversation_id, "current", 3); + + assert_eq!( + super::provider_finished_action_disposition(&run_id, Some(&work_id), None, ¤t), + super::ProviderFinishedActionDisposition::AwaitBatchCommit, + ); + assert_eq!( + super::provider_finished_action_disposition(&run_id, None, Some(&work_id), ¤t), + super::ProviderFinishedActionDisposition::Resume, + ); + assert_eq!( + super::provider_finished_action_disposition( + &run_id, + None, + Some(&work_id), + &provider_execution_ref(conversation_id, "current", 2), + ), + super::ProviderFinishedActionDisposition::Ignore, + ); + assert_eq!( + super::provider_finished_action_disposition( + &run_id, + None, + Some(&work_id), + &provider_execution_ref(conversation_id, "old", 3), + ), + super::ProviderFinishedActionDisposition::Ignore, + ); +} + +#[test] +fn provider_boundary_prioritizes_completion_and_waits_for_committed_results() { + assert_eq!( + super::provider_boundary_intent( + super::ProviderBoundaryPhase::Ready, + true, + true, + true, + true, + true, + 0, + ), + super::ProviderBoundaryIntent::Park + ); + assert_eq!( + super::provider_boundary_intent( + super::ProviderBoundaryPhase::Unsafe, + false, + true, + true, + true, + true, + 0, + ), + super::ProviderBoundaryIntent::Park + ); + assert_eq!( + super::provider_boundary_intent( + super::ProviderBoundaryPhase::Ready, + false, + true, + true, + true, + true, + 0, + ), + super::ProviderBoundaryIntent::ApplyCompletion + ); + assert_eq!( + super::provider_boundary_intent( + super::ProviderBoundaryPhase::Ready, + false, + false, + true, + false, + true, + 0, + ), + super::ProviderBoundaryIntent::ApplyMonitorObservation + ); +} + +#[test] +fn provider_monitor_prose_retry_is_bounded_then_parks() { + assert_eq!( + super::provider_boundary_intent( + super::ProviderBoundaryPhase::AwaitingDriver, + false, + false, + false, + true, + true, + 0, + ), + super::ProviderBoundaryIntent::RetryMonitor + ); + assert_eq!( + super::provider_boundary_intent( + super::ProviderBoundaryPhase::AwaitingDriver, + false, + false, + false, + true, + true, + super::MAX_PROVIDER_MONITOR_PROSE_CONTINUATIONS, + ), + super::ProviderBoundaryIntent::Park + ); + assert_eq!( + super::provider_boundary_intent( + super::ProviderBoundaryPhase::AwaitingDriver, + false, + false, + false, + false, + false, + 0, + ), + super::ProviderBoundaryIntent::CompleteRun + ); +} + +#[test] +fn provider_completion_requires_current_run_and_exact_monitor_identity() { + let conversation_id = AIConversationId::new(); + let run_id = ProviderRunId::new("current"); + let block_id = BlockId::new(); + let other_block_id = BlockId::new(); + let action_id = AIAgentActionId::from("command-1".to_owned()); + let other_action_id = AIAgentActionId::from("command-2".to_owned()); + let execution_ref = provider_execution_ref(conversation_id, "current", 3); + let command_action_refs = HashMap::from([(action_id.clone(), execution_ref.clone())]); + + assert!(super::provider_command_completion_matches( + &run_id, + &command_action_refs, + None, + &block_id, + Some(&action_id), + )); + assert!(!super::provider_command_completion_matches( + &run_id, + &command_action_refs, + None, + &block_id, + Some(&other_action_id), + )); + + let monitor = super::ProviderCommandMonitorState { + run_id: run_id.clone(), + originating_work_id: execution_ref.work_id(), + originating_call_id: action_id.to_string(), + initial_requested_command_action_id: action_id.clone(), + block_id: block_id.clone(), + command: "sleep 10".to_owned(), + cli_task_id: TaskId::new("cli-task".to_owned()), + }; + assert!(super::provider_command_completion_matches( + &run_id, + &command_action_refs, + Some(&monitor), + &block_id, + Some(&action_id), + )); + assert!(super::provider_command_completion_matches( + &run_id, + &command_action_refs, + Some(&monitor), + &block_id, + None, + )); + assert!(!super::provider_command_completion_matches( + &run_id, + &command_action_refs, + Some(&monitor), + &other_block_id, + Some(&action_id), + )); + assert!(!super::provider_command_completion_matches( + &ProviderRunId::new("replacement"), + &command_action_refs, + Some(&monitor), + &block_id, + Some(&action_id), + )); +} + +#[test] +fn pending_provider_completion_reconciles_with_its_committed_snapshot() { + let block_id = BlockId::new(); + let action_id = AIAgentActionId::from("command-1".to_owned()); + let mut completion = super::PendingProviderCommandCompletion { + block_id: block_id.clone(), + initial_requested_command_action_id: Some(action_id.clone()), + command: String::new(), + output: "done".to_owned(), + exit_code: 0, + }; + + assert!(super::reconcile_provider_completion_with_snapshot( + Some(&mut completion), + &block_id, + &action_id, + Some("sleep 10"), + None, + ) + .unwrap()); + assert_eq!(completion.command, "sleep 10"); + assert!(super::reconcile_provider_completion_with_snapshot( + None, + &block_id, + &action_id, + Some("sleep 10"), + None, + ) + .is_ok_and(|matched| !matched)); +} + +#[test] +fn pending_provider_completion_rejects_a_different_snapshot() { + let block_id = BlockId::new(); + let action_id = AIAgentActionId::from("command-1".to_owned()); + let mut completion = super::PendingProviderCommandCompletion { + block_id, + initial_requested_command_action_id: Some(action_id.clone()), + command: String::new(), + output: "done".to_owned(), + exit_code: 0, + }; + + assert!(super::reconcile_provider_completion_with_snapshot( + Some(&mut completion), + &BlockId::new(), + &action_id, + Some("sleep 10"), + None, + ) + .is_err()); + let completion_block_id = completion.block_id.clone(); + assert!(super::reconcile_provider_completion_with_snapshot( + Some(&mut completion), + &completion_block_id, + &AIAgentActionId::from("command-2".to_owned()), + Some("sleep 10"), + None, + ) + .is_err()); +} + +#[test] +fn nonzero_provider_completion_is_continuation_evidence() { + let completion = super::PendingProviderCommandCompletion { + block_id: BlockId::new(), + initial_requested_command_action_id: None, + command: "cargo test".to_owned(), + output: "one test failed".to_owned(), + exit_code: 17, + }; + let galaxy_agent_core::MessageContent::Text(observation) = completion.observation() else { + panic!("command completion must be text evidence"); + }; + assert!(observation.contains("exit code 17")); + assert!(observation.contains("nonzero exit is not automatic run completion")); + assert!(observation.contains("Continue the original objective")); +} + +#[test] +fn provider_command_result_classifier_covers_snapshot_and_finished_variants() { + let block_id = BlockId::new(); + let exit_code = ExitCode::from(17); + let expected_snapshot = super::ProviderCommandResult::Snapshot { + block_id: block_id.clone(), + command: Some("sleep 10".to_owned()), + }; + assert_eq!( + super::classify_provider_command_result(&AIAgentActionResultType::RequestCommandOutput( + RequestCommandOutputResult::LongRunningCommandSnapshot { + block_id: block_id.clone(), + command: "sleep 10".to_owned(), + grid_contents: "running".to_owned(), + cursor: String::new(), + is_alt_screen_active: false, + }, + ),), + Some(expected_snapshot.clone()) + ); + assert_eq!( + super::classify_provider_command_result(&AIAgentActionResultType::ReadShellCommandOutput( + ReadShellCommandOutputResult::LongRunningCommandSnapshot { + command: "sleep 10".to_owned(), + block_id: block_id.clone(), + grid_contents: "running".to_owned(), + cursor: String::new(), + is_alt_screen_active: false, + is_preempted: false, + }, + ),), + Some(expected_snapshot) + ); + for result in [ + AIAgentActionResultType::WriteToLongRunningShellCommand( + WriteToLongRunningShellCommandResult::Snapshot { + block_id: block_id.clone(), + grid_contents: "running".to_owned(), + cursor: String::new(), + is_alt_screen_active: false, + is_preempted: false, + }, + ), + AIAgentActionResultType::TransferShellCommandControlToUser( + TransferShellCommandControlToUserResult::Snapshot { + block_id: block_id.clone(), + grid_contents: "running".to_owned(), + cursor: String::new(), + is_alt_screen_active: false, + is_preempted: false, + }, + ), + ] { + assert_eq!( + super::classify_provider_command_result(&result), + Some(super::ProviderCommandResult::Snapshot { + block_id: block_id.clone(), + command: None, + }) + ); + } + + let expected_finished = super::ProviderCommandResult::Finished { + block_id: block_id.clone(), + command: Some("sleep 10".to_owned()), + output: "failed".to_owned(), + exit_code: 17, + }; + assert_eq!( + super::classify_provider_command_result(&AIAgentActionResultType::RequestCommandOutput( + RequestCommandOutputResult::Completed { + block_id: block_id.clone(), + command: "sleep 10".to_owned(), + output: "failed".to_owned(), + exit_code, + start_ts: None, + completed_ts: None, + }, + ),), + Some(expected_finished.clone()) + ); + assert_eq!( + super::classify_provider_command_result(&AIAgentActionResultType::ReadShellCommandOutput( + ReadShellCommandOutputResult::CommandFinished { + command: "sleep 10".to_owned(), + block_id: block_id.clone(), + output: "failed".to_owned(), + exit_code, + start_ts: None, + completed_ts: None, + }, + ),), + Some(expected_finished) + ); + for result in [ + AIAgentActionResultType::WriteToLongRunningShellCommand( + WriteToLongRunningShellCommandResult::CommandFinished { + block_id: block_id.clone(), + output: "failed".to_owned(), + exit_code, + start_ts: None, + completed_ts: None, + }, + ), + AIAgentActionResultType::TransferShellCommandControlToUser( + TransferShellCommandControlToUserResult::CommandFinished { + block_id: block_id.clone(), + output: "failed".to_owned(), + exit_code, + start_ts: None, + completed_ts: None, + }, + ), + ] { + assert_eq!( + super::classify_provider_command_result(&result), + Some(super::ProviderCommandResult::Finished { + block_id: block_id.clone(), + command: None, + output: "failed".to_owned(), + exit_code: 17, + }) + ); + } +} + +#[test] +fn provider_command_result_classifier_ignores_cancelled_and_error_variants() { + let results = [ + AIAgentActionResultType::RequestCommandOutput( + RequestCommandOutputResult::CancelledBeforeExecution, + ), + AIAgentActionResultType::RequestCommandOutput(RequestCommandOutputResult::Denylisted { + command: "blocked".to_owned(), + }), + AIAgentActionResultType::WriteToLongRunningShellCommand( + WriteToLongRunningShellCommandResult::Cancelled, + ), + AIAgentActionResultType::WriteToLongRunningShellCommand( + WriteToLongRunningShellCommandResult::Error(ShellCommandError::BlockNotFound), + ), + AIAgentActionResultType::ReadShellCommandOutput(ReadShellCommandOutputResult::Cancelled), + AIAgentActionResultType::ReadShellCommandOutput(ReadShellCommandOutputResult::Error( + ShellCommandError::BlockNotFound, + )), + AIAgentActionResultType::TransferShellCommandControlToUser( + TransferShellCommandControlToUserResult::Cancelled, + ), + AIAgentActionResultType::TransferShellCommandControlToUser( + TransferShellCommandControlToUserResult::Error(ShellCommandError::BlockNotFound), + ), + ]; + assert!(results + .iter() + .all(|result| super::classify_provider_command_result(result).is_none())); +} + #[test] fn no_action_tool_error_recovery_detects_unfulfilled_tool_intent() { assert_eq!( @@ -121,18 +1031,10 @@ fn no_action_tool_error_recovery_ignores_normal_answers_and_non_failed_tools() { ); } -#[test] -fn tool_queue_decision_blocks_failed_tool_proposal_before_snapshot_fallback() { - assert_eq!( - super::tool_queue_decision(false, false, true, false, 1, 1,), - super::ToolQueueDecision::BlockedFailedToolProposal - ); -} - #[test] fn tool_queue_decision_blocks_parent_tools_while_child_agents_are_active() { assert_eq!( - super::tool_queue_decision(false, false, false, true, 2, 0,), + super::tool_queue_decision(false, false, true, 2), super::ToolQueueDecision::BlockedActiveChildAgents ); } @@ -140,23 +1042,20 @@ fn tool_queue_decision_blocks_parent_tools_while_child_agents_are_active() { #[test] fn tool_queue_decision_preserves_existing_terminal_precedence() { assert_eq!( - super::tool_queue_decision(true, false, true, true, 1, 1,), + super::tool_queue_decision(true, false, true, 1), super::ToolQueueDecision::Cancelled ); assert_eq!( - super::tool_queue_decision(false, true, true, true, 1, 1,), + super::tool_queue_decision(false, true, true, 1), super::ToolQueueDecision::UnfinishedExchange ); } #[test] -fn tool_queue_decision_uses_snapshot_fallback_only_when_unblocked() { - let decision = super::tool_queue_decision(false, false, false, false, 1, 1); +fn tool_queue_decision_queues_actions_when_unblocked() { + let decision = super::tool_queue_decision(false, false, false, 1); - assert_eq!( - decision, - super::ToolQueueDecision::QueueActionsWithStreamSnapshotFallback - ); + assert_eq!(decision, super::ToolQueueDecision::QueueActions); assert!(decision.will_queue_actions()); } @@ -493,36 +1392,6 @@ fn input_for_query_converts_prompt_attachments_and_ignores_live_staging() { }); } -#[test] -fn cancelling_conversation_aborts_pending_auto_resume() { - App::test((), |mut app| async move { - initialize_app_for_terminal_view(&mut app); - let terminal = add_window_with_terminal(&mut app, None); - - // An ID with no backing conversation: if the scheduled wait ever - // completes, the resume is a harmless no-op. - let conversation_id = AIConversationId::new(); - - terminal.update(&mut app, |terminal, ctx| { - terminal.ai_controller().update(ctx, |controller, ctx| { - controller.schedule_auto_resume_after_error(conversation_id, ctx); - assert!(controller - .pending_auto_resume_handles - .contains_key(&conversation_id)); - - controller.cancel_conversation_progress( - conversation_id, - CancellationReason::ManuallyCancelled, - ctx, - ); - assert!(!controller - .pending_auto_resume_handles - .contains_key(&conversation_id)); - }); - }); - }); -} - #[test] fn user_follow_up_does_not_cancel_unresolved_ask_user_question() { App::test((), |mut app| async move { diff --git a/app/src/ai/blocklist/history_model.rs b/app/src/ai/blocklist/history_model.rs index 716c0d9e..644a91bb 100644 --- a/app/src/ai/blocklist/history_model.rs +++ b/app/src/ai/blocklist/history_model.rs @@ -563,6 +563,44 @@ impl BlocklistAIHistoryModel { conversation.write_updated_conversation_state(ctx); } + pub(crate) fn persist_active_provider_run_json( + &mut self, + conversation_id: AIConversationId, + snapshot: Option, + ctx: &mut ModelContext, + ) -> Result<(), UpdateHistoryError> { + let conversation = self + .conversations_by_id + .get_mut(&conversation_id) + .ok_or(UpdateHistoryError::ConversationNotFound(conversation_id))?; + conversation.set_active_provider_run_json(snapshot); + conversation.write_updated_conversation_state(ctx); + Ok(()) + } + + pub(crate) fn rebind_provider_projection( + &mut self, + conversation_id: AIConversationId, + task_id: &TaskId, + exchange_id: AIAgentExchangeId, + response_stream_id: ResponseStreamId, + terminal_surface_id: EntityId, + ctx: &mut ModelContext, + ) -> Result<(), UpdateHistoryError> { + let conversation = self + .conversations_by_id + .get_mut(&conversation_id) + .ok_or(UpdateHistoryError::ConversationNotFound(conversation_id))?; + conversation.rebind_provider_projection( + task_id, + exchange_id, + response_stream_id, + terminal_surface_id, + ctx, + )?; + Ok(()) + } + fn update_cached_metadata_for_conversation(&mut self, conversation_id: AIConversationId) { let Some(conversation) = self.conversations_by_id.get(&conversation_id) else { return; @@ -1652,6 +1690,7 @@ impl BlocklistAIHistoryModel { let conversation_data = AgentConversationData { agent_backend: source_conversation.agent_backend().for_fork(), + active_provider_run_json: None, server_conversation_token: None, conversation_usage_metadata: Some(source_conversation.usage_metadata()), reverted_action_ids, @@ -1816,6 +1855,7 @@ impl BlocklistAIHistoryModel { // be recomputed based on the retained exchanges in a follow-up. let conversation_data = AgentConversationData { agent_backend: conversation.agent_backend().for_fork(), + active_provider_run_json: None, server_conversation_token: None, conversation_usage_metadata: None, reverted_action_ids, @@ -2824,6 +2864,9 @@ fn merged_remote_child_placeholder_conversation_data( // Placeholder authoritative. agent_backend: placeholder.agent_backend().clone(), + // Active process-local provider runs cannot be merged from a cloud transcript. + active_provider_run_json: None, + // Cloud authoritative. server_conversation_token: cloud_conversation .server_conversation_token() diff --git a/app/src/ai/blocklist/inline_action/requested_command.rs b/app/src/ai/blocklist/inline_action/requested_command.rs index c3696783..0dd3ef87 100644 --- a/app/src/ai/blocklist/inline_action/requested_command.rs +++ b/app/src/ai/blocklist/inline_action/requested_command.rs @@ -418,21 +418,21 @@ impl RequestedCommandView { if !is_finished { ctx.subscribe_to_model(action_model, |me, _, event, ctx| { match event { - BlocklistAIActionEvent::QueuedAction(action_id) + BlocklistAIActionEvent::QueuedAction { action_id, .. } if *action_id == me.action_id => { ctx.notify(); } - BlocklistAIActionEvent::ActionBlockedOnUserConfirmation(action_id) - if *action_id == me.action_id => - { + BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { + action_id, .. + } if *action_id == me.action_id => { if me.action_type.is_requested_command() { me.ensure_editor(ctx); } me.set_is_header_expanded(true, ctx); ctx.notify(); } - BlocklistAIActionEvent::ExecutingAction(action_id) + BlocklistAIActionEvent::ExecutingAction { action_id, .. } if *action_id == me.action_id => { // For shared-session viewers, sync the command text from the action when it starts executing. diff --git a/app/src/ai/blocklist/inline_action/run_agents_card_view.rs b/app/src/ai/blocklist/inline_action/run_agents_card_view.rs index 847468b8..64df1ebc 100644 --- a/app/src/ai/blocklist/inline_action/run_agents_card_view.rs +++ b/app/src/ai/blocklist/inline_action/run_agents_card_view.rs @@ -376,7 +376,7 @@ impl RunAgentsCardView { { ctx.notify(); } - BlocklistAIActionEvent::ActionBlockedOnUserConfirmation(action_id) + BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { action_id, .. } if action_id == &action_id_for_action_events => { // Normal case: streaming is complete and the action is diff --git a/app/src/ai/blocklist/passive_suggestions/maa.rs b/app/src/ai/blocklist/passive_suggestions/maa.rs index 4622cf1d..d697686c 100644 --- a/app/src/ai/blocklist/passive_suggestions/maa.rs +++ b/app/src/ai/blocklist/passive_suggestions/maa.rs @@ -11,7 +11,6 @@ use warpui::r#async::SpawnedFutureHandle; use warpui::{Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; use super::super::controller::{BlocklistAIController, BlocklistAIControllerEvent}; -use crate::ai::agent::api::generate_multi_agent_output; use crate::ai::agent::conversation::AIConversationId; use crate::ai::agent::{ AIIdentifiers, FileContext, PassiveCodeDiffEntry, PassiveSuggestionTrigger, diff --git a/app/src/ai/openai/mod.rs b/app/src/ai/openai/mod.rs index 91efa128..db5993b6 100644 --- a/app/src/ai/openai/mod.rs +++ b/app/src/ai/openai/mod.rs @@ -2,7 +2,6 @@ pub mod client; pub mod convert; pub mod request_translator; pub mod response_translator; -pub mod translator; #[cfg(test)] #[path = "convert_tests.rs"] diff --git a/app/src/ai/openai/translator.rs b/app/src/ai/openai/translator.rs deleted file mode 100644 index 0a720961..00000000 --- a/app/src/ai/openai/translator.rs +++ /dev/null @@ -1,190 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use warp_multi_agent_api as api; - -use super::client::{OpenAIClient, OpenAIClientConfig, OpenAIError}; -use super::convert::build_openai_request; -use super::request_translator::sanitize_messages_for_openai; -use super::response_translator::{openai_stream_to_response_events, OpenAIStreamContext}; -use crate::ai::agent::api::LegacyResponseStream; -use crate::ai::bedrock::request_translator; -use crate::ai::provider::types::{ - flatten_tool_history_for_no_tools_turn, ConversationMessage, MessageContent, MessageRole, -}; - -const DEFAULT_MAX_OUTPUT_TOKENS: u32 = 64_000; - -pub struct TranslatorRequest { - pub config: OpenAIClientConfig, - pub model_id: String, - pub root_task_id: Option, - pub message_history: Vec, - pub tool_result_archive: Vec, - pub progressive_summary: Option, - pub messages_sent: Arc>>, - /// Global rules (name, content) from the local CloudModel. - pub global_rules: Vec<(String, String)>, - /// Whether the native input should be emitted as a transcript-visible user query. - pub emit_user_query_message: bool, -} - -pub(crate) struct PreparedTurn { - pub(crate) task_id: String, - pub(crate) needs_create_task: bool, - pub(crate) user_query: Option, - pub(crate) messages: Vec, - pub(crate) system_prompt: Option, - pub(crate) tools: Vec, - pub(crate) model_id: String, - pub(crate) persistent_message_count: usize, -} - -pub(crate) fn prepare_turn(params: &TranslatorRequest, request: &mut api::Request) -> PreparedTurn { - let task_id = params.root_task_id.clone().unwrap_or_else(|| { - request - .task_context - .as_ref() - .and_then(|tc| tc.tasks.first()) - .map(|task| task.id.clone()) - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()) - }); - - let needs_create_task = request - .task_context - .as_ref() - .map(|task_context| task_context.tasks.is_empty()) - .unwrap_or(true); - let model_id = if params.model_id.is_empty() || params.model_id == "auto" { - params - .config - .model - .clone() - .unwrap_or_else(|| "anthropic/claude-sonnet-4-6".to_string()) - } else { - params - .config - .model - .clone() - .unwrap_or_else(|| params.model_id.clone()) - }; - - request_translator::inject_input_messages_into_task(request); - let new_input_messages = request_translator::extract_new_input_messages(request); - let persistent_message_count = params.message_history.len() + new_input_messages.len(); - let mut messages = Vec::new(); - - if let Some(summary) = ¶ms.progressive_summary { - messages.push(ConversationMessage { - role: MessageRole::User, - content: MessageContent::Text(format!( - "\n{summary}\n\n\n\ - The above summarizes earlier conversation history. The detailed messages below are the most recent exchanges." - )), - }); - messages.push(ConversationMessage { - role: MessageRole::Assistant, - content: MessageContent::Text( - "Understood, I have the prior context. Continuing with the recent conversation." - .to_string(), - ), - }); - } - - messages.extend(params.message_history.clone()); - messages.extend(new_input_messages); - for message in &mut messages { - message.truncate_tool_results_for_provider_request(); - } - sanitize_messages_for_openai(&mut messages); - let tools = request_translator::extract_tools(request); - if tools_are_inline_only(&tools) { - flatten_tool_history_for_no_tools_turn(&mut messages); - } - - PreparedTurn { - task_id, - needs_create_task, - user_query: params - .emit_user_query_message - .then(|| request_translator::extract_user_query_text(request)) - .flatten(), - messages, - system_prompt: request_translator::extract_system_prompt(request, ¶ms.global_rules), - tools, - model_id, - persistent_message_count, - } -} - -fn tools_are_inline_only(tools: &[crate::ai::provider::types::ToolDefinition]) -> bool { - tools.iter().all(|tool| tool.name == "recall_tool_history") -} - -pub async fn execute( - params: TranslatorRequest, - request: &mut api::Request, -) -> Result { - let client = OpenAIClient::from_config(params.config.clone()); - let PreparedTurn { - task_id, - needs_create_task, - user_query, - mut messages, - system_prompt, - tools, - model_id, - persistent_message_count, - } = prepare_turn(¶ms, request); - - log::info!( - "[openai] Translator: model={model_id}, task_id={task_id}, needs_create_task={needs_create_task}" - ); - - log::info!( - "[openai] Sending {} messages, system_prompt={}, tools={}", - messages.len(), - system_prompt.is_some(), - tools.len() - ); - - let max_output_tokens = params - .config - .max_output_tokens - .unwrap_or(DEFAULT_MAX_OUTPUT_TOKENS) - .min(i32::MAX as u32) as i32; - - let request_body = build_openai_request( - messages.clone(), - system_prompt, - tools, - max_output_tokens, - None, - &model_id, - ); - - let byte_stream = client.chat_completions_stream(request_body).await?; - - // Store the message history for the controller - if let Ok(mut sent) = params.messages_sent.lock() { - if persistent_message_count > 0 && messages.len() >= persistent_message_count { - *sent = messages.split_off(messages.len() - persistent_message_count); - } else { - *sent = messages; - } - } - - let stream = openai_stream_to_response_events( - byte_stream, - OpenAIStreamContext { - task_id, - needs_create_task, - user_query, - messages_sent: params.messages_sent.clone(), - model_id, - max_context_tokens: params.config.max_input_tokens, - tool_result_archive: params.tool_result_archive, - }, - ); - - Ok(stream) -} diff --git a/app/src/ai/runtime/event_translator.rs b/app/src/ai/runtime/event_translator.rs index 0666defe..a1a3b2d8 100644 --- a/app/src/ai/runtime/event_translator.rs +++ b/app/src/ai/runtime/event_translator.rs @@ -1,12 +1,14 @@ use std::collections::HashMap; use galaxy_agent_core::{ - AgentEvent, RuntimeActivity, RuntimeActivityStatus, RuntimeCapabilities, StopReason, Usage, + AgentEvent, ProviderRunOutcome, RuntimeActivity, RuntimeActivityStatus, RuntimeCapabilities, + StopReason, Usage, }; use uuid::Uuid; use warp_multi_agent_api::response_event::stream_finished; 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::{ build_add_agent_output_message, build_append_text, build_create_task, build_stream_init, @@ -14,6 +16,7 @@ use crate::ai::bedrock::response_translator::{ }; use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage}; +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] pub(crate) struct RuntimeResponseConfig { pub(crate) task_id: String, pub(crate) conversation_id: String, @@ -41,17 +44,104 @@ pub(crate) struct RuntimeResponseTranslator { context_usage: Option<(u64, u64)>, } +/// Projects a multi-turn provider run into one existing Galaxy response stream. +/// Intermediate model stops remain coordinator-internal; only the run outcome +/// emits the UI's terminal `Finished` event. +pub(crate) struct ProviderRunResponseProjector { + translator: RuntimeResponseTranslator, + has_started_model_turn: bool, + finished: bool, +} + +impl ProviderRunResponseProjector { + pub(crate) fn new(config: RuntimeResponseConfig) -> Self { + Self { + translator: RuntimeResponseTranslator::new(config), + has_started_model_turn: false, + finished: false, + } + } + + pub(crate) fn restored(config: RuntimeResponseConfig) -> Self { + Self { + translator: RuntimeResponseTranslator::restored(config), + has_started_model_turn: false, + finished: false, + } + } + + pub(crate) fn project( + &mut self, + projection: ProviderRunProjection, + ) -> Result, String> { + if self.finished { + return Err("provider run projection is already finished".to_string()); + } + match projection { + ProviderRunProjection::ModelTurnStarted { .. } => { + if self.has_started_model_turn { + self.translator.begin_followup_turn(); + } + self.has_started_model_turn = true; + self.translator.translate(AgentEvent::TurnStarted { + runtime_request_id: String::new(), + }) + } + ProviderRunProjection::ModelEvent { event, .. } => self.translator.translate(event), + ProviderRunProjection::ModelRetry { .. } + | ProviderRunProjection::ToolBatchReady { .. } => Ok(Vec::new()), + } + } + + pub(crate) fn set_task_id(&mut self, task_id: impl Into) { + self.translator.set_task_id(task_id); + } + + pub(crate) fn finish( + &mut self, + outcome: &ProviderRunOutcome, + ) -> Result, String> { + if self.finished { + return Err("provider run projection is already finished".to_string()); + } + self.finished = true; + match outcome { + ProviderRunOutcome::Completed(completion) => { + self.translator.translate(AgentEvent::TurnStopped { + reason: completion.stop_reason.clone(), + }) + } + ProviderRunOutcome::Failed(failure) => { + Ok(self.translator.provider_failure(&failure.message)) + } + ProviderRunOutcome::Cancelled { .. } => { + self.translator.translate(AgentEvent::TurnStopped { + reason: StopReason::Cancelled, + }) + } + } + } +} + impl RuntimeResponseTranslator { pub(crate) fn new(config: RuntimeResponseConfig) -> Self { + Self::with_initialization(config, false) + } + + pub(crate) fn restored(config: RuntimeResponseConfig) -> Self { + Self::with_initialization(config, true) + } + + fn with_initialization(config: RuntimeResponseConfig, initialized: bool) -> Self { Self { config, request_id: Uuid::new_v4().to_string(), - initialized: false, + initialized, text_message_id: None, reasoning_message_id: None, activity_message_ids: HashMap::new(), activities: HashMap::new(), - has_visible_output: false, + has_visible_output: initialized, usage: Usage::default(), context_usage: None, } @@ -152,6 +242,18 @@ impl RuntimeResponseTranslator { self.reasoning_message_id = None; } + pub(crate) fn set_task_id(&mut self, task_id: impl Into) { + let task_id = task_id.into(); + if self.config.task_id == task_id { + return; + } + self.config.task_id = task_id; + self.text_message_id = None; + self.reasoning_message_id = None; + self.activity_message_ids.clear(); + self.activities.clear(); + } + fn initialize(&mut self, events: &mut Vec) { if self.initialized { return; @@ -256,7 +358,23 @@ impl RuntimeResponseTranslator { } fn finished(&self, reason: StopReason) -> ResponseEvent { - let reason = map_stop_reason(reason); + self.finished_with_reason(map_stop_reason(reason)) + } + + fn provider_failure(&mut self, message: &str) -> Vec { + let mut events = Vec::new(); + self.initialize(&mut events); + events.push( + self.finished_with_reason(stream_finished::Reason::InternalError( + stream_finished::InternalError { + message: message.to_owned(), + }, + )), + ); + events + } + + fn finished_with_reason(&self, reason: stream_finished::Reason) -> ResponseEvent { if !self.config.capabilities.host_managed_history { let (used_tokens, context_size) = self.context_usage.unwrap_or_default(); return build_context_finished( diff --git a/app/src/ai/runtime/event_translator_tests.rs b/app/src/ai/runtime/event_translator_tests.rs index c2c9e260..44bd01e6 100644 --- a/app/src/ai/runtime/event_translator_tests.rs +++ b/app/src/ai/runtime/event_translator_tests.rs @@ -3,8 +3,9 @@ use galaxy_agent_core::{ }; use warp_multi_agent_api::{client_action, message, response_event}; -use super::{RuntimeResponseConfig, RuntimeResponseTranslator}; +use super::{ProviderRunResponseProjector, RuntimeResponseConfig, RuntimeResponseTranslator}; use crate::ai::agent::runtime_activity; +use crate::ai::runtime::provider_run_coordinator::ProviderRunProjection; fn provider_translator() -> RuntimeResponseTranslator { RuntimeResponseTranslator::new(RuntimeResponseConfig { @@ -32,6 +33,51 @@ fn session_translator() -> RuntimeResponseTranslator { }) } +#[test] +fn restored_provider_projection_skips_stream_initialization() { + let config = RuntimeResponseConfig { + task_id: "task".to_owned(), + conversation_id: "conversation".to_owned(), + needs_create_task: true, + user_query: Some("do not duplicate".to_owned()), + model_id: "model".to_owned(), + max_context_tokens: Some(1_000), + capabilities: RuntimeCapabilities::provider(), + empty_output_message: None, + }; + let mut projector = ProviderRunResponseProjector::restored(config); + let work_id = galaxy_agent_core::ExternalWorkId { + run_id: galaxy_agent_core::ProviderRunId::new("run"), + epoch: galaxy_agent_core::RunEpoch::new(2), + }; + + assert!(projector + .project(ProviderRunProjection::ModelTurnStarted { + work_id: work_id.clone(), + runtime_request_id: "request".to_owned(), + retry_attempt: 0, + }) + .unwrap() + .is_empty()); + let events = projector + .project(ProviderRunProjection::ModelEvent { + work_id, + event: AgentEvent::TextDelta { + text: "continued".to_owned(), + }, + }) + .unwrap(); + + assert_eq!(events.len(), 1); + let Some(response_event::Type::ClientActions(actions)) = &events[0].r#type else { + panic!("restored output should append through a client action"); + }; + assert!(matches!( + actions.actions[0].action, + Some(client_action::Action::AddMessagesToTask(_)) + )); +} + #[test] fn provider_and_session_runtimes_share_text_translation() { for mut translator in [provider_translator(), session_translator()] { @@ -63,6 +109,41 @@ fn provider_and_session_runtimes_share_text_translation() { } } +#[test] +fn retargeting_starts_new_text_and_reasoning_messages_on_the_new_task() { + let mut translator = provider_translator(); + translator + .translate(AgentEvent::TextDelta { + text: "root text".to_owned(), + }) + .expect("root text"); + translator + .translate(AgentEvent::ReasoningDelta { + text: "root reasoning".to_owned(), + }) + .expect("root reasoning"); + + translator.set_task_id("cli-task"); + for event in [ + AgentEvent::TextDelta { + text: "cli text".to_owned(), + }, + AgentEvent::ReasoningDelta { + text: "cli reasoning".to_owned(), + }, + ] { + let translated = translator.translate(event).expect("retargeted output"); + let Some(response_event::Type::ClientActions(actions)) = &translated[0].r#type else { + panic!("expected retargeted client action"); + }; + let Some(client_action::Action::AddMessagesToTask(add)) = &actions.actions[0].action else { + panic!("retargeted output must start a new message"); + }; + assert_eq!(add.task_id, "cli-task"); + assert_eq!(add.messages[0].task_id, "cli-task"); + } +} + #[test] fn reasoning_uses_the_native_reasoning_message_contract() { let mut translator = provider_translator(); diff --git a/app/src/ai/runtime/mod.rs b/app/src/ai/runtime/mod.rs index 7ddf3487..58c89634 100644 --- a/app/src/ai/runtime/mod.rs +++ b/app/src/ai/runtime/mod.rs @@ -1,9 +1,16 @@ mod event_translator; -mod provider; +mod provider_run_coordinator; mod rig; mod rig_request; mod rig_tool; -pub(crate) use event_translator::{RuntimeResponseConfig, RuntimeResponseTranslator}; -pub(crate) use provider::ProviderRuntime; -pub(crate) use rig::{rig_bedrock_response_stream, rig_openai_response_stream}; +pub(crate) use event_translator::{ + ProviderRunResponseProjector, RuntimeResponseConfig, RuntimeResponseTranslator, +}; +pub(crate) use provider_run_coordinator::{ + ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderToolExecutionRef, + ProviderToolLifecycleOutcome, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE, +}; +pub(crate) use rig::{ + prepare_provider_run, provider_runtime_for_request, PreparedProviderRun, ProviderActionContext, +}; diff --git a/app/src/ai/runtime/provider.rs b/app/src/ai/runtime/provider.rs deleted file mode 100644 index 77e569a6..00000000 --- a/app/src/ai/runtime/provider.rs +++ /dev/null @@ -1,27 +0,0 @@ -use futures::channel::oneshot; - -use crate::ai::agent::api::{self, ConvertToAPITypeError}; -use crate::ai::provider::ProviderConfig; - -/// Application-facing provider runtime dispatcher. -/// -/// OpenAI-compatible models can opt into the provider-neutral Rig runtime; -/// other models continue through their current translators while migration is -/// in progress. Both paths preserve the existing UI response stream contract. -pub(crate) struct ProviderRuntime { - provider_config: ProviderConfig, -} - -impl ProviderRuntime { - pub(crate) fn new(provider_config: ProviderConfig) -> Self { - Self { provider_config } - } - - pub(crate) async fn start_turn( - self, - params: api::RequestParams, - cancellation_rx: oneshot::Receiver<()>, - ) -> Result { - api::generate_multi_agent_output(self.provider_config, params, cancellation_rx).await - } -} diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs new file mode 100644 index 00000000..93c60528 --- /dev/null +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -0,0 +1,743 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::error::Error; +use std::fmt; +use std::sync::Arc; + +use futures::future::BoxFuture; +use futures::StreamExt; +use galaxy_agent_core::{ + AgentError, AgentErrorKind, AgentEvent, AgentRuntime, CompletedModelTurn, ContentPart, + ExternalWorkId, ModelFailureDisposition, PendingToolBatch, ProviderModelCall, + ProviderRequestProfile, ProviderRun, ProviderRunFailureKind, ProviderRunId, ProviderRunLimits, + ProviderRunOutcome, ProviderRunPhase, ProviderRunProtocolError, ProviderRunState, + ProviderRunStep, RunEpoch, RuntimeKind, StopReason, ToolEvent, TurnControl, TurnRequest, Usage, +}; + +use crate::ai::agent::conversation::AIConversationId; + +pub(crate) const BASE_PROVIDER_PROFILE: &str = "base"; +pub(crate) const CLI_MONITOR_PROVIDER_PROFILE: &str = "cli-monitor"; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum ProviderRunProjection { + ModelTurnStarted { + work_id: ExternalWorkId, + runtime_request_id: String, + retry_attempt: u32, + }, + ModelEvent { + work_id: ExternalWorkId, + event: AgentEvent, + }, + ModelRetry { + work_id: ExternalWorkId, + retry_attempt: u32, + error: AgentError, + }, + ToolBatchReady { + batch: PendingToolBatch, + }, +} + +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum ProviderRunBlock { + Tools(PendingToolBatch), + AwaitingDriver { + work_id: ExternalWorkId, + stop_reason: StopReason, + }, + Done(ProviderRunOutcome), +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] +pub struct ProviderToolExecutionRef { + pub(crate) conversation_id: AIConversationId, + pub(crate) run_id: ProviderRunId, + pub(crate) epoch: RunEpoch, + pub(crate) call_id: String, +} + +impl ProviderToolExecutionRef { + pub(crate) fn new( + conversation_id: AIConversationId, + work_id: &ExternalWorkId, + call_id: impl Into, + ) -> Self { + Self { + conversation_id, + run_id: work_id.run_id.clone(), + epoch: work_id.epoch, + call_id: call_id.into(), + } + } + + pub(crate) fn work_id(&self) -> ExternalWorkId { + ExternalWorkId { + run_id: self.run_id.clone(), + epoch: self.epoch, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum ProviderToolLifecycleOutcome { + Pending, + BatchCommitted, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum ProviderRunCoordinatorError { + InvalidRuntime(String), + InvalidToolLifecycle(String), + Core(ProviderRunProtocolError), +} + +impl fmt::Display for ProviderRunCoordinatorError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::InvalidRuntime(message) | Self::InvalidToolLifecycle(message) => { + f.write_str(message) + } + Self::Core(error) => error.fmt(f), + } + } +} + +impl Error for ProviderRunCoordinatorError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + Self::InvalidRuntime(_) | Self::InvalidToolLifecycle(_) => None, + Self::Core(error) => Some(error), + } + } +} + +impl From for ProviderRunCoordinatorError { + fn from(value: ProviderRunProtocolError) -> Self { + Self::Core(value) + } +} + +#[derive(Clone)] +pub(crate) struct ProviderRunProfile { + pub(crate) runtime: Arc, + pub(crate) request: TurnRequest, +} + +impl ProviderRunProfile { + pub(crate) fn new(runtime: Arc, request: TurnRequest) -> Self { + Self { runtime, request } + } +} + +pub(crate) struct ProviderRunCoordinator { + run: ProviderRun, + profiles: BTreeMap, +} + +impl ProviderRunCoordinator { + pub(crate) fn from_request( + run_id: impl Into, + runtime: Arc, + request: TurnRequest, + tool_result_archive: Vec, + limits: ProviderRunLimits, + ) -> Result { + let mut run = ProviderRun::new( + run_id, + request.messages.clone(), + BASE_PROVIDER_PROFILE, + limits, + ); + run.replace_tool_result_archive(tool_result_archive); + let mut profiles = BTreeMap::new(); + profiles.insert( + BASE_PROVIDER_PROFILE.to_string(), + ProviderRunProfile::new(runtime, request), + ); + Self::new(run, profiles) + } + + pub(crate) fn new( + run: ProviderRun, + profiles: BTreeMap, + ) -> Result { + if !profiles.contains_key(run.profile().as_str()) { + return Err(ProviderRunCoordinatorError::InvalidRuntime(format!( + "provider request profile '{}' is not configured", + run.profile().as_str() + ))); + } + for (profile, config) in &profiles { + validate_profile_runtime(profile, config.runtime.as_ref())?; + } + Ok(Self { run, profiles }) + } + + pub(crate) fn run(&self) -> &ProviderRun { + &self.run + } + + pub(crate) fn run_mut(&mut self) -> &mut ProviderRun { + &mut self.run + } + + pub(crate) fn profile_request(&self, profile: &str) -> Option<&TurnRequest> { + self.profiles.get(profile).map(|profile| &profile.request) + } + + pub(crate) fn insert_profile( + &mut self, + profile: impl Into, + runtime: Arc, + request: TurnRequest, + ) -> Result<(), ProviderRunCoordinatorError> { + let profile = profile.into(); + validate_profile_runtime(&profile, runtime.as_ref())?; + self.profiles + .insert(profile, ProviderRunProfile::new(runtime, request)); + Ok(()) + } + + pub(crate) fn apply_tool_lifecycle( + &mut self, + execution_ref: &ProviderToolExecutionRef, + event: &ToolEvent, + ) -> Result { + let event_call_id = tool_event_call_id(event)?; + if event_call_id != execution_ref.call_id { + return Err(ProviderRunCoordinatorError::InvalidToolLifecycle(format!( + "provider tool lifecycle call ID mismatch: expected '{}', received '{}'", + execution_ref.call_id, event_call_id + ))); + } + + let work_id = execution_ref.work_id(); + match event { + ToolEvent::Proposed { .. } => { + return Err(ProviderRunCoordinatorError::InvalidToolLifecycle( + "tool proposals must be committed by the model turn before action execution" + .to_string(), + )); + } + ToolEvent::PermissionRequested { request } => { + self.run + .request_tool_permission(&work_id, request.clone())?; + } + ToolEvent::PermissionResolved { + request_id, + call_id, + decision, + } => { + self.run.resolve_tool_permission( + &work_id, + call_id, + request_id, + decision.clone(), + )?; + } + ToolEvent::Started { call_id } => { + self.run.start_tool(&work_id, call_id)?; + } + ToolEvent::Completed { result } => { + self.run.complete_tool(&work_id, result.clone())?; + } + } + + let batch_is_complete = match self.run.state() { + ProviderRunState::AwaitingTools { batch } => batch.is_complete(), + ProviderRunState::ReadyToCallModel + | ProviderRunState::AwaitingModel { .. } + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingDriver { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => false, + }; + if batch_is_complete { + self.run.commit_tool_batch(&work_id)?; + Ok(ProviderToolLifecycleOutcome::BatchCommitted) + } else { + Ok(ProviderToolLifecycleOutcome::Pending) + } + } + + pub(crate) async fn drive_until_blocked( + &mut self, + control: TurnControl, + project: F, + ) -> Result + where + F: FnMut(ProviderRunProjection) -> Result<(), String>, + { + self.drive_until_blocked_with_checkpoint(control, project, |_| Box::pin(async { Ok(()) })) + .await + } + + pub(crate) async fn drive_until_blocked_with_checkpoint( + &mut self, + control: TurnControl, + mut project: F, + mut checkpoint: C, + ) -> Result + where + F: FnMut(ProviderRunProjection) -> Result<(), String>, + C: FnMut(ProviderRun) -> BoxFuture<'static, Result<(), String>>, + { + loop { + match self.run.next_step()? { + Some(ProviderRunStep::CallModel(call)) => { + if !self.checkpoint_or_fail(&mut checkpoint).await? { + continue; + } + self.drive_model_call(call, control.clone(), &mut project) + .await?; + } + Some(ProviderRunStep::DispatchTools(batch)) => { + if !self.checkpoint_or_fail(&mut checkpoint).await? { + continue; + } + if !self.project_or_fail( + ProviderRunProjection::ToolBatchReady { + batch: batch.clone(), + }, + &mut project, + )? { + continue; + } + if batch.is_complete() { + self.run.commit_tool_batch(&batch.work_id)?; + continue; + } + return Ok(ProviderRunBlock::Tools(batch)); + } + Some(ProviderRunStep::Done(outcome)) => { + if !self.checkpoint_or_fail(&mut checkpoint).await? { + continue; + } + return Ok(ProviderRunBlock::Done(outcome)); + } + None => match self.run.state() { + ProviderRunState::AwaitingDriver { + work_id, + stop_reason, + } => { + let work_id = work_id.clone(); + let stop_reason = stop_reason.clone(); + if !self.checkpoint_or_fail(&mut checkpoint).await? { + continue; + } + return Ok(ProviderRunBlock::AwaitingDriver { + work_id, + stop_reason, + }); + } + ProviderRunState::ReadyToCallModel + | ProviderRunState::AwaitingModel { .. } + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingTools { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => { + return Err(ProviderRunCoordinatorError::Core( + ProviderRunProtocolError::UnexpectedState { + expected: ProviderRunPhase::AwaitingDriver, + actual: self.run.state().phase(), + }, + )); + } + }, + } + } + } + + async fn checkpoint_or_fail( + &mut self, + checkpoint: &mut C, + ) -> Result + where + C: FnMut(ProviderRun) -> BoxFuture<'static, Result<(), String>>, + { + match checkpoint(self.run.clone()).await { + Ok(()) => Ok(true), + Err(message) => { + self.run.fail( + ProviderRunFailureKind::ExternalWork, + format!("provider run checkpoint failed: {message}"), + )?; + Ok(false) + } + } + } + + async fn drive_model_call( + &mut self, + call: ProviderModelCall, + control: TurnControl, + project: &mut F, + ) -> Result<(), ProviderRunCoordinatorError> + where + F: FnMut(ProviderRunProjection) -> Result<(), String>, + { + let Some(profile) = self.profiles.get(call.profile.as_str()).cloned() else { + self.run.fail( + ProviderRunFailureKind::ExternalWork, + format!( + "provider request profile '{}' is not configured", + call.profile.as_str() + ), + )?; + return Ok(()); + }; + let advertised_tools = profile + .request + .tools + .iter() + .map(|tool| tool.name.clone()) + .collect::>(); + let request = request_for_model_call(profile.request, &call); + let stream = match profile.runtime.start_turn(request, control).await { + Ok(stream) => stream, + Err(error) => { + self.handle_model_failure(&call.work_id, error, project)?; + return Ok(()); + } + }; + futures::pin_mut!(stream); + let mut buffer = ModelTurnBuffer::default(); + + while let Some(event) = stream.next().await { + let event = match event { + Ok(event) => event, + Err(error) => { + self.handle_model_failure(&call.work_id, error, project)?; + return Ok(()); + } + }; + match event { + AgentEvent::TurnStarted { runtime_request_id } => { + if buffer.started { + self.handle_model_failure( + &call.work_id, + protocol_error("provider emitted more than one TurnStarted event"), + project, + )?; + return Ok(()); + } + if runtime_request_id.is_empty() { + self.handle_model_failure( + &call.work_id, + protocol_error("provider emitted an empty runtime request ID"), + project, + )?; + return Ok(()); + } + buffer.started = true; + if !self.project_or_fail( + ProviderRunProjection::ModelTurnStarted { + work_id: call.work_id.clone(), + runtime_request_id, + retry_attempt: call.retry_attempt, + }, + project, + )? { + return Ok(()); + } + } + AgentEvent::TextDelta { text } => { + if !self.ensure_model_started(&call.work_id, &buffer, project)? { + return Ok(()); + } + buffer.text.push_str(&text); + if !self.project_or_fail( + ProviderRunProjection::ModelEvent { + work_id: call.work_id.clone(), + event: AgentEvent::TextDelta { text }, + }, + project, + )? { + return Ok(()); + } + } + AgentEvent::ReasoningDelta { text } => { + if !self.ensure_model_started(&call.work_id, &buffer, project)? { + return Ok(()); + } + buffer.reasoning.push_str(&text); + if !self.project_or_fail( + ProviderRunProjection::ModelEvent { + work_id: call.work_id.clone(), + event: AgentEvent::ReasoningDelta { text }, + }, + project, + )? { + return Ok(()); + } + } + AgentEvent::ReasoningCompleted { text, signature } => { + if !self.ensure_model_started(&call.work_id, &buffer, project)? { + return Ok(()); + } + if !text.is_empty() { + buffer.reasoning.clone_from(&text); + } + buffer.reasoning_signature.clone_from(&signature); + if !self.project_or_fail( + ProviderRunProjection::ModelEvent { + work_id: call.work_id.clone(), + event: AgentEvent::ReasoningCompleted { text, signature }, + }, + project, + )? { + return Ok(()); + } + } + AgentEvent::Tool { + event: ToolEvent::Proposed { call: tool_call }, + } => { + if !self.ensure_model_started(&call.work_id, &buffer, project)? { + return Ok(()); + } + buffer.tool_calls.push(tool_call); + } + AgentEvent::UsageUpdated { usage } => { + if !self.ensure_model_started(&call.work_id, &buffer, project)? { + return Ok(()); + } + buffer.usage.clone_from(&usage); + let cumulative_usage = combined_usage(self.run.usage(), &usage); + if !self.project_or_fail( + ProviderRunProjection::ModelEvent { + work_id: call.work_id.clone(), + event: AgentEvent::UsageUpdated { + usage: cumulative_usage, + }, + }, + project, + )? { + return Ok(()); + } + } + AgentEvent::TurnStopped { reason } => { + if !self.ensure_model_started(&call.work_id, &buffer, project)? { + return Ok(()); + } + if reason == StopReason::Cancelled { + self.run.cancel("provider model call was cancelled")?; + 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( + ProviderRunFailureKind::Protocol, + format!("provider returned an invalid completed turn: {error}"), + )?; + } + return Ok(()); + } + AgentEvent::Tool { + event: + ToolEvent::PermissionRequested { .. } + | ToolEvent::PermissionResolved { .. } + | ToolEvent::Started { .. } + | ToolEvent::Completed { .. }, + } + | AgentEvent::RuntimeActivityUpdated { .. } + | AgentEvent::ContextUsageUpdated { .. } + | AgentEvent::UserInputAccepted { .. } + | AgentEvent::RuntimeNotice { .. } => { + self.handle_model_failure( + &call.work_id, + protocol_error( + "direct-provider transport emitted a non-model lifecycle event", + ), + project, + )?; + return Ok(()); + } + } + } + + self.handle_model_failure( + &call.work_id, + protocol_error("provider stream ended before TurnStopped"), + project, + )?; + Ok(()) + } + + fn ensure_model_started( + &mut self, + work_id: &ExternalWorkId, + buffer: &ModelTurnBuffer, + project: &mut F, + ) -> Result + where + F: FnMut(ProviderRunProjection) -> Result<(), String>, + { + if buffer.started { + return Ok(true); + } + self.handle_model_failure( + work_id, + protocol_error("provider emitted model output before TurnStarted"), + project, + )?; + Ok(false) + } + + fn handle_model_failure( + &mut self, + work_id: &ExternalWorkId, + error: AgentError, + project: &mut F, + ) -> Result<(), ProviderRunCoordinatorError> + where + F: FnMut(ProviderRunProjection) -> Result<(), String>, + { + let disposition = self.run.register_model_failure(work_id, error.clone())?; + if disposition == ModelFailureDisposition::RetryScheduled { + let retry_attempt = match self.run.state() { + ProviderRunState::AwaitingModel { call } => call.retry_attempt, + ProviderRunState::ReadyToCallModel + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingTools { .. } + | ProviderRunState::AwaitingDriver { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => { + return Err(ProviderRunCoordinatorError::Core( + ProviderRunProtocolError::UnexpectedState { + expected: ProviderRunPhase::AwaitingModel, + actual: self.run.state().phase(), + }, + )); + } + }; + self.project_or_fail( + ProviderRunProjection::ModelRetry { + work_id: work_id.clone(), + retry_attempt, + error, + }, + project, + )?; + } + Ok(()) + } + + fn project_or_fail( + &mut self, + event: ProviderRunProjection, + project: &mut F, + ) -> Result + where + F: FnMut(ProviderRunProjection) -> Result<(), String>, + { + match project(event) { + Ok(()) => Ok(true), + Err(message) => { + self.run.fail( + ProviderRunFailureKind::Projection, + format!("provider run projection failed: {message}"), + )?; + Ok(false) + } + } + } +} + +#[derive(Default)] +struct ModelTurnBuffer { + started: bool, + text: String, + reasoning: String, + reasoning_signature: Option, + tool_calls: Vec, + usage: Usage, +} + +impl ModelTurnBuffer { + fn complete( + self, + stop_reason: StopReason, + advertised_tools: BTreeSet, + ) -> CompletedModelTurn { + let mut assistant_content = Vec::new(); + if !self.reasoning.is_empty() || self.reasoning_signature.is_some() { + assistant_content.push(ContentPart::Reasoning { + text: self.reasoning, + signature: self.reasoning_signature, + }); + } + if !self.text.is_empty() { + assistant_content.push(ContentPart::Text(self.text)); + } + CompletedModelTurn { + assistant_content, + tool_calls: self.tool_calls, + usage: self.usage, + stop_reason, + advertised_tools, + } + } +} + +fn validate_profile_runtime( + profile: &str, + runtime: &dyn AgentRuntime, +) -> Result<(), ProviderRunCoordinatorError> { + let descriptor = runtime.descriptor(); + if descriptor.kind != RuntimeKind::Provider { + return Err(ProviderRunCoordinatorError::InvalidRuntime(format!( + "runtime '{}' for provider request profile '{profile}' is not a direct-provider transport", + descriptor.id + ))); + } + if !descriptor.capabilities.host_managed_history || !descriptor.capabilities.host_tool_execution + { + return Err(ProviderRunCoordinatorError::InvalidRuntime(format!( + "runtime '{}' for provider request profile '{profile}' does not expose Galaxy-owned history and tools", + descriptor.id + ))); + } + Ok(()) +} + +fn request_for_model_call(mut template: TurnRequest, call: &ProviderModelCall) -> TurnRequest { + template.prompt = None; + template.messages = call.messages.clone(); + template +} + +fn combined_usage(previous: &Usage, current: &Usage) -> Usage { + Usage { + input_tokens: previous.input_tokens.saturating_add(current.input_tokens), + output_tokens: previous.output_tokens.saturating_add(current.output_tokens), + cached_input_tokens: previous + .cached_input_tokens + .saturating_add(current.cached_input_tokens), + cache_creation_input_tokens: previous + .cache_creation_input_tokens + .saturating_add(current.cache_creation_input_tokens), + } +} + +fn tool_event_call_id(event: &ToolEvent) -> Result<&str, ProviderRunCoordinatorError> { + match event { + ToolEvent::Proposed { call } => Ok(&call.id), + ToolEvent::PermissionRequested { request } => Ok(&request.call_id), + ToolEvent::PermissionResolved { call_id, .. } | ToolEvent::Started { call_id } => { + Ok(call_id) + } + ToolEvent::Completed { result } => Ok(&result.call_id), + } +} + +fn protocol_error(message: impl Into) -> AgentError { + AgentError::new(AgentErrorKind::Protocol, message) +} + +#[cfg(test)] +#[path = "provider_run_coordinator_tests.rs"] +mod tests; diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs new file mode 100644 index 00000000..091dd33c --- /dev/null +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -0,0 +1,963 @@ +use std::collections::VecDeque; +use std::sync::Mutex; + +use async_trait::async_trait; +use galaxy_agent_core::{ + turn_control, AgentEventStream, ConversationMessage, MessageContent, MessageRole, + PermissionDecision, PermissionKind, PermissionRequest, ProviderRunFailure, RuntimeCapabilities, + RuntimeDescriptor, ToolCall, ToolDefinition, ToolResult, ToolResultStatus, +}; +use warp_multi_agent_api::response_event; + +use super::*; +use crate::ai::runtime::event_translator::{ProviderRunResponseProjector, RuntimeResponseConfig}; + +type ScriptedEvent = Result; +type ScriptedTurn = Result, AgentError>; + +struct ScriptedRuntime { + descriptor: RuntimeDescriptor, + turns: Mutex>, + requests: Mutex>, +} + +impl ScriptedRuntime { + fn new(turns: Vec) -> Self { + Self::with_id("scripted", turns) + } + + fn with_id(id: &str, turns: Vec) -> Self { + Self { + descriptor: RuntimeDescriptor { + id: id.to_string(), + display_name: "Scripted provider".to_string(), + kind: RuntimeKind::Provider, + capabilities: RuntimeCapabilities::provider(), + }, + turns: Mutex::new(turns.into()), + requests: Mutex::new(Vec::new()), + } + } + + fn requests(&self) -> Vec { + self.requests.lock().unwrap().clone() + } +} + +#[async_trait] +impl AgentRuntime for ScriptedRuntime { + fn descriptor(&self) -> &RuntimeDescriptor { + &self.descriptor + } + + async fn start_turn( + &self, + request: TurnRequest, + _control: TurnControl, + ) -> Result { + self.requests.lock().unwrap().push(request); + let turn = self + .turns + .lock() + .unwrap() + .pop_front() + .expect("scripted runtime ran out of turns")?; + Ok(Box::pin(futures::stream::iter(turn))) + } +} + +fn request() -> TurnRequest { + let mut request = TurnRequest::new( + "test-model", + vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("Inspect and fix the issue".to_string()), + }], + ); + request.conversation_id = Some("conversation".to_string()); + request.system_prompt = Some("Use tools and finish the task.".to_string()); + request.tools = vec![ + ToolDefinition { + name: "read_files".to_string(), + description: "Read files".to_string(), + input_schema: serde_json::json!({"type": "object"}), + }, + ToolDefinition { + name: galaxy_agent_core::RECALL_TOOL_HISTORY_NAME.to_string(), + description: "Recall tool results".to_string(), + input_schema: serde_json::json!({"type": "object"}), + }, + ]; + request +} + +fn started(id: &str) -> ScriptedEvent { + Ok(AgentEvent::TurnStarted { + runtime_request_id: id.to_string(), + }) +} + +fn usage(input_tokens: u64, output_tokens: u64) -> ScriptedEvent { + Ok(AgentEvent::UsageUpdated { + usage: Usage { + input_tokens, + output_tokens, + ..Usage::default() + }, + }) +} + +fn stopped(reason: StopReason) -> ScriptedEvent { + Ok(AgentEvent::TurnStopped { reason }) +} + +fn tool_call(id: &str, name: &str) -> ToolCall { + ToolCall { + id: id.to_string(), + name: name.to_string(), + arguments: serde_json::json!({"path": "Cargo.toml"}), + } +} + +fn tool_turn() -> ScriptedTurn { + Ok(vec![ + started("request-tools"), + Ok(AgentEvent::TextDelta { + text: "I will inspect it.".to_string(), + }), + Ok(AgentEvent::Tool { + event: ToolEvent::Proposed { + call: tool_call("read-1", "read_files"), + }, + }), + usage(10, 4), + stopped(StopReason::Completed), + ]) +} + +fn tool_turn_with_calls(calls: Vec) -> ScriptedTurn { + let mut events = vec![started("request-tools")]; + events.extend(calls.into_iter().map(|call| { + Ok(AgentEvent::Tool { + event: ToolEvent::Proposed { call }, + }) + })); + events.push(usage(10, 4)); + events.push(stopped(StopReason::Completed)); + Ok(events) +} + +fn answer_turn() -> ScriptedTurn { + Ok(vec![ + started("request-answer"), + Ok(AgentEvent::TextDelta { + text: "The issue is fixed.".to_string(), + }), + usage(20, 5), + stopped(StopReason::Completed), + ]) +} + +fn coordinator(runtime: Arc) -> ProviderRunCoordinator { + ProviderRunCoordinator::from_request( + "run-1", + runtime, + request(), + Vec::new(), + ProviderRunLimits::default(), + ) + .unwrap() +} + +async fn coordinator_awaiting_tools( + calls: Vec, +) -> (ProviderRunCoordinator, PendingToolBatch) { + let runtime = Arc::new(ScriptedRuntime::new(vec![tool_turn_with_calls(calls)])); + let mut coordinator = coordinator(runtime); + let (_sender, control) = turn_control(); + let block = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(); + let ProviderRunBlock::Tools(batch) = block else { + panic!("expected external tool work"); + }; + (coordinator, batch) +} + +fn tool_result(call_id: &str, content: &str, status: ToolResultStatus) -> ToolEvent { + ToolEvent::Completed { + result: ToolResult { + call_id: call_id.to_string(), + content: content.to_string(), + status, + }, + } +} + +fn collect_projection( + events: &mut Vec, +) -> impl FnMut(ProviderRunProjection) -> Result<(), String> + '_ { + |event| { + events.push(event); + Ok(()) + } +} + +#[tokio::test] +async fn one_run_drives_model_tool_and_followup_turns_with_atomic_history() { + let runtime = Arc::new(ScriptedRuntime::new(vec![tool_turn(), answer_turn()])); + let mut coordinator = coordinator(runtime.clone()); + let mut projections = Vec::new(); + let (_sender, control) = turn_control(); + + let first_block = coordinator + .drive_until_blocked(control, collect_projection(&mut projections)) + .await + .unwrap(); + let ProviderRunBlock::Tools(batch) = first_block else { + panic!("expected external tool work"); + }; + assert_eq!(batch.work_id.epoch.get(), 1); + assert_eq!(batch.calls.len(), 1); + assert!(projections.iter().all(|projection| !matches!( + projection, + ProviderRunProjection::ModelEvent { + event: AgentEvent::Tool { .. }, + .. + } + ))); + assert!(matches!( + projections.last(), + Some(ProviderRunProjection::ToolBatchReady { .. }) + )); + + coordinator + .run_mut() + .start_tool(&batch.work_id, "read-1") + .unwrap(); + coordinator + .run_mut() + .complete_tool( + &batch.work_id, + ToolResult { + call_id: "read-1".to_string(), + content: "manifest contents".to_string(), + status: ToolResultStatus::Success, + }, + ) + .unwrap(); + coordinator + .run_mut() + .commit_tool_batch(&batch.work_id) + .unwrap(); + + projections.clear(); + let (_sender, control) = turn_control(); + let second_block = coordinator + .drive_until_blocked(control, collect_projection(&mut projections)) + .await + .unwrap(); + let ProviderRunBlock::AwaitingDriver { + work_id, + stop_reason, + } = second_block + else { + panic!("expected driver decision"); + }; + assert_eq!(stop_reason, StopReason::Completed); + assert_eq!(work_id.epoch.get(), 3); + assert_eq!(coordinator.run().usage().input_tokens, 30); + assert_eq!(coordinator.run().usage().output_tokens, 9); + + let requests = runtime.requests(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].messages.len(), 1); + assert_eq!(requests[1].messages.len(), 3); + let MessageContent::MultiPart(results) = &requests[1].messages[2].content else { + panic!("expected atomic tool result message"); + }; + assert!(matches!( + results.as_slice(), + [ContentPart::ToolResult { + tool_use_id, + content, + is_error: false, + }] if tool_use_id == "read-1" && content == "manifest contents" + )); + assert!(projections.iter().any(|projection| matches!( + projection, + ProviderRunProjection::ModelEvent { + event: AgentEvent::UsageUpdated { usage }, + .. + } if usage.input_tokens == 30 && usage.output_tokens == 9 + ))); + + coordinator.run_mut().complete(&work_id).unwrap(); + let (_sender, control) = turn_control(); + assert_eq!( + coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(), + ProviderRunBlock::Done(ProviderRunOutcome::Completed( + galaxy_agent_core::ProviderRunCompletion { + stop_reason: StopReason::Completed, + } + )) + ); +} + +#[tokio::test] +async fn request_profiles_route_through_their_own_runtime_and_template() { + let base_runtime = Arc::new(ScriptedRuntime::with_id( + "base-runtime", + vec![answer_turn()], + )); + let cli_runtime = Arc::new(ScriptedRuntime::with_id("cli-runtime", vec![answer_turn()])); + let mut base_request = request(); + base_request.model = "base-model".into(); + base_request.system_prompt = Some("base prompt".to_string()); + let mut cli_request = request(); + cli_request.model = "cli-model".into(); + cli_request.system_prompt = Some("cli prompt".to_string()); + let run = ProviderRun::new( + "run-1", + base_request.messages.clone(), + BASE_PROVIDER_PROFILE, + ProviderRunLimits::default(), + ); + let mut profiles = BTreeMap::new(); + profiles.insert( + BASE_PROVIDER_PROFILE.to_string(), + ProviderRunProfile::new(base_runtime.clone(), base_request), + ); + profiles.insert( + CLI_MONITOR_PROVIDER_PROFILE.to_string(), + ProviderRunProfile::new(cli_runtime.clone(), cli_request), + ); + let mut coordinator = ProviderRunCoordinator::new(run, profiles).unwrap(); + + let (_sender, control) = turn_control(); + let ProviderRunBlock::AwaitingDriver { work_id, .. } = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap() + else { + panic!("expected base driver decision"); + }; + coordinator + .run_mut() + .continue_with_observation( + &work_id, + MessageContent::Text("command is still running".to_string()), + CLI_MONITOR_PROVIDER_PROFILE, + ) + .unwrap(); + let (_sender, control) = turn_control(); + assert!(matches!( + coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(), + ProviderRunBlock::AwaitingDriver { .. } + )); + + let base_requests = base_runtime.requests(); + let cli_requests = cli_runtime.requests(); + assert_eq!(base_requests.len(), 1); + assert_eq!(base_requests[0].model.as_str(), "base-model"); + assert_eq!( + base_requests[0].system_prompt.as_deref(), + Some("base prompt") + ); + assert_eq!(cli_requests.len(), 1); + assert_eq!(cli_requests[0].model.as_str(), "cli-model"); + assert_eq!(cli_requests[0].system_prompt.as_deref(), Some("cli prompt")); + assert!(matches!( + cli_requests[0].messages.last(), + Some(ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text(text), + }) if text == "command is still running" + )); +} + +#[tokio::test] +async fn correlated_parallel_lifecycle_commits_results_in_original_call_order() { + let (mut coordinator, batch) = coordinator_awaiting_tools(vec![ + tool_call("first", "read_files"), + tool_call("second", "read_files"), + ]) + .await; + let conversation_id = AIConversationId::new(); + let first = ProviderToolExecutionRef::new(conversation_id, &batch.work_id, "first"); + let second = ProviderToolExecutionRef::new(conversation_id, &batch.work_id, "second"); + + assert_eq!( + coordinator + .apply_tool_lifecycle( + &first, + &ToolEvent::Started { + call_id: "first".to_string(), + }, + ) + .unwrap(), + ProviderToolLifecycleOutcome::Pending + ); + coordinator + .apply_tool_lifecycle( + &second, + &ToolEvent::Started { + call_id: "second".to_string(), + }, + ) + .unwrap(); + assert_eq!( + coordinator + .apply_tool_lifecycle( + &second, + &tool_result("second", "second result", ToolResultStatus::Success), + ) + .unwrap(), + ProviderToolLifecycleOutcome::Pending + ); + assert_eq!(coordinator.run().transcript().len(), 2); + assert_eq!( + coordinator + .apply_tool_lifecycle( + &first, + &tool_result("first", "first result", ToolResultStatus::Success), + ) + .unwrap(), + ProviderToolLifecycleOutcome::BatchCommitted + ); + + let MessageContent::MultiPart(parts) = &coordinator.run().transcript().last().unwrap().content + else { + panic!("expected atomic tool result message"); + }; + let ids = parts + .iter() + .map(|part| match part { + ContentPart::ToolResult { tool_use_id, .. } => tool_use_id.as_str(), + ContentPart::Text(_) + | ContentPart::Reasoning { .. } + | ContentPart::Image { .. } + | ContentPart::ToolUse { .. } => panic!("expected only tool results"), + }) + .collect::>(); + assert_eq!(ids, vec!["first", "second"]); +} + +#[tokio::test] +async fn correlated_lifecycle_rejects_wrong_run_epoch_call_and_duplicate_without_mutation() { + let (mut coordinator, batch) = coordinator_awaiting_tools(vec![ + tool_call("first", "read_files"), + tool_call("second", "read_files"), + ]) + .await; + let conversation_id = AIConversationId::new(); + let first = ProviderToolExecutionRef::new(conversation_id, &batch.work_id, "first"); + let second = ProviderToolExecutionRef::new(conversation_id, &batch.work_id, "second"); + let initial = coordinator.run().clone(); + + let mut wrong_run = first.clone(); + wrong_run.run_id = ProviderRunId::new("wrong-run"); + assert!(matches!( + coordinator + .apply_tool_lifecycle( + &wrong_run, + &ToolEvent::Started { + call_id: "first".to_string(), + }, + ) + .unwrap_err(), + ProviderRunCoordinatorError::Core(ProviderRunProtocolError::WorkMismatch { .. }) + )); + assert_eq!(coordinator.run(), &initial); + + let mut stale = first.clone(); + stale.epoch = RunEpoch::new(stale.epoch.get() + 1); + assert!(matches!( + coordinator + .apply_tool_lifecycle( + &stale, + &ToolEvent::Started { + call_id: "first".to_string(), + }, + ) + .unwrap_err(), + ProviderRunCoordinatorError::Core(ProviderRunProtocolError::WorkMismatch { .. }) + )); + assert_eq!(coordinator.run(), &initial); + + assert!(matches!( + coordinator + .apply_tool_lifecycle( + &first, + &ToolEvent::Started { + call_id: "other".to_string(), + }, + ) + .unwrap_err(), + ProviderRunCoordinatorError::InvalidToolLifecycle(_) + )); + assert_eq!(coordinator.run(), &initial); + + coordinator + .apply_tool_lifecycle( + &second, + &tool_result("second", "done", ToolResultStatus::Success), + ) + .unwrap(); + let after_completion = coordinator.run().clone(); + assert!(matches!( + coordinator + .apply_tool_lifecycle( + &second, + &tool_result("second", "duplicate", ToolResultStatus::Success), + ) + .unwrap_err(), + ProviderRunCoordinatorError::Core(ProviderRunProtocolError::DuplicateToolUpdate { .. }) + )); + assert_eq!(coordinator.run(), &after_completion); +} + +#[tokio::test] +async fn permission_denial_is_the_only_correlated_terminal_result() { + let (mut coordinator, batch) = + coordinator_awaiting_tools(vec![tool_call("shell", "read_files")]).await; + let execution_ref = + ProviderToolExecutionRef::new(AIConversationId::new(), &batch.work_id, "shell"); + let request = PermissionRequest { + id: "permission-shell".to_string(), + call_id: "shell".to_string(), + kind: PermissionKind::Execute, + reason: Some("run a command".to_string()), + }; + + assert_eq!( + coordinator + .apply_tool_lifecycle( + &execution_ref, + &ToolEvent::PermissionRequested { + request: request.clone(), + }, + ) + .unwrap(), + ProviderToolLifecycleOutcome::Pending + ); + assert_eq!( + coordinator + .apply_tool_lifecycle( + &execution_ref, + &ToolEvent::PermissionResolved { + request_id: request.id, + call_id: request.call_id, + decision: PermissionDecision::Denied { + reason: Some("not allowed".to_string()), + }, + }, + ) + .unwrap(), + ProviderToolLifecycleOutcome::BatchCommitted + ); + let committed = coordinator.run().clone(); + assert!(matches!( + &committed.transcript().last().unwrap().content, + MessageContent::MultiPart(parts) + if matches!(parts.as_slice(), [ContentPart::ToolResult { + tool_use_id, + content, + is_error: true, + }] if tool_use_id == "shell" && content == "not allowed") + )); + + assert!(coordinator + .apply_tool_lifecycle( + &execution_ref, + &tool_result("shell", "duplicate denial", ToolResultStatus::Denied), + ) + .is_err()); + assert_eq!(coordinator.run(), &committed); +} + +#[tokio::test] +async fn execution_failure_commits_one_correlated_error_result() { + let (mut coordinator, batch) = + coordinator_awaiting_tools(vec![tool_call("read", "read_files")]).await; + let execution_ref = + ProviderToolExecutionRef::new(AIConversationId::new(), &batch.work_id, "read"); + + coordinator + .apply_tool_lifecycle( + &execution_ref, + &ToolEvent::Started { + call_id: "read".to_string(), + }, + ) + .unwrap(); + assert_eq!( + coordinator + .apply_tool_lifecycle( + &execution_ref, + &tool_result("read", "file missing", ToolResultStatus::Error), + ) + .unwrap(), + ProviderToolLifecycleOutcome::BatchCommitted + ); + assert!(matches!( + &coordinator.run().transcript().last().unwrap().content, + MessageContent::MultiPart(parts) + if matches!(parts.as_slice(), [ContentPart::ToolResult { + tool_use_id, + content, + is_error: true, + }] if tool_use_id == "read" && content == "file missing") + )); +} + +#[tokio::test] +async fn inline_tool_batches_continue_without_leaving_the_coordinator() { + let recall_turn = Ok(vec![ + started("request-recall"), + Ok(AgentEvent::Tool { + event: ToolEvent::Proposed { + call: ToolCall { + id: "recall-1".to_string(), + name: galaxy_agent_core::RECALL_TOOL_HISTORY_NAME.to_string(), + arguments: serde_json::json!({"search_query": "missing"}), + }, + }, + }), + usage(10, 1), + stopped(StopReason::Completed), + ]); + let runtime = Arc::new(ScriptedRuntime::new(vec![recall_turn, answer_turn()])); + let mut coordinator = coordinator(runtime.clone()); + let mut projections = Vec::new(); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked(control, collect_projection(&mut projections)) + .await + .unwrap(); + + assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. })); + assert_eq!(runtime.requests().len(), 2); + assert!(projections.iter().any(|projection| matches!( + projection, + ProviderRunProjection::ToolBatchReady { batch } if batch.is_complete() + ))); + let requests = runtime.requests(); + assert_eq!(requests[1].messages.len(), 3); + let MessageContent::MultiPart(results) = &requests[1].messages[2].content else { + panic!("expected inline result batch"); + }; + assert!(matches!( + results.as_slice(), + [ContentPart::ToolResult { + tool_use_id, + is_error: false, + .. + }] if tool_use_id == "recall-1" + )); +} + +#[tokio::test] +async fn recoverable_start_failure_retries_the_same_work_identity() { + let mut recoverable = AgentError::new(AgentErrorKind::Transport, "temporary network error"); + recoverable.recoverable = true; + let runtime = Arc::new(ScriptedRuntime::new(vec![Err(recoverable), answer_turn()])); + let mut coordinator = coordinator(runtime.clone()); + let mut projections = Vec::new(); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked(control, collect_projection(&mut projections)) + .await + .unwrap(); + + assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. })); + assert_eq!(runtime.requests().len(), 2); + assert_eq!(coordinator.run().model_retries(), 1); + let retry = projections + .iter() + .find_map(|projection| match projection { + ProviderRunProjection::ModelRetry { + work_id, + retry_attempt, + .. + } => Some((work_id.clone(), *retry_attempt)), + ProviderRunProjection::ModelTurnStarted { .. } + | ProviderRunProjection::ModelEvent { .. } + | ProviderRunProjection::ToolBatchReady { .. } => None, + }) + .expect("retry projection"); + let started = projections + .iter() + .find_map(|projection| match projection { + ProviderRunProjection::ModelTurnStarted { + work_id, + retry_attempt: 1, + .. + } => Some(work_id.clone()), + ProviderRunProjection::ModelTurnStarted { .. } + | ProviderRunProjection::ModelEvent { .. } + | ProviderRunProjection::ModelRetry { .. } + | ProviderRunProjection::ToolBatchReady { .. } => None, + }) + .expect("retried model start"); + assert_eq!(retry.0, started); + assert_eq!(retry.1, 1); +} + +#[tokio::test] +async fn tool_projection_failure_synthesizes_a_result_and_fails_the_run() { + let runtime = Arc::new(ScriptedRuntime::new(vec![tool_turn()])); + let mut coordinator = coordinator(runtime); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked(control, |event| match event { + ProviderRunProjection::ToolBatchReady { .. } => { + Err("task projection disappeared".to_string()) + } + ProviderRunProjection::ModelTurnStarted { .. } + | ProviderRunProjection::ModelEvent { .. } + | ProviderRunProjection::ModelRetry { .. } => Ok(()), + }) + .await + .unwrap(); + + let ProviderRunBlock::Done(ProviderRunOutcome::Failed(failure)) = block else { + panic!("projection failure must terminate the run visibly"); + }; + assert_eq!(failure.kind, ProviderRunFailureKind::Projection); + let MessageContent::MultiPart(results) = + &coordinator.run().transcript().last().unwrap().content + else { + panic!("pending tool must receive a synthetic error result"); + }; + assert!(matches!( + results.as_slice(), + [ContentPart::ToolResult { + tool_use_id, + is_error: true, + .. + }] if tool_use_id == "read-1" + )); +} + +#[tokio::test] +async fn stream_without_terminal_event_fails_instead_of_committing_partial_output() { + let runtime = Arc::new(ScriptedRuntime::new(vec![Ok(vec![ + started("request-eof"), + Ok(AgentEvent::TextDelta { + text: "partial".to_string(), + }), + ])])); + let mut coordinator = coordinator(runtime); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(); + + let ProviderRunBlock::Done(ProviderRunOutcome::Failed(failure)) = block else { + panic!("unexpected EOF must fail the run"); + }; + assert_eq!(failure.kind, ProviderRunFailureKind::ModelCall); + assert_eq!(coordinator.run().transcript().len(), 1); +} + +#[tokio::test] +async fn provider_cancellation_does_not_commit_partial_assistant_content() { + let expected_transcript = request().messages; + let runtime = Arc::new(ScriptedRuntime::new(vec![Ok(vec![ + started("request-cancelled"), + Ok(AgentEvent::TextDelta { + text: "partial".to_string(), + }), + stopped(StopReason::Cancelled), + ])])); + let mut coordinator = coordinator(runtime); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(); + + assert_eq!( + block, + ProviderRunBlock::Done(ProviderRunOutcome::Cancelled { + reason: "provider model call was cancelled".to_string(), + }) + ); + assert_eq!(coordinator.run().transcript(), expected_transcript); +} + +#[tokio::test] +async fn transcript_projector_emits_one_ui_stream_for_the_whole_run() { + let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); + let mut coordinator = coordinator(runtime); + let mut projector = ProviderRunResponseProjector::new(RuntimeResponseConfig { + task_id: "task".to_string(), + conversation_id: "conversation".to_string(), + needs_create_task: true, + user_query: Some("Inspect and fix the issue".to_string()), + model_id: "test-model".to_string(), + max_context_tokens: Some(100_000), + capabilities: RuntimeCapabilities::provider(), + empty_output_message: None, + }); + let mut ui_events = Vec::new(); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked(control, |projection| { + ui_events.extend(projector.project(projection)?); + Ok(()) + }) + .await + .unwrap(); + let ProviderRunBlock::AwaitingDriver { work_id, .. } = block else { + panic!("expected driver decision"); + }; + + assert_eq!( + count_response_events(&ui_events, ResponseEventKind::Init), + 1 + ); + assert_eq!( + count_response_events(&ui_events, ResponseEventKind::Finished), + 0 + ); + assert!(ui_events + .iter() + .any(|event| matches!(event.r#type, Some(response_event::Type::ClientActions(_))))); + + coordinator.run_mut().complete(&work_id).unwrap(); + let (_sender, control) = turn_control(); + let ProviderRunBlock::Done(outcome) = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap() + else { + panic!("expected terminal run"); + }; + ui_events.extend(projector.finish(&outcome).unwrap()); + assert_eq!( + count_response_events(&ui_events, ResponseEventKind::Finished), + 1 + ); +} + +#[test] +fn transcript_projector_preserves_provider_failure_message() { + let mut projector = ProviderRunResponseProjector::new(RuntimeResponseConfig { + task_id: "task".to_string(), + conversation_id: "conversation".to_string(), + needs_create_task: false, + user_query: None, + model_id: "test-model".to_string(), + max_context_tokens: Some(100_000), + capabilities: RuntimeCapabilities::provider(), + empty_output_message: None, + }); + let events = projector + .finish(&ProviderRunOutcome::Failed(ProviderRunFailure { + kind: ProviderRunFailureKind::ModelCall, + message: "upstream provider rejected the request".to_string(), + source: None, + })) + .unwrap(); + + let finished = events + .iter() + .find_map(|event| match &event.r#type { + Some(response_event::Type::Finished(finished)) => Some(finished), + Some(response_event::Type::Init(_)) + | Some(response_event::Type::ClientActions(_)) + | None => None, + }) + .expect("terminal provider failure"); + assert!(matches!( + &finished.reason, + Some(response_event::stream_finished::Reason::InternalError(error)) + if error.message == "upstream provider rejected the request" + )); +} + +#[derive(Clone, Copy)] +enum ResponseEventKind { + Init, + Finished, +} + +fn count_response_events( + events: &[warp_multi_agent_api::ResponseEvent], + kind: ResponseEventKind, +) -> usize { + events + .iter() + .filter(|event| match (&event.r#type, kind) { + (Some(response_event::Type::Init(_)), ResponseEventKind::Init) + | (Some(response_event::Type::Finished(_)), ResponseEventKind::Finished) => true, + (Some(response_event::Type::ClientActions(_)), ResponseEventKind::Init) + | (Some(response_event::Type::ClientActions(_)), ResponseEventKind::Finished) + | (Some(response_event::Type::Init(_)), ResponseEventKind::Finished) + | (Some(response_event::Type::Finished(_)), ResponseEventKind::Init) + | (None, ResponseEventKind::Init) + | (None, ResponseEventKind::Finished) => false, + }) + .count() +} + +#[test] +fn session_runtime_is_rejected_before_any_turn_can_start() { + struct SessionRuntime { + descriptor: RuntimeDescriptor, + } + + #[async_trait] + impl AgentRuntime for SessionRuntime { + fn descriptor(&self) -> &RuntimeDescriptor { + &self.descriptor + } + + async fn start_turn( + &self, + _request: TurnRequest, + _control: TurnControl, + ) -> Result { + unreachable!() + } + } + + let runtime = Arc::new(SessionRuntime { + descriptor: RuntimeDescriptor { + id: "session".to_string(), + display_name: "Session runtime".to_string(), + kind: RuntimeKind::Acp, + capabilities: RuntimeCapabilities::session_runtime(), + }, + }); + let error = ProviderRunCoordinator::from_request( + "run-1", + runtime, + request(), + Vec::new(), + ProviderRunLimits::default(), + ) + .err() + .expect("session runtime must be rejected"); + + assert!(matches!( + error, + ProviderRunCoordinatorError::InvalidRuntime(_) + )); +} diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index 5bed5df4..b78b334d 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -1,12 +1,7 @@ use std::collections::HashMap; use std::sync::Arc; -use futures::channel::oneshot; -use futures::{FutureExt, StreamExt}; -use galaxy_agent_core::{ - turn_control, AgentError, AgentEvent, AgentRuntime, MessageContent, MessageRole, StopReason, - ToolCall, ToolCallDecision, ToolEvent, ToolPolicy, ToolResult, TurnCommand, Usage, -}; +use galaxy_agent_core::{AgentRuntime, ToolCall, TurnRequest}; use galaxy_agent_rig::{ AnthropicRuntime, AnthropicRuntimeConfig, ChatGPTSubscriptionRuntime, ChatGPTSubscriptionRuntimeConfig, GeminiRuntime, GeminiRuntimeConfig, OpenAICompatibleRuntime, @@ -16,597 +11,242 @@ use uuid::Uuid; use warp_multi_agent_api::ToolType; use super::rig_request::{ - prepare_bedrock_rig_turn, prepare_rig_turn, MCPToolTarget, PreparedRigTurn, + prepare_bedrock_rig_turn_for_mode, prepare_rig_turn, prepare_rig_turn_for_mode, MCPToolTarget, + PreparedRigTurn, RigRequestMode, }; use super::rig_tool::action_from_tool_call; -use crate::ai::agent::api::{Event, RequestParams, ResponseStream, StreamEvent}; +use super::ProviderRunProfile; +use crate::ai::agent::api::RequestParams; use crate::ai::agent::AIAgentAction; -use crate::ai::bedrock::client::{BedrockClient, BedrockClientConfig}; +use crate::ai::bedrock::client::BedrockClient; +use crate::ai::bedrock::convert::CachingConfig; use crate::ai::bedrock::external_config::ExternalBedrockConfig; -use crate::ai::bedrock::response_translator::build_add_agent_output_message; -use crate::ai::openai::client::OpenAIClientConfig; -use crate::ai::provider::types::{ContentPart, ConversationMessage}; -use crate::ai::runtime::{RuntimeResponseConfig, RuntimeResponseTranslator}; -use crate::server::server_api::AIApiError; +use crate::ai::provider::types::ConversationMessage; +use crate::ai::runtime::RuntimeResponseConfig; use crate::settings::OpenAIProviderKind; -const MAX_INLINE_TOOL_CONTINUATIONS: usize = 3; -const INLINE_TOOL_LOOP_MESSAGE: &str = - "I couldn't continue because the model repeatedly searched prior tool history without making progress. Please retry with a more specific instruction."; +pub(crate) struct PreparedProviderRun { + pub(crate) base_profile: ProviderRunProfile, + pub(crate) cli_monitor_profile: Option, + pub(crate) tool_result_archive: Vec, + pub(crate) messages_sent: Arc>>, + pub(crate) persistence_offset: usize, + pub(crate) response_config: RuntimeResponseConfig, + pub(crate) action_context: ProviderActionContext, +} -pub(crate) fn rig_openai_response_stream( - config: OpenAIClientConfig, - params: RequestParams, - supported_tools: Vec, - supported_cli_agent_tools: Vec, - cancellation_rx: oneshot::Receiver<()>, -) -> ResponseStream { - let skill_path_origin = params.session_context.skill_path_origin(); - let prepared = prepare_rig_turn(&config, params, supported_tools, supported_cli_agent_tools); - let model_id = prepared.request.model.as_str().to_string(); - match config.kind { - OpenAIProviderKind::OpenAI | OpenAIProviderKind::LiteLLM => { - let runtime = OpenAICompatibleRuntime::new(OpenAICompatibleRuntimeConfig { - base_url: config.base_url, - api_key: config.api_key, - model: model_id.clone(), - max_output_tokens: config.max_output_tokens.map(u64::from), - supports_system_messages: config.supports_system_messages, - }); - rig_response_stream( - runtime, - prepared, - skill_path_origin, - config.max_input_tokens, - "rig_openai_compatible", - cancellation_rx, - ) - } - OpenAIProviderKind::ChatGPTSubscription => { - let runtime = ChatGPTSubscriptionRuntime::new(ChatGPTSubscriptionRuntimeConfig { - model: model_id, - reasoning_effort: config.reasoning_effort, - max_output_tokens: config.max_output_tokens.map(u64::from), - auth_file: None, - }); - rig_response_stream( - runtime, - prepared, - skill_path_origin, - config.max_input_tokens, - "rig_chatgpt_subscription", - cancellation_rx, - ) - } - OpenAIProviderKind::Anthropic => { - let runtime = AnthropicRuntime::new(AnthropicRuntimeConfig { - api_key: config.api_key.unwrap_or_default(), - model: model_id, - max_output_tokens: config.max_output_tokens.map(u64::from), - }); - rig_response_stream( - runtime, - prepared, - skill_path_origin, - config.max_input_tokens, - "rig_anthropic", - cancellation_rx, - ) - } - OpenAIProviderKind::Gemini => { - let runtime = GeminiRuntime::new(GeminiRuntimeConfig { - api_key: config.api_key.unwrap_or_default(), - model: model_id, - max_output_tokens: config.max_output_tokens.map(u64::from), - }); - rig_response_stream( - runtime, - prepared, - skill_path_origin, - config.max_input_tokens, - "rig_gemini", - cancellation_rx, - ) - } - OpenAIProviderKind::VertexAI => { - let runtime = VertexAiRuntime::new(VertexAiRuntimeConfig { - project_id: config.project_id.unwrap_or_default(), - location: config.location.unwrap_or_else(|| "global".to_string()), - model: model_id, - max_output_tokens: config.max_output_tokens.map(u64::from), - }); - rig_response_stream( - runtime, - prepared, - skill_path_origin, - config.max_input_tokens, - "rig_vertex_ai", - cancellation_rx, - ) +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub(crate) struct ProviderActionContext { + task_id: String, + skill_path_origin: ai::skills::SkillPathOrigin, + mcp_tool_aliases: HashMap, +} + +impl ProviderActionContext { + pub(crate) fn task_id(&self) -> &str { + &self.task_id + } + + pub(crate) fn set_task_id(&mut self, task_id: impl Into) { + self.task_id = task_id.into(); + } + + #[cfg(test)] + pub(crate) fn new_for_test(task_id: impl Into) -> Self { + Self { + task_id: task_id.into(), + skill_path_origin: ai::skills::SkillPathOrigin::Local, + mcp_tool_aliases: HashMap::new(), } } + + pub(crate) fn action_from_tool_call(&self, call: &ToolCall) -> Result { + action_from_tool_call( + &self.task_id, + call, + &self.skill_path_origin, + &self.mcp_tool_aliases, + ) + } } -pub(crate) async fn rig_bedrock_response_stream( - config: BedrockClientConfig, - params: RequestParams, - supported_tools: Vec, - supported_cli_agent_tools: Vec, - cancellation_rx: oneshot::Receiver<()>, -) -> anyhow::Result { +pub(crate) async fn prepare_provider_run( + base_provider_config: crate::ai::provider::ProviderConfig, + cli_provider_config: crate::ai::provider::ProviderConfig, + mut params: RequestParams, +) -> anyhow::Result { + let (supported_tools, supported_cli_agent_tools) = + crate::ai::agent::api::prepare_direct_provider_params(&mut params); let skill_path_origin = params.session_context.skill_path_origin(); let max_context_tokens = params.context_window_limit; - let model = params.model.as_str().to_string(); - let max_output_tokens = Some(64_000); - let cross_region_inference = config.cross_region_inference; - let external_config = ExternalBedrockConfig::load(); - let prompt_caching = !external_config.disable_prompt_caching; - let client = BedrockClient::from_config(config).await?; - let runtime = client.rig_runtime( - model.clone(), - cross_region_inference, - prompt_caching, - max_output_tokens, - )?; - let prepared = prepare_bedrock_rig_turn( - model, - max_output_tokens, + let mut cli_params = params.clone(); + cli_params.model = params.cli_agent_model.clone(); + + let (base_runtime, prepared) = prepare_provider_profile( + base_provider_config, params, - supported_tools, - supported_cli_agent_tools, - ); + supported_tools.clone(), + supported_cli_agent_tools.clone(), + None, + ) + .await?; + let cli_monitor_profile = match cli_provider_config { + crate::ai::provider::ProviderConfig::None => None, + provider_config => { + let (runtime, prepared) = prepare_provider_profile( + provider_config, + cli_params, + supported_tools, + supported_cli_agent_tools, + Some(RigRequestMode::Cli), + ) + .await?; + Some(ProviderRunProfile::new(runtime, prepared.request)) + } + }; - Ok(rig_response_stream( - runtime, - prepared, - skill_path_origin, - max_context_tokens, - "rig_bedrock", - cancellation_rx, - )) -} - -fn rig_response_stream( - runtime: R, - prepared: PreparedRigTurn, - skill_path_origin: ai::skills::SkillPathOrigin, - max_context_tokens: Option, - stream_type: &'static str, - cancellation_rx: oneshot::Receiver<()>, -) -> ResponseStream -where - R: AgentRuntime + Send + Sync + 'static, -{ - let runtime_capabilities = runtime.descriptor().capabilities.clone(); let PreparedRigTurn { task_id, needs_create_task, user_query, - request: turn_request, + request, persistent_messages, tool_result_archive, messages_sent, mcp_tool_aliases, } = prepared; - store_messages_sent(&messages_sent, &persistent_messages); - - let conversation_id = turn_request.conversation_id.clone(); - let model_id = turn_request.model.as_str().to_string(); - let tool_policy = ToolPolicy::new(&turn_request.tools); - let stream = async_stream::stream! { - let cancel_future = cancellation_rx.fuse(); - futures::pin_mut!(cancel_future); - - let conversation_id = conversation_id.unwrap_or_else(|| Uuid::new_v4().to_string()); - let mut translator = RuntimeResponseTranslator::new(RuntimeResponseConfig { - task_id: task_id.clone(), - conversation_id, - needs_create_task, - user_query, - model_id, - max_context_tokens, - capabilities: runtime_capabilities, - empty_output_message: None, - }); - let mut turn_request = turn_request; - let mut cumulative_usage = Usage::default(); - let mut inline_continuation_count = 0; - - 'provider_turns: loop { - let (control_sender, control) = turn_control(); - let start_future = runtime.start_turn(turn_request.clone(), control).fuse(); - futures::pin_mut!(start_future); - - let mut agent_events = futures::select_biased! { - _ = cancel_future => { - let _ = control_sender.try_send(TurnCommand::Cancel); - match start_future.await { - Ok(stream) => stream, - Err(error) => { - yield Err(agent_error(error, stream_type)); - return; - } - } - } - result = start_future => match result { - Ok(stream) => stream, - Err(error) => { - yield Err(agent_error(error, stream_type)); - return; - } - }, - }; - - let mut full_text = String::new(); - let mut full_reasoning = String::new(); - let mut reasoning_signature = None; - let mut proposed_tools = Vec::new(); - let mut assistant_history_index = None; - let mut handled_inline_tool = false; - let mut proposed_client_tool = false; - - loop { - let next_event = agent_events.next().fuse(); - futures::pin_mut!(next_event); - futures::select_biased! { - _ = cancel_future => { - let _ = control_sender.try_send(TurnCommand::Cancel); - } - event = next_event => { - let Some(event) = event else { - yield Err(Arc::new(AIApiError::UnexpectedEof)); - return; - }; - let event = match event { - Ok(event) => event, - Err(error) => { - yield Err(agent_error(error, stream_type)); - return; - } - }; - - match event { - AgentEvent::Tool { - event: ToolEvent::Proposed { call }, - } => { - proposed_tools.push(call.clone()); - sync_assistant_turn( - &messages_sent, - &full_reasoning, - reasoning_signature.as_deref(), - &full_text, - &proposed_tools, - &mut assistant_history_index, - ); - let history = messages_sent - .lock() - .map(|sent| sent.clone()) - .unwrap_or_default(); - match tool_policy.decide(&call, &history, &tool_result_archive) { - ToolCallDecision::Execute => { - proposed_client_tool = true; - match build_tool_proposed( - &task_id, - &call, - &skill_path_origin, - &mcp_tool_aliases, - ) { - Ok(action) => yield Ok(StreamEvent::ToolProposed(action)), - Err(message) => { - yield Err(agent_error(AgentError::new( - galaxy_agent_core::AgentErrorKind::Protocol, - message, - ), stream_type)); - return; - } - } - } - ToolCallDecision::Inline(result) => { - handled_inline_tool = true; - append_tool_result(&messages_sent, result); - } - ToolCallDecision::Reject(result) => { - log::warn!( - "Rig model called unavailable tool '{}' (id={})", - call.name, - call.id - ); - let error_display = format!( - "Failed tool call: `{}`\n\n{}", - call.name, result.content - ); - append_tool_result(&messages_sent, result); - let message_id = Uuid::new_v4().to_string(); - yield Ok(StreamEvent::Response(build_add_agent_output_message( - &task_id, - &message_id, - &error_display, - ))); - } - } - } - AgentEvent::UsageUpdated { usage } => { - accumulate_usage(&mut cumulative_usage, &usage); - let response_events = match translator.translate( - AgentEvent::UsageUpdated { - usage: cumulative_usage.clone(), - }, - ) { - Ok(response_events) => response_events, - Err(message) => { - yield Err(agent_error(AgentError::new( - galaxy_agent_core::AgentErrorKind::Protocol, - message, - ), stream_type)); - return; - } - }; - for response_event in response_events { - yield Ok(StreamEvent::Response(response_event)); - } - } - AgentEvent::TurnStopped { mut reason } => { - sync_assistant_turn( - &messages_sent, - &full_reasoning, - reasoning_signature.as_deref(), - &full_text, - &proposed_tools, - &mut assistant_history_index, - ); - if reason == StopReason::Completed - && handled_inline_tool - && !proposed_client_tool - { - if inline_continuation_count < MAX_INLINE_TOOL_CONTINUATIONS { - inline_continuation_count += 1; - turn_request.messages = match copy_messages(&messages_sent) { - Ok(messages) => messages, - Err(()) => { - yield Err(agent_error(AgentError::new( - galaxy_agent_core::AgentErrorKind::Protocol, - "could not access Rig conversation history for inline tool continuation", - ), stream_type)); - return; - } - }; - translator.begin_followup_turn(); - log::info!( - "Continuing Rig provider turn after inline tool result ({inline_continuation_count}/{MAX_INLINE_TOOL_CONTINUATIONS})" - ); - continue 'provider_turns; - } - - log::warn!( - "Rig provider exceeded {MAX_INLINE_TOOL_CONTINUATIONS} inline tool continuations" - ); - append_assistant_text(&messages_sent, INLINE_TOOL_LOOP_MESSAGE); - let response_events = match translator.translate( - AgentEvent::RuntimeNotice { - message: INLINE_TOOL_LOOP_MESSAGE.to_string(), - }, - ) { - Ok(response_events) => response_events, - Err(message) => { - yield Err(agent_error(AgentError::new( - galaxy_agent_core::AgentErrorKind::Protocol, - message, - ), stream_type)); - return; - } - }; - for response_event in response_events { - yield Ok(StreamEvent::Response(response_event)); - } - reason = StopReason::ToolLoopLimit; - } - let response_events = match translator - .translate(AgentEvent::TurnStopped { reason }) - { - Ok(response_events) => response_events, - Err(message) => { - yield Err(agent_error(AgentError::new( - galaxy_agent_core::AgentErrorKind::Protocol, - message, - ), stream_type)); - return; - } - }; - for response_event in response_events { - yield Ok(StreamEvent::Response(response_event)); - } - return; - } - event => { - match &event { - AgentEvent::TextDelta { text } => full_text.push_str(text), - AgentEvent::ReasoningDelta { text } => { - full_reasoning.push_str(text); - } - AgentEvent::ReasoningCompleted { text, signature } => { - if !text.is_empty() { - full_reasoning.clone_from(text); - } - reasoning_signature.clone_from(signature); - } - AgentEvent::TurnStarted { .. } - | AgentEvent::Tool { .. } - | AgentEvent::UsageUpdated { .. } - | AgentEvent::RuntimeActivityUpdated { .. } - | AgentEvent::ContextUsageUpdated { .. } - | AgentEvent::UserInputAccepted { .. } - | AgentEvent::RuntimeNotice { .. } - | AgentEvent::TurnStopped { .. } => {} - } - let response_events = match translator.translate(event) { - Ok(response_events) => response_events, - Err(message) => { - yield Err(agent_error(AgentError::new( - galaxy_agent_core::AgentErrorKind::Protocol, - message, - ), stream_type)); - return; - } - }; - for response_event in response_events { - yield Ok(StreamEvent::Response(response_event)); - } - } - } - } - } - } - } + let persistence_offset = request + .messages + .len() + .saturating_sub(persistent_messages.len()); + let response_config = RuntimeResponseConfig { + task_id: task_id.clone(), + conversation_id: request + .conversation_id + .clone() + .unwrap_or_else(|| Uuid::new_v4().to_string()), + needs_create_task, + user_query, + model_id: request.model.as_str().to_string(), + max_context_tokens, + capabilities: base_runtime.descriptor().capabilities.clone(), + empty_output_message: None, }; - - Box::pin(stream) -} - -fn store_messages_sent( - messages_sent: &std::sync::Arc>>, - messages: &[ConversationMessage], -) { - let Ok(mut sent) = messages_sent.lock() else { - return; - }; - *sent = messages.to_vec(); -} - -fn copy_messages( - messages_sent: &std::sync::Arc>>, -) -> Result, ()> { - messages_sent - .lock() - .map(|sent| sent.clone()) - .map_err(|_| ()) -} - -fn append_tool_result( - messages_sent: &std::sync::Arc>>, - result: ToolResult, -) { - let is_error = result.is_error(); - let message = ConversationMessage { - role: MessageRole::User, - content: MessageContent::ToolResult { - tool_use_id: result.call_id, - content: result.content, - is_error, + Ok(PreparedProviderRun { + base_profile: ProviderRunProfile::new(base_runtime, request), + cli_monitor_profile, + tool_result_archive, + messages_sent, + persistence_offset, + response_config, + action_context: ProviderActionContext { + task_id, + skill_path_origin, + mcp_tool_aliases, }, - }; - if let Ok(mut sent) = messages_sent.lock() { - sent.push(message); - } + }) } -fn append_assistant_text( - messages_sent: &std::sync::Arc>>, - text: &str, -) { - if let Ok(mut sent) = messages_sent.lock() { - sent.push(ConversationMessage { - role: MessageRole::Assistant, - content: MessageContent::Text(text.to_string()), - }); - } -} - -fn accumulate_usage(total: &mut Usage, usage: &Usage) { - total.input_tokens = total.input_tokens.saturating_add(usage.input_tokens); - total.output_tokens = total.output_tokens.saturating_add(usage.output_tokens); - total.cached_input_tokens = total - .cached_input_tokens - .saturating_add(usage.cached_input_tokens); - total.cache_creation_input_tokens = total - .cache_creation_input_tokens - .saturating_add(usage.cache_creation_input_tokens); -} - -fn sync_assistant_turn( - messages_sent: &std::sync::Arc>>, - reasoning_text: &str, - reasoning_signature: Option<&str>, - text: &str, - tool_calls: &[ToolCall], - history_index: &mut Option, -) { - let has_reasoning = !reasoning_text.is_empty() || reasoning_signature.is_some(); - let mut parts = Vec::with_capacity( - usize::from(has_reasoning) + usize::from(!text.is_empty()) + tool_calls.len(), - ); - if has_reasoning { - parts.push(ContentPart::Reasoning { - text: reasoning_text.to_string(), - signature: reasoning_signature.map(str::to_string), - }); - } - if !text.is_empty() { - parts.push(ContentPart::Text(text.to_string())); - } - parts.extend(tool_calls.iter().map(|call| ContentPart::ToolUse { - tool_use_id: call.id.clone(), - name: call.name.clone(), - input: call.arguments.clone(), - })); - if parts.is_empty() { - return; - } - - let content = if parts.len() == 1 { - match parts.pop().unwrap() { - ContentPart::Text(text) => MessageContent::Text(text), - ContentPart::ToolUse { - tool_use_id, - name, - input, - } => MessageContent::ToolUse { - tool_use_id, - name, - input, - }, - reasoning @ ContentPart::Reasoning { .. } => MessageContent::MultiPart(vec![reasoning]), - ContentPart::Image { .. } | ContentPart::ToolResult { .. } => unreachable!(), +async fn prepare_provider_profile( + provider_config: crate::ai::provider::ProviderConfig, + params: RequestParams, + supported_tools: Vec, + supported_cli_agent_tools: Vec, + mode: Option, +) -> anyhow::Result<(Arc, PreparedRigTurn)> { + let model = params.model.as_str().to_string(); + let prepared = match &provider_config { + crate::ai::provider::ProviderConfig::OpenAI(config) => match mode { + Some(mode) => prepare_rig_turn_for_mode( + config, + params, + supported_tools, + supported_cli_agent_tools, + mode, + ), + None => prepare_rig_turn(config, params, supported_tools, supported_cli_agent_tools), + }, + crate::ai::provider::ProviderConfig::Bedrock(_) => prepare_bedrock_rig_turn_for_mode( + model, + Some(64_000), + params, + supported_tools, + supported_cli_agent_tools, + mode, + ), + crate::ai::provider::ProviderConfig::None => { + anyhow::bail!( + "No AI runtime configured. Enable an agent runtime or model provider in settings." + ); } - } else { - MessageContent::MultiPart(parts) - }; - let message = ConversationMessage { - role: MessageRole::Assistant, - content, }; + let runtime = provider_runtime_for_request(provider_config, &prepared.request).await?; + Ok((runtime, prepared)) +} - let Ok(mut sent) = messages_sent.lock() else { - return; - }; - if let Some(index) = *history_index { - if index < sent.len() { - sent[index] = message; - return; +/// Rebuilds a one-turn provider transport from current settings and a persisted request. +/// Credentials remain in the live provider config and never enter the run snapshot. +pub(crate) async fn provider_runtime_for_request( + provider_config: crate::ai::provider::ProviderConfig, + request: &TurnRequest, +) -> anyhow::Result> { + let model = request.model.as_str().to_string(); + let runtime: Arc = match provider_config { + crate::ai::provider::ProviderConfig::OpenAI(config) => match config.kind { + OpenAIProviderKind::OpenAI | OpenAIProviderKind::LiteLLM => Arc::new( + OpenAICompatibleRuntime::new(OpenAICompatibleRuntimeConfig { + base_url: config.base_url, + api_key: config.api_key, + model, + max_output_tokens: config.max_output_tokens.map(u64::from), + supports_system_messages: config.supports_system_messages, + }), + ), + OpenAIProviderKind::ChatGPTSubscription => Arc::new(ChatGPTSubscriptionRuntime::new( + ChatGPTSubscriptionRuntimeConfig { + model, + reasoning_effort: config.reasoning_effort, + max_output_tokens: config.max_output_tokens.map(u64::from), + auth_file: None, + }, + )), + OpenAIProviderKind::Anthropic => { + Arc::new(AnthropicRuntime::new(AnthropicRuntimeConfig { + api_key: config.api_key.unwrap_or_default(), + model, + max_output_tokens: config.max_output_tokens.map(u64::from), + })) + } + OpenAIProviderKind::Gemini => Arc::new(GeminiRuntime::new(GeminiRuntimeConfig { + api_key: config.api_key.unwrap_or_default(), + model, + max_output_tokens: config.max_output_tokens.map(u64::from), + })), + OpenAIProviderKind::VertexAI => Arc::new(VertexAiRuntime::new(VertexAiRuntimeConfig { + project_id: config.project_id.unwrap_or_default(), + location: config.location.unwrap_or_else(|| "global".to_string()), + model, + max_output_tokens: config.max_output_tokens.map(u64::from), + })), + }, + crate::ai::provider::ProviderConfig::Bedrock(config) => { + let max_output_tokens = Some(64_000); + let cross_region_inference = config.cross_region_inference; + let caching_config = + CachingConfig::from_external_config(&ExternalBedrockConfig::load()); + let client = BedrockClient::from_config(config).await?; + Arc::new(client.agent_runtime( + model, + cross_region_inference, + max_output_tokens, + caching_config, + )?) } - } - *history_index = Some(sent.len()); - sent.push(message); -} - -fn build_tool_proposed( - task_id: &str, - call: &ToolCall, - skill_path_origin: &ai::skills::SkillPathOrigin, - mcp_tool_aliases: &HashMap, -) -> Result { - action_from_tool_call(task_id, call, skill_path_origin, mcp_tool_aliases) -} - -fn agent_error(error: AgentError, stream_type: &'static str) -> Arc { - Arc::new( - AIApiError::Stream { - stream_type, - source: anyhow::anyhow!(error), + crate::ai::provider::ProviderConfig::None => { + anyhow::bail!( + "No AI runtime configured. Enable an agent runtime or model provider in settings." + ); } - .into_quota_limit_if_provider_budget_exhausted(), - ) + }; + Ok(runtime) } - -#[cfg(test)] -#[path = "rig_tests.rs"] -mod tests; diff --git a/app/src/ai/runtime/rig_request.rs b/app/src/ai/runtime/rig_request.rs index f32bffae..a44f11a6 100644 --- a/app/src/ai/runtime/rig_request.rs +++ b/app/src/ai/runtime/rig_request.rs @@ -33,8 +33,8 @@ pub(crate) struct PreparedRigTurn { pub mcp_tool_aliases: HashMap, } -#[derive(Clone, Debug, Eq, PartialEq)] -pub(super) struct MCPToolTarget { +#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] +pub(crate) struct MCPToolTarget { pub server_id: Option, pub name: String, } @@ -52,6 +52,25 @@ pub(crate) fn prepare_rig_turn( params, supported_tools, supported_cli_agent_tools, + None, + ) +} + +pub(crate) fn prepare_rig_turn_for_mode( + config: &OpenAIClientConfig, + params: RequestParams, + supported_tools: Vec, + supported_cli_agent_tools: Vec, + mode: RigRequestMode, +) -> PreparedRigTurn { + prepare_rig_turn_for_provider( + config.model.clone(), + config.max_output_tokens.map(u64::from), + RigRequestSanitizer::OpenAICompatible, + params, + supported_tools, + supported_cli_agent_tools, + Some(mode), ) } @@ -61,6 +80,24 @@ pub(crate) fn prepare_bedrock_rig_turn( params: RequestParams, supported_tools: Vec, supported_cli_agent_tools: Vec, +) -> PreparedRigTurn { + prepare_bedrock_rig_turn_for_mode( + model, + max_output_tokens, + params, + supported_tools, + supported_cli_agent_tools, + None, + ) +} + +pub(crate) fn prepare_bedrock_rig_turn_for_mode( + model: String, + max_output_tokens: Option, + params: RequestParams, + supported_tools: Vec, + supported_cli_agent_tools: Vec, + mode: Option, ) -> PreparedRigTurn { prepare_rig_turn_for_provider( Some(model), @@ -69,6 +106,7 @@ pub(crate) fn prepare_bedrock_rig_turn( params, supported_tools, supported_cli_agent_tools, + mode, ) } @@ -85,6 +123,7 @@ fn prepare_rig_turn_for_provider( params: RequestParams, supported_tools: Vec, supported_cli_agent_tools: Vec, + mode_override: Option, ) -> PreparedRigTurn { let RequestParams { input, @@ -107,7 +146,7 @@ fn prepare_rig_turn_for_provider( .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); let needs_create_task = tasks.is_empty(); let user_query = input.iter().find_map(input_user_query); - let mode = request_mode(&input); + let mode = mode_override.unwrap_or_else(|| request_mode(&input)); let available_tools = match mode { RigRequestMode::Cli => supported_cli_agent_tools, RigRequestMode::CompletedCommandAssessment => Vec::new(), @@ -119,11 +158,9 @@ fn prepare_rig_turn_for_provider( tool_definitions(&available_tools, mcp_context.as_ref()); match mode { RigRequestMode::Cli => { - // History recall cannot advance a running command and is handled inline by the Rig - // adapter (without producing a client action that can trigger another turn). Keeping it - // in the CLI tool list lets the model spend its entire monitor turn recalling the prior - // snapshot instead of scheduling `read_shell_command_output`, so make polling the only - // way to inspect the active command here. + // History recall cannot advance a running command. Keeping it in the CLI tool list lets + // the model spend its monitor turn recalling a prior snapshot instead of scheduling + // `read_shell_command_output`, so make polling the only inspection path here. tools.retain(|tool| tool.name != "recall_tool_history"); } RigRequestMode::CompletedCommandAssessment => { @@ -435,7 +472,7 @@ fn input_user_query(input: &AIAgentInput) -> Option { } #[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum RigRequestMode { +pub(crate) enum RigRequestMode { Normal, Plan, Orchestrate, diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index c673d8d4..8a725a1f 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -7,7 +7,10 @@ use galaxy_agent_core::{ContentPart, MessageContent, MessageRole, ToolResult, To use galaxy_util::local_or_remote_path::LocalOrRemotePath; use warp_multi_agent_api::ToolType; -use super::{input_messages, prepare_bedrock_rig_turn, prepare_rig_turn, tool_definitions}; +use super::{ + input_messages, prepare_bedrock_rig_turn, prepare_rig_turn, prepare_rig_turn_for_mode, + tool_definitions, RigRequestMode, +}; use crate::ai::agent::api::RequestParams; use crate::ai::agent::task::TaskId; use crate::ai::agent::{ @@ -344,6 +347,38 @@ fn completed_command_assessment_uses_root_history_without_tools_or_monitor_instr assert_eq!(prepared.request.messages, prepared.persistent_messages); } +#[test] +fn forced_cli_profile_uses_monitor_prompt_and_tools_for_an_initial_query() { + let mut params = RequestParams::new_for_test(); + params.input = vec![user_query("Run the checks")]; + + let prepared = prepare_rig_turn_for_mode( + &config(), + params, + vec![ToolType::RunShellCommand], + vec![ToolType::ReadShellCommandOutput], + RigRequestMode::Cli, + ); + let prompt = prepared.request.system_prompt.expect("system prompt"); + + assert!(prompt.contains("## Running Command Monitor")); + assert!(prepared + .request + .tools + .iter() + .any(|tool| tool.name == "read_shell_command_output")); + assert!(!prepared + .request + .tools + .iter() + .any(|tool| tool.name == "run_shell_command")); + assert!(!prepared + .request + .tools + .iter() + .any(|tool| tool.name == "recall_tool_history")); +} + #[test] fn lrc_snapshot_follow_up_uses_the_cli_monitor_prompt_and_tools() { let block_id: galaxy_terminal::model::BlockId = "precmd-lrc-test".to_string().into(); diff --git a/app/src/ai/runtime/rig_tests.rs b/app/src/ai/runtime/rig_tests.rs deleted file mode 100644 index a1a22ebe..00000000 --- a/app/src/ai/runtime/rig_tests.rs +++ /dev/null @@ -1,468 +0,0 @@ -use std::collections::{HashMap, VecDeque}; -use std::sync::{Arc, Mutex}; - -use ai::skills::SkillPathOrigin; -use async_trait::async_trait; -use futures::channel::oneshot; -use futures::StreamExt; -use galaxy_agent_core::{ - AgentError, AgentErrorKind, AgentEvent, AgentEventStream, AgentRuntime, ContentPart, - ConversationMessage, MessageContent, MessageRole, RuntimeCapabilities, RuntimeDescriptor, - RuntimeKind, StopReason, ToolCall, ToolDefinition, ToolEvent, ToolResult, ToolResultStatus, - TurnControl, TurnRequest, Usage, RECALL_TOOL_HISTORY_NAME, -}; -use warp_multi_agent_api::{client_action, message, response_event}; - -use super::{ - append_tool_result, build_tool_proposed, rig_response_stream, sync_assistant_turn, - PreparedRigTurn, INLINE_TOOL_LOOP_MESSAGE, MAX_INLINE_TOOL_CONTINUATIONS, -}; -use crate::ai::agent::api::StreamEvent; - -#[test] -fn tool_proposal_matches_the_domain_permission_contract() { - let action = build_tool_proposed( - "task", - &ToolCall { - id: "call-1".to_string(), - name: "run_shell_command".to_string(), - arguments: serde_json::json!({ - "command": "cargo test", - "is_read_only": true - }), - }, - &SkillPathOrigin::Local, - &HashMap::new(), - ) - .unwrap(); - - assert_eq!(action.id.to_string(), "call-1"); - assert!(matches!( - action.action, - crate::ai::agent::AIAgentActionType::RequestCommandOutput { - command, - is_read_only: Some(true), - .. - } if command == "cargo test" - )); -} - -#[test] -fn mcp_tool_proposal_routes_directly_to_the_mcp_executor_contract() { - let action = build_tool_proposed( - "task", - &ToolCall { - id: "call-mcp".to_string(), - name: "mcp__11111111-1111-4111-8111-111111111111__read_file".to_string(), - arguments: serde_json::json!({"path": "Cargo.toml"}), - }, - &SkillPathOrigin::Local, - &HashMap::new(), - ) - .unwrap(); - - assert!(matches!( - action.action, - crate::ai::agent::AIAgentActionType::CallMCPTool { - server_id: Some(server_id), - name, - .. - } if server_id.to_string() == "11111111-1111-4111-8111-111111111111" - && name == "read_file" - )); -} - -#[test] -fn assistant_history_is_updated_before_fast_tool_execution_can_continue() { - let messages = Arc::new(Mutex::new(Vec::new())); - let mut history_index = None; - let first_call = ToolCall { - id: "call-1".to_string(), - name: "read_files".to_string(), - arguments: serde_json::json!({"files": ["Cargo.toml"]}), - }; - let second_call = ToolCall { - id: "call-2".to_string(), - name: "grep".to_string(), - arguments: serde_json::json!({"queries": ["rig"]}), - }; - - sync_assistant_turn( - &messages, - "", - None, - "I'll inspect both.", - std::slice::from_ref(&first_call), - &mut history_index, - ); - sync_assistant_turn( - &messages, - "", - None, - "I'll inspect both.", - &[first_call, second_call], - &mut history_index, - ); - - let messages = messages.lock().unwrap(); - assert_eq!(messages.len(), 1); - let MessageContent::MultiPart(parts) = &messages[0].content else { - panic!("expected combined assistant content"); - }; - assert_eq!(parts.len(), 3); - assert!( - matches!(&parts[0], galaxy_agent_core::ContentPart::Text(text) if text == "I'll inspect both.") - ); - assert!( - matches!(&parts[1], galaxy_agent_core::ContentPart::ToolUse { tool_use_id, .. } if tool_use_id == "call-1") - ); - assert!( - matches!(&parts[2], galaxy_agent_core::ContentPart::ToolUse { tool_use_id, .. } if tool_use_id == "call-2") - ); -} - -#[test] -fn signed_reasoning_is_persisted_before_the_tool_call() { - let messages = Arc::new(Mutex::new(Vec::new())); - let mut history_index = None; - let call = ToolCall { - id: "call-1".to_string(), - name: "read_files".to_string(), - arguments: serde_json::json!({"files": ["Cargo.toml"]}), - }; - - sync_assistant_turn( - &messages, - "I should inspect the manifest.", - Some("signed-reasoning"), - "", - std::slice::from_ref(&call), - &mut history_index, - ); - - let messages = messages.lock().unwrap(); - let MessageContent::MultiPart(parts) = &messages[0].content else { - panic!("expected reasoning and tool call parts"); - }; - assert!(matches!( - parts.as_slice(), - [ - ContentPart::Reasoning { - text, - signature: Some(signature), - }, - ContentPart::ToolUse { tool_use_id, .. }, - ] if text == "I should inspect the manifest." - && signature == "signed-reasoning" - && tool_use_id == "call-1" - )); -} - -#[test] -fn rejected_tool_result_is_paired_with_the_assistant_call_in_history() { - let messages = Arc::new(Mutex::new(Vec::new())); - let mut history_index = None; - let call = ToolCall { - id: "call-unknown".to_string(), - name: "invented_tool".to_string(), - arguments: serde_json::json!({}), - }; - sync_assistant_turn( - &messages, - "", - None, - "", - std::slice::from_ref(&call), - &mut history_index, - ); - append_tool_result( - &messages, - ToolResult { - call_id: call.id.clone(), - content: "tool is unavailable".to_string(), - status: ToolResultStatus::Error, - }, - ); - - let messages = messages.lock().unwrap(); - assert_eq!(messages.len(), 2); - assert_eq!(messages[0].role, MessageRole::Assistant); - assert!(matches!( - &messages[0].content, - MessageContent::ToolUse { - tool_use_id, - name, - .. - } if tool_use_id == "call-unknown" && name == "invented_tool" - )); - assert_eq!(messages[1].role, MessageRole::User); - assert!(matches!( - &messages[1].content, - MessageContent::ToolResult { - tool_use_id, - content, - is_error: true, - } if tool_use_id == "call-unknown" && content == "tool is unavailable" - )); -} - -struct ScriptedRuntime { - descriptor: RuntimeDescriptor, - turns: Mutex>>, - requests: Arc>>, -} - -impl ScriptedRuntime { - fn new(turns: Vec>, requests: Arc>>) -> Self { - Self { - descriptor: RuntimeDescriptor { - id: "scripted-provider".to_string(), - display_name: "Scripted provider".to_string(), - kind: RuntimeKind::Provider, - capabilities: RuntimeCapabilities::provider(), - }, - turns: Mutex::new(turns.into()), - requests, - } - } -} - -#[async_trait] -impl AgentRuntime for ScriptedRuntime { - fn descriptor(&self) -> &RuntimeDescriptor { - &self.descriptor - } - - async fn start_turn( - &self, - request: TurnRequest, - _control: TurnControl, - ) -> Result { - self.requests.lock().unwrap().push(request); - let events = self.turns.lock().unwrap().pop_front().ok_or_else(|| { - AgentError::new( - AgentErrorKind::Protocol, - "scripted provider ran out of turns", - ) - })?; - Ok(Box::pin(futures::stream::iter(events.into_iter().map(Ok)))) - } -} - -fn recall_turn(index: usize) -> Vec { - vec![ - AgentEvent::TurnStarted { - runtime_request_id: format!("request-{index}"), - }, - AgentEvent::Tool { - event: ToolEvent::Proposed { - call: ToolCall { - id: format!("recall-{index}"), - name: RECALL_TOOL_HISTORY_NAME.to_string(), - arguments: serde_json::json!({"search_query": "missing"}), - }, - }, - }, - AgentEvent::UsageUpdated { - usage: Usage { - input_tokens: 10, - output_tokens: 1, - ..Usage::default() - }, - }, - AgentEvent::TurnStopped { - reason: StopReason::Completed, - }, - ] -} - -fn answer_turn() -> Vec { - vec![ - AgentEvent::TurnStarted { - runtime_request_id: "request-answer".to_string(), - }, - AgentEvent::TextDelta { - text: "Continuing with the answer.".to_string(), - }, - AgentEvent::UsageUpdated { - usage: Usage { - input_tokens: 20, - output_tokens: 3, - ..Usage::default() - }, - }, - AgentEvent::TurnStopped { - reason: StopReason::Completed, - }, - ] -} - -fn prepared_turn(messages_sent: Arc>>) -> PreparedRigTurn { - let initial_messages = vec![ConversationMessage { - role: MessageRole::User, - content: MessageContent::Text("Inspect the issue.".to_string()), - }]; - let mut request = TurnRequest::new("test-model", initial_messages.clone()); - request.conversation_id = Some("conversation".to_string()); - request.tools = vec![ToolDefinition { - name: RECALL_TOOL_HISTORY_NAME.to_string(), - description: "Recall prior tool output".to_string(), - input_schema: serde_json::json!({"type": "object"}), - }]; - PreparedRigTurn { - task_id: "task".to_string(), - needs_create_task: false, - user_query: None, - request, - persistent_messages: initial_messages, - tool_result_archive: Vec::new(), - messages_sent, - mcp_tool_aliases: HashMap::new(), - } -} - -async fn run_scripted_turn( - turns: Vec>, -) -> (Vec, Vec, Vec) { - let requests = Arc::new(Mutex::new(Vec::new())); - let messages_sent = Arc::new(Mutex::new(Vec::new())); - let runtime = ScriptedRuntime::new(turns, requests.clone()); - let (cancel_tx, cancellation_rx) = oneshot::channel(); - let events = rig_response_stream( - runtime, - prepared_turn(messages_sent.clone()), - SkillPathOrigin::Local, - Some(100_000), - "scripted", - cancellation_rx, - ) - .collect::>() - .await - .into_iter() - .collect::, _>>() - .expect("scripted response should succeed"); - drop(cancel_tx); - let requests = requests.lock().unwrap().clone(); - let messages_sent = messages_sent.lock().unwrap().clone(); - (events, requests, messages_sent) -} - -fn agent_output_texts(events: &[StreamEvent]) -> Vec<&str> { - let mut texts = Vec::new(); - for event in events { - let StreamEvent::Response(response) = event else { - continue; - }; - let Some(response_event::Type::ClientActions(actions)) = &response.r#type else { - continue; - }; - for action in &actions.actions { - let Some(client_action::Action::AddMessagesToTask(add)) = &action.action else { - continue; - }; - for message in &add.messages { - if let Some(message::Message::AgentOutput(output)) = &message.message { - texts.push(output.text.as_str()); - } - } - } - } - texts -} - -#[tokio::test] -async fn inline_recall_starts_a_followup_provider_turn_with_the_paired_result() { - let (events, requests, messages_sent) = - run_scripted_turn(vec![recall_turn(1), answer_turn()]).await; - - assert_eq!(requests.len(), 2); - assert_eq!(requests[1].messages.len(), 3); - assert!(matches!( - &requests[1].messages[1].content, - MessageContent::ToolUse { - tool_use_id, - name, - .. - } if tool_use_id == "recall-1" && name == RECALL_TOOL_HISTORY_NAME - )); - assert!(matches!( - &requests[1].messages[2].content, - MessageContent::ToolResult { - tool_use_id, - content, - is_error: false, - } if tool_use_id == "recall-1" - && content == "No matching tool calls found in conversation history." - )); - assert_eq!( - events - .iter() - .filter(|event| matches!( - event, - StreamEvent::Response(response) - if matches!(response.r#type, Some(response_event::Type::Init(_))) - )) - .count(), - 1 - ); - assert_eq!( - events - .iter() - .filter(|event| matches!( - event, - StreamEvent::Response(response) - if matches!(response.r#type, Some(response_event::Type::Finished(_))) - )) - .count(), - 1 - ); - assert!(events - .iter() - .all(|event| !matches!(event, StreamEvent::ToolProposed(_)))); - assert_eq!(agent_output_texts(&events), ["Continuing with the answer."]); - - let finished = events.iter().find_map(|event| { - let StreamEvent::Response(response) = event else { - return None; - }; - let Some(response_event::Type::Finished(finished)) = &response.r#type else { - return None; - }; - Some(finished) - }); - let finished = finished.expect("stream should finish"); - assert_eq!(finished.token_usage[0].total_input, 30); - assert_eq!(finished.token_usage[0].output, 4); - assert!(matches!( - messages_sent.last().map(|message| &message.content), - Some(MessageContent::Text(text)) if text == "Continuing with the answer." - )); -} - -#[tokio::test] -async fn repeated_inline_recall_stops_with_a_visible_loop_limit_message() { - let turns = (0..=MAX_INLINE_TOOL_CONTINUATIONS) - .map(recall_turn) - .collect(); - let (events, requests, messages_sent) = run_scripted_turn(turns).await; - - assert_eq!(requests.len(), MAX_INLINE_TOOL_CONTINUATIONS + 1); - assert!(agent_output_texts(&events).contains(&INLINE_TOOL_LOOP_MESSAGE)); - assert!(matches!( - messages_sent.last().map(|message| &message.content), - Some(MessageContent::Text(text)) if text == INLINE_TOOL_LOOP_MESSAGE - )); - let finished = events.iter().find_map(|event| { - let StreamEvent::Response(response) = event else { - return None; - }; - let Some(response_event::Type::Finished(finished)) = &response.r#type else { - return None; - }; - Some(finished) - }); - assert!(matches!( - finished.and_then(|finished| finished.reason.as_ref()), - Some(response_event::stream_finished::Reason::Other(_)) - )); -} diff --git a/app/src/ai/runtime/rig_tool.rs b/app/src/ai/runtime/rig_tool.rs index 73ce56cc..37c1a31b 100644 --- a/app/src/ai/runtime/rig_tool.rs +++ b/app/src/ai/runtime/rig_tool.rs @@ -18,7 +18,7 @@ use crate::ai::agent::{ }; use crate::ai::document::ai_document_model::AIDocumentId; -pub(super) fn action_from_tool_call( +pub(crate) fn action_from_tool_call( task_id: &str, call: &ToolCall, skill_path_origin: &SkillPathOrigin, diff --git a/app/src/pane_group/mod_tests.rs b/app/src/pane_group/mod_tests.rs index e325fd67..9a7d9424 100644 --- a/app/src/pane_group/mod_tests.rs +++ b/app/src/pane_group/mod_tests.rs @@ -386,6 +386,7 @@ fn persisted_remote_child_conversation( conversation_id: conversation_id.to_string(), conversation_data: serde_json::to_string(&AgentConversationData { agent_backend: Default::default(), + active_provider_run_json: None, server_conversation_token: Some("restored-child-token".to_string()), conversation_usage_metadata: None, reverted_action_ids: None, diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index d43d4408..f5e6675b 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -6930,13 +6930,7 @@ impl TerminalView { } self.ai_controller.update(ctx, |controller, ctx| { - controller.resume_conversation( - *conversation_id, - /*can_attempt_resume_on_error*/ true, - /*is_auto_resume_after_error*/ false, - vec![], - ctx, - ); + controller.resume_conversation(*conversation_id, vec![], ctx); }); } @@ -7261,7 +7255,7 @@ impl TerminalView { ctx: &mut ViewContext, ) { match event { - BlocklistAIActionEvent::ActionBlockedOnUserConfirmation(_) => { + BlocklistAIActionEvent::ActionBlockedOnUserConfirmation { .. } => { let is_agent_in_control = self .model .lock() @@ -7272,7 +7266,7 @@ impl TerminalView { self.redetermine_terminal_focus(ctx); } } - BlocklistAIActionEvent::ExecutingAction(..) => { + BlocklistAIActionEvent::ExecutingAction { .. } => { self.redetermine_terminal_focus(ctx); ctx.notify(); } @@ -7378,7 +7372,7 @@ impl TerminalView { ); } } - BlocklistAIActionEvent::QueuedAction(_) + BlocklistAIActionEvent::QueuedAction { .. } | BlocklistAIActionEvent::ToolLifecycle { .. } => {} } } @@ -11168,13 +11162,7 @@ impl TerminalView { }; self.ai_controller.update(ctx, |controller, ctx| { - controller.resume_conversation( - conversation_id, - /*can_attempt_resume_on_error*/ true, - /*is_auto_resume_after_error*/ false, - resume_context, - ctx, - ); + controller.resume_conversation(conversation_id, resume_context, ctx); }); } diff --git a/app/src/terminal/view/load_ai_conversation.rs b/app/src/terminal/view/load_ai_conversation.rs index 29cbeaa1..6b46cd7a 100644 --- a/app/src/terminal/view/load_ai_conversation.rs +++ b/app/src/terminal/view/load_ai_conversation.rs @@ -937,6 +937,7 @@ impl TerminalView { let conversation_data = AgentConversationData { agent_backend: Default::default(), + active_provider_run_json: None, server_conversation_token: None, conversation_usage_metadata: None, reverted_action_ids: None, diff --git a/crates/ai/src/skills/conversion.rs b/crates/ai/src/skills/conversion.rs index 4bb8310e..8aeb8fbd 100644 --- a/crates/ai/src/skills/conversion.rs +++ b/crates/ai/src/skills/conversion.rs @@ -34,7 +34,7 @@ pub enum SkillConversionError { /// Live agent responses can be decoded from the active session's location. Restored payloads do /// not carry enough session identity to safely reconstruct path-based skill locations, so callers /// must use [`SkillPathOrigin::Unavailable`] rather than silently assuming the local filesystem. -#[derive(Debug, Clone, Eq, PartialEq)] +#[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize)] pub enum SkillPathOrigin { Local, Remote { diff --git a/crates/galaxy_agent_core/src/lib.rs b/crates/galaxy_agent_core/src/lib.rs index 1c7565a5..231b6d69 100644 --- a/crates/galaxy_agent_core/src/lib.rs +++ b/crates/galaxy_agent_core/src/lib.rs @@ -4,10 +4,12 @@ //! concrete runtimes such as Rig-backed providers or ACP agents. It must not //! depend on GalaxyUI, provider SDKs, persistence, or Warp wire protocols. +mod provider_run; mod runtime; mod tool_policy; mod types; +pub use provider_run::*; pub use runtime::*; pub use tool_policy::*; pub use types::*; diff --git a/crates/galaxy_agent_core/src/provider_run.rs b/crates/galaxy_agent_core/src/provider_run.rs new file mode 100644 index 00000000..340da19c --- /dev/null +++ b/crates/galaxy_agent_core/src/provider_run.rs @@ -0,0 +1,1306 @@ +use std::collections::{BTreeMap, BTreeSet, HashSet}; +use std::error::Error; +use std::fmt; + +use serde::{Deserialize, Serialize}; + +use crate::{ + AgentError, ContentPart, ConversationMessage, MessageContent, MessageRole, PermissionDecision, + PermissionRequest, StopReason, ToolCall, ToolCallDecision, ToolPolicy, ToolResult, + ToolResultStatus, Usage, +}; + +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct ProviderRunId(String); + +impl ProviderRunId { + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl From for ProviderRunId { + fn from(value: String) -> Self { + Self(value) + } +} + +impl From<&str> for ProviderRunId { + fn from(value: &str) -> Self { + Self(value.to_string()) + } +} + +#[derive( + Clone, Copy, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize, +)] +#[serde(transparent)] +pub struct RunEpoch(u64); + +impl RunEpoch { + pub const fn new(value: u64) -> Self { + Self(value) + } + + pub const fn get(self) -> u64 { + self.0 + } + + fn checked_next(self) -> Option { + self.0.checked_add(1).map(Self) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct ExternalWorkId { + pub run_id: ProviderRunId, + pub epoch: RunEpoch, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct ProviderRequestProfile(String); + +impl ProviderRequestProfile { + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl From for ProviderRequestProfile { + fn from(value: String) -> Self { + Self(value) + } +} + +impl From<&str> for ProviderRequestProfile { + fn from(value: &str) -> Self { + Self(value.to_string()) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ProviderRunLimits { + pub max_model_turns: u32, + pub max_model_retries_per_turn: u32, +} + +impl Default for ProviderRunLimits { + fn default() -> Self { + Self { + max_model_turns: 100, + max_model_retries_per_turn: 2, + } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ProviderModelCall { + pub work_id: ExternalWorkId, + pub profile: ProviderRequestProfile, + pub messages: Vec, + pub retry_attempt: u32, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CompletedModelTurn { + /// Canonical non-tool assistant content. Tool calls are appended in `tool_calls` order. + pub assistant_content: Vec, + pub tool_calls: Vec, + pub usage: Usage, + pub stop_reason: StopReason, + pub advertised_tools: BTreeSet, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct PendingModelCall { + pub work_id: ExternalWorkId, + pub retry_attempt: u32, + pub last_error: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub enum PendingToolCallState { + Proposed, + PermissionPending { + request: PermissionRequest, + }, + Approved { + request_id: String, + decision: PermissionDecision, + }, + Executing, + Resolved { + result: ToolResult, + }, +} + +impl PendingToolCallState { + pub fn result(&self) -> Option<&ToolResult> { + match self { + Self::Resolved { result } => Some(result), + Self::Proposed + | Self::PermissionPending { .. } + | Self::Approved { .. } + | Self::Executing => None, + } + } + + fn name(&self) -> &'static str { + match self { + Self::Proposed => "proposed", + Self::PermissionPending { .. } => "permission_pending", + Self::Approved { .. } => "approved", + Self::Executing => "executing", + Self::Resolved { .. } => "resolved", + } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct PendingToolCall { + pub call: ToolCall, + pub state: PendingToolCallState, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct PendingToolBatch { + pub work_id: ExternalWorkId, + pub calls: Vec, +} + +impl PendingToolBatch { + pub fn is_complete(&self) -> bool { + self.calls.iter().all(|call| call.state.result().is_some()) + } + + pub fn unresolved_call_ids(&self) -> Vec { + self.calls + .iter() + .filter(|call| call.state.result().is_none()) + .map(|call| call.call.id.clone()) + .collect() + } + + fn ordered_results(&self) -> Option> { + self.calls + .iter() + .map(|call| call.state.result().cloned()) + .collect() + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub enum ProviderRunState { + ReadyToCallModel, + AwaitingModel { + call: PendingModelCall, + }, + ResolvingModel { + turn: CompletedModelTurn, + }, + AwaitingTools { + batch: PendingToolBatch, + }, + AwaitingDriver { + work_id: ExternalWorkId, + stop_reason: StopReason, + }, + Done { + completion: ProviderRunCompletion, + }, + Failed { + failure: ProviderRunFailure, + }, + Cancelled { + reason: String, + }, +} + +impl ProviderRunState { + pub fn phase(&self) -> ProviderRunPhase { + match self { + Self::ReadyToCallModel => ProviderRunPhase::ReadyToCallModel, + Self::AwaitingModel { .. } => ProviderRunPhase::AwaitingModel, + Self::ResolvingModel { .. } => ProviderRunPhase::ResolvingModel, + Self::AwaitingTools { .. } => ProviderRunPhase::AwaitingTools, + Self::AwaitingDriver { .. } => ProviderRunPhase::AwaitingDriver, + Self::Done { .. } => ProviderRunPhase::Done, + Self::Failed { .. } => ProviderRunPhase::Failed, + Self::Cancelled { .. } => ProviderRunPhase::Cancelled, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum ProviderRunPhase { + ReadyToCallModel, + AwaitingModel, + ResolvingModel, + AwaitingTools, + AwaitingDriver, + Done, + Failed, + Cancelled, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ProviderRunCompletion { + pub stop_reason: StopReason, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum ProviderRunFailureKind { + ModelCall, + RetryLimitExceeded, + TurnLimitExceeded, + Protocol, + Projection, + Restore, + ExternalWork, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ProviderRunFailure { + pub kind: ProviderRunFailureKind, + pub message: String, + pub source: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum ProviderRunOutcome { + Completed(ProviderRunCompletion), + Failed(ProviderRunFailure), + Cancelled { reason: String }, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub enum ProviderRunStep { + CallModel(ProviderModelCall), + DispatchTools(PendingToolBatch), + Done(ProviderRunOutcome), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ModelFailureDisposition { + RetryScheduled, + RunFailed, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct ProviderRunRestoreNormalization { + pub permission_call_ids_reset: Vec, + pub interrupted_call_ids: Vec, + pub committed_tool_batch: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum ProviderRunProtocolError { + WorkMismatch { + expected: ExternalWorkId, + received: ExternalWorkId, + }, + UnexpectedState { + expected: ProviderRunPhase, + actual: ProviderRunPhase, + }, + InvalidModelTurn { + message: String, + }, + UnknownToolCall { + call_id: String, + }, + DuplicateToolUpdate { + call_id: String, + }, + InvalidToolTransition { + call_id: String, + state: String, + update: String, + }, + PermissionRequestMismatch { + call_id: String, + expected_request_id: String, + received_request_id: String, + }, + IncompleteToolBatch { + missing_call_ids: Vec, + }, + ToolResultSetMismatch { + expected_call_ids: Vec, + received_call_ids: Vec, + }, + DuplicateToolResult { + call_id: String, + }, + InvalidDriverObservation { + message: String, + }, + EpochExhausted, + Terminal, +} + +impl fmt::Display for ProviderRunProtocolError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::WorkMismatch { expected, received } => write!( + f, + "external work mismatch: expected {}:{}, received {}:{}", + expected.run_id.as_str(), + expected.epoch.get(), + received.run_id.as_str(), + received.epoch.get() + ), + Self::UnexpectedState { expected, actual } => { + write!(f, "expected run phase {expected:?}, found {actual:?}") + } + Self::InvalidModelTurn { message } => write!(f, "invalid model turn: {message}"), + Self::UnknownToolCall { call_id } => write!(f, "unknown tool call '{call_id}'"), + Self::DuplicateToolUpdate { call_id } => { + write!(f, "tool call '{call_id}' already has a result") + } + Self::InvalidToolTransition { + call_id, + state, + update, + } => write!( + f, + "cannot apply {update} to tool call '{call_id}' while it is {state}" + ), + Self::PermissionRequestMismatch { + call_id, + expected_request_id, + received_request_id, + } => write!( + f, + "permission request mismatch for '{call_id}': expected '{expected_request_id}', received '{received_request_id}'" + ), + Self::IncompleteToolBatch { missing_call_ids } => write!( + f, + "tool batch is incomplete; missing results for {}", + missing_call_ids.join(", ") + ), + Self::ToolResultSetMismatch { + expected_call_ids, + received_call_ids, + } => write!( + f, + "tool result set mismatch: expected [{}], received [{}]", + expected_call_ids.join(", "), + received_call_ids.join(", ") + ), + Self::DuplicateToolResult { call_id } => { + write!(f, "duplicate tool result for '{call_id}'") + } + Self::InvalidDriverObservation { message } => { + write!(f, "invalid driver observation: {message}") + } + Self::EpochExhausted => f.write_str("provider run epoch is exhausted"), + Self::Terminal => f.write_str("provider run is already terminal"), + } + } +} + +impl Error for ProviderRunProtocolError {} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ProviderRun { + id: ProviderRunId, + epoch: RunEpoch, + transcript: Vec, + tool_result_archive: Vec, + usage: Usage, + model_turns: u32, + model_retries: u32, + limits: ProviderRunLimits, + profile: ProviderRequestProfile, + state: ProviderRunState, +} + +impl ProviderRun { + pub fn new( + id: impl Into, + transcript: Vec, + profile: impl Into, + mut limits: ProviderRunLimits, + ) -> Self { + limits.max_model_turns = limits.max_model_turns.max(1); + Self { + id: id.into(), + epoch: RunEpoch::default(), + transcript, + tool_result_archive: Vec::new(), + usage: Usage::default(), + model_turns: 0, + model_retries: 0, + limits, + profile: profile.into(), + state: ProviderRunState::ReadyToCallModel, + } + } + + pub fn id(&self) -> &ProviderRunId { + &self.id + } + + pub fn epoch(&self) -> RunEpoch { + self.epoch + } + + pub fn profile(&self) -> &ProviderRequestProfile { + &self.profile + } + + pub fn state(&self) -> &ProviderRunState { + &self.state + } + + pub fn transcript(&self) -> &[ConversationMessage] { + &self.transcript + } + + pub fn tool_result_archive(&self) -> &[ConversationMessage] { + &self.tool_result_archive + } + + pub fn replace_tool_result_archive(&mut self, archive: Vec) { + self.tool_result_archive = archive; + } + + pub fn usage(&self) -> &Usage { + &self.usage + } + + pub fn model_turns(&self) -> u32 { + self.model_turns + } + + pub fn model_retries(&self) -> u32 { + self.model_retries + } + + pub fn is_terminal(&self) -> bool { + matches!( + self.state, + ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } + ) + } + + pub fn normalize_after_restore( + &mut self, + ) -> Result { + let ProviderRunState::AwaitingTools { batch } = &mut self.state else { + return Ok(ProviderRunRestoreNormalization::default()); + }; + + let mut normalization = ProviderRunRestoreNormalization::default(); + for pending in &mut batch.calls { + match &pending.state { + PendingToolCallState::PermissionPending { .. } => { + normalization + .permission_call_ids_reset + .push(pending.call.id.clone()); + pending.state = PendingToolCallState::Proposed; + } + PendingToolCallState::Approved { .. } | PendingToolCallState::Executing => { + normalization + .interrupted_call_ids + .push(pending.call.id.clone()); + pending.state = PendingToolCallState::Resolved { + result: ToolResult { + call_id: pending.call.id.clone(), + content: "Tool execution was interrupted by application restart and was not replayed." + .to_string(), + status: ToolResultStatus::Error, + }, + }; + } + PendingToolCallState::Proposed | PendingToolCallState::Resolved { .. } => {} + } + } + + if batch.is_complete() { + let work_id = batch.work_id.clone(); + self.commit_tool_batch(&work_id)?; + normalization.committed_tool_batch = true; + } + Ok(normalization) + } + + pub fn active_work_id(&self) -> Option<&ExternalWorkId> { + match &self.state { + ProviderRunState::AwaitingModel { call } => Some(&call.work_id), + ProviderRunState::AwaitingTools { batch } => Some(&batch.work_id), + ProviderRunState::AwaitingDriver { work_id, .. } => Some(work_id), + ProviderRunState::ReadyToCallModel + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => None, + } + } + + pub fn ready_work_id(&self) -> Option { + matches!(self.state, ProviderRunState::ReadyToCallModel).then(|| self.current_work_id()) + } + + pub fn next_step(&mut self) -> Result, ProviderRunProtocolError> { + loop { + match self.state.clone() { + ProviderRunState::ReadyToCallModel => { + if self.model_turns >= self.limits.max_model_turns { + self.state = ProviderRunState::Failed { + failure: ProviderRunFailure { + kind: ProviderRunFailureKind::TurnLimitExceeded, + message: format!( + "provider run reached its {} model-turn limit", + self.limits.max_model_turns + ), + source: None, + }, + }; + continue; + } + self.state = ProviderRunState::AwaitingModel { + call: PendingModelCall { + work_id: self.current_work_id(), + retry_attempt: 0, + last_error: None, + }, + }; + } + ProviderRunState::AwaitingModel { call } => { + return Ok(Some(ProviderRunStep::CallModel(ProviderModelCall { + work_id: call.work_id, + profile: self.profile.clone(), + messages: self.transcript.clone(), + retry_attempt: call.retry_attempt, + }))); + } + ProviderRunState::ResolvingModel { turn } => { + if turn.tool_calls.is_empty() { + self.state = ProviderRunState::AwaitingDriver { + work_id: self.current_work_id(), + stop_reason: turn.stop_reason, + }; + } else { + let batch = self.build_tool_batch(&turn); + self.state = ProviderRunState::AwaitingTools { batch }; + } + } + ProviderRunState::AwaitingTools { batch } => { + return Ok(Some(ProviderRunStep::DispatchTools(batch))); + } + ProviderRunState::AwaitingDriver { .. } => return Ok(None), + ProviderRunState::Done { completion } => { + return Ok(Some(ProviderRunStep::Done(ProviderRunOutcome::Completed( + completion, + )))); + } + ProviderRunState::Failed { failure } => { + return Ok(Some(ProviderRunStep::Done(ProviderRunOutcome::Failed( + failure, + )))); + } + ProviderRunState::Cancelled { reason } => { + return Ok(Some(ProviderRunStep::Done(ProviderRunOutcome::Cancelled { + reason, + }))); + } + } + } + } + + pub fn accept_model_turn( + &mut self, + work_id: &ExternalWorkId, + turn: CompletedModelTurn, + ) -> Result<(), ProviderRunProtocolError> { + let expected = match &self.state { + ProviderRunState::AwaitingModel { call } => &call.work_id, + ProviderRunState::ReadyToCallModel + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingTools { .. } + | ProviderRunState::AwaitingDriver { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => { + return Err(self.unexpected_state(ProviderRunPhase::AwaitingModel)); + } + }; + validate_work_id(expected, work_id)?; + validate_model_turn(&turn)?; + let next_epoch = self.next_epoch()?; + let assistant_message = assistant_message(&turn); + + self.transcript.push(assistant_message); + add_usage(&mut self.usage, &turn.usage); + self.model_turns = self.model_turns.saturating_add(1); + self.epoch = next_epoch; + self.state = ProviderRunState::ResolvingModel { turn }; + Ok(()) + } + + pub fn register_model_failure( + &mut self, + work_id: &ExternalWorkId, + error: AgentError, + ) -> Result { + let call = match &mut self.state { + ProviderRunState::AwaitingModel { call } => call, + ProviderRunState::ReadyToCallModel + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingTools { .. } + | ProviderRunState::AwaitingDriver { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => { + return Err(self.unexpected_state(ProviderRunPhase::AwaitingModel)); + } + }; + validate_work_id(&call.work_id, work_id)?; + + if error.recoverable && call.retry_attempt < self.limits.max_model_retries_per_turn { + call.retry_attempt = call.retry_attempt.saturating_add(1); + call.last_error = Some(error); + self.model_retries = self.model_retries.saturating_add(1); + return Ok(ModelFailureDisposition::RetryScheduled); + } + + let kind = if error.recoverable { + ProviderRunFailureKind::RetryLimitExceeded + } else { + ProviderRunFailureKind::ModelCall + }; + let message = if error.recoverable { + format!( + "provider model call failed after {} retries: {}", + call.retry_attempt, error.message + ) + } else { + error.message.clone() + }; + self.state = ProviderRunState::Failed { + failure: ProviderRunFailure { + kind, + message, + source: Some(error), + }, + }; + Ok(ModelFailureDisposition::RunFailed) + } + + pub fn request_tool_permission( + &mut self, + work_id: &ExternalWorkId, + request: PermissionRequest, + ) -> Result<(), ProviderRunProtocolError> { + let call = self.pending_tool_call_mut(work_id, &request.call_id)?; + match &call.state { + PendingToolCallState::Proposed => { + call.state = PendingToolCallState::PermissionPending { request }; + Ok(()) + } + PendingToolCallState::Resolved { .. } => { + Err(ProviderRunProtocolError::DuplicateToolUpdate { + call_id: call.call.id.clone(), + }) + } + PendingToolCallState::PermissionPending { .. } + | PendingToolCallState::Approved { .. } + | PendingToolCallState::Executing => { + Err(invalid_tool_transition(call, "permission request")) + } + } + } + + pub fn resolve_tool_permission( + &mut self, + work_id: &ExternalWorkId, + call_id: &str, + request_id: &str, + decision: PermissionDecision, + ) -> Result<(), ProviderRunProtocolError> { + let call = self.pending_tool_call_mut(work_id, call_id)?; + let pending_request_id = match &call.state { + PendingToolCallState::PermissionPending { request } => request.id.clone(), + PendingToolCallState::Resolved { .. } => { + return Err(ProviderRunProtocolError::DuplicateToolUpdate { + call_id: call.call.id.clone(), + }); + } + PendingToolCallState::Proposed + | PendingToolCallState::Approved { .. } + | PendingToolCallState::Executing => { + return Err(invalid_tool_transition(call, "permission resolution")); + } + }; + if pending_request_id != request_id { + return Err(ProviderRunProtocolError::PermissionRequestMismatch { + call_id: call_id.to_string(), + expected_request_id: pending_request_id, + received_request_id: request_id.to_string(), + }); + } + + match decision { + PermissionDecision::AllowOnce | PermissionDecision::AlwaysAllow => { + call.state = PendingToolCallState::Approved { + request_id: request_id.to_string(), + decision, + }; + } + PermissionDecision::Denied { reason } => { + let content = reason.unwrap_or_else(|| "Tool permission was denied.".to_string()); + call.state = PendingToolCallState::Resolved { + result: ToolResult { + call_id: call_id.to_string(), + content, + status: ToolResultStatus::Denied, + }, + }; + } + } + Ok(()) + } + + pub fn start_tool( + &mut self, + work_id: &ExternalWorkId, + call_id: &str, + ) -> Result<(), ProviderRunProtocolError> { + let call = self.pending_tool_call_mut(work_id, call_id)?; + match &call.state { + PendingToolCallState::Proposed | PendingToolCallState::Approved { .. } => { + call.state = PendingToolCallState::Executing; + Ok(()) + } + PendingToolCallState::Resolved { .. } => { + Err(ProviderRunProtocolError::DuplicateToolUpdate { + call_id: call.call.id.clone(), + }) + } + PendingToolCallState::PermissionPending { .. } | PendingToolCallState::Executing => { + Err(invalid_tool_transition(call, "tool start")) + } + } + } + + pub fn complete_tool( + &mut self, + work_id: &ExternalWorkId, + result: ToolResult, + ) -> Result<(), ProviderRunProtocolError> { + let call = self.pending_tool_call_mut(work_id, &result.call_id)?; + match &call.state { + PendingToolCallState::Proposed + | PendingToolCallState::Approved { .. } + | PendingToolCallState::Executing => { + call.state = PendingToolCallState::Resolved { result }; + Ok(()) + } + PendingToolCallState::Resolved { .. } => { + Err(ProviderRunProtocolError::DuplicateToolUpdate { + call_id: call.call.id.clone(), + }) + } + PendingToolCallState::PermissionPending { .. } => { + Err(invalid_tool_transition(call, "tool completion")) + } + } + } + + pub fn cancel_tool( + &mut self, + work_id: &ExternalWorkId, + call_id: &str, + reason: impl Into, + ) -> Result<(), ProviderRunProtocolError> { + let call = self.pending_tool_call_mut(work_id, call_id)?; + match &call.state { + PendingToolCallState::Resolved { .. } => { + Err(ProviderRunProtocolError::DuplicateToolUpdate { + call_id: call.call.id.clone(), + }) + } + PendingToolCallState::Proposed + | PendingToolCallState::PermissionPending { .. } + | PendingToolCallState::Approved { .. } + | PendingToolCallState::Executing => { + call.state = PendingToolCallState::Resolved { + result: ToolResult { + call_id: call_id.to_string(), + content: reason.into(), + status: ToolResultStatus::Cancelled, + }, + }; + Ok(()) + } + } + } + + pub fn complete_tool_batch( + &mut self, + work_id: &ExternalWorkId, + results: Vec, + ) -> Result<(), ProviderRunProtocolError> { + let batch = self.pending_tool_batch(work_id)?; + let expected = batch.unresolved_call_ids(); + let mut by_call_id = BTreeMap::new(); + for result in results { + let call_id = result.call_id.clone(); + if by_call_id.insert(call_id.clone(), result).is_some() { + return Err(ProviderRunProtocolError::DuplicateToolResult { call_id }); + } + } + let received = by_call_id.keys().cloned().collect::>(); + let expected_set = expected.iter().cloned().collect::>(); + let received_set = received.iter().cloned().collect::>(); + if expected_set != received_set { + return Err(ProviderRunProtocolError::ToolResultSetMismatch { + expected_call_ids: expected, + received_call_ids: received, + }); + } + + let batch = self.pending_tool_batch_mut(work_id)?; + for call in &mut batch.calls { + if call.state.result().is_some() { + continue; + } + let result = by_call_id + .remove(&call.call.id) + .expect("validated tool result set must contain every unresolved call"); + call.state = PendingToolCallState::Resolved { result }; + } + Ok(()) + } + + pub fn commit_tool_batch( + &mut self, + work_id: &ExternalWorkId, + ) -> Result<(), ProviderRunProtocolError> { + let batch = self.pending_tool_batch(work_id)?.clone(); + let Some(results) = batch.ordered_results() else { + return Err(ProviderRunProtocolError::IncompleteToolBatch { + missing_call_ids: batch.unresolved_call_ids(), + }); + }; + let next_epoch = self.next_epoch()?; + self.transcript.push(tool_result_message(results)); + self.epoch = next_epoch; + self.state = ProviderRunState::ReadyToCallModel; + Ok(()) + } + + pub fn complete(&mut self, work_id: &ExternalWorkId) -> Result<(), ProviderRunProtocolError> { + let (expected, stop_reason) = match &self.state { + ProviderRunState::AwaitingDriver { + work_id, + stop_reason, + } => (work_id, stop_reason.clone()), + ProviderRunState::ReadyToCallModel + | ProviderRunState::AwaitingModel { .. } + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingTools { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => { + return Err(self.unexpected_state(ProviderRunPhase::AwaitingDriver)); + } + }; + validate_work_id(expected, work_id)?; + self.state = ProviderRunState::Done { + completion: ProviderRunCompletion { stop_reason }, + }; + Ok(()) + } + + pub fn continue_with_observation( + &mut self, + work_id: &ExternalWorkId, + observation: MessageContent, + next_profile: impl Into, + ) -> Result<(), ProviderRunProtocolError> { + let expected = match &self.state { + ProviderRunState::AwaitingDriver { work_id, .. } => work_id, + ProviderRunState::ReadyToCallModel + | ProviderRunState::AwaitingModel { .. } + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingTools { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => { + return Err(self.unexpected_state(ProviderRunPhase::AwaitingDriver)); + } + }; + validate_work_id(expected, work_id)?; + validate_driver_observation(&observation)?; + let next_epoch = self.next_epoch()?; + self.transcript.push(ConversationMessage { + role: MessageRole::User, + content: observation, + }); + self.profile = next_profile.into(); + self.epoch = next_epoch; + self.state = ProviderRunState::ReadyToCallModel; + Ok(()) + } + + pub fn continue_ready_with_observation( + &mut self, + work_id: &ExternalWorkId, + observation: MessageContent, + next_profile: impl Into, + ) -> Result<(), ProviderRunProtocolError> { + if !matches!(self.state, ProviderRunState::ReadyToCallModel) { + return Err(self.unexpected_state(ProviderRunPhase::ReadyToCallModel)); + } + validate_work_id(&self.current_work_id(), work_id)?; + validate_driver_observation(&observation)?; + let next_epoch = self.next_epoch()?; + self.transcript.push(ConversationMessage { + role: MessageRole::User, + content: observation, + }); + self.profile = next_profile.into(); + self.epoch = next_epoch; + Ok(()) + } + + pub fn cancel(&mut self, reason: impl Into) -> Result<(), ProviderRunProtocolError> { + let reason = reason.into(); + self.finish_run_with_pending_tools( + ToolResultStatus::Cancelled, + reason.clone(), + ProviderRunState::Cancelled { reason }, + ) + } + + pub fn fail( + &mut self, + kind: ProviderRunFailureKind, + message: impl Into, + ) -> Result<(), ProviderRunProtocolError> { + let message = message.into(); + let terminal = ProviderRunState::Failed { + failure: ProviderRunFailure { + kind, + message: message.clone(), + source: None, + }, + }; + self.finish_run_with_pending_tools(ToolResultStatus::Error, message, terminal) + } + + fn current_work_id(&self) -> ExternalWorkId { + ExternalWorkId { + run_id: self.id.clone(), + epoch: self.epoch, + } + } + + fn next_epoch(&self) -> Result { + self.epoch + .checked_next() + .ok_or(ProviderRunProtocolError::EpochExhausted) + } + + fn build_tool_batch(&self, turn: &CompletedModelTurn) -> PendingToolBatch { + let policy = ToolPolicy::from_names(turn.advertised_tools.iter().cloned()); + let calls = turn + .tool_calls + .iter() + .cloned() + .map(|call| { + let state = if !call.arguments.is_object() { + PendingToolCallState::Resolved { + result: ToolResult { + call_id: call.id.clone(), + content: format!( + "Error: '{}' received malformed arguments; expected a JSON object.", + call.name + ), + status: ToolResultStatus::Error, + }, + } + } else { + match policy.decide(&call, &self.transcript, &self.tool_result_archive) { + ToolCallDecision::Execute => PendingToolCallState::Proposed, + ToolCallDecision::Inline(result) | ToolCallDecision::Reject(result) => { + PendingToolCallState::Resolved { result } + } + } + }; + PendingToolCall { call, state } + }) + .collect(); + PendingToolBatch { + work_id: self.current_work_id(), + calls, + } + } + + fn pending_tool_batch( + &self, + work_id: &ExternalWorkId, + ) -> Result<&PendingToolBatch, ProviderRunProtocolError> { + let batch = match &self.state { + ProviderRunState::AwaitingTools { batch } => batch, + ProviderRunState::ReadyToCallModel + | ProviderRunState::AwaitingModel { .. } + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingDriver { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => { + return Err(self.unexpected_state(ProviderRunPhase::AwaitingTools)); + } + }; + validate_work_id(&batch.work_id, work_id)?; + Ok(batch) + } + + fn pending_tool_batch_mut( + &mut self, + work_id: &ExternalWorkId, + ) -> Result<&mut PendingToolBatch, ProviderRunProtocolError> { + let actual = self.state.phase(); + let batch = match &mut self.state { + ProviderRunState::AwaitingTools { batch } => batch, + ProviderRunState::ReadyToCallModel + | ProviderRunState::AwaitingModel { .. } + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingDriver { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => { + return Err(ProviderRunProtocolError::UnexpectedState { + expected: ProviderRunPhase::AwaitingTools, + actual, + }); + } + }; + validate_work_id(&batch.work_id, work_id)?; + Ok(batch) + } + + fn pending_tool_call_mut( + &mut self, + work_id: &ExternalWorkId, + call_id: &str, + ) -> Result<&mut PendingToolCall, ProviderRunProtocolError> { + let batch = self.pending_tool_batch_mut(work_id)?; + batch + .calls + .iter_mut() + .find(|call| call.call.id == call_id) + .ok_or_else(|| ProviderRunProtocolError::UnknownToolCall { + call_id: call_id.to_string(), + }) + } + + fn finish_run_with_pending_tools( + &mut self, + status: ToolResultStatus, + result_content: String, + terminal: ProviderRunState, + ) -> Result<(), ProviderRunProtocolError> { + match self.state.clone() { + ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => Err(ProviderRunProtocolError::Terminal), + ProviderRunState::ResolvingModel { turn } => { + if turn.tool_calls.is_empty() { + self.state = ProviderRunState::AwaitingDriver { + work_id: self.current_work_id(), + stop_reason: turn.stop_reason, + }; + } else { + let batch = self.build_tool_batch(&turn); + self.state = ProviderRunState::AwaitingTools { batch }; + } + self.finish_run_with_pending_tools(status, result_content, terminal) + } + ProviderRunState::AwaitingTools { mut batch } => { + for call in &mut batch.calls { + if call.state.result().is_none() { + call.state = PendingToolCallState::Resolved { + result: ToolResult { + call_id: call.call.id.clone(), + content: result_content.clone(), + status, + }, + }; + } + } + let results = batch + .ordered_results() + .expect("all unresolved tool calls were assigned a terminal result"); + let next_epoch = self.next_epoch()?; + self.transcript.push(tool_result_message(results)); + self.epoch = next_epoch; + self.state = terminal; + Ok(()) + } + ProviderRunState::ReadyToCallModel + | ProviderRunState::AwaitingModel { .. } + | ProviderRunState::AwaitingDriver { .. } => { + self.state = terminal; + Ok(()) + } + } + } + + fn unexpected_state(&self, expected: ProviderRunPhase) -> ProviderRunProtocolError { + ProviderRunProtocolError::UnexpectedState { + expected, + actual: self.state.phase(), + } + } +} + +fn validate_work_id( + expected: &ExternalWorkId, + received: &ExternalWorkId, +) -> Result<(), ProviderRunProtocolError> { + if expected == received { + Ok(()) + } else { + Err(ProviderRunProtocolError::WorkMismatch { + expected: expected.clone(), + received: received.clone(), + }) + } +} + +fn validate_model_turn(turn: &CompletedModelTurn) -> Result<(), ProviderRunProtocolError> { + for part in &turn.assistant_content { + match part { + ContentPart::Text(_) | ContentPart::Reasoning { .. } | ContentPart::Image { .. } => {} + ContentPart::ToolUse { .. } | ContentPart::ToolResult { .. } => { + return Err(ProviderRunProtocolError::InvalidModelTurn { + message: "assistant_content must not contain tool calls or results".to_string(), + }); + } + } + } + + let mut call_ids = HashSet::new(); + for call in &turn.tool_calls { + if call.id.is_empty() { + return Err(ProviderRunProtocolError::InvalidModelTurn { + message: "tool call IDs must not be empty".to_string(), + }); + } + if !call_ids.insert(call.id.as_str()) { + return Err(ProviderRunProtocolError::InvalidModelTurn { + message: format!("duplicate tool call ID '{}'", call.id), + }); + } + } + Ok(()) +} + +fn assistant_message(turn: &CompletedModelTurn) -> ConversationMessage { + let mut parts = turn.assistant_content.clone(); + parts.extend(turn.tool_calls.iter().map(|call| ContentPart::ToolUse { + tool_use_id: call.id.clone(), + name: call.name.clone(), + input: call.arguments.clone(), + })); + let content = match parts.as_slice() { + [ContentPart::Text(text)] => MessageContent::Text(text.clone()), + [] => MessageContent::Text(String::new()), + [ContentPart::Reasoning { .. }] + | [ContentPart::Image { .. }] + | [ContentPart::ToolUse { .. }] + | [ContentPart::ToolResult { .. }] + | [_, _, ..] => MessageContent::MultiPart(parts), + }; + ConversationMessage { + role: MessageRole::Assistant, + content, + } +} + +fn tool_result_message(results: Vec) -> ConversationMessage { + ConversationMessage { + role: MessageRole::User, + content: MessageContent::MultiPart( + results + .into_iter() + .map(|result| ContentPart::ToolResult { + tool_use_id: result.call_id, + content: result.content, + is_error: !matches!(result.status, ToolResultStatus::Success), + }) + .collect(), + ), + } +} + +fn validate_driver_observation( + observation: &MessageContent, +) -> Result<(), ProviderRunProtocolError> { + match observation { + MessageContent::Text(_) => Ok(()), + MessageContent::MultiPart(parts) => { + for part in parts { + match part { + ContentPart::Text(_) + | ContentPart::Reasoning { .. } + | ContentPart::Image { .. } => {} + ContentPart::ToolUse { .. } | ContentPart::ToolResult { .. } => { + return Err(ProviderRunProtocolError::InvalidDriverObservation { + message: + "continuation observations cannot inject tool calls or results" + .to_string(), + }); + } + } + } + Ok(()) + } + MessageContent::ToolUse { .. } | MessageContent::ToolResult { .. } => { + Err(ProviderRunProtocolError::InvalidDriverObservation { + message: "continuation observations cannot inject tool calls or results" + .to_string(), + }) + } + } +} + +fn invalid_tool_transition(call: &PendingToolCall, update: &str) -> ProviderRunProtocolError { + ProviderRunProtocolError::InvalidToolTransition { + call_id: call.call.id.clone(), + state: call.state.name().to_string(), + update: update.to_string(), + } +} + +fn add_usage(total: &mut Usage, turn: &Usage) { + total.input_tokens = total.input_tokens.saturating_add(turn.input_tokens); + total.output_tokens = total.output_tokens.saturating_add(turn.output_tokens); + total.cached_input_tokens = total + .cached_input_tokens + .saturating_add(turn.cached_input_tokens); + total.cache_creation_input_tokens = total + .cache_creation_input_tokens + .saturating_add(turn.cache_creation_input_tokens); +} + +#[cfg(test)] +#[path = "provider_run_tests.rs"] +mod tests; diff --git a/crates/galaxy_agent_core/src/provider_run_tests.rs b/crates/galaxy_agent_core/src/provider_run_tests.rs new file mode 100644 index 00000000..d2be435c --- /dev/null +++ b/crates/galaxy_agent_core/src/provider_run_tests.rs @@ -0,0 +1,811 @@ +use serde_json::json; + +use super::*; +use crate::{AgentErrorKind, PermissionKind}; + +fn initial_messages() -> Vec { + vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("Finish the task".to_string()), + }] +} + +fn run_with_limits(limits: ProviderRunLimits) -> ProviderRun { + ProviderRun::new("run-1", initial_messages(), "base", limits) +} + +fn run() -> ProviderRun { + run_with_limits(ProviderRunLimits::default()) +} + +fn next_model_call(run: &mut ProviderRun) -> ProviderModelCall { + let Some(ProviderRunStep::CallModel(call)) = run.next_step().unwrap() else { + panic!("expected model call"); + }; + call +} + +fn text_turn(text: &str) -> CompletedModelTurn { + CompletedModelTurn { + assistant_content: vec![ContentPart::Text(text.to_string())], + tool_calls: Vec::new(), + usage: Usage { + input_tokens: 10, + output_tokens: 3, + ..Usage::default() + }, + stop_reason: StopReason::Completed, + advertised_tools: BTreeSet::new(), + } +} + +fn tool_call(id: &str, name: &str) -> ToolCall { + ToolCall { + id: id.to_string(), + name: name.to_string(), + arguments: json!({"id": id}), + } +} + +fn tool_turn(calls: Vec, advertised_tools: &[&str]) -> CompletedModelTurn { + CompletedModelTurn { + assistant_content: vec![ContentPart::Text("I will use tools.".to_string())], + tool_calls: calls, + usage: Usage { + input_tokens: 20, + output_tokens: 5, + cached_input_tokens: 4, + ..Usage::default() + }, + stop_reason: StopReason::Completed, + advertised_tools: advertised_tools + .iter() + .map(|name| (*name).to_string()) + .collect(), + } +} + +fn accept_tool_turn(run: &mut ProviderRun, turn: CompletedModelTurn) -> PendingToolBatch { + let model_call = next_model_call(run); + run.accept_model_turn(&model_call.work_id, turn).unwrap(); + let Some(ProviderRunStep::DispatchTools(batch)) = run.next_step().unwrap() else { + panic!("expected tool dispatch"); + }; + batch +} + +fn successful_result(call_id: &str, content: &str) -> ToolResult { + ToolResult { + call_id: call_id.to_string(), + content: content.to_string(), + status: ToolResultStatus::Success, + } +} + +fn assert_serialization_round_trip(run: &ProviderRun) { + let json = serde_json::to_string(run).unwrap(); + let restored: ProviderRun = serde_json::from_str(&json).unwrap(); + assert_eq!(&restored, run); +} + +#[test] +fn next_step_reemits_identical_pending_model_work() { + let mut run = run(); + + let first = run.next_step().unwrap(); + let second = run.next_step().unwrap(); + + assert_eq!(first, second); + assert_eq!(run.epoch(), RunEpoch::new(0)); + assert_serialization_round_trip(&run); +} + +#[test] +fn stale_model_completion_is_rejected_without_mutation() { + let mut run = run(); + let call = next_model_call(&mut run); + let stale = ExternalWorkId { + run_id: call.work_id.run_id.clone(), + epoch: RunEpoch::new(call.work_id.epoch.get() + 1), + }; + let before = run.clone(); + + let error = run + .accept_model_turn(&stale, text_turn("done")) + .unwrap_err(); + + assert!(matches!( + error, + ProviderRunProtocolError::WorkMismatch { .. } + )); + assert_eq!(run, before); +} + +#[test] +fn model_retries_reuse_work_identity_and_stop_at_the_budget() { + let mut run = run_with_limits(ProviderRunLimits { + max_model_turns: 5, + max_model_retries_per_turn: 1, + }); + let call = next_model_call(&mut run); + let mut recoverable = AgentError::new(AgentErrorKind::Transport, "network failed"); + recoverable.recoverable = true; + + assert_eq!( + run.register_model_failure(&call.work_id, recoverable.clone()) + .unwrap(), + ModelFailureDisposition::RetryScheduled + ); + let retry = next_model_call(&mut run); + assert_eq!(retry.work_id, call.work_id); + assert_eq!(retry.retry_attempt, 1); + assert_eq!(retry.messages, call.messages); + assert_eq!(run.model_retries(), 1); + + assert_eq!( + run.register_model_failure(&retry.work_id, recoverable) + .unwrap(), + ModelFailureDisposition::RunFailed + ); + let Some(ProviderRunStep::Done(ProviderRunOutcome::Failed(failure))) = run.next_step().unwrap() + else { + panic!("expected failed run"); + }; + assert_eq!(failure.kind, ProviderRunFailureKind::RetryLimitExceeded); +} + +#[test] +fn text_only_turn_requires_an_explicit_driver_decision() { + let mut run = run(); + let call = next_model_call(&mut run); + run.accept_model_turn(&call.work_id, text_turn("finished")) + .unwrap(); + + assert_eq!(run.next_step().unwrap(), None); + let driver_work = run.active_work_id().unwrap().clone(); + assert_eq!(driver_work.epoch, RunEpoch::new(1)); + assert!(!run.is_terminal()); + + run.complete(&driver_work).unwrap(); + assert_eq!( + run.next_step().unwrap(), + Some(ProviderRunStep::Done(ProviderRunOutcome::Completed( + ProviderRunCompletion { + stop_reason: StopReason::Completed, + } + ))) + ); +} + +#[test] +fn driver_continuation_appends_observation_switches_profile_and_advances_epoch() { + let mut run = run(); + let call = next_model_call(&mut run); + run.accept_model_turn(&call.work_id, text_turn("command is still running")) + .unwrap(); + assert_eq!(run.next_step().unwrap(), None); + let driver_work = run.active_work_id().unwrap().clone(); + + run.continue_with_observation( + &driver_work, + MessageContent::Text("command exited with code 1".to_string()), + "cli-monitor", + ) + .unwrap(); + + assert_eq!(run.epoch(), RunEpoch::new(2)); + assert_eq!(run.profile().as_str(), "cli-monitor"); + let next = next_model_call(&mut run); + assert_eq!(next.work_id.epoch, RunEpoch::new(2)); + assert_eq!( + next.messages.last(), + Some(&ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("command exited with code 1".to_string()), + }) + ); +} + +#[test] +fn ready_continuation_appends_observation_switches_profile_and_rejects_reuse() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![tool_call("shell", "run_shell_command")], + &["run_shell_command"], + ), + ); + run.complete_tool( + &batch.work_id, + successful_result("shell", "command is still running"), + ) + .unwrap(); + run.commit_tool_batch(&batch.work_id).unwrap(); + let ready_work = run.ready_work_id().expect("ready work identity"); + + run.continue_ready_with_observation( + &ready_work, + MessageContent::Text("Monitor command block-1.".to_string()), + "cli-monitor", + ) + .unwrap(); + + assert_eq!(run.epoch(), RunEpoch::new(3)); + assert_eq!(run.profile().as_str(), "cli-monitor"); + assert!(matches!( + run.transcript().last(), + Some(ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text(text), + }) if text == "Monitor command block-1." + )); + let continued = run.clone(); + assert!(matches!( + run.continue_ready_with_observation( + &ready_work, + MessageContent::Text("duplicate".to_string()), + "cli-monitor", + ) + .unwrap_err(), + ProviderRunProtocolError::WorkMismatch { .. } + )); + assert_eq!(run, continued); +} + +#[test] +fn unknown_malformed_and_inline_tools_are_pre_resolved_in_the_same_batch() { + let mut run = ProviderRun::new( + "run-1", + vec![ + ConversationMessage { + role: MessageRole::Assistant, + content: MessageContent::ToolUse { + tool_use_id: "old-read".to_string(), + name: "read_files".to_string(), + input: json!({"path": "old.txt"}), + }, + }, + ConversationMessage { + role: MessageRole::User, + content: MessageContent::ToolResult { + tool_use_id: "old-read".to_string(), + content: "old contents".to_string(), + is_error: false, + }, + }, + ], + "base", + ProviderRunLimits::default(), + ); + let mut malformed = tool_call("malformed", "read_files"); + malformed.arguments = json!("not an object"); + let recall = ToolCall { + id: "recall".to_string(), + name: crate::RECALL_TOOL_HISTORY_NAME.to_string(), + arguments: json!({"tool_use_id": "old-read"}), + }; + let turn = tool_turn( + vec![ + tool_call("external", "read_files"), + tool_call("unknown", "invented_tool"), + malformed, + recall, + ], + &["read_files", crate::RECALL_TOOL_HISTORY_NAME], + ); + + let batch = accept_tool_turn(&mut run, turn); + + assert!(matches!( + batch.calls[0].state, + PendingToolCallState::Proposed + )); + for index in [1, 2, 3] { + assert!(matches!( + batch.calls[index].state, + PendingToolCallState::Resolved { .. } + )); + } + let recall_result = batch.calls[3].state.result().unwrap(); + assert!(recall_result.content.contains("old contents")); + assert_eq!(batch.unresolved_call_ids(), vec!["external"]); +} + +#[test] +fn parallel_tool_results_commit_atomically_in_original_call_order() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![ + tool_call("first", "read_files"), + tool_call("second", "grep"), + ], + &["read_files", "grep"], + ), + ); + + run.start_tool(&batch.work_id, "first").unwrap(); + run.start_tool(&batch.work_id, "second").unwrap(); + run.complete_tool(&batch.work_id, successful_result("second", "second result")) + .unwrap(); + run.complete_tool(&batch.work_id, successful_result("first", "first result")) + .unwrap(); + + assert_eq!(run.transcript().len(), 2); + let Some(ProviderRunStep::DispatchTools(completed)) = run.next_step().unwrap() else { + panic!("completed batch must remain recoverable until commit"); + }; + assert!(completed.is_complete()); + assert_eq!(run.epoch(), RunEpoch::new(1)); + run.commit_tool_batch(&batch.work_id).unwrap(); + assert_eq!(run.epoch(), RunEpoch::new(2)); + + let MessageContent::MultiPart(parts) = &run.transcript().last().unwrap().content else { + panic!("expected atomic multi-part result message"); + }; + let ids = parts + .iter() + .map(|part| match part { + ContentPart::ToolResult { tool_use_id, .. } => tool_use_id.as_str(), + ContentPart::Text(_) + | ContentPart::Reasoning { .. } + | ContentPart::Image { .. } + | ContentPart::ToolUse { .. } => panic!("expected only tool results"), + }) + .collect::>(); + assert_eq!(ids, vec!["first", "second"]); +} + +#[test] +fn exact_batch_submission_rejects_missing_duplicate_and_unknown_results_without_mutation() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![ + tool_call("first", "read_files"), + tool_call("second", "grep"), + ], + &["read_files", "grep"], + ), + ); + let before = run.clone(); + + let missing = run + .complete_tool_batch(&batch.work_id, vec![successful_result("first", "one")]) + .unwrap_err(); + assert!(matches!( + missing, + ProviderRunProtocolError::ToolResultSetMismatch { .. } + )); + assert_eq!(run, before); + + let duplicate = run + .complete_tool_batch( + &batch.work_id, + vec![ + successful_result("first", "one"), + successful_result("first", "again"), + ], + ) + .unwrap_err(); + assert_eq!( + duplicate, + ProviderRunProtocolError::DuplicateToolResult { + call_id: "first".to_string(), + } + ); + assert_eq!(run, before); + + let unknown = run + .complete_tool_batch( + &batch.work_id, + vec![ + successful_result("first", "one"), + successful_result("unknown", "bad"), + ], + ) + .unwrap_err(); + assert!(matches!( + unknown, + ProviderRunProtocolError::ToolResultSetMismatch { .. } + )); + assert_eq!(run, before); +} + +#[test] +fn permission_denial_becomes_one_correlated_result() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![tool_call("shell", "run_shell_command")], + &["run_shell_command"], + ), + ); + let request = PermissionRequest { + id: "permission-shell".to_string(), + call_id: "shell".to_string(), + kind: PermissionKind::Execute, + reason: Some("run a command".to_string()), + }; + + run.request_tool_permission(&batch.work_id, request) + .unwrap(); + let wrong_request = run + .resolve_tool_permission( + &batch.work_id, + "shell", + "wrong", + PermissionDecision::AllowOnce, + ) + .unwrap_err(); + assert!(matches!( + wrong_request, + ProviderRunProtocolError::PermissionRequestMismatch { .. } + )); + run.resolve_tool_permission( + &batch.work_id, + "shell", + "permission-shell", + PermissionDecision::Denied { + reason: Some("not allowed".to_string()), + }, + ) + .unwrap(); + + let Some(ProviderRunStep::DispatchTools(completed)) = run.next_step().unwrap() else { + panic!("expected completed tool batch"); + }; + let result = completed.calls[0].state.result().unwrap(); + assert_eq!(result.status, ToolResultStatus::Denied); + assert_eq!(result.content, "not allowed"); + assert!(completed.is_complete()); +} + +#[test] +fn stale_unknown_and_duplicate_tool_updates_do_not_mutate_the_batch() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn(vec![tool_call("read", "read_files")], &["read_files"]), + ); + let stale = ExternalWorkId { + run_id: batch.work_id.run_id.clone(), + epoch: RunEpoch::new(batch.work_id.epoch.get() + 1), + }; + let before = run.clone(); + assert!(matches!( + run.start_tool(&stale, "read").unwrap_err(), + ProviderRunProtocolError::WorkMismatch { .. } + )); + assert_eq!(run, before); + assert_eq!( + run.start_tool(&batch.work_id, "missing").unwrap_err(), + ProviderRunProtocolError::UnknownToolCall { + call_id: "missing".to_string(), + } + ); + assert_eq!(run, before); + + run.complete_tool(&batch.work_id, successful_result("read", "ok")) + .unwrap(); + let completed = run.clone(); + assert_eq!( + run.complete_tool(&batch.work_id, successful_result("read", "again")) + .unwrap_err(), + ProviderRunProtocolError::DuplicateToolUpdate { + call_id: "read".to_string(), + } + ); + assert_eq!(run, completed); +} + +#[test] +fn run_cancellation_preserves_completed_results_and_synthesizes_the_rest() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![tool_call("done", "read_files"), tool_call("active", "grep")], + &["read_files", "grep"], + ), + ); + run.complete_tool(&batch.work_id, successful_result("done", "contents")) + .unwrap(); + run.start_tool(&batch.work_id, "active").unwrap(); + + run.cancel("user cancelled").unwrap(); + + assert_eq!(run.epoch(), RunEpoch::new(2)); + let MessageContent::MultiPart(parts) = &run.transcript().last().unwrap().content else { + panic!("expected tool results"); + }; + assert!(matches!( + &parts[0], + ContentPart::ToolResult { + tool_use_id, + content, + is_error: false, + } if tool_use_id == "done" && content == "contents" + )); + assert!(matches!( + &parts[1], + ContentPart::ToolResult { + tool_use_id, + content, + is_error: true, + } if tool_use_id == "active" && content == "user cancelled" + )); + assert_eq!( + run.next_step().unwrap(), + Some(ProviderRunStep::Done(ProviderRunOutcome::Cancelled { + reason: "user cancelled".to_string(), + })) + ); +} + +#[test] +fn failure_while_tools_are_pending_records_correlated_errors_before_terminal_state() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn(vec![tool_call("project", "read_files")], &["read_files"]), + ); + assert_eq!(batch.work_id.epoch, RunEpoch::new(1)); + + run.fail( + ProviderRunFailureKind::Projection, + "proposal could not be projected", + ) + .unwrap(); + + assert_eq!(run.epoch(), RunEpoch::new(2)); + let MessageContent::MultiPart(parts) = &run.transcript().last().unwrap().content else { + panic!("expected synthesized tool result"); + }; + assert!(matches!( + &parts[0], + ContentPart::ToolResult { + tool_use_id, + is_error: true, + .. + } if tool_use_id == "project" + )); + let Some(ProviderRunStep::Done(ProviderRunOutcome::Failed(failure))) = run.next_step().unwrap() + else { + panic!("expected failed run"); + }; + assert_eq!(failure.kind, ProviderRunFailureKind::Projection); +} + +#[test] +fn turn_limit_cannot_finish_successfully_after_tools_require_another_model_turn() { + let mut run = run_with_limits(ProviderRunLimits { + max_model_turns: 1, + max_model_retries_per_turn: 0, + }); + let batch = accept_tool_turn( + &mut run, + tool_turn(vec![tool_call("read", "read_files")], &["read_files"]), + ); + run.complete_tool(&batch.work_id, successful_result("read", "ok")) + .unwrap(); + run.commit_tool_batch(&batch.work_id).unwrap(); + + let Some(ProviderRunStep::Done(ProviderRunOutcome::Failed(failure))) = run.next_step().unwrap() + else { + panic!("expected turn-limit failure"); + }; + assert_eq!(failure.kind, ProviderRunFailureKind::TurnLimitExceeded); +} + +#[test] +fn invalid_model_turn_does_not_commit_partial_content_usage_or_epoch() { + let mut run = run(); + let call = next_model_call(&mut run); + let duplicate = tool_call("duplicate", "read_files"); + let turn = tool_turn(vec![duplicate.clone(), duplicate], &["read_files"]); + let before = run.clone(); + + assert!(matches!( + run.accept_model_turn(&call.work_id, turn).unwrap_err(), + ProviderRunProtocolError::InvalidModelTurn { .. } + )); + assert_eq!(run, before); +} + +#[test] +fn every_nonterminal_phase_round_trips_through_json() { + let mut ready = run(); + assert_serialization_round_trip(&ready); + + let call = next_model_call(&mut ready); + assert_serialization_round_trip(&ready); + + let mut resolving_tools = ready.clone(); + resolving_tools + .accept_model_turn( + &call.work_id, + tool_turn(vec![tool_call("read", "read_files")], &["read_files"]), + ) + .unwrap(); + assert_eq!( + resolving_tools.state().phase(), + ProviderRunPhase::ResolvingModel + ); + assert_serialization_round_trip(&resolving_tools); + + let Some(ProviderRunStep::DispatchTools(_)) = resolving_tools.next_step().unwrap() else { + panic!("expected tool phase"); + }; + assert_serialization_round_trip(&resolving_tools); + + let mut resolving_text = ready; + resolving_text + .accept_model_turn(&call.work_id, text_turn("done")) + .unwrap(); + assert_serialization_round_trip(&resolving_text); + assert_eq!(resolving_text.next_step().unwrap(), None); + assert_eq!( + resolving_text.state().phase(), + ProviderRunPhase::AwaitingDriver + ); + assert_serialization_round_trip(&resolving_text); +} + +#[test] +fn restore_normalization_preserves_safe_nonterminal_states() { + let ready = run(); + let mut awaiting_model = ready.clone(); + let call = next_model_call(&mut awaiting_model); + + let mut resolving = awaiting_model.clone(); + resolving + .accept_model_turn(&call.work_id, text_turn("done")) + .unwrap(); + let mut awaiting_driver = resolving.clone(); + assert_eq!(awaiting_driver.next_step().unwrap(), None); + + for mut candidate in [ready, awaiting_model, resolving, awaiting_driver] { + let before = candidate.clone(); + assert_eq!( + candidate.normalize_after_restore().unwrap(), + ProviderRunRestoreNormalization::default() + ); + assert_eq!(candidate, before); + } +} + +#[test] +fn restore_normalization_reproposes_permissions_and_interrupts_unsafe_tools() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![ + tool_call("proposed", "read_files"), + tool_call("permission", "grep"), + tool_call("approved", "run_shell_command"), + tool_call("executing", "read_files"), + tool_call("resolved", "grep"), + ], + &["read_files", "grep", "run_shell_command"], + ), + ); + let permission = PermissionRequest { + id: "permission-1".to_string(), + call_id: "permission".to_string(), + kind: PermissionKind::Read, + reason: None, + }; + run.request_tool_permission(&batch.work_id, permission) + .unwrap(); + let approved = PermissionRequest { + id: "permission-2".to_string(), + call_id: "approved".to_string(), + kind: PermissionKind::Execute, + reason: None, + }; + run.request_tool_permission(&batch.work_id, approved) + .unwrap(); + run.resolve_tool_permission( + &batch.work_id, + "approved", + "permission-2", + PermissionDecision::AllowOnce, + ) + .unwrap(); + run.start_tool(&batch.work_id, "executing").unwrap(); + run.complete_tool( + &batch.work_id, + successful_result("resolved", "already finished"), + ) + .unwrap(); + + let normalization = run.normalize_after_restore().unwrap(); + + assert_eq!( + normalization.permission_call_ids_reset, + vec!["permission".to_string()] + ); + assert_eq!( + normalization.interrupted_call_ids, + vec!["approved".to_string(), "executing".to_string()] + ); + assert!(!normalization.committed_tool_batch); + let ProviderRunState::AwaitingTools { batch } = run.state() else { + panic!("partially resolved batch must remain pending"); + }; + assert!(matches!( + batch.calls[0].state, + PendingToolCallState::Proposed + )); + assert!(matches!( + batch.calls[1].state, + PendingToolCallState::Proposed + )); + for index in [2, 3] { + let result = batch.calls[index].state.result().unwrap(); + assert_eq!(result.status, ToolResultStatus::Error); + assert!(result.content.contains("was not replayed")); + } + assert_eq!( + batch.calls[4].state.result(), + Some(&successful_result("resolved", "already finished")) + ); +} + +#[test] +fn restore_normalization_commits_a_fully_resolved_batch() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![ + tool_call("executing", "read_files"), + tool_call("resolved", "grep"), + ], + &["read_files", "grep"], + ), + ); + run.start_tool(&batch.work_id, "executing").unwrap(); + run.complete_tool( + &batch.work_id, + successful_result("resolved", "already finished"), + ) + .unwrap(); + + let normalization = run.normalize_after_restore().unwrap(); + + assert!(normalization.committed_tool_batch); + assert_eq!( + normalization.interrupted_call_ids, + vec!["executing".to_string()] + ); + assert_eq!(run.state().phase(), ProviderRunPhase::ReadyToCallModel); + assert_eq!(run.epoch(), RunEpoch::new(2)); + let next = next_model_call(&mut run); + assert_eq!(next.work_id.epoch, RunEpoch::new(2)); + let MessageContent::MultiPart(parts) = &next.messages.last().unwrap().content else { + panic!("expected committed tool results"); + }; + assert!(matches!( + &parts[0], + ContentPart::ToolResult { + tool_use_id, + is_error: true, + .. + } if tool_use_id == "executing" + )); + assert!(matches!( + &parts[1], + ContentPart::ToolResult { + tool_use_id, + is_error: false, + .. + } if tool_use_id == "resolved" + )); +} diff --git a/crates/galaxy_agent_core/src/tool_policy.rs b/crates/galaxy_agent_core/src/tool_policy.rs index 76052944..61b520ad 100644 --- a/crates/galaxy_agent_core/src/tool_policy.rs +++ b/crates/galaxy_agent_core/src/tool_policy.rs @@ -95,8 +95,12 @@ impl ToolLoopGuard { impl ToolPolicy { pub fn new(tools: &[ToolDefinition]) -> Self { + Self::from_names(tools.iter().map(|tool| tool.name.clone())) + } + + pub fn from_names(names: impl IntoIterator) -> Self { Self { - advertised_tools: tools.iter().map(|tool| tool.name.clone()).collect(), + advertised_tools: names.into_iter().collect(), } } diff --git a/crates/persistence/src/model.rs b/crates/persistence/src/model.rs index 03a82ea1..1783ff7a 100644 --- a/crates/persistence/src/model.rs +++ b/crates/persistence/src/model.rs @@ -1079,6 +1079,9 @@ pub struct AcpConversationData { pub struct AgentConversationData { #[serde(default, skip_serializing_if = "AgentBackend::is_provider")] pub agent_backend: AgentBackend, + /// Versioned application-owned snapshot for an active direct-provider run. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub active_provider_run_json: Option, pub server_conversation_token: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub conversation_usage_metadata: Option, diff --git a/crates/persistence/src/model_tests.rs b/crates/persistence/src/model_tests.rs index 8b31b4c7..6adf00f4 100644 --- a/crates/persistence/src/model_tests.rs +++ b/crates/persistence/src/model_tests.rs @@ -137,6 +137,7 @@ fn is_restorable_accepts_empty_and_single_task_conversations() { fn agent_conversation_data_roundtrips_last_event_sequence() { let data = AgentConversationData { agent_backend: Default::default(), + active_provider_run_json: None, server_conversation_token: None, conversation_usage_metadata: None, reverted_action_ids: None, @@ -181,6 +182,31 @@ fn agent_conversation_data_defaults_legacy_rows_to_provider_backend() { assert_eq!(data.agent_backend, AgentBackend::Provider); } +#[test] +fn agent_conversation_data_defaults_legacy_rows_without_provider_run_snapshot() { + let data: AgentConversationData = serde_json::from_str(r#"{"server_conversation_token":null}"#) + .expect("legacy rows must deserialize"); + + assert_eq!(data.active_provider_run_json, None); +} + +#[test] +fn agent_conversation_data_roundtrips_provider_run_snapshot() { + let snapshot = r#"{"version":1,"run_id":"run-1"}"#; + let data = AgentConversationData { + active_provider_run_json: Some(snapshot.to_string()), + ..Default::default() + }; + + let json = serde_json::to_string(&data).expect("serialize"); + let roundtripped: AgentConversationData = serde_json::from_str(&json).expect("deserialize"); + + assert_eq!( + roundtripped.active_provider_run_json.as_deref(), + Some(snapshot) + ); +} + #[test] fn agent_conversation_data_roundtrips_acp_backend() { let data = AgentConversationData { @@ -214,6 +240,7 @@ fn agent_conversation_data_omits_default_provider_backend() { fn agent_conversation_data_roundtrips_remote_child_marker() { let data = AgentConversationData { agent_backend: Default::default(), + active_provider_run_json: None, server_conversation_token: None, conversation_usage_metadata: None, reverted_action_ids: None, @@ -241,6 +268,7 @@ fn agent_conversation_data_roundtrips_remote_child_marker() { fn agent_conversation_data_roundtrips_optimistic_root_marker() { let data = AgentConversationData { agent_backend: Default::default(), + active_provider_run_json: None, server_conversation_token: None, conversation_usage_metadata: None, reverted_action_ids: None, @@ -280,6 +308,7 @@ fn agent_conversation_data_deserializes_legacy_payload_without_last_event_sequen fn agent_conversation_data_skips_serializing_none_last_event_sequence() { let data = AgentConversationData { agent_backend: Default::default(), + active_provider_run_json: None, server_conversation_token: None, conversation_usage_metadata: None, reverted_action_ids: None, @@ -309,6 +338,7 @@ fn agent_conversation_data_skips_serializing_none_last_event_sequence() { fn agent_conversation_data_roundtrips_pinned() { let data = AgentConversationData { agent_backend: Default::default(), + active_provider_run_json: None, server_conversation_token: None, conversation_usage_metadata: None, reverted_action_ids: None, @@ -336,6 +366,7 @@ fn agent_conversation_data_roundtrips_pinned() { fn agent_conversation_data_skips_serializing_unpinned() { let data = AgentConversationData { agent_backend: Default::default(), + active_provider_run_json: None, server_conversation_token: None, conversation_usage_metadata: None, reverted_action_ids: None,