From 15cfad80fd2b9a82aea631b49623fccd201d3cd2 Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Tue, 28 Apr 2026 05:59:07 -0400 Subject: [PATCH 01/13] Wake up remote Claude Code agents on new events --- app/src/ai/agent/conversation_tests.rs | 11 + app/src/ai/agent_sdk/common.rs | 52 +- app/src/ai/agent_sdk/driver.rs | 87 ++- .../agent_sdk/driver/error_classification.rs | 21 +- .../driver/error_classification_tests.rs | 29 + .../agent_sdk/driver/harness/claude_code.rs | 558 +++++++++++++--- .../harness/claude_code/parent_bridge.rs | 92 ++- .../driver/harness/claude_code_tests.rs | 179 ++++++ app/src/ai/agent_sdk/driver/harness/gemini.rs | 16 +- app/src/ai/agent_sdk/driver/harness/mod.rs | 81 ++- app/src/ai/agent_sdk/driver/snapshot_tests.rs | 4 + app/src/ai/agent_sdk/mod.rs | 175 ++--- app/src/ai/blocklist/controller.rs | 283 ++++++++- app/src/ai/blocklist/history_model.rs | 4 - app/src/ai/blocklist/history_model_test.rs | 44 ++ .../blocklist/orchestration_event_poller.rs | 215 ++++++- .../orchestration_event_poller_tests.rs | 598 ++++++------------ app/src/ai/blocklist/orchestration_events.rs | 68 ++ .../blocklist/orchestration_events_tests.rs | 114 ++-- app/src/server/server_api.rs | 10 + app/src/server/server_api/ai.rs | 7 +- app/src/server/server_api/harness_support.rs | 153 ++++- app/src/terminal/view.rs | 3 + crates/http_client/src/lib.rs | 7 + crates/warp_cli/src/agent.rs | 6 +- crates/warp_cli/src/lib_tests.rs | 18 + 26 files changed, 1974 insertions(+), 861 deletions(-) diff --git a/app/src/ai/agent/conversation_tests.rs b/app/src/ai/agent/conversation_tests.rs index d330f264c..535286b3f 100644 --- a/app/src/ai/agent/conversation_tests.rs +++ b/app/src/ai/agent/conversation_tests.rs @@ -121,6 +121,17 @@ fn restored_conversation_defaults_autoexecute_override_when_not_persisted() { ); } +#[test] +fn restored_conversation_uses_persisted_last_event_sequence() { + let conversation_data: AgentConversationData = + serde_json::from_str(r#"{"server_conversation_token":null,"last_event_sequence":42}"#) + .unwrap(); + + let conversation = restored_conversation(Some(conversation_data)); + + assert_eq!(conversation.last_event_sequence(), Some(42)); +} + #[test] fn restored_conversation_defaults_unknown_persisted_autoexecute_override() { let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true); diff --git a/app/src/ai/agent_sdk/common.rs b/app/src/ai/agent_sdk/common.rs index 22ac02192..5683d8f16 100644 --- a/app/src/ai/agent_sdk/common.rs +++ b/app/src/ai/agent_sdk/common.rs @@ -5,14 +5,7 @@ use std::future::Future; use std::sync::Arc; use std::time::Duration; -use futures::TryFutureExt; -use inquire::{InquireError, Select}; -use warp_cli::agent::Harness; -use warp_cli::environment::{EnvironmentCreateArgs, EnvironmentUpdateArgs}; -use warpui::r#async::FutureExt; -use warpui::{AppContext, GetSingletonModelHandle, SingletonEntity as _, UpdateModel}; - -use crate::ai::agent::conversation::ServerAIConversationMetadata; +use crate::ai::agent::conversation::{AIAgentHarness, ServerAIConversationMetadata}; use crate::ai::agent_sdk::driver::{AgentDriverError, WARP_DRIVE_SYNC_TIMEOUT}; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::cloud_environments::CloudAmbientAgentEnvironment; @@ -25,6 +18,12 @@ use crate::server::server_api::ai::AIClient; use crate::server::server_api::ServerApiProvider; use crate::workspaces::update_manager::TeamUpdateManager; use crate::workspaces::user_workspaces::UserWorkspaces; +use futures::TryFutureExt; +use inquire::{InquireError, Select}; +use warp_cli::agent::Harness; +use warp_cli::environment::{EnvironmentCreateArgs, EnvironmentUpdateArgs}; +use warpui::r#async::FutureExt; +use warpui::{AppContext, GetSingletonModelHandle, SingletonEntity as _, UpdateModel}; /// How long to wait for workspace metadata to refresh. pub const WORKSPACE_METADATA_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); @@ -144,14 +143,8 @@ pub fn refresh_warp_drive( .map_err(|_| anyhow::anyhow!("Timed out waiting for Warp Drive to sync")) } -/// Fetch the conversation's server metadata and validate that its harness matches the caller's -/// `--harness` choice. Returns the metadata on success so the caller can reuse it (e.g. for the -/// server conversation token). -/// -/// Called up-front before any task/config-build logic consumes `args.harness`, so a mismatch -/// error surfaces before side effects like task creation. We deliberately do NOT auto-upgrade -/// the harness: `Harness::Oz` default with a Claude conversation id is treated as a mismatch -/// and errors out. +/// Fetch the conversation's server metadata and validate that its harness +/// matches the caller's `--harness` choice. pub(super) async fn fetch_and_validate_conversation_harness( ai_client: Arc, conversation_id: &str, @@ -160,7 +153,7 @@ pub(super) async fn fetch_and_validate_conversation_harness( let metadata = ai_client .list_ai_conversation_metadata(Some(vec![conversation_id.to_string()])) .await - .map_err(|e| AgentDriverError::ConversationLoadFailed(format!("{e:#}")))? + .map_err(|error| AgentDriverError::ConversationLoadFailed(format!("{error:#}")))? .into_iter() .next() .ok_or_else(|| { @@ -169,17 +162,36 @@ pub(super) async fn fetch_and_validate_conversation_harness( )) })?; - if metadata.harness != args_harness { + let expected = harness_label(metadata.harness); + let got = harness_label_from_cli(args_harness); + if expected != got { return Err(AgentDriverError::ConversationHarnessMismatch { conversation_id: conversation_id.to_string(), - expected: Harness::from(metadata.harness).to_string(), - got: args_harness.to_string(), + expected: expected.to_string(), + got: got.to_string(), }); } Ok(metadata) } +fn harness_label(harness: AIAgentHarness) -> &'static str { + match harness { + AIAgentHarness::Oz => "oz", + AIAgentHarness::ClaudeCode => "claude", + AIAgentHarness::Gemini => "gemini", + } +} + +fn harness_label_from_cli(harness: Harness) -> &'static str { + match harness { + Harness::Oz => "oz", + Harness::Claude => "claude", + Harness::OpenCode => "opencode", + Harness::Gemini => "gemini", + } +} + /// Format an object owner for display in the CLI. pub fn format_owner(owner: &Owner) -> &'static str { // TODO: For potentially-shared objects, consider looking up the particular user/team name. diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index 84ad49ab8..f792fcc0d 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -22,7 +22,8 @@ use crate::ai::skills::{SkillManager, SkillWatcher}; use crate::ai::{ agent::conversation::AIConversationId, agent_sdk::driver::harness::{ - task_env_vars, HarnessKind, HarnessRunner, ResumePayload, SavePoint, ThirdPartyHarness, + task_env_vars, HarnessCleanupDisposition, HarnessKind, HarnessRunner, ResumePayload, + SavePoint, ThirdPartyHarness, }, }; use crate::terminal::cli_agent_sessions::plugin_manager::{ @@ -196,16 +197,6 @@ impl IdleTimeoutSender { } } -/// How to resume an existing conversation when starting an agent run. -/// -/// The Oz harness restores the full conversation transcript into the terminal pane and treats -/// any new prompt as a follow-up; third-party harnesses round-trip a harness-specific payload -/// (see [`ResumePayload`]) instead. -pub enum ResumeOptions { - Oz(Box), - ThirdParty(Box), -} - /// Options for initializing the agent driver. pub struct AgentDriverOptions { /// Initial working directory for the agent's terminal session. @@ -220,10 +211,12 @@ pub struct AgentDriverOptions { pub should_share: bool, /// How long to keep the session alive after the agent run completes, if at all. pub idle_on_complete: Option, - /// If set, resume an existing conversation instead of starting fresh. The variant - /// determines which harness-specific path is taken (Oz transcript restore vs. - /// third-party-harness payload rehydration). - pub resume: Option, + /// Conversation to restore when creating the terminal. + /// Any ambient agent prompt will be sent as a follow-up to this conversation. + pub conversation_restoration: Option, + /// If set, resume an existing third-party-harness conversation instead of + /// starting fresh. + pub resume_payload: Option, /// Cloud providers to configure within the agent's session. pub cloud_providers: Vec>, /// Resolved environment configuration, if any. @@ -268,6 +261,10 @@ pub struct AgentDriver { // The conversation ID to continue (if provided). restored_conversation_id: Option, + /// If set, a third-party-harness conversation to resume. Consumed when + /// preparing the harness runner and cleared afterward. + resume_payload: Option, + /// Cloud providers set up within this driver session. cloud_providers: Vec>, @@ -278,11 +275,6 @@ pub struct AgentDriver { snapshot_disabled: bool, snapshot_upload_timeout: Duration, snapshot_script_timeout: Duration, - - /// If set, a third-party-harness conversation to resume. Consumed by `prepare_harness` - /// when building the runner and taken back to `None` after use so subsequent runs start - /// fresh. - resume_payload: Option, } pub(crate) enum SDKConversationOutputStatus { @@ -412,15 +404,6 @@ pub enum AgentDriverError { expected: String, got: String, }, - #[error( - "Task {task_id} was created with the {expected} harness, but --harness {got} was requested. \ - Re-run with --harness {expected} (or omit --harness to match) to continue this task." - )] - TaskHarnessMismatch { - task_id: String, - expected: String, - got: String, - }, #[error( "Conversation {conversation_id} has no stored transcript for the {harness} harness. \ The prior run may have crashed before saving any state." @@ -469,7 +452,8 @@ impl AgentDriver { should_share, idle_on_complete, secrets, - resume, + conversation_restoration, + resume_payload, cloud_providers, environment, selected_harness, @@ -478,15 +462,6 @@ impl AgentDriver { snapshot_script_timeout, } = options; - // Split the unified resume option into the two internal slots that the rest of - // the driver consumes: terminal-driven Oz transcript restoration vs. third-party - // harness payload rehydration. - let (conversation_restoration, resume_payload) = match resume { - Some(ResumeOptions::Oz(restoration)) => (Some(*restoration), None), - Some(ResumeOptions::ThirdParty(payload)) => (None, Some(*payload)), - None => (None, None), - }; - safe_info!( safe: ("Initializing agent driver: share={should_share}, idle_on_complete={idle_on_complete:?}"), full: ( @@ -625,6 +600,7 @@ impl AgentDriver { harness: None, idle_on_complete, restored_conversation_id, + resume_payload, cloud_providers, environment, snapshot_disabled: snapshot_disabled.unwrap_or(false), @@ -632,7 +608,6 @@ impl AgentDriver { .unwrap_or(snapshot::DEFAULT_SNAPSHOT_UPLOAD_TIMEOUT), snapshot_script_timeout: snapshot_script_timeout .unwrap_or(snapshot::DEFAULT_DECLARATIONS_SCRIPT_TIMEOUT), - resume_payload, }) } @@ -1520,11 +1495,6 @@ impl AgentDriver { .await .map_err(|_| AgentDriverError::InvalidRuntimeState)?; harness.prepare_environment_config(&working_dir, system_prompt.as_deref(), &secrets)?; - - // Pull the resume payload off the driver so the harness runner can rehydrate any - // existing session/conversation state before launching its CLI. The payload variant - // is harness-specific; harnesses match on their own [`ResumePayload`] variant and - // ignore others. let resume = foreground .spawn(|me, _| me.resume_payload.take()) .await @@ -1588,14 +1558,31 @@ impl AgentDriver { // Final save after the command finishes. log::debug!("Triggering final save of harness conversation data"); - report_if_error!(runner + let final_save_succeeded = match runner .save_conversation(SavePoint::Final, foreground) .await - .context("Failed to save harness conversation (final)")); - report_if_error!(runner - .cleanup(foreground) + .context("Failed to save harness conversation (final)") + { + Ok(()) => true, + Err(err) => { + report_error!(err); + false + } + }; + let cleanup_disposition = if final_save_succeeded + && matches!(command_result.as_ref(), Ok(exit_code) if exit_code.was_successful()) + { + HarnessCleanupDisposition::PreserveResumptionStateIfSupported + } else { + HarnessCleanupDisposition::DropResumptionState + }; + if let Err(err) = runner + .cleanup(cleanup_disposition, foreground) .await - .context("Failed to clean up harness runtime state")); + .context("Failed to clean up harness runtime state") + { + report_error!(err); + } let exit_code = command_result?; log::debug!("Agent harness exited with status {exit_code}"); diff --git a/app/src/ai/agent_sdk/driver/error_classification.rs b/app/src/ai/agent_sdk/driver/error_classification.rs index 47d0f23d1..3353936d7 100644 --- a/app/src/ai/agent_sdk/driver/error_classification.rs +++ b/app/src/ai/agent_sdk/driver/error_classification.rs @@ -248,7 +248,11 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS PlatformErrorCode::InternalError, ), ), - AgentDriverError::ConversationHarnessMismatch { conversation_id, expected, got } => ( + AgentDriverError::ConversationHarnessMismatch { + conversation_id, + expected, + got, + } => ( AgentTaskState::Failed, TaskStatusUpdate::with_error_code( format!( @@ -258,17 +262,10 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS PlatformErrorCode::EnvironmentSetupFailed, ), ), - AgentDriverError::TaskHarnessMismatch { task_id, expected, got } => ( - AgentTaskState::Failed, - TaskStatusUpdate::with_error_code( - format!( - "Task {task_id} was created with the {expected} harness, but --harness {got} was requested. \ - Re-run with --harness {expected} (or omit --harness) to continue this task." - ), - PlatformErrorCode::EnvironmentSetupFailed, - ), - ), - AgentDriverError::ConversationResumeStateMissing { harness, conversation_id } => ( + AgentDriverError::ConversationResumeStateMissing { + harness, + conversation_id, + } => ( AgentTaskState::Failed, TaskStatusUpdate::with_error_code( format!( diff --git a/app/src/ai/agent_sdk/driver/error_classification_tests.rs b/app/src/ai/agent_sdk/driver/error_classification_tests.rs index 2f2f85e19..1d0769d48 100644 --- a/app/src/ai/agent_sdk/driver/error_classification_tests.rs +++ b/app/src/ai/agent_sdk/driver/error_classification_tests.rs @@ -99,6 +99,35 @@ fn environment_not_found_is_failed_with_resource_not_found() { ); } +#[test] +fn conversation_harness_mismatch_is_failed_with_env_setup() { + let (state, update) = classify_driver_error(&AgentDriverError::ConversationHarnessMismatch { + conversation_id: "conv-123".into(), + expected: "claude".into(), + got: "oz".into(), + }); + assert_eq!(state, AgentTaskState::Failed); + assert_eq!( + update.error_code, + Some(PlatformErrorCode::EnvironmentSetupFailed) + ); + assert!(update.message.contains("conv-123")); + assert!(update.message.contains("--harness claude")); +} + +#[test] +fn conversation_resume_state_missing_is_failed_with_resource_not_found() { + let (state, update) = + classify_driver_error(&AgentDriverError::ConversationResumeStateMissing { + harness: "claude".into(), + conversation_id: "conv-123".into(), + }); + assert_eq!(state, AgentTaskState::Failed); + assert_eq!(update.error_code, Some(PlatformErrorCode::ResourceNotFound)); + assert!(update.message.contains("conv-123")); + assert!(update.message.contains("claude")); +} + // --- ShareSessionFailed variants --- #[test] diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code.rs b/app/src/ai/agent_sdk/driver/harness/claude_code.rs index b6761eb61..c56e0029f 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::fmt::Write as _; +use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -11,11 +12,15 @@ use serde_json::{Map, Value}; use tempfile::NamedTempFile; use uuid::Uuid; use warp_cli::agent::Harness; +use warp_core::safe_warn; use warpui::{ModelHandle, ModelSpawner}; use crate::ai::agent::conversation::AIConversationId; +use crate::ai::agent_events::MessageHydrator; use crate::ai::ambient_agents::AmbientAgentTaskId; -use crate::server::server_api::harness_support::{upload_to_target, HarnessSupportClient}; +use crate::server::server_api::harness_support::{ + upload_to_target, HarnessSupportClient, ResolvePromptRequest, +}; use crate::server::server_api::ServerApi; use crate::terminal::model::block::BlockId; use crate::terminal::model::session::ExecuteCommandOptions; @@ -23,31 +28,155 @@ use crate::terminal::CLIAgent; use super::super::terminal::{CommandHandle, TerminalDriver}; use super::super::{AgentDriver, AgentDriverError}; -use super::claude_transcript::{ - claude_config_dir, read_envelope, write_envelope, write_session_index_entry, ClaudeResumeInfo, - ClaudeTranscriptEnvelope, -}; use super::json_utils::{read_json_file_or_default, write_json_file}; use super::{ - write_temp_file, HarnessRunner, ManagedSecretValue, ResumePayload, SavePoint, ThirdPartyHarness, + cli_agent_session_status, write_temp_file, HarnessCleanupDisposition, HarnessRunner, + ManagedSecretValue, ResumePayload, SavePoint, ThirdPartyHarness, }; mod parent_bridge; #[cfg(test)] use super::super::OZ_MESSAGE_LISTENER_STATE_ROOT_ENV; -use parent_bridge::MessageBridge; -#[cfg(test)] use parent_bridge::{ acknowledge_parent_bridge_hook_output, ensure_parent_bridge_state_dir, - parent_bridge_char_count, parent_bridge_hook_output_ack_file, parent_bridge_hook_output_file, - parent_bridge_root, parent_bridge_staged_message_path, parent_bridge_surfaced_message_path, - prepare_parent_bridge_hook_output, render_parent_bridge_message_block, - stage_parent_bridge_message, MessageBridgeHookOutput, MessageBridgeMessageRecord, + parent_bridge_max_context_chars, parent_bridge_root, prepare_parent_bridge_hook_output, + stage_parent_bridge_message, MessageBridge, MessageBridgeCleanupDisposition, + MessageBridgeMessageRecord, +}; +#[cfg(test)] +use parent_bridge::{ + parent_bridge_char_count, parent_bridge_event_cursor_file, parent_bridge_hook_output_ack_file, + parent_bridge_hook_output_file, parent_bridge_staged_message_path, + parent_bridge_surfaced_message_path, read_parent_bridge_event_cursor, + render_parent_bridge_message_block, write_parent_bridge_event_cursor, MessageBridgeHookOutput, MESSAGE_BRIDGE_CONTEXT_PREAMBLE, }; +#[derive(Debug)] +pub(crate) struct ClaudeResumeInfo { + pub(crate) conversation_id: AIConversationId, + pub(crate) session_id: Uuid, + pub(crate) envelope: ClaudeTranscriptEnvelope, +} + +#[derive(Debug, Clone)] +pub(crate) struct ClaudeWakeMessage { + pub(crate) sequence: i64, + pub(crate) message_id: String, + pub(crate) sender_run_id: String, + pub(crate) subject: String, + pub(crate) body: String, + pub(crate) occurred_at: String, +} + +impl From for MessageBridgeMessageRecord { + fn from(value: ClaudeWakeMessage) -> Self { + Self { + sequence: value.sequence, + message_id: value.message_id, + sender_run_id: value.sender_run_id, + subject: value.subject, + body: value.body, + occurred_at: value.occurred_at, + } + } +} + +#[derive(Debug)] +pub(crate) struct ClaudeWakeRemoteContext { + session_id: Uuid, + envelope: ClaudeTranscriptEnvelope, + wake_prompt: String, +} + pub(crate) struct ClaudeHarness; +impl ClaudeHarness { + pub(crate) async fn fetch_local_wake_remote_context( + task_id: AmbientAgentTaskId, + server_api: Arc, + ) -> Result { + let resolved = server_api + .resolve_prompt_for_task( + &task_id, + ResolvePromptRequest { + skill: None, + attachments_dir: None, + }, + ) + .await + .with_context(|| format!("Failed to resolve Claude wake prompt for task {task_id}"))?; + let bytes = server_api + .fetch_transcript_for_task(&task_id) + .await + .with_context(|| format!("Failed to fetch Claude transcript for task {task_id}"))?; + let envelope: ClaudeTranscriptEnvelope = + serde_json::from_slice(&bytes).with_context(|| { + format!("Failed to deserialize Claude transcript for wake task {task_id}") + })?; + let wake_prompt = match resolved.resumption_prompt { + Some(resumption_prompt) if !resumption_prompt.is_empty() => { + format!( + "{resumption_prompt} + +{CLAUDE_WAKE_PROMPT}" + ) + } + _ => CLAUDE_WAKE_PROMPT.to_string(), + }; + Ok(ClaudeWakeRemoteContext { + session_id: envelope.uuid, + envelope, + wake_prompt, + }) + } + + pub(crate) async fn prepare_local_wake_command( + server_api: Arc, + working_dir: Option, + mut remote: ClaudeWakeRemoteContext, + pending_messages: Vec, + ) -> Result { + let working_dir = working_dir.unwrap_or_else(|| remote.envelope.cwd.clone()); + prepare_claude_environment_config(&working_dir, &HashMap::new()) + .context("Failed to prepare Claude environment for wake")?; + + remote.envelope.cwd = working_dir.clone(); + let config_root = claude_config_dir().context("Failed to resolve Claude config dir")?; + write_envelope(&remote.envelope, &config_root) + .context("Failed to rehydrate Claude transcript for wake")?; + if let Err(error) = write_session_index_entry(remote.session_id, &working_dir, &config_root) + { + log::warn!("Failed to update Claude sessions-index.json for wake: {error:#}"); + } + + let state_dir = parent_bridge_root()?.join(remote.session_id.to_string()); + ensure_parent_bridge_state_dir(&state_dir)?; + let hydrator = MessageHydrator::new(server_api); + acknowledge_parent_bridge_hook_output(&hydrator, &state_dir).await?; + for record in pending_messages + .into_iter() + .map(MessageBridgeMessageRecord::from) + { + stage_parent_bridge_message(&state_dir, &record)?; + } + prepare_parent_bridge_hook_output(&hydrator, &state_dir, parent_bridge_max_context_chars()) + .await?; + + let prompt_path = state_dir.join(CLAUDE_WAKE_PROMPT_FILE_NAME); + std::fs::write(&prompt_path, remote.wake_prompt.as_bytes()) + .with_context(|| format!("Failed to write {}", prompt_path.display()))?; + + Ok(claude_command( + CLIAgent::Claude.command_prefix(), + &remote.session_id, + &prompt_path.display().to_string(), + None, + true, + )) + } +} + #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] impl ThirdPartyHarness for ClaudeHarness { @@ -77,10 +206,6 @@ impl ThirdPartyHarness for ClaudeHarness { }) } - /// Fetch the Claude Code transcript for the current task's conversation and wrap it - /// into a [`ResumePayload::Claude`]. Maps a server 404 to - /// [`AgentDriverError::ConversationResumeStateMissing`] tagged as the `claude` harness - /// so the user sees a resume-specific error rather than a generic load failure. async fn fetch_resume_payload( &self, conversation_id: &AIConversationId, @@ -91,8 +216,6 @@ impl ThirdPartyHarness for ClaudeHarness { .fetch_transcript() .await .map_err(|err| { - // A 404 from the server maps to "no stored transcript" so the CLI can tell - // the user the prior run never saved state. let message = format!("{err:#}").to_lowercase(); if message.contains("status 404") { AgentDriverError::ConversationResumeStateMissing { @@ -108,10 +231,9 @@ impl ThirdPartyHarness for ClaudeHarness { "Failed to deserialize Claude transcript for {conversation_id_str}: {err:#}" )) })?; - let session_id = envelope.uuid; Ok(Some(ResumePayload::Claude(ClaudeResumeInfo { conversation_id: *conversation_id, - session_id, + session_id: envelope.uuid, envelope, }))) } @@ -127,14 +249,9 @@ impl ThirdPartyHarness for ClaudeHarness { terminal_driver: ModelHandle, resume: Option, ) -> Result, AgentDriverError> { - // Extract the Claude variant; any other variant is ignored since it belongs to a - // different harness. Today there are no other variants, but this keeps the shape - // ready for future CLI-specific payloads. let claude_resume = resume.map(|payload| match payload { ResumePayload::Claude(info) => info, }); - // Claude treats the user-turn message as immediate intent, so the resumption preamble - // is most reliable when prepended directly to the prompt that gets piped into the CLI. let owned_prompt = match resumption_prompt { Some(preamble) if !preamble.is_empty() => format!("{preamble}\n\n{prompt}"), _ => prompt.to_string(), @@ -156,14 +273,12 @@ impl ThirdPartyHarness for ClaudeHarness { const CLAUDE_CODE_FORMAT: &str = "claude_code_cli"; /// Command used to exit claude. const CLAUDE_EXIT_COMMAND: &str = "/exit"; +const CLAUDE_WAKE_PROMPT: &str = + "New lead-agent messages are available. Read the latest lead-agent updates and continue the task accordingly."; +const CLAUDE_WAKE_PROMPT_FILE_NAME: &str = "wake-turn-prompt.txt"; /// Build the shell command that launches the Claude CLI for a given session and /// prompt file. -/// -/// When `resuming` is true we pass `--resume ` so Claude picks up the -/// existing on-disk session; otherwise we pass `--session-id ` to pin a -/// fresh session to that id. If `system_prompt_path` is provided, the CLI is -/// told to append its contents to the base system prompt. fn claude_command( cli_name: &str, session_id: &Uuid, @@ -207,14 +322,10 @@ struct ClaudeHarnessRunner { parent_bridge: Option, /// Lazily cached output of `claude --version`. claude_version: Mutex>, - /// When resuming an existing conversation, we pin the runner's server conversation id - /// up front instead of calling `create_external_conversation` in [`HarnessRunner::start`]. - /// Subsequent saves overwrite the same GCS objects keyed by this id. preexisting_conversation_id: Option, } impl ClaudeHarnessRunner { - #[allow(clippy::too_many_arguments)] fn new( cli_command: &str, prompt: &str, @@ -229,32 +340,26 @@ impl ClaudeHarnessRunner { // avoiding shell-quoting issues with complex content (e.g. skill instructions). let temp_file = write_temp_file("oz_prompt_", prompt)?; let prompt_path = temp_file.path().display().to_string(); - let (session_id, preexisting_conversation_id, resuming) = match resume { Some(ClaudeResumeInfo { conversation_id, session_id, mut envelope, }) => { - // Rehydrate the stored envelope under the current working directory so - // `claude --resume ` finds the jsonl under ~/.claude/projects//. - // The original envelope's cwd usually points at the cloud sandbox path, which - // doesn't exist locally. envelope.cwd = working_dir.to_path_buf(); - let config_root = claude_config_dir().map_err(|e| { + let config_root = claude_config_dir().map_err(|error| { AgentDriverError::ConfigBuildFailed( - e.context("Failed to resolve Claude config dir"), + error.context("Failed to resolve Claude config dir"), ) })?; - write_envelope(&envelope, &config_root).map_err(|e| { + write_envelope(&envelope, &config_root).map_err(|error| { AgentDriverError::ConfigBuildFailed( - e.context("Failed to rehydrate Claude transcript"), + error.context("Failed to rehydrate Claude transcript"), ) })?; - // Index write is best-effort: upstream Claude versions vary in how they use - // `sessions-index.json`, so losing the index entry shouldn't abort the run. - if let Err(e) = write_session_index_entry(session_id, working_dir, &config_root) { - log::warn!("Failed to update Claude sessions-index.json: {e:#}"); + if let Err(error) = write_session_index_entry(session_id, working_dir, &config_root) + { + log::warn!("Failed to update Claude sessions-index.json: {error:#}"); } (session_id, Some(conversation_id), true) } @@ -362,9 +467,33 @@ impl ClaudeHarnessRunner { .await } - fn cleanup_parent_bridge(&self) -> Result<()> { + async fn should_preserve_parent_bridge( + &self, + cleanup_disposition: HarnessCleanupDisposition, + foreground: &ModelSpawner, + ) -> bool { + if !matches!( + cleanup_disposition, + HarnessCleanupDisposition::PreserveResumptionStateIfSupported + ) { + return false; + } + + !matches!( + cli_agent_session_status(&self.terminal_driver, foreground).await, + Some(crate::terminal::cli_agent_sessions::CLIAgentSessionStatus::Blocked { .. }) + | Some(crate::terminal::cli_agent_sessions::CLIAgentSessionStatus::InProgress) + ) + } + + fn cleanup_parent_bridge(&self, preserve_state: bool) -> Result<()> { if let Some(parent_bridge) = self.parent_bridge.as_ref() { - parent_bridge.cleanup()?; + let cleanup_disposition = if preserve_state { + MessageBridgeCleanupDisposition::PreserveState + } else { + MessageBridgeCleanupDisposition::RemoveState + }; + parent_bridge.cleanup(cleanup_disposition)?; } Ok(()) } @@ -377,28 +506,23 @@ impl HarnessRunner for ClaudeHarnessRunner { &self, foreground: &ModelSpawner, ) -> Result { - // When resuming, we already have a server conversation id from the prior run. - // Otherwise create a fresh external conversation record for this run. - // TODO(REMOTE-1149): `create_external_conversation` currently won't work for local CLI - // runs. We should either support it or have a fallback. let conversation_id = match self.preexisting_conversation_id { - Some(id) => { - log::info!("Resuming external conversation {id}"); - id - } - None => { - let id = self - .client - .create_external_conversation(CLAUDE_CODE_FORMAT) - .await - .map_err(|e| { - log::error!("Failed to create external conversation: {e}"); - AgentDriverError::ConfigBuildFailed(e) - })?; - log::info!("Created external conversation {id}"); - id + Some(conversation_id) => { + log::info!("Resuming external conversation {conversation_id}"); + conversation_id } + None => self + .client + .create_external_conversation(CLAUDE_CODE_FORMAT) + .await + .map_err(|e| { + log::error!("Failed to create external conversation: {e}"); + AgentDriverError::ConfigBuildFailed(e) + })?, }; + if self.preexisting_conversation_id.is_none() { + log::info!("Created external conversation {conversation_id}"); + } self.start_parent_bridge(foreground) .await .map_err(AgentDriverError::ConfigBuildFailed)?; @@ -414,7 +538,7 @@ impl HarnessRunner for ClaudeHarnessRunner { { Ok(command_handle) => command_handle, Err(err) => { - self.cleanup_parent_bridge() + self.cleanup_parent_bridge(false) .map_err(AgentDriverError::ConfigBuildFailed)?; return Err(err); } @@ -494,9 +618,16 @@ impl HarnessRunner for ClaudeHarnessRunner { Ok(()) } - async fn cleanup(&self, _foreground: &ModelSpawner) -> Result<()> { + async fn cleanup( + &self, + cleanup_disposition: HarnessCleanupDisposition, + foreground: &ModelSpawner, + ) -> Result<()> { self.flush_parent_bridge_acks().await?; - self.cleanup_parent_bridge() + let preserve_state = self + .should_preserve_parent_bridge(cleanup_disposition, foreground) + .await; + self.cleanup_parent_bridge(preserve_state) } } @@ -527,6 +658,283 @@ async fn upload_transcript( upload_to_target(client.http_client(), &target, body).await } +// ─── Transcript envelope ────────────────────────────────────────────────────── + +/// JSON envelope sent to the server representing a complete Claude Code session. +/// +/// Bundles the main session transcript, any subagent transcripts, and +/// per-agent TODO lists assembled from the Claude state directory. +#[derive(Debug, PartialEq, Serialize, Deserialize)] +pub(crate) struct ClaudeTranscriptEnvelope { + /// The directory that the Claude Code session started in. + cwd: PathBuf, + /// Unique session identifier. + uuid: Uuid, + /// Claude Code version, if available. + #[serde(default, skip_serializing_if = "Option::is_none")] + claude_version: Option, + /// List of messages in the main agent conversation. + entries: Vec, + /// Messages in each subagent conversation, keyed by the agent filename (e.g. `"agent-aac0b7f3db6bccfaf"`). + subagents: HashMap>, + /// TODO lists for each agent, keyed on the session and agent (e.g. `"-agent-"`). + todos: HashMap, +} + +/// Encode a filesystem path as a Claude config directory name, matching the +/// Claude CLI convention of replacing every `/` with `-`. +/// +/// Example: `/Users/ben/src/foo` → `-Users-ben-src-foo` +fn encode_cwd(cwd: &Path) -> String { + cwd.to_string_lossy().replace(['/', '.'], "-") +} + +/// Resolve the Claude config directory. +/// +/// Reads `$CLAUDE_CONFIG_DIR` if set, otherwise falls back to `~/.claude`. +// +/// TODO(REMOTE-1209): Use the transcript path reported by our hook. +fn claude_config_dir() -> Result { + if let Ok(dir) = std::env::var("CLAUDE_CONFIG_DIR") { + return Ok(PathBuf::from(dir)); + } + dirs::home_dir() + .map(|h| h.join(".claude")) + .ok_or_else(|| anyhow::anyhow!("could not determine home directory")) +} + +/// Assemble a [`ClaudeTranscriptEnvelope`] from the Claude config directory. +/// +/// Reads: +/// - `/projects//.jsonl` - main transcript +/// - `/projects///subagents/*.jsonl` - subagents +/// - `/todos/-agent-*.json` - per-agent todo lists +/// +/// If the main JSONL does not exist yet (e.g. during an early periodic save) +/// the envelope is returned with an empty `entries` list rather than an error. +fn read_envelope( + session_uuid: Uuid, + cwd: &Path, + config_root: &Path, +) -> Result { + let encoded = encode_cwd(cwd); + let projects_dir = config_root.join("projects").join(&encoded); + + // Main session transcript. + let session_file = projects_dir.join(format!("{session_uuid}.jsonl")); + let entries = read_jsonl(&session_file)?; + + // Subagents are stored in a directory named after the session UUID. + let mut subagents: HashMap> = HashMap::new(); + let subagents_dir = projects_dir + .join(session_uuid.to_string()) + .join("subagents"); + if subagents_dir.is_dir() { + for entry in std::fs::read_dir(&subagents_dir) + .with_context(|| format!("Failed to read subagents dir {}", subagents_dir.display()))? + { + let entry = entry?; + let path = entry.path(); + if path.extension().and_then(|e| e.to_str()) != Some("jsonl") { + continue; + } + let Some(stem) = path.file_stem().and_then(|s| s.to_str()) else { + continue; + }; + subagents.insert(stem.to_owned(), read_jsonl(&path)?); + } + } + + // Per-agent todo lists. + let mut todos: HashMap = HashMap::new(); + let todos_dir = config_root.join("todos"); + let todos_prefix = format!("{session_uuid}-agent-"); + if todos_dir.is_dir() { + for entry in std::fs::read_dir(&todos_dir) + .with_context(|| format!("Failed to read todos dir {}", todos_dir.display()))? + { + let entry = entry?; + let path = entry.path(); + if path.extension().and_then(|e| e.to_str()) != Some("json") { + continue; + } + let Some(stem) = path.file_stem().and_then(|s| s.to_str()) else { + continue; + }; + if !stem.starts_with(&todos_prefix) { + continue; + } + match std::fs::read_to_string(&path) { + Ok(content) => match serde_json::from_str(&content) { + Ok(value) => { + todos.insert(stem.to_owned(), value); + } + Err(e) => log::warn!("Failed to parse todos file {}: {e}", path.display()), + }, + Err(e) => log::warn!("Failed to read todos file {}: {e}", path.display()), + } + } + } + + Ok(ClaudeTranscriptEnvelope { + cwd: cwd.to_path_buf(), + uuid: session_uuid, + claude_version: None, + entries, + subagents, + todos, + }) +} + +/// Write a [`ClaudeTranscriptEnvelope`] back to disk using the same layout +/// that Claude Code uses. +/// +/// Creates: +/// - `/projects//.jsonl` - main transcript +/// - `/projects///subagents/.jsonl` - subagents +/// - `/todos/.json` - per-agent todo lists +fn write_envelope(envelope: &ClaudeTranscriptEnvelope, config_root: &Path) -> Result<()> { + let encoded = encode_cwd(&envelope.cwd); + let projects_dir = config_root.join("projects").join(&encoded); + std::fs::create_dir_all(&projects_dir) + .with_context(|| format!("Failed to create {}", projects_dir.display()))?; + + // Main session JSONL. + let session_file = projects_dir.join(format!("{}.jsonl", envelope.uuid)); + std::fs::write(&session_file, entries_to_jsonl(&envelope.entries)?) + .with_context(|| format!("Failed to write {}", session_file.display()))?; + + // Subagent JSONLs. + if !envelope.subagents.is_empty() { + let subagents_dir = projects_dir + .join(envelope.uuid.to_string()) + .join("subagents"); + std::fs::create_dir_all(&subagents_dir) + .with_context(|| format!("Failed to create {}", subagents_dir.display()))?; + for (stem, entries) in &envelope.subagents { + let path = subagents_dir.join(format!("{stem}.jsonl")); + std::fs::write(&path, entries_to_jsonl(entries)?) + .with_context(|| format!("Failed to write {}", path.display()))?; + } + } + + // Per-agent todo lists. + if !envelope.todos.is_empty() { + let todos_dir = config_root.join("todos"); + std::fs::create_dir_all(&todos_dir) + .with_context(|| format!("Failed to create {}", todos_dir.display()))?; + for (stem, value) in &envelope.todos { + let path = todos_dir.join(format!("{stem}.json")); + std::fs::write(&path, serde_json::to_vec(value)?) + .with_context(|| format!("Failed to write {}", path.display()))?; + } + } + + Ok(()) +} + +/// Filename of Claude's global session index. +const SESSIONS_INDEX_FILENAME: &str = "sessions-index.json"; + +/// Upsert an entry for `session_uuid` into `/sessions-index.json` so Claude's +/// `claude --resume ` lookup can find the rehydrated jsonl. +/// +/// Best-effort: callers should log a warning on failure rather than aborting the run. +fn write_session_index_entry(session_uuid: Uuid, cwd: &Path, config_root: &Path) -> Result<()> { + let index_path = config_root.join(SESSIONS_INDEX_FILENAME); + + let mut index: serde_json::Map = match std::fs::read_to_string(&index_path) { + Ok(content) => match serde_json::from_str::(&content) { + Ok(Value::Object(map)) => map, + Ok(_) => { + safe_warn!( + safe: ("sessions-index.json is not a JSON object; overwriting"), + full: ("sessions-index.json at {} is not a JSON object; overwriting", index_path.display()) + ); + serde_json::Map::new() + } + Err(error) => { + safe_warn!( + safe: ("Failed to parse sessions-index.json; overwriting"), + full: ("Failed to parse sessions-index.json at {}: {error}; overwriting", index_path.display()) + ); + serde_json::Map::new() + } + }, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => serde_json::Map::new(), + Err(error) => { + return Err(anyhow::Error::from(error) + .context(format!("Failed to read {}", index_path.display()))); + } + }; + + let encoded = encode_cwd(cwd); + let transcript_path = format!("projects/{encoded}/{session_uuid}.jsonl"); + let entry = serde_json::json!({ + "sessionId": session_uuid.to_string(), + "cwd": cwd.to_string_lossy(), + "projectPath": encoded, + "transcriptPath": transcript_path, + }); + index.insert(session_uuid.to_string(), entry); + + if let Some(parent) = index_path.parent() { + std::fs::create_dir_all(parent) + .with_context(|| format!("Failed to create {}", parent.display()))?; + } + std::fs::write( + &index_path, + serde_json::to_vec_pretty(&Value::Object(index)) + .context("Failed to serialize sessions-index.json")?, + ) + .with_context(|| format!("Failed to write {}", index_path.display()))?; + Ok(()) +} +/// Serialize a slice of JSON values as a JSONL byte string (one value per line). +fn entries_to_jsonl(entries: &[Value]) -> Result> { + let mut buf = Vec::new(); + for entry in entries { + serde_json::to_writer(&mut buf, entry)?; + buf.push(b'\n'); + } + Ok(buf) +} + +/// Read a JSONL file, returning one parsed [`Value`] per non-blank line. +/// +/// Lines that fail to parse as JSON are skipped with a warning rather than +/// causing the entire read to fail. A missing file returns an empty [`Vec`]. +fn read_jsonl(path: &Path) -> Result> { + let file = match std::fs::File::open(path) { + Ok(f) => f, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(e) => { + return Err( + anyhow::Error::from(e).context(format!("Failed to open {}", path.display())) + ); + } + }; + let reader = BufReader::new(file); + let mut entries = Vec::new(); + for line in reader.lines() { + let line = line.with_context(|| format!("Failed to read line from {}", path.display()))?; + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + match serde_json::from_str(trimmed) { + Ok(value) => entries.push(value), + Err(e) => { + safe_warn!( + safe: ("Skipping malformed JSONL entry"), + full: ("Skipping malformed JSONL entry in {}: {e}", path.display()) + ); + } + } + } + Ok(entries) +} + fn prepare_claude_environment_config( working_dir: &Path, secrets: &HashMap, @@ -586,7 +994,7 @@ const CLAUDE_JSON_FILE_NAME: &str = ".claude.json"; const CLAUDE_SETTINGS_FILE_NAME: &str = "settings.json"; const ANTHROPIC_API_KEY_SUFFIX_LEN: usize = 20; -#[derive(Default, Deserialize, Serialize, Debug)] +#[derive(Default, Deserialize, Serialize)] #[serde(rename_all = "camelCase")] struct ClaudeConfig { #[serde(default)] @@ -601,7 +1009,7 @@ struct ClaudeConfig { extra: Map, } -#[derive(Default, Deserialize, Serialize, Debug)] +#[derive(Default, Deserialize, Serialize)] #[serde(rename_all = "camelCase")] struct CustomApiKeyResponses { #[serde(default)] @@ -610,7 +1018,7 @@ struct CustomApiKeyResponses { extra: Map, } -#[derive(Default, Deserialize, Serialize, Debug)] +#[derive(Default, Deserialize, Serialize)] #[serde(rename_all = "camelCase")] struct ClaudeProjectConfig { #[serde(default)] @@ -619,7 +1027,7 @@ struct ClaudeProjectConfig { extra: Map, } -#[derive(Default, Deserialize, Serialize, Debug)] +#[derive(Default, Deserialize, Serialize)] #[serde(rename_all = "camelCase")] struct ClaudeSettings { #[serde(default)] diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs b/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs index fe05e8771..6f20b2176 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs @@ -28,12 +28,14 @@ use crate::ai::agent_events::{ AgentEventDriverConfig, MessageHydrator, ServerApiAgentEventSource, }; use crate::ai::agent_sdk::driver::{AgentDriver, OZ_MESSAGE_LISTENER_STATE_ROOT_ENV}; +use crate::server::server_api::ai::AIClient; use crate::server::server_api::ai::AgentRunEvent; use crate::server::server_api::ServerApi; const LEGACY_MESSAGE_LISTENER_STATE_ROOT_ENV: &str = "OZ_PARENT_STATE_ROOT"; const PARENT_BRIDGE_DEFAULT_STATE_ROOT: &str = ".claude-code/oz-parent-bridge"; const PARENT_BRIDGE_SURFACED_DIR_NAME: &str = "surfaced"; +const PARENT_BRIDGE_EVENT_CURSOR_FILE_NAME: &str = "event-cursor.json"; const PARENT_BRIDGE_HOOK_OUTPUT_FILE_NAME: &str = "pending-hook-output.json"; const PARENT_BRIDGE_HOOK_OUTPUT_ACK_FILE_NAME: &str = "pending-hook-output.ack"; const PARENT_BRIDGE_MAX_CONTEXT_CHARS_ENV: &str = "OZ_PARENT_MAX_CONTEXT_CHARS"; @@ -48,6 +50,16 @@ pub(super) struct MessageBridge { runtime: Mutex>, state_lock: AsyncMutex<()>, } + +pub(super) enum MessageBridgeCleanupDisposition { + RemoveState, + PreserveState, +} + +#[derive(Debug, PartialEq, Serialize, Deserialize)] +struct MessageBridgeEventCursor { + since_sequence: i64, +} struct MessageBridgeRuntime { task: SpawnedFutureHandle, } @@ -55,6 +67,7 @@ struct MessageBridgeRuntime { struct MessageBridgeEventConsumer { run_id: String, state_dir: PathBuf, + server_api: Arc, } #[cfg_attr(target_family = "wasm", async_trait(?Send))] @@ -91,6 +104,13 @@ impl AgentEventConsumer for MessageBridgeEventConsumer { Ok(AgentEventConsumerControlFlow::Continue) } + + async fn persist_cursor(&mut self, sequence: i64) -> anyhow::Result<()> { + write_parent_bridge_event_cursor(&self.state_dir, sequence)?; + self.server_api + .update_event_sequence_on_server(&self.run_id, sequence) + .await + } } #[derive(Clone, Debug, Serialize, Deserialize)] @@ -199,10 +219,13 @@ impl MessageBridge { acknowledge_parent_bridge_hook_output(&hydrator, &self.state_dir).await } - pub(super) fn cleanup(&self) -> Result<()> { + pub(super) fn cleanup(&self, disposition: MessageBridgeCleanupDisposition) -> Result<()> { if let Some(runtime) = self.runtime.lock().take() { runtime.task.abort(); } + if matches!(disposition, MessageBridgeCleanupDisposition::PreserveState) { + return Ok(()); + } match fs::remove_dir_all(&self.state_dir) { Ok(()) => {} Err(err) if err.kind() == std::io::ErrorKind::NotFound => {} @@ -249,6 +272,10 @@ pub(super) fn parent_bridge_hook_output_ack_file(state_dir: &Path) -> PathBuf { state_dir.join(PARENT_BRIDGE_HOOK_OUTPUT_ACK_FILE_NAME) } +pub(super) fn parent_bridge_event_cursor_file(state_dir: &Path) -> PathBuf { + state_dir.join(PARENT_BRIDGE_EVENT_CURSOR_FILE_NAME) +} + fn parent_bridge_message_path(dir: &Path, sequence: i64, message_id: &str) -> PathBuf { dir.join(format!("{sequence:020}-{message_id}.json")) } @@ -277,6 +304,28 @@ pub(super) fn ensure_parent_bridge_state_dir(state_dir: &Path) -> Result<()> { Ok(()) } +pub(super) fn read_parent_bridge_event_cursor(state_dir: &Path) -> Result { + let path = parent_bridge_event_cursor_file(state_dir); + if !path.exists() { + return Ok(0); + } + + let cursor = serde_json::from_slice::( + &fs::read(&path).with_context(|| format!("Failed to read {}", path.display()))?, + ) + .with_context(|| format!("Failed to parse {}", path.display()))?; + Ok(cursor.since_sequence) +} + +pub(super) fn write_parent_bridge_event_cursor(state_dir: &Path, sequence: i64) -> Result<()> { + write_parent_bridge_json_atomically( + &parent_bridge_event_cursor_file(state_dir), + &MessageBridgeEventCursor { + since_sequence: sequence, + }, + ) +} + pub(super) fn stage_parent_bridge_message( state_dir: &Path, record: &MessageBridgeMessageRecord, @@ -288,7 +337,7 @@ pub(super) fn stage_parent_bridge_message( Ok(()) } -fn parent_bridge_max_context_chars() -> usize { +pub(super) fn parent_bridge_max_context_chars() -> usize { std::env::var(PARENT_BRIDGE_MAX_CONTEXT_CHARS_ENV) .ok() .and_then(|value| value.trim().parse::().ok()) @@ -579,15 +628,44 @@ async fn run_parent_bridge_forever( state_dir: PathBuf, ) -> Result<()> { ensure_parent_bridge_state_dir(&state_dir)?; + let since_sequence = + read_parent_bridge_resume_cursor(server_api.as_ref(), &run_id, &state_dir).await?; // The shared driver keeps `since_sequence` in memory across its own retry - // loop, which is all this per-session bridge needs because the state dir is - // not reused across sessions. - let config = AgentEventDriverConfig::retry_forever(vec![run_id.clone()], 0); - let source = ServerApiAgentEventSource::new(server_api); - let mut consumer = MessageBridgeEventConsumer { run_id, state_dir }; + // loop and we also persist it inside the session state dir so dormant runs + // can resume without replaying already handled events. + let config = AgentEventDriverConfig::retry_forever(vec![run_id.clone()], since_sequence); + let source = ServerApiAgentEventSource::new(server_api.clone()); + let mut consumer = MessageBridgeEventConsumer { + run_id, + state_dir, + server_api, + }; run_agent_event_driver(source, config, &mut consumer).await } +async fn read_parent_bridge_resume_cursor( + server_api: &ServerApi, + run_id: &str, + state_dir: &Path, +) -> Result { + let local_sequence = read_parent_bridge_event_cursor(state_dir)?; + let Ok(task_id) = run_id.parse() else { + return Ok(local_sequence); + }; + + let server_sequence = match server_api.get_ambient_agent_task(&task_id).await { + Ok(task) => task.last_event_sequence.unwrap_or(0), + Err(err) => { + log::warn!( + "Failed to read server-backed event cursor for Claude message bridge run {run_id}: {err:#}" + ); + 0 + } + }; + + Ok(local_sequence.max(server_sequence)) +} + fn write_parent_bridge_json_atomically(path: &Path, value: &T) -> Result<()> { write_parent_bridge_bytes_atomically(path, &serde_json::to_vec(value)?) } diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs b/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs index 71342cf28..5775e9870 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs @@ -10,6 +10,7 @@ use uuid::Uuid; use super::*; use crate::ai::agent_events::MessageHydrator; use crate::server::server_api::ai::{MockAIClient, ReadAgentMessageResponse}; +use crate::server::server_api::ServerApiProvider; fn sample_parent_bridge_message( sequence: i64, @@ -91,6 +92,32 @@ fn claude_command_pipes_prompt_path() { ); } +#[test] +fn write_session_index_entry_creates_expected_entry() { + let tmp = TempDir::new().unwrap(); + let cwd = Path::new("/my/project"); + let session_id = Uuid::new_v4(); + + write_session_index_entry(session_id, cwd, tmp.path()).unwrap(); + + let index_path = tmp.path().join("sessions-index.json"); + let index: Value = serde_json::from_slice(&fs::read(index_path).unwrap()).unwrap(); + let session_key = session_id.to_string(); + let entry = &index[&session_key]; + let encoded = encode_cwd(cwd); + + assert_eq!(entry["sessionId"], Value::String(session_key.clone())); + assert_eq!( + entry["cwd"], + Value::String(cwd.to_string_lossy().into_owned()) + ); + assert_eq!(entry["projectPath"], Value::String(encoded.clone())); + assert_eq!( + entry["transcriptPath"], + Value::String(format!("projects/{encoded}/{session_id}.jsonl")) + ); +} + #[test] #[serial_test::serial] fn parent_bridge_root_prefers_environment_override() { @@ -119,6 +146,77 @@ fn stage_parent_bridge_message_writes_message_record() { assert!(staged_record.sender_run_id.is_empty()); } +#[tokio::test] +async fn parent_bridge_event_cursor_defaults_to_zero_when_missing() { + let tmp = TempDir::new().unwrap(); + let state_dir = tmp.path().join("session-123"); + ensure_parent_bridge_state_dir(&state_dir).unwrap(); + + assert_eq!(read_parent_bridge_event_cursor(&state_dir).unwrap(), 0); + assert!(!parent_bridge_event_cursor_file(&state_dir).exists()); +} + +#[tokio::test] +async fn parent_bridge_event_cursor_round_trips() { + let tmp = TempDir::new().unwrap(); + let state_dir = tmp.path().join("session-123"); + ensure_parent_bridge_state_dir(&state_dir).unwrap(); + + write_parent_bridge_event_cursor(&state_dir, 42).unwrap(); + + assert_eq!(read_parent_bridge_event_cursor(&state_dir).unwrap(), 42); + assert!(parent_bridge_event_cursor_file(&state_dir).exists()); +} + +#[test] +#[serial_test::serial] +fn message_bridge_cleanup_preserves_state_for_wakeable_runs() { + let tmp = TempDir::new().unwrap(); + std::env::set_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV, tmp.path()); + + let session_id = Uuid::new_v4(); + let bridge = MessageBridge::new("run-123".to_string(), session_id).unwrap(); + let state_dir = tmp.path().join(session_id.to_string()); + ensure_parent_bridge_state_dir(&state_dir).unwrap(); + let record = sample_staged_parent_bridge_message(42, "msg-123"); + stage_parent_bridge_message(&state_dir, &record).unwrap(); + write_parent_bridge_event_cursor(&state_dir, 42).unwrap(); + + bridge + .cleanup(MessageBridgeCleanupDisposition::PreserveState) + .unwrap(); + std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV); + + assert!(state_dir.exists()); + assert!(parent_bridge_staged_message_path(&state_dir, 42, "msg-123").exists()); + assert_eq!(read_parent_bridge_event_cursor(&state_dir).unwrap(), 42); +} + +#[test] +#[serial_test::serial] +fn message_bridge_cleanup_removes_state_for_non_wakeable_runs() { + let tmp = TempDir::new().unwrap(); + std::env::set_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV, tmp.path()); + + let session_id = Uuid::new_v4(); + let bridge = MessageBridge::new("run-123".to_string(), session_id).unwrap(); + let state_dir = tmp.path().join(session_id.to_string()); + ensure_parent_bridge_state_dir(&state_dir).unwrap(); + stage_parent_bridge_message( + &state_dir, + &sample_staged_parent_bridge_message(42, "msg-123"), + ) + .unwrap(); + write_parent_bridge_event_cursor(&state_dir, 42).unwrap(); + + bridge + .cleanup(MessageBridgeCleanupDisposition::RemoveState) + .unwrap(); + std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV); + + assert!(!state_dir.exists()); +} + #[tokio::test] async fn prepare_parent_bridge_hook_output_moves_selected_messages_to_surfaced_dir() { let tmp = TempDir::new().unwrap(); @@ -543,6 +641,87 @@ fn resolve_suffix_returns_none_for_short_key() { assert_eq!(resolve_anthropic_api_key_suffix(&secrets), None); } +#[test] +#[serial_test::serial] +fn prepare_local_wake_command_rehydrates_transcript_and_stages_messages() { + let home_dir = TempDir::new().unwrap(); + let claude_config_dir = TempDir::new().unwrap(); + let bridge_state_root = TempDir::new().unwrap(); + let working_dir = home_dir.path().join("workspace/project"); + fs::create_dir_all(&working_dir).unwrap(); + + std::env::set_var("HOME", home_dir.path()); + std::env::set_var("CLAUDE_CONFIG_DIR", claude_config_dir.path()); + std::env::set_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV, bridge_state_root.path()); + + let session_id = Uuid::new_v4(); + let remote = ClaudeWakeRemoteContext { + session_id, + envelope: ClaudeTranscriptEnvelope { + cwd: Path::new("/stale/cwd").to_path_buf(), + uuid: session_id, + claude_version: None, + entries: vec![serde_json::json!({"type": "assistant", "text": "done"})], + subagents: HashMap::new(), + todos: HashMap::new(), + }, + wake_prompt: "resume prompt\n\nwake prompt".to_string(), + }; + let message = ClaudeWakeMessage { + sequence: 42, + message_id: "message-1".to_string(), + sender_run_id: "parent-run-123".to_string(), + subject: "Please pivot".to_string(), + body: "Inspect the failing tests first.".to_string(), + occurred_at: "2026-04-17T15:46:00Z".to_string(), + }; + + let command = futures::executor::block_on(ClaudeHarness::prepare_local_wake_command( + ServerApiProvider::new_for_test().get(), + Some(working_dir.clone()), + remote, + vec![message.clone()], + )) + .unwrap(); + + let state_dir = bridge_state_root.path().join(session_id.to_string()); + let prompt_path = state_dir.join(CLAUDE_WAKE_PROMPT_FILE_NAME); + let surfaced_path = + parent_bridge_surfaced_message_path(&state_dir, message.sequence, &message.message_id); + + assert!(command.contains("--resume")); + assert!(command.contains(&session_id.to_string())); + assert!(command.contains(CLAUDE_WAKE_PROMPT_FILE_NAME)); + assert_eq!( + fs::read_to_string(&prompt_path).unwrap(), + "resume prompt\n\nwake prompt" + ); + assert!(surfaced_path.exists()); + assert!(parent_bridge_hook_output_file(&state_dir).exists()); + let surfaced_record: MessageBridgeMessageRecord = + serde_json::from_slice(&fs::read(&surfaced_path).unwrap()).unwrap(); + assert_eq!(surfaced_record.sender_run_id, message.sender_run_id); + assert_eq!(surfaced_record.subject, message.subject); + assert_eq!(surfaced_record.body, message.body); + + let restored_envelope = + read_envelope(session_id, &working_dir, claude_config_dir.path()).unwrap(); + assert_eq!(restored_envelope.cwd, working_dir); + assert_eq!( + restored_envelope.entries, + vec![serde_json::json!({"type": "assistant", "text": "done"})] + ); + assert!(home_dir.path().join(".claude.json").exists()); + assert!(claude_config_dir + .path() + .join(CLAUDE_SETTINGS_FILE_NAME) + .exists()); + + std::env::remove_var("HOME"); + std::env::remove_var("CLAUDE_CONFIG_DIR"); + std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV); +} + #[test] fn resolve_suffix_returns_none_for_short_anthropic_api_key() { let secrets = HashMap::from([( diff --git a/app/src/ai/agent_sdk/driver/harness/gemini.rs b/app/src/ai/agent_sdk/driver/harness/gemini.rs index c3ecda57e..88a1a0021 100644 --- a/app/src/ai/agent_sdk/driver/harness/gemini.rs +++ b/app/src/ai/agent_sdk/driver/harness/gemini.rs @@ -22,7 +22,10 @@ use crate::terminal::CLIAgent; use super::super::terminal::{CommandHandle, TerminalDriver}; use super::super::{AgentDriver, AgentDriverError}; use super::json_utils::{read_json_file_or_default, write_json_file}; -use super::{write_temp_file, HarnessRunner, ResumePayload, SavePoint, ThirdPartyHarness}; +use super::{ + write_temp_file, HarnessCleanupDisposition, HarnessRunner, ResumePayload, SavePoint, + ThirdPartyHarness, +}; pub(crate) struct GeminiHarness; @@ -71,9 +74,6 @@ impl ThirdPartyHarness for GeminiHarness { terminal_driver: ModelHandle, _resume: Option, ) -> Result, AgentDriverError> { - // Gemini does not support conversation resume yet. When it does, it will add its - // own `ResumePayload::Gemini(..)` variant and override `fetch_resume_payload`, - // and decide how to surface the user-turn resumption preamble. let client: Arc = server_api; Ok(Box::new(GeminiHarnessRunner::new( self.cli_agent().command_prefix(), @@ -215,6 +215,14 @@ impl HarnessRunner for GeminiHarnessRunner { ) .await } + + async fn cleanup( + &self, + _cleanup_disposition: HarnessCleanupDisposition, + _foreground: &ModelSpawner, + ) -> Result<()> { + Ok(()) + } } fn prepare_gemini_environment_config( diff --git a/app/src/ai/agent_sdk/driver/harness/mod.rs b/app/src/ai/agent_sdk/driver/harness/mod.rs index 34816a361..787f9e3e4 100644 --- a/app/src/ai/agent_sdk/driver/harness/mod.rs +++ b/app/src/ai/agent_sdk/driver/harness/mod.rs @@ -35,21 +35,15 @@ use super::{ }; mod claude_code; -pub(crate) mod claude_transcript; mod gemini; mod json_utils; -pub(crate) use claude_code::ClaudeHarness; -use claude_transcript::ClaudeResumeInfo; +use claude_code::ClaudeResumeInfo; +pub(crate) use claude_code::{ClaudeHarness, ClaudeWakeMessage, ClaudeWakeRemoteContext}; use gemini::GeminiHarness; /// Harness-agnostic payload describing how to resume an existing conversation. -/// -/// Each variant carries the data a specific harness needs to rehydrate state before its CLI -/// launches. Harnesses match on the variant they produce and ignore others; new CLIs that -/// want resume support add a new variant and override [`ThirdPartyHarness::fetch_resume_payload`]. pub(crate) enum ResumePayload { - /// Claude Code session state fetched from the server's transcript endpoint. Claude(ClaudeResumeInfo), } @@ -88,15 +82,6 @@ pub(crate) trait ThirdPartyHarness: Send + Sync { } /// Fetch the harness-specific resume payload for an existing conversation. - /// - /// The driver calls this when the user passes `--conversation ` and the harness - /// matches the stored conversation's harness. Harnesses that don't support resume - /// use the default impl, which returns `Ok(None)` and causes the run to start fresh. - /// - /// Implementations download the raw transcript via [`HarnessSupportClient::fetch_transcript`] - /// (which derives the conversation from the current task's `agent_conversation_id`) and - /// own all harness-specific deserialization and error mapping (e.g. a 404 maps to - /// [`AgentDriverError::ConversationResumeStateMissing`] tagged with the harness label). async fn fetch_resume_payload( &self, _conversation_id: &AIConversationId, @@ -106,15 +91,6 @@ pub(crate) trait ThirdPartyHarness: Send + Sync { } /// Build a runner for executing this harness with the given prompt. - /// - /// If `resume` is `Some`, the harness matches on its own [`ResumePayload`] variant and - /// reuses the stored session/conversation ids instead of minting fresh ones. Variants - /// belonging to other harnesses are ignored. - /// - /// `resumption_prompt`, when non-empty, is a short user-turn preamble the server emits - /// during a resumed session. Each harness decides exactly how to surface it (e.g. Claude - /// prepends it to the user-turn prompt that gets piped into the CLI). Harnesses that - /// don't yet support resumption can ignore it. #[allow(clippy::too_many_arguments)] fn build_runner( &self, @@ -158,16 +134,12 @@ impl fmt::Debug for HarnessKind { } /// Build a [`HarnessKind`] for the given [`Harness`]. -/// -/// We shouldn't ever get a `--harness unknown` here because clap should handle -/// it. -pub(crate) fn harness_kind(harness: Harness) -> Result { +pub(crate) fn harness_kind(harness: Harness) -> HarnessKind { match harness { - Harness::Oz => Ok(HarnessKind::Oz), - Harness::Claude => Ok(HarnessKind::ThirdParty(Box::new(ClaudeHarness))), - Harness::OpenCode => Ok(HarnessKind::Unsupported(Harness::OpenCode)), - Harness::Gemini => Ok(HarnessKind::ThirdParty(Box::new(GeminiHarness))), - Harness::Unknown => Err(AgentDriverError::InvalidRuntimeState), + Harness::Oz => HarnessKind::Oz, + Harness::Claude => HarnessKind::ThirdParty(Box::new(ClaudeHarness)), + Harness::OpenCode => HarnessKind::Unsupported(Harness::OpenCode), + Harness::Gemini => HarnessKind::ThirdParty(Box::new(GeminiHarness)), } } @@ -322,6 +294,18 @@ pub(crate) enum SavePoint { PostTurn, } +/// Controls how much harness-owned state should survive cleanup after the CLI +/// exits. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum HarnessCleanupDisposition { + /// Tear down all harness-owned resume and wake state. + DropResumptionState, + /// The harness exited cleanly and its final save completed, so wake/resume + /// state may be preserved if the harness-specific runtime also considers + /// the run complete. + PreserveResumptionStateIfSupported, +} + /// Stateful per-run representation of an external harness produced /// by [`ThirdPartyHarness::build_runner`]. /// @@ -359,7 +343,11 @@ pub(crate) trait HarnessRunner: Send + Sync { } /// Clean up any harness-owned background state after the harness exits. - async fn cleanup(&self, _foreground: &ModelSpawner) -> Result<()> { + async fn cleanup( + &self, + _cleanup_disposition: HarnessCleanupDisposition, + _foreground: &ModelSpawner, + ) -> Result<()> { Ok(()) } } @@ -370,20 +358,29 @@ pub(crate) async fn has_running_cli_agent( terminal_driver: &ModelHandle, foreground: &ModelSpawner, ) -> bool { + matches!( + cli_agent_session_status(terminal_driver, foreground).await, + Some(CLIAgentSessionStatus::InProgress) + ) +} + +/// Returns the tracked CLI agent session status for the terminal, if any. +pub(crate) async fn cli_agent_session_status( + terminal_driver: &ModelHandle, + foreground: &ModelSpawner, +) -> Option { let driver = terminal_driver.clone(); - let Ok(running) = foreground + foreground .spawn(move |_, ctx| { let terminal_view_id = driver.as_ref(ctx).terminal_view().id(); CLIAgentSessionsModel::handle(ctx) .as_ref(ctx) .session(terminal_view_id) - .is_some_and(|s| s.status == CLIAgentSessionStatus::InProgress) + .map(|session| session.status.clone()) }) .await - else { - return false; - }; - running + .ok() + .flatten() } /// Create a [`NamedTempFile`] with the given prefix and write `content` into it. diff --git a/app/src/ai/agent_sdk/driver/snapshot_tests.rs b/app/src/ai/agent_sdk/driver/snapshot_tests.rs index 322fb6b44..31c0cb4d5 100644 --- a/app/src/ai/agent_sdk/driver/snapshot_tests.rs +++ b/app/src/ai/agent_sdk/driver/snapshot_tests.rs @@ -86,6 +86,10 @@ impl HarnessSupportClient for TestClient { unimplemented!("not used by upload_snapshot_from_declarations_file") } + async fn fetch_transcript(&self) -> Result { + unimplemented!("not used by upload_snapshot_from_declarations_file") + } + async fn get_block_snapshot_upload_target( &self, _conversation_id: &AIConversationId, diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index 391e3c2de..6b9c34711 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -68,7 +68,8 @@ use crate::ai::skills::{ }; pub(crate) use driver::harness::{ - task_env_vars, validate_cli_installed, ClaudeHarness, ThirdPartyHarness, + task_env_vars, validate_cli_installed, ClaudeHarness, ClaudeWakeMessage, + ClaudeWakeRemoteContext, ThirdPartyHarness, }; pub use driver::AgentDriver; use telemetry::CliTelemetryEvent; @@ -400,7 +401,7 @@ fn build_merged_config_and_task( model: model_override, profile: args.profile.clone(), mcp_specs: runtime_mcp_specs, - harness: harness_kind(args.harness)?, + harness: harness_kind(args.harness), }; Ok((merged_config, task)) @@ -460,7 +461,7 @@ fn build_server_side_task( model: model_override, profile, mcp_specs: runtime_mcp_specs, - harness: harness_kind(args.harness)?, + harness: harness_kind(args.harness), }; Ok((config, task)) @@ -567,43 +568,47 @@ impl AgentDriverRunner { // Set up and run the driver, reporting any errors back to the server. let result: Result<(), AgentDriverError> = async { - // Pull relevant variables out of args before moving it into the closure. + // Pull relevant variables out of args before moving it into the + // build_driver_options_and_task future. let share_requests = args.share.share.clone(); let bedrock_inference_role = args.bedrock_inference_role.clone(); let args_harness = args.harness; - // `--conversation` path (user-invoked local resume): validate before any task side - // effects so mismatches fail fast. The `--task-id` path derives its conversation id - // from the server-side task metadata inside `build_driver_options_and_task`. Both - // can currently be passed together (the worker server-side appends `--conversation` - // alongside `--task-id` for Slack/Linear followups); when both are set, the explicit - // `--conversation` value wins via the merge below. - if let Some(conversation_id) = args.conversation.as_deref() { - common::fetch_and_validate_conversation_harness( + let resume_context = if let Some(conversation_id) = args.conversation.clone() { + let metadata = common::fetch_and_validate_conversation_harness( server_api.clone(), - conversation_id, + &conversation_id, args_harness, ) .await?; - } - let resume_conversation_id = args.conversation.clone(); + Some((conversation_id, metadata)) + } else { + None + }; - // Build driver options and task, handling task creation or existing task setup. - // For the `--task-id` path, `task_conversation_id` is the `conversation_id` read off - // the fetched `AmbientAgentTask` (set by the server when linking the task to an - // existing conversation, e.g. via `run-cloud --conversation`). - let (mut driver_options, task, task_conversation_id) = + let (mut driver_options, task, environment_id, task_conversation_id) = Self::build_driver_options_and_task(&foreground, args, &server_api).await?; // Update the effective task ID so errors are reported correctly. // This only matters if we created a task ID locally. task_id = driver_options.task_id.or(task_id); + let resume_context = match (resume_context, task_conversation_id) { + (Some(context), _) => Some(context), + (None, Some(conversation_id)) => { + let metadata = common::fetch_and_validate_conversation_harness( + server_api.clone(), + &conversation_id, + args_harness, + ) + .await?; + Some((conversation_id, metadata)) + } + (None, None) => None, + }; - // The `--task-id` branch already validated `args_harness` against the task's harness - // setting inside `build_driver_options_and_task`; the conversation that the task spawned - // necessarily uses the same harness, so no extra conversation-metadata roundtrip is - // needed here. Just merge the task's linked conversation id into the resume target. - let resume_conversation_id = resume_conversation_id.or(task_conversation_id); + // Resolve the environment. We make sure this happens after resolving the task ID + // so that errors are reported. + Self::resolve_environment(&foreground, environment_id, &mut driver_options).await?; let bedrock_task_id = driver_options.task_id.map(|id| id.to_string()); @@ -653,11 +658,13 @@ impl AgentDriverRunner { } // Pull conversation information, if we have it - if let Some(conversation_id) = resume_conversation_id { - driver_options.resume = Self::load_conversation_information( + if let Some((conversation_id, resume_metadata)) = resume_context { + Self::load_conversation_information( &foreground, conversation_id, + resume_metadata, &task.harness, + &mut driver_options, ) .await?; } @@ -748,15 +755,13 @@ impl AgentDriverRunner { /// Build the AgentDriverOptions and Task, handling task creation or existing task setup. /// - /// The third tuple element is the conversation id read off the server-side task metadata - /// on the `--task-id` branch. It's `None` when no task id was passed or when the task is - /// not linked to a conversation; callers use it to drive `--task-id`-implied resume - /// without requiring the caller to also pass `--conversation`. + /// Returns the driver options, the task, the unresolved environment ID (if any), and any + /// conversation id read off server-side task metadata. async fn build_driver_options_and_task( foreground: &ModelSpawner, args: RunAgentArgs, server_api: &Arc, - ) -> Result<(AgentDriverOptions, Task, Option), AgentDriverError> { + ) -> Result<(AgentDriverOptions, Task, Option, Option), AgentDriverError> { // Get the working directory let working_dir = match args.cwd.as_ref() { Some(dir) => dunce::canonicalize(dir) @@ -791,7 +796,8 @@ impl AgentDriverRunner { should_share, idle_on_complete: args.idle_on_complete.map(|d| d.into()), secrets: Default::default(), - resume: None, + conversation_restoration: None, + resume_payload: None, cloud_providers: Vec::new(), environment: None, selected_harness: args.harness, @@ -813,9 +819,7 @@ impl AgentDriverRunner { let environment_id = merged_config.environment_id.clone(); - // Handle secrets/attachments fetch (existing task) or task creation (new run). - // The existing-task branch also surfaces the task's `conversation_id` (if any) so - // the caller can wire up resume without a separate `--conversation` arg. + // Handle secrets/attachments fetch (existing task) or task creation (new run) let task_conversation_id = if let Some(task_id_str) = task_id_str { Self::fetch_secrets_and_attachments( foreground, @@ -848,10 +852,7 @@ impl AgentDriverRunner { None }; - // Resolve environment and cloud providers. - Self::resolve_environment(foreground, environment_id, &mut driver_options).await?; - - Ok((driver_options, task, task_conversation_id)) + Ok((driver_options, task, environment_id, task_conversation_id)) } /// Creates a new task on the server for this agent run, sets the task ID on the driver @@ -905,10 +906,6 @@ impl AgentDriverRunner { /// When starting an agent run from an existing task_id, fetch secrets, task metadata, /// and task attachments (images and files) from the server and update the driver options. - /// - /// Returns the task's `conversation_id` when the server has linked the task to an existing - /// AI conversation (e.g. a `run-cloud --conversation` spawn). The caller uses this to drive - /// transcript rehydration without a separate `--conversation` CLI arg. async fn fetch_secrets_and_attachments( foreground: &ModelSpawner, task_id_str: String, @@ -1028,42 +1025,15 @@ impl AgentDriverRunner { } } }; - let (parent_run_id, task_conversation_id, task_harness) = match task_metadata_result { - Ok(Some(task_metadata)) => { - // The task's harness is stored on the snapshot; if absent, it's the default Oz. - let task_harness = task_metadata - .agent_config_snapshot - .as_ref() - .and_then(|c| c.harness.as_ref()) - .map(|h| h.harness_type) - .unwrap_or(Harness::Oz); - ( - task_metadata.parent_run_id, - task_metadata.conversation_id, - Some(task_harness), - ) - } - Ok(None) => (None, None, None), + let (parent_run_id, task_conversation_id) = match task_metadata_result { + Ok(Some(task_metadata)) => (task_metadata.parent_run_id, task_metadata.conversation_id), + Ok(None) => (None, None), Err(err) => { log::warn!("Failed to fetch task metadata: {err:#}"); - (None, None, None) + (None, None) } }; - // Validate the requested `--harness` against the task's harness setting. This avoids the - // extra conversation-metadata roundtrip that would otherwise be needed downstream when the - // task is linked to an existing conversation, since task harness and conversation harness - // always match (the task spawned the conversation). - if let Some(task_harness) = task_harness { - if task_harness != driver_options.selected_harness { - return Err(AgentDriverError::TaskHarnessMismatch { - task_id: task_id_str, - expected: task_harness.to_string(), - got: driver_options.selected_harness.to_string(), - }); - } - } - // Set the task ID on the ServerApi so it's sent with all subsequent requests. foreground .spawn(move |_, ctx| { @@ -1091,24 +1061,17 @@ impl AgentDriverRunner { } /// If we are starting this agent run from an existing conversation, load the conversation - /// data from the server and return the harness-specific [`ResumeOptions`] payload that the - /// caller plugs onto [`AgentDriverOptions::resume`]. - /// - /// `harness` is the resolved harness from the task config (already validated against the - /// conversation's metadata up-front by [`common::fetch_and_validate_conversation_harness`]). - /// - /// For the Oz harness, fetches the full conversation and returns a [`driver::ResumeOptions::Oz`]. - /// For third-party harnesses, delegates to [`ThirdPartyHarness::fetch_resume_payload`] and - /// wraps the returned payload (if any) in [`driver::ResumeOptions::ThirdParty`]; each harness - /// owns its server call and error mapping. Returns `None` if a third-party harness has no - /// resume payload to surface. + /// data from the server and set the relevant driver options. async fn load_conversation_information( foreground: &ModelSpawner, conversation_id: String, + resume_metadata: crate::ai::agent::conversation::ServerAIConversationMetadata, harness: &HarnessKind, - ) -> Result, AgentDriverError> { + driver_options: &mut AgentDriverOptions, + ) -> Result<(), AgentDriverError> { match harness { HarnessKind::Oz => { + let _ = resume_metadata; let server_api = foreground .spawn(|_, ctx| { ServerApiProvider::handle(ctx) @@ -1134,33 +1097,35 @@ impl AgentDriverRunner { ) })?; - Ok(Some(driver::ResumeOptions::Oz(Box::new( - ConversationRestorationInNewPaneType::Historical { + driver_options.conversation_restoration = + Some(ConversationRestorationInNewPaneType::Historical { conversation, should_use_live_appearance: false, ambient_agent_task_id: None, - }, - )))) + }); } - HarnessKind::ThirdParty(h) => { + HarnessKind::ThirdParty(harness) => { + let _ = resume_metadata; let harness_support_client = foreground .spawn(|_, ctx| ServerApiProvider::as_ref(ctx).get_harness_support_client()) .await?; - let resume_conversation_id = AIConversationId::try_from(conversation_id.clone()) + let resume_conversation_id = AIConversationId::try_from(conversation_id) .map_err(|err| AgentDriverError::ConversationLoadFailed(format!("{err:#}")))?; - Ok( - h.fetch_resume_payload(&resume_conversation_id, harness_support_client) - .await? - .map(|payload| driver::ResumeOptions::ThirdParty(Box::new(payload))), - ) + driver_options.resume_payload = harness + .fetch_resume_payload(&resume_conversation_id, harness_support_client) + .await?; + } + HarnessKind::Unsupported(harness) => { + return Err(AgentDriverError::HarnessSetupFailed { + harness: harness.to_string(), + reason: format!( + "The {harness} harness is only supported for local child agent launches." + ), + }); } - HarnessKind::Unsupported(harness) => Err(AgentDriverError::HarnessSetupFailed { - harness: harness.to_string(), - reason: format!( - "The {harness} harness is only supported for local child agent launches." - ), - }), } + + Ok(()) } /// Resolve the environment and store into `driver_options`. @@ -1389,7 +1354,7 @@ fn resolve_orchestration_harness_label() -> &'static str { Some(Harness::Claude) => "claude", Some(Harness::OpenCode) => "opencode", Some(Harness::Gemini) => "gemini", - Some(Harness::Unknown) | None => "unknown", + _ => "unknown", } } diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 75809b953..8934869f2 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -30,7 +30,7 @@ use crate::ai::agent::{ PassiveSuggestionTriggerType, RunningCommand, }; use crate::ai::agent::{DocumentContentAttachmentSource, FileContext}; -use crate::ai::ambient_agents::AmbientAgentTaskId; +use crate::ai::ambient_agents::{AmbientAgentTaskId, AmbientAgentTaskState}; use crate::ai::document::ai_document_model::{ AIDocumentId, AIDocumentModel, AIDocumentUserEditStatus, }; @@ -42,6 +42,7 @@ use crate::ai::{ FinishedAIAgentOutput, RenderableAIError, RequestCost, RequestMetadata, StaticQueryType, UserQueryMode, }, + agent_sdk::{ClaudeHarness, ClaudeWakeMessage, ClaudeWakeRemoteContext}, llms::LLMPreferences, AIRequestUsageModel, }; @@ -52,7 +53,8 @@ use crate::network::NetworkStatus; use crate::notebooks::editor::model::FileLinkResolutionContext; use crate::persistence::ModelEvent; use crate::search::slash_command_menu::static_commands::commands; -use crate::server::server_api::AIApiError; +use crate::server::server_api::ai::AIClient; +use crate::server::server_api::{AIApiError, ServerApiProvider}; use crate::terminal::model::block::{ formatted_terminal_contents_for_input, BlockId, CURSOR_MARKER, }; @@ -72,13 +74,17 @@ use parking_lot::FairMutex; use pending_response_streams::PendingResponseStreams; use session_sharing_protocol::common::ParticipantId; use std::collections::{HashMap, HashSet}; +use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; +use warp_cli::agent::Harness; use warp_core::assertions::safe_assert; use warp_multi_agent_api::{message, Task, ToolType}; use warpui::r#async::{SpawnedFutureHandle, Timer}; -use super::orchestration_events::{OrchestrationEventService, OrchestrationEventServiceEvent}; +use super::orchestration_events::{ + OrchestrationEventService, OrchestrationEventServiceEvent, PendingEvent, PendingEventDetail, +}; use warpui::{AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; #[derive(Debug, Clone)] @@ -163,6 +169,10 @@ pub enum BlocklistAIControllerEvent { filename: Option, }, + ExecuteLocalHarnessCommand { + command: String, + }, + FreeTierLimitCheckTriggered, } @@ -313,6 +323,8 @@ pub struct BlocklistAIController { /// 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. + pending_local_claude_wakes: HashMap, /// Passive conversations explicitly requested to follow up after actions complete. pending_passive_follow_ups: HashSet, /// Passive suggestion results that should be included with the next request @@ -551,6 +563,7 @@ impl BlocklistAIController { 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_passive_suggestion_results: HashMap::new(), } @@ -1493,32 +1506,94 @@ impl BlocklistAIController { self.pending_passive_follow_ups.remove(&conversation_id); } - /// Handles the EventsReady signal. Checks readiness, drains - /// pending events from the service, and injects them into the conversation. - fn handle_pending_events_ready( - &mut self, + fn conversation_ready_for_pending_events( + &self, conversation_id: AIConversationId, - ctx: &mut ModelContext, - ) { + ctx: &ModelContext, + ) -> bool { let owns = BlocklistAIHistoryModel::as_ref(ctx) .all_live_conversations_for_terminal_view(self.terminal_view_id) - .any(|c| c.id() == conversation_id); + .any(|conversation| conversation.id() == conversation_id); if !owns { - return; + return false; } if self .in_flight_response_streams .has_active_stream_for_conversation(conversation_id, ctx) { - return; + return false; } - // Only drain when the conversation is actually idle. - let is_success = BlocklistAIHistoryModel::as_ref(ctx) + BlocklistAIHistoryModel::as_ref(ctx) .conversation(&conversation_id) - .is_some_and(|c| matches!(c.status(), ConversationStatus::Success)); - if !is_success { + .is_some_and(|conversation| { + matches!(conversation.status(), ConversationStatus::Success) + }) + } + + fn local_claude_wake_candidate( + &self, + conversation_id: AIConversationId, + ctx: &ModelContext, + ) -> Option<(AmbientAgentTaskId, Option)> { + let conversation = BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id)?; + if !conversation.is_child_agent_conversation() || conversation.is_remote_child() { + return None; + } + + Some(( + conversation.task_id()?, + self.active_session + .as_ref(ctx) + .current_working_directory() + .cloned() + .map(PathBuf::from), + )) + } + + fn pending_event_to_claude_wake_message(event: &PendingEvent) -> Option { + let PendingEventDetail::Message { + sequence, + message_id, + subject, + message_body, + occurred_at, + .. + } = &event.detail + else { + return None; + }; + + Some(ClaudeWakeMessage { + sequence: *sequence, + message_id: message_id.clone(), + sender_run_id: event.source_agent_id.clone(), + subject: subject.clone(), + body: message_body.clone(), + occurred_at: occurred_at.clone(), + }) + } + + fn schedule_pending_events_ready_retry( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + ctx.spawn( + async move { Timer::after(Duration::from_secs(2)).await }, + move |me, _, ctx| { + me.handle_pending_events_ready(conversation_id, ctx); + }, + ); + } + + fn inject_pending_events_for_request( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + if !self.conversation_ready_for_pending_events(conversation_id, ctx) { return; } @@ -1555,6 +1630,182 @@ impl BlocklistAIController { } } + fn maybe_prepare_local_claude_wake( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) -> bool { + if self + .pending_local_claude_wakes + .contains_key(&conversation_id) + { + return true; + } + + let Some((task_id, working_dir)) = self.local_claude_wake_candidate(conversation_id, ctx) + else { + return false; + }; + + let pending_message_events = OrchestrationEventService::handle(ctx) + .update(ctx, |svc, _| { + svc.peek_pending_message_events(conversation_id) + }); + if pending_message_events.is_empty() { + return false; + } + + let pending_message_event_ids = pending_message_events + .iter() + .map(|event| event.event_id.clone()) + .collect_vec(); + let server_api = ServerApiProvider::as_ref(ctx).get(); + let handle = ctx.spawn( + async move { + let task = server_api.get_ambient_agent_task(&task_id).await?; + let harness = task + .agent_config_snapshot + .as_ref() + .and_then(|snapshot| snapshot.harness.as_ref()) + .map(|config| config.harness_type); + if task.state != AmbientAgentTaskState::Succeeded + || harness != Some(Harness::Claude) + { + return Ok::, anyhow::Error>(None); + } + + ClaudeHarness::fetch_local_wake_remote_context(task_id, server_api.clone()) + .await + .map(Some) + }, + move |me, result, ctx| { + me.pending_local_claude_wakes.remove(&conversation_id); + + let remote = match result { + Ok(Some(remote)) => remote, + Ok(None) => { + me.inject_pending_events_for_request(conversation_id, ctx); + return; + } + Err(err) => { + log::warn!( + "Failed to prepare dormant Claude wake context for {conversation_id:?}: {err:#}" + ); + me.schedule_pending_events_ready_retry(conversation_id, ctx); + return; + } + }; + + if !me.conversation_ready_for_pending_events(conversation_id, ctx) { + return; + } + + let removed_events = OrchestrationEventService::handle(ctx).update(ctx, |svc, _| { + svc.take_pending_events_by_id(conversation_id, &pending_message_event_ids) + }); + if removed_events.len() != pending_message_event_ids.len() { + if !removed_events.is_empty() { + OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { + svc.prepend_pending_events(conversation_id, removed_events, ctx); + }); + } + return; + } + + let wake_messages = removed_events + .iter() + .filter_map(Self::pending_event_to_claude_wake_message) + .collect_vec(); + if wake_messages.len() != removed_events.len() { + OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { + svc.prepend_pending_events(conversation_id, removed_events, ctx); + }); + return; + } + + let removed_events_for_retry = removed_events.clone(); + let server_api = ServerApiProvider::as_ref(ctx).get(); + let handle = ctx.spawn( + async move { + ClaudeHarness::prepare_local_wake_command( + server_api, + working_dir, + remote, + wake_messages, + ) + .await + }, + move |me, result, ctx| { + me.pending_local_claude_wakes.remove(&conversation_id); + match result { + Ok(command) => { + if !me.conversation_ready_for_pending_events(conversation_id, ctx) { + OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { + svc.prepend_pending_events( + conversation_id, + removed_events_for_retry.clone(), + ctx, + ); + }); + return; + } + + BlocklistAIHistoryModel::handle(ctx).update( + ctx, + |history_model, ctx| { + history_model.update_conversation_status( + me.terminal_view_id, + conversation_id, + ConversationStatus::InProgress, + ctx, + ); + }, + ); + ctx.emit(BlocklistAIControllerEvent::ExecuteLocalHarnessCommand { + command, + }); + } + Err(err) => { + log::warn!( + "Failed to finalize dormant Claude wake for {conversation_id:?}: {err:#}" + ); + OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { + svc.prepend_pending_events( + conversation_id, + removed_events_for_retry.clone(), + ctx, + ); + }); + } + } + }, + ); + me.pending_local_claude_wakes.insert(conversation_id, handle); + }, + ); + self.pending_local_claude_wakes + .insert(conversation_id, handle); + true + } + + /// Handles the EventsReady signal. Checks readiness, drains + /// pending events from the service, and injects them into the conversation. + fn handle_pending_events_ready( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + if !self.conversation_ready_for_pending_events(conversation_id, ctx) { + return; + } + + if self.maybe_prepare_local_claude_wake(conversation_id, ctx) { + return; + } + + self.inject_pending_events_for_request(conversation_id, ctx); + } + pub fn resume_conversation( &mut self, conversation_id: AIConversationId, diff --git a/app/src/ai/blocklist/history_model.rs b/app/src/ai/blocklist/history_model.rs index 3beeb8b4e..e5e5149a6 100644 --- a/app/src/ai/blocklist/history_model.rs +++ b/app/src/ai/blocklist/history_model.rs @@ -1092,8 +1092,6 @@ impl BlocklistAIHistoryModel { parent_conversation_id: None, run_id: None, autoexecute_override: Some(source_conversation.autoexecute_override().into()), - // The event cursor belongs to the source conversation's run; the - // forked conversation will establish its own cursor. last_event_sequence: None, }; let forked_conversation_id = AIConversationId::new(); @@ -1247,8 +1245,6 @@ impl BlocklistAIHistoryModel { parent_conversation_id: None, run_id: None, autoexecute_override: Some(conversation.autoexecute_override().into()), - // The event cursor belongs to the source conversation's run; the - // forked conversation will establish its own cursor. last_event_sequence: None, }; diff --git a/app/src/ai/blocklist/history_model_test.rs b/app/src/ai/blocklist/history_model_test.rs index 6dc2a4f48..6dd4d2d06 100644 --- a/app/src/ai/blocklist/history_model_test.rs +++ b/app/src/ai/blocklist/history_model_test.rs @@ -879,6 +879,50 @@ fn test_toggle_autoexecute_override_persists_updated_conversation_state() { }); } +#[test] +fn test_update_event_sequence_persists_updated_conversation_state() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + + let (sender, receiver) = std::sync::mpsc::sync_channel(1); + let mut global_resource_handles = GlobalResourceHandles::mock(&mut app); + global_resource_handles.model_event_sender = Some(sender); + app.add_singleton_model(|_| GlobalResourceHandlesProvider::new(global_resource_handles)); + + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); + let terminal_view_id = EntityId::new(); + + let conversation_id = history_model.update(&mut app, |history_model, ctx| { + history_model.start_new_conversation(terminal_view_id, false, false, ctx) + }); + + history_model.update(&mut app, |history_model, ctx| { + history_model.update_event_sequence(conversation_id, 42, ctx); + }); + + let event = receiver.recv_timeout(Duration::from_secs(1)).unwrap(); + + let ModelEvent::UpdateMultiAgentConversation { + conversation_id: persisted_conversation_id, + conversation_data, + .. + } = event + else { + panic!("expected UpdateMultiAgentConversation event"); + }; + + assert_eq!(persisted_conversation_id, conversation_id.to_string()); + assert_eq!(conversation_data.last_event_sequence, Some(42)); + + history_model.read(&app, |history_model, _| { + let conversation = history_model + .conversation(&conversation_id) + .expect("conversation should exist"); + assert_eq!(conversation.last_event_sequence(), Some(42)); + }); + }); +} + #[test] fn test_find_by_token_after_merge_cloud_metadata() { App::test((), |mut app| async move { diff --git a/app/src/ai/blocklist/orchestration_event_poller.rs b/app/src/ai/blocklist/orchestration_event_poller.rs index b8ecd8197..104a2d616 100644 --- a/app/src/ai/blocklist/orchestration_event_poller.rs +++ b/app/src/ai/blocklist/orchestration_event_poller.rs @@ -8,6 +8,7 @@ use crate::ai::agent_events::{ run_agent_event_driver, AgentEventConsumer, AgentEventConsumerControlFlow, AgentEventDriverConfig, MessageHydrator, ServerApiAgentEventSource, }; +use crate::ai::ambient_agents::{AmbientAgentTask, AmbientAgentTaskId}; use crate::server::server_api::ai::{AIClient, AgentRunEvent}; use crate::server::server_api::{ServerApi, ServerApiProvider}; use anyhow::anyhow; @@ -90,6 +91,7 @@ pub struct OrchestrationEventPoller { server_api: Arc, watched_run_ids: HashMap>, event_cursor: HashMap, + restore_fetch_failures: HashMap, poll_backoff_index: HashMap, pending_delivery: HashMap, conversation_statuses: HashMap, @@ -123,6 +125,24 @@ impl OrchestrationEventPoller { server_api, watched_run_ids: HashMap::new(), event_cursor: HashMap::new(), + restore_fetch_failures: HashMap::new(), + poll_backoff_index: HashMap::new(), + pending_delivery: HashMap::new(), + conversation_statuses: HashMap::new(), + poll_in_flight: HashSet::new(), + sse_connections: HashMap::new(), + next_sse_generation: 0, + } + } + + #[cfg(test)] + fn new_with_clients_for_test(ai_client: Arc, server_api: Arc) -> Self { + Self { + ai_client, + server_api, + watched_run_ids: HashMap::new(), + event_cursor: HashMap::new(), + restore_fetch_failures: HashMap::new(), poll_backoff_index: HashMap::new(), pending_delivery: HashMap::new(), conversation_statuses: HashMap::new(), @@ -199,6 +219,9 @@ impl OrchestrationEventPoller { self.on_conversation_status_updated(*conversation_id, ctx); } } + BlocklistAIHistoryEvent::RestoredConversations { + conversation_ids, .. + } => self.on_restored_conversations(conversation_ids, ctx), BlocklistAIHistoryEvent::ConversationServerTokenAssigned { conversation_id, .. } => self.on_server_token_assigned(*conversation_id, ctx), @@ -215,6 +238,7 @@ impl OrchestrationEventPoller { } => { self.watched_run_ids.remove(conversation_id); self.event_cursor.remove(conversation_id); + self.restore_fetch_failures.remove(conversation_id); self.poll_backoff_index.remove(conversation_id); self.pending_delivery.remove(conversation_id); self.conversation_statuses.remove(conversation_id); @@ -491,11 +515,56 @@ impl OrchestrationEventPoller { // Trigger event delivery when a conversation with watched run_ids // becomes idle. With the event-push flag this opens an SSE stream; // otherwise it falls back to the existing polling loop. - if became_success && self.watched_run_ids.contains_key(&conversation_id) { + if became_success + && self.watched_run_ids.contains_key(&conversation_id) + && !self.restore_fetch_failures.contains_key(&conversation_id) + { self.start_event_delivery(conversation_id, ctx); } } + fn on_restored_conversations( + &mut self, + conversation_ids: &[AIConversationId], + ctx: &mut ModelContext, + ) { + for conversation_id in conversation_ids { + let Some(conversation) = + BlocklistAIHistoryModel::as_ref(ctx).conversation(conversation_id) + else { + continue; + }; + if conversation.is_viewing_shared_session() { + continue; + } + + self.conversation_statuses + .insert(*conversation_id, conversation.status().clone()); + + let sqlite_cursor = conversation.last_event_sequence().unwrap_or(0); + self.event_cursor.insert(*conversation_id, sqlite_cursor); + + let Some(run_id) = conversation.run_id() else { + continue; + }; + + self.watched_run_ids + .entry(*conversation_id) + .or_default() + .insert(run_id.clone()); + + let Ok(task_id) = run_id.parse::() else { + self.maybe_start_delivery_after_restore(*conversation_id, ctx); + continue; + }; + + self.restore_fetch_failures + .entry(*conversation_id) + .or_insert(0); + self.spawn_restore_fetch(*conversation_id, task_id, sqlite_cursor, ctx); + } + } + fn on_server_token_assigned( &mut self, conversation_id: AIConversationId, @@ -521,6 +590,116 @@ impl OrchestrationEventPoller { self.register_watched_run_id(conversation_id, run_id, ctx); } + fn spawn_restore_fetch( + &mut self, + conversation_id: AIConversationId, + task_id: AmbientAgentTaskId, + sqlite_cursor: i64, + ctx: &mut ModelContext, + ) { + let ai_client = self.ai_client.clone(); + let task_id_string = task_id.to_string(); + let retry_task_id_string = task_id_string.clone(); + ctx.spawn( + async move { + let task_id: AmbientAgentTaskId = task_id_string.parse()?; + ai_client.get_ambient_agent_task(&task_id).await + }, + move |me, result, ctx| match result { + Ok(task) => me.finish_restore_fetch(conversation_id, sqlite_cursor, task, ctx), + Err(err) => { + log::warn!( + "Failed to restore orchestration event state for {conversation_id:?}: {err:#}" + ); + me.start_restore_fetch_retry_timer( + conversation_id, + retry_task_id_string.clone(), + ctx, + ); + } + }, + ); + } + + fn finish_restore_fetch( + &mut self, + conversation_id: AIConversationId, + sqlite_cursor: i64, + task: AmbientAgentTask, + ctx: &mut ModelContext, + ) { + self.restore_fetch_failures.remove(&conversation_id); + + let merged_cursor = sqlite_cursor.max(task.last_event_sequence.unwrap_or(0)); + self.event_cursor.insert(conversation_id, merged_cursor); + + let watched = self.watched_run_ids.entry(conversation_id).or_default(); + let mut added_child = false; + for child_run_id in task.children { + if watched.insert(child_run_id) { + added_child = true; + } + } + + if added_child && self.sse_connections.contains_key(&conversation_id) { + self.reconnect_sse(conversation_id, ctx); + } + + self.maybe_start_delivery_after_restore(conversation_id, ctx); + } + + fn start_restore_fetch_retry_timer( + &mut self, + conversation_id: AIConversationId, + task_id: String, + ctx: &mut ModelContext, + ) { + let failures = self + .restore_fetch_failures + .entry(conversation_id) + .or_insert(0); + *failures += 1; + let delay_secs = POLL_BACKOFF_STEPS[(*failures - 1).min(POLL_BACKOFF_STEPS.len() - 1)]; + + ctx.spawn( + async move { Timer::after(Duration::from_secs(delay_secs)).await }, + move |me, _, ctx| { + if !me.restore_fetch_failures.contains_key(&conversation_id) { + return; + } + let Some(_) = BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) + else { + me.restore_fetch_failures.remove(&conversation_id); + return; + }; + + let Ok(task_id) = task_id.parse::() else { + me.restore_fetch_failures.remove(&conversation_id); + me.maybe_start_delivery_after_restore(conversation_id, ctx); + return; + }; + + let sqlite_cursor = me.event_cursor.get(&conversation_id).copied().unwrap_or(0); + me.spawn_restore_fetch(conversation_id, task_id, sqlite_cursor, ctx); + }, + ); + } + + fn maybe_start_delivery_after_restore( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + let status = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .map(|conversation| conversation.status().clone()) + .or_else(|| self.conversation_statuses.get(&conversation_id).cloned()); + + if matches!(status, Some(ConversationStatus::Success)) { + self.start_event_delivery(conversation_id, ctx); + } + } + fn on_streaming_exchange_updated( &mut self, conversation_id: AIConversationId, @@ -687,6 +866,26 @@ impl OrchestrationEventPoller { .max() .unwrap_or(previous_cursor); self.event_cursor.insert(conversation_id, max_seq); + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.update_event_sequence(conversation_id, max_seq, ctx); + }); + + let own_run_id = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .and_then(|conversation| conversation.run_id()); + if let Some(run_id) = own_run_id { + let ai_client = self.ai_client.clone(); + ctx.spawn( + async move { ai_client.update_event_sequence_on_server(&run_id, max_seq).await }, + move |_, result, _| { + if let Err(err) = result { + log::warn!( + "Failed to persist server-backed event cursor for {conversation_id:?}: {err:#}" + ); + } + }, + ); + } // Persist the cursor to SQLite so that after a restart we can resume // event delivery from this sequence number without re-delivering @@ -746,7 +945,7 @@ impl OrchestrationEventPoller { // Only reset backoff when events actually produce pending items. self.poll_backoff_index.remove(&conversation_id); - let pending = build_pending_events(messages, lifecycle_events); + let pending = build_pending_events(&events, messages, lifecycle_events); OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { svc.enqueue_polled_events(conversation_id, pending, ctx); }); @@ -1042,20 +1241,32 @@ fn convert_lifecycle_events( } fn build_pending_events( + events: &[AgentRunEvent], messages: Vec, lifecycle_events: Vec, ) -> Vec { let mut pending = Vec::with_capacity(messages.len() + lifecycle_events.len()); for msg in &messages { + let metadata = events + .iter() + .find(|event| { + event.event_type == "new_message" + && event.ref_id.as_deref() == Some(msg.message_id.as_str()) + }) + .map(|event| (event.sequence, event.occurred_at.clone())); + let (sequence, occurred_at) = + metadata.unwrap_or_else(|| (0, chrono::Utc::now().to_rfc3339())); pending.push(PendingEvent { event_id: msg.message_id.clone(), source_agent_id: msg.sender_agent_id.clone(), attempt_count: 0, detail: PendingEventDetail::Message { + sequence, message_id: msg.message_id.clone(), addresses: msg.addresses.clone(), subject: msg.subject.clone(), message_body: msg.message_body.clone(), + occurred_at, }, }); } diff --git a/app/src/ai/blocklist/orchestration_event_poller_tests.rs b/app/src/ai/blocklist/orchestration_event_poller_tests.rs index 06d5a209e..c9775dacb 100644 --- a/app/src/ai/blocklist/orchestration_event_poller_tests.rs +++ b/app/src/ai/blocklist/orchestration_event_poller_tests.rs @@ -1,8 +1,20 @@ use super::*; +use crate::ai::agent::conversation::{AIConversation, AIConversationId, ConversationStatus}; use crate::ai::agent_events::{ agent_event_backoff, agent_event_failures_exceeded_threshold, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS, }; +use crate::ai::ambient_agents::{AmbientAgentTask, AmbientAgentTaskState}; +use crate::persistence::{model::AgentConversationData, ModelEvent}; +use crate::server::server_api::ai::MockAIClient; +use crate::server::server_api::ServerApiProvider; +use crate::test_util::settings::initialize_settings_for_tests; +use crate::{GlobalResourceHandles, GlobalResourceHandlesProvider}; +use chrono::Utc; +use std::collections::HashSet; +use std::sync::Arc; +use warp_multi_agent_api as api; +use warpui::{App, EntityId}; #[test] fn sse_backoff_escalates_then_caps() { @@ -105,75 +117,51 @@ fn convert_lifecycle_events_maps_run_restarted() { )); } -#[test] -fn ai_conversation_new_restored_preserves_last_event_sequence() { - // Guards against regressions that drop the field when wiring the restore - // path: a conversation restored with `last_event_sequence: Some(N)` - // should expose it via `conversation.last_event_sequence()`. - use crate::ai::agent::conversation::{AIConversation, AIConversationId}; - use crate::persistence::model::AgentConversationData; - - let task = api::Task { - id: "root".to_string(), - messages: vec![api::Message { - id: "m1".to_string(), - task_id: "root".to_string(), - server_message_data: String::new(), - citations: vec![], - message: Some(api::message::Message::AgentOutput( - api::message::AgentOutput { - text: "hi".to_string(), - }, - )), - request_id: String::new(), - timestamp: None, +fn restored_conversation( + conversation_id: AIConversationId, + run_id: String, + last_event_sequence: Option, +) -> AIConversation { + AIConversation::new_restored( + conversation_id, + vec![api::Task { + id: "root-task".to_string(), + messages: vec![], + dependencies: None, + description: String::new(), + summary: String::new(), + server_data: String::new(), }], - dependencies: None, - description: String::new(), - summary: String::new(), - server_data: String::new(), - }; - let data = AgentConversationData { - server_conversation_token: None, - conversation_usage_metadata: None, - reverted_action_ids: None, - forked_from_server_conversation_token: None, - artifacts_json: None, - parent_agent_id: None, - agent_name: None, - parent_conversation_id: None, - run_id: None, - autoexecute_override: None, - last_event_sequence: Some(42), - }; - let conversation = - AIConversation::new_restored(AIConversationId::new(), vec![task], Some(data)) - .expect("should restore"); - assert_eq!(conversation.last_event_sequence(), Some(42)); -} - -// ---- Helpers for App-based poller tests ---- - -fn make_ambient_task_with_children( - children: Vec, -) -> crate::ai::ambient_agents::AmbientAgentTask { - let mut task = make_ambient_task_with_event_seq(None); - task.children = children; - task + Some(AgentConversationData { + server_conversation_token: None, + conversation_usage_metadata: None, + reverted_action_ids: None, + forked_from_server_conversation_token: None, + artifacts_json: None, + parent_agent_id: None, + agent_name: None, + parent_conversation_id: None, + run_id: Some(run_id), + autoexecute_override: None, + last_event_sequence, + }), + ) + .expect("restored conversation should build") } -fn make_ambient_task_with_event_seq( +fn ambient_agent_task( + task_id: crate::ai::ambient_agents::AmbientAgentTaskId, last_event_sequence: Option, -) -> crate::ai::ambient_agents::AmbientAgentTask { - use chrono::Utc; - crate::ai::ambient_agents::AmbientAgentTask { - task_id: "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(), + children: Vec, +) -> AmbientAgentTask { + AmbientAgentTask { + task_id, parent_run_id: None, - title: "test".to_string(), - state: crate::ai::ambient_agents::AmbientAgentTaskState::Succeeded, - prompt: "prompt".to_string(), + title: "Task".to_string(), + state: AmbientAgentTaskState::Succeeded, + prompt: String::new(), created_at: Utc::now(), - started_at: Some(Utc::now()), + started_at: None, updated_at: Utc::now(), status_message: None, source: None, @@ -182,95 +170,29 @@ fn make_ambient_task_with_event_seq( creator: None, conversation_id: None, request_usage: None, + is_sandbox_running: false, agent_config_snapshot: None, artifacts: vec![], - is_sandbox_running: false, last_event_sequence, - children: vec![], + children, } } #[test] -fn finish_restore_fetch_uses_server_cursor_when_sqlite_is_absent() { - use crate::ai::agent::conversation::AIConversation; - use crate::server::server_api::ai::MockAIClient; - use crate::server::server_api::ServerApiProvider; - use std::sync::Arc; - use warpui::App; - +fn finish_restore_fetch_merges_server_cursor_and_child_runs() { App::test((), |mut app| async move { - let _v2_guard = FeatureFlag::OrchestrationV2.override_enabled(true); - let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); + let terminal_view_id = EntityId::new(); + let conversation_id = AIConversationId::new(); + let run_id = uuid::Uuid::new_v4().to_string(); - // Restore a conversation with no SQLite cursor (`last_event_sequence: - // None`). After the server fetch completes with `Some(42)` we expect - // the in-memory cursor to be 42 (max(0, 42)). - let conversation = AIConversation::new(false); - let conversation_id = conversation.id(); - let terminal_view_id = warpui::EntityId::new(); - history_model.update(&mut app, |model, ctx| { - model.restore_conversations(terminal_view_id, vec![conversation], ctx); - }); - - let mock = MockAIClient::new(); - let ai_client: Arc = Arc::new(mock); - let server_api = ServerApiProvider::new_for_test().get(); - - let poller = app.add_singleton_model(|ctx| { - OrchestrationEventPoller::new_with_clients_for_test(ai_client, server_api, ctx) - }); - - // Seed event_cursor as on_restored_conversations would before spawning - // the async fetch. Without this the guard that detects mid-flight - // conversation deletion would fire and return early. - poller.update(&mut app, |me, _| { - me.event_cursor.insert(conversation_id, 0); - }); - - let task_id: crate::ai::ambient_agents::AmbientAgentTaskId = - "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(); - poller.update(&mut app, |me, ctx| { - me.finish_restore_fetch( - conversation_id, - task_id, - /* sqlite_cursor */ 0, - Ok(make_ambient_task_with_event_seq(Some(42))), + history_model.update(&mut app, |history_model, ctx| { + history_model.restore_conversations( + terminal_view_id, + vec![restored_conversation(conversation_id, run_id.clone(), None)], ctx, ); - }); - - poller.read(&app, |me, _| { - assert_eq!(me.event_cursor.get(&conversation_id).copied(), Some(42)); - }); - }); -} - -#[test] -fn restored_inprogress_parent_defers_delivery_until_success() { - use crate::ai::agent::conversation::{AIConversation, AIConversationId, ConversationStatus}; - use crate::server::server_api::ai::MockAIClient; - use crate::server::server_api::ServerApiProvider; - use std::sync::Arc; - use warpui::App; - - App::test((), |mut app| async move { - let _v2_guard = FeatureFlag::OrchestrationV2.override_enabled(true); - - let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); - - let mut conversation = AIConversation::new(false); - // Use a parsable UUID-shaped run_id so the poller can construct - // an `AmbientAgentTaskId` for the (mocked) server fetch. - conversation.set_run_id("550e8400-e29b-41d4-a716-446655440100".to_string()); - let conversation_id: AIConversationId = conversation.id(); - let terminal_view_id = warpui::EntityId::new(); - history_model.update(&mut app, |model, ctx| { - model.restore_conversations(terminal_view_id, vec![conversation], ctx); - // The default status after restore is `InProgress` for live - // conversations, but assert it explicitly to make the test - // self-documenting. - model.update_conversation_status( + history_model.update_conversation_status( terminal_view_id, conversation_id, ConversationStatus::InProgress, @@ -278,258 +200,98 @@ fn restored_inprogress_parent_defers_delivery_until_success() { ); }); - let mut mock = MockAIClient::new(); - // The async restore fetch may or may not complete during the test; - // a permissive expectation prevents spurious panics either way. - mock.expect_get_ambient_agent_task() - .returning(|_| Ok(make_ambient_task_with_event_seq(None))); - mock.expect_poll_agent_events() - .returning(|_, _, _| Ok(vec![])); - mock.expect_update_event_sequence_on_server() - .returning(|_, _| Ok(())); - let ai_client: Arc = Arc::new(mock); let server_api = ServerApiProvider::new_for_test().get(); - - let poller = app.add_singleton_model(|ctx| { - OrchestrationEventPoller::new_with_clients_for_test(ai_client, server_api, ctx) + let poller = app.add_singleton_model(move |_| { + OrchestrationEventPoller::new_with_clients_for_test( + Arc::new(MockAIClient::new()), + server_api, + ) }); - // Synchronous part of `on_restored_conversations`: cursor seeded, - // own run_id watched. No event delivery yet because parent is - // InProgress. - poller.update(&mut app, |me, ctx| { - me.on_restored_conversations(vec![conversation_id], ctx); - }); - poller.read(&app, |me, _| { - assert_eq!(me.event_cursor.get(&conversation_id).copied(), Some(0)); - assert!( - me.watched_run_ids - .get(&conversation_id) - .is_some_and(|w| !w.is_empty()), - "own run_id should have been registered as watched" - ); - assert!( - !me.poll_in_flight.contains(&conversation_id), - "InProgress parent must not start polling" - ); - assert!( - me.sse_connections.is_empty(), - "InProgress parent must not open SSE" - ); + let child_run_id = uuid::Uuid::new_v4().to_string(); + let task_id = run_id.parse().expect("run_id should parse"); + let task = ambient_agent_task(task_id, Some(27), vec![child_run_id.clone()]); + poller.update(&mut app, |poller, ctx| { + poller + .watched_run_ids + .insert(conversation_id, HashSet::from([run_id.clone()])); + poller.restore_fetch_failures.insert(conversation_id, 1); + poller.finish_restore_fetch(conversation_id, 0, task, ctx); }); + poller.read(&app, |poller, _| { + assert_eq!(poller.event_cursor.get(&conversation_id), Some(&27)); + assert!(!poller.restore_fetch_failures.contains_key(&conversation_id)); - // Transitioning the conversation to Success should trigger event - // delivery (poll_and_inject in the non-SSE default path). - history_model.update(&mut app, |model, ctx| { - model.update_conversation_status( - terminal_view_id, - conversation_id, - ConversationStatus::Success, - ctx, - ); - }); - poller.read(&app, |me, _| { - assert!( - me.poll_in_flight.contains(&conversation_id), - "Success transition with watched run_ids should start delivery" - ); + let watched = poller + .watched_run_ids + .get(&conversation_id) + .expect("watched runs should exist"); + assert!(watched.contains(&run_id)); + assert!(watched.contains(&child_run_id)); }); }); } #[test] -fn handle_poll_result_persists_max_seq_to_history_model() { - use crate::ai::agent::conversation::{AIConversation, AIConversationId}; - use crate::persistence::ModelEvent; - use crate::server::server_api::ai::MockAIClient; - use crate::server::server_api::ServerApiProvider; - use crate::test_util::settings::initialize_settings_for_tests; - use crate::{GlobalResourceHandles, GlobalResourceHandlesProvider}; - use std::sync::Arc; - use warpui::App; +fn build_pending_events_preserves_message_sequence_and_timestamp() { + let pending = build_pending_events( + &[AgentRunEvent { + event_type: "new_message".to_string(), + run_id: "recipient-run".to_string(), + ref_id: Some("message-1".to_string()), + execution_id: None, + occurred_at: "2026-02-01T12:34:56Z".to_string(), + sequence: 42, + }], + vec![crate::ai::agent::ReceivedMessageInput { + message_id: "message-1".to_string(), + sender_agent_id: "sender-run".to_string(), + addresses: vec!["recipient-run".to_string()], + subject: "subject".to_string(), + message_body: "body".to_string(), + }], + vec![], + ); - App::test((), |mut app| async move { - let _v2_guard = FeatureFlag::OrchestrationV2.override_enabled(true); + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].event_id, "message-1"); + match &pending[0].detail { + PendingEventDetail::Message { + sequence, + message_id, + occurred_at, + .. + } => { + assert_eq!(*sequence, 42); + assert_eq!(message_id, "message-1"); + assert_eq!(occurred_at, "2026-02-01T12:34:56Z"); + } + PendingEventDetail::Lifecycle { .. } => panic!("expected message event"), + } +} - // `update_event_sequence` calls `write_updated_conversation_state`, - // which reads `GeneralSettings`, `AppExecutionMode`, and the global - // resource sender. Wire all of these up so the SQLite write can run. +#[test] +fn handle_poll_result_updates_persisted_and_server_cursor() { + App::test((), |mut app| async move { initialize_settings_for_tests(&mut app); - let (sender, receiver) = std::sync::mpsc::sync_channel::(4); + + let (sender, receiver) = std::sync::mpsc::sync_channel(1); let mut global_resource_handles = GlobalResourceHandles::mock(&mut app); global_resource_handles.model_event_sender = Some(sender); app.add_singleton_model(|_| GlobalResourceHandlesProvider::new(global_resource_handles)); let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); + let terminal_view_id = EntityId::new(); + let conversation_id = AIConversationId::new(); + let run_id = uuid::Uuid::new_v4().to_string(); - let mut conversation = AIConversation::new(false); - conversation.set_run_id("550e8400-e29b-41d4-a716-446655440200".to_string()); - let conversation_id: AIConversationId = conversation.id(); - let terminal_view_id = warpui::EntityId::new(); - history_model.update(&mut app, |model, ctx| { - model.restore_conversations(terminal_view_id, vec![conversation], ctx); - }); - - let mut mock = MockAIClient::new(); - // The fire-and-forget server PATCH should be issued; permissive Ok. - mock.expect_update_event_sequence_on_server() - .returning(|_, _| Ok(())); - let ai_client: Arc = Arc::new(mock); - let server_api = ServerApiProvider::new_for_test().get(); - - let poller = app.add_singleton_model(|ctx| { - OrchestrationEventPoller::new_with_clients_for_test(ai_client, server_api, ctx) - }); - - // Build a poll batch with max sequence = 42. Use an unrecognized - // event_type so `convert_lifecycle_events` returns empty and the - // function early-exits before touching `OrchestrationEventService` - // (which we did not register in this test App). - let events = vec![ - AgentRunEvent { - event_type: "unrecognized_event_type".to_string(), - run_id: "some-other-run".to_string(), - ref_id: None, - execution_id: None, - occurred_at: "2026-01-01T00:00:00Z".to_string(), - sequence: 17, - }, - AgentRunEvent { - event_type: "unrecognized_event_type".to_string(), - run_id: "some-other-run".to_string(), - ref_id: None, - execution_id: None, - occurred_at: "2026-01-01T00:00:00Z".to_string(), - sequence: 42, - }, - ]; - - poller.update(&mut app, |me, ctx| { - me.handle_poll_result( - conversation_id, - /* self_run_id */ "some-other-run", - /* previous_cursor */ 0, - events, - /* messages */ vec![], - ctx, - ); - }); - - history_model.read(&app, |model, _| { - let last_seq = model - .conversation(&conversation_id) - .and_then(|c| c.last_event_sequence()); - assert_eq!( - last_seq, - Some(42), - "BlocklistAIHistoryModel.update_event_sequence must be called with max_seq" - ); - }); - - // Drain at least one persistence event to confirm the SQLite write - // path was triggered (sanity check for the side effect, not the - // primary assertion). - let _ = receiver.recv_timeout(std::time::Duration::from_secs(1)); - }); -} - -#[test] -fn finish_restore_fetch_no_ops_when_conversation_deleted_mid_flight() { - // If the conversation is removed while the async fetch is in-flight, the - // RemoveConversation handler clears event_cursor. finish_restore_fetch - // uses the missing cursor as a sentinel and must not re-populate - // watched_run_ids or event_cursor for the deleted conversation. - use crate::ai::agent::conversation::AIConversation; - use crate::server::server_api::ai::MockAIClient; - use crate::server::server_api::ServerApiProvider; - use std::sync::Arc; - use warpui::App; - - App::test((), |mut app| async move { - let _v2_guard = FeatureFlag::OrchestrationV2.override_enabled(true); - - let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); - - let mut conversation = AIConversation::new(false); - conversation.set_run_id("550e8400-e29b-41d4-a716-446655440300".to_string()); - let conversation_id = conversation.id(); - let terminal_view_id = warpui::EntityId::new(); - history_model.update(&mut app, |model, ctx| { - model.restore_conversations(terminal_view_id, vec![conversation], ctx); - }); - - let mock = MockAIClient::new(); - let ai_client: Arc = Arc::new(mock); - let server_api = ServerApiProvider::new_for_test().get(); - - let poller = app.add_singleton_model(|ctx| { - OrchestrationEventPoller::new_with_clients_for_test(ai_client, server_api, ctx) - }); - - // Seed cursor as on_restored_conversations would. - poller.update(&mut app, |me, _| { - me.event_cursor.insert(conversation_id, 0); - }); - - // Simulate the RemoveConversation handler firing while the fetch is - // in-flight: it clears event_cursor (and all other state). - poller.update(&mut app, |me, _| { - me.watched_run_ids.remove(&conversation_id); - me.event_cursor.remove(&conversation_id); - }); - - // The in-flight fetch now completes — with children. - let task_id: crate::ai::ambient_agents::AmbientAgentTaskId = - "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(); - poller.update(&mut app, |me, ctx| { - me.finish_restore_fetch( - conversation_id, - task_id, - /* sqlite_cursor */ 0, - Ok(make_ambient_task_with_children(vec![ - "child-run-1".to_string() - ])), + history_model.update(&mut app, |history_model, ctx| { + history_model.restore_conversations( + terminal_view_id, + vec![restored_conversation(conversation_id, run_id.clone(), None)], ctx, ); - }); - - poller.read(&app, |me, _| { - assert!( - !me.watched_run_ids.contains_key(&conversation_id), - "watched_run_ids must not be repopulated for a deleted conversation" - ); - assert!( - !me.event_cursor.contains_key(&conversation_id), - "event_cursor must not be repopulated for a deleted conversation" - ); - }); - }); -} - -#[test] -fn finish_restore_fetch_reconnects_sse_when_children_added_to_open_connection() { - // When a status transition races with the restore fetch and opens SSE - // before children are known, finish_restore_fetch must reconnect SSE - // with the updated run_id set rather than leaving children unwatched. - use crate::ai::agent::conversation::{AIConversation, ConversationStatus}; - use crate::server::server_api::ai::MockAIClient; - use crate::server::server_api::ServerApiProvider; - use std::sync::Arc; - use warpui::App; - - App::test((), |mut app| async move { - let _v2_guard = FeatureFlag::OrchestrationV2.override_enabled(true); - - let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); - - let own_run_id = "550e8400-e29b-41d4-a716-446655440400"; - let mut conversation = AIConversation::new(false); - conversation.set_run_id(own_run_id.to_string()); - let conversation_id = conversation.id(); - let terminal_view_id = warpui::EntityId::new(); - history_model.update(&mut app, |model, ctx| { - model.restore_conversations(terminal_view_id, vec![conversation], ctx); - model.update_conversation_status( + history_model.update_conversation_status( terminal_view_id, conversation_id, ConversationStatus::InProgress, @@ -537,66 +299,58 @@ fn finish_restore_fetch_reconnects_sse_when_children_added_to_open_connection() ); }); - let mock = MockAIClient::new(); - let ai_client: Arc = Arc::new(mock); - let server_api = ServerApiProvider::new_for_test().get(); + let expected_run_id = run_id.clone(); + let mut ai_client = MockAIClient::new(); + ai_client + .expect_update_event_sequence_on_server() + .withf(move |run_id, sequence| run_id == expected_run_id.as_str() && *sequence == 17) + .times(1) + .returning(|_, _| Ok(())); - let poller = app.add_singleton_model(|ctx| { - OrchestrationEventPoller::new_with_clients_for_test(ai_client, server_api, ctx) + let server_api = ServerApiProvider::new_for_test().get(); + let poller = app.add_singleton_model(move |_| { + OrchestrationEventPoller::new_with_clients_for_test(Arc::new(ai_client), server_api) }); - // Seed the state on_restored_conversations would have set up, then - // inject a fake open SSE connection (generation 0) simulating the - // race: a status transition fired before the restore fetch completed. - let (_, rx) = futures::channel::mpsc::unbounded::(); - poller.update(&mut app, |me, _| { - me.event_cursor.insert(conversation_id, 0); - me.watched_run_ids - .entry(conversation_id) - .or_default() - .insert(own_run_id.to_string()); - me.sse_connections.insert( + poller.update(&mut app, |poller, ctx| { + poller.handle_poll_result( conversation_id, - SseConnectionState { - event_receiver: rx, - generation: 0, - }, + &run_id, + 0, + vec![AgentRunEvent { + event_type: "new_message".to_string(), + run_id: run_id.clone(), + ref_id: Some("message-1".to_string()), + execution_id: None, + occurred_at: "2026-01-01T00:00:00Z".to_string(), + sequence: 17, + }], + vec![], + ctx, ); - me.next_sse_generation = 1; }); - // The restore fetch returns with a child run_id. - let task_id: crate::ai::ambient_agents::AmbientAgentTaskId = - "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(); - poller.update(&mut app, |me, ctx| { - me.finish_restore_fetch( - conversation_id, - task_id, - /* sqlite_cursor */ 0, - Ok(make_ambient_task_with_children(vec![ - "child-run-1".to_string() - ])), - ctx, - ); + let event = receiver.recv_timeout(Duration::from_secs(1)).unwrap(); + let ModelEvent::UpdateMultiAgentConversation { + conversation_id: persisted_conversation_id, + conversation_data, + .. + } = event + else { + panic!("expected UpdateMultiAgentConversation event"); + }; + + assert_eq!(persisted_conversation_id, conversation_id.to_string()); + assert_eq!(conversation_data.last_event_sequence, Some(17)); + poller.read(&app, |poller, _| { + assert_eq!(poller.event_cursor.get(&conversation_id), Some(&17)); }); - poller.read(&app, |me, _| { - assert!( - me.watched_run_ids - .get(&conversation_id) - .is_some_and(|w| w.contains("child-run-1")), - "child run_id must be in watched set" - ); - // The old generation-0 connection must have been replaced by a - // new one with a higher generation, proving SSE was reconnected. - let gen = me - .sse_connections - .get(&conversation_id) - .map(|s| s.generation); - assert!( - gen.is_some_and(|g| g > 0), - "SSE must be reconnected (new generation) after children are discovered; got gen={gen:?}" - ); + history_model.read(&app, |history_model, _| { + let conversation = history_model + .conversation(&conversation_id) + .expect("conversation should exist"); + assert_eq!(conversation.last_event_sequence(), Some(17)); }); }); } diff --git a/app/src/ai/blocklist/orchestration_events.rs b/app/src/ai/blocklist/orchestration_events.rs index 381906196..febb47c63 100644 --- a/app/src/ai/blocklist/orchestration_events.rs +++ b/app/src/ai/blocklist/orchestration_events.rs @@ -56,10 +56,12 @@ impl LifecycleEventDetailStage { #[derive(Debug, Clone)] pub enum PendingEventDetail { Message { + sequence: i64, message_id: String, addresses: Vec, subject: String, message_body: String, + occurred_at: String, }, Lifecycle { event: api::AgentEvent, @@ -747,6 +749,7 @@ impl OrchestrationEventService { // We keep `message_id` stable across targets so dedupe/threading can reason // about a single message delivered to multiple recipients. let message_id = Uuid::new_v4().to_string(); + let occurred_at = chrono::Utc::now().to_rfc3339(); for (_, target_conversation_id) in resolved_targets { let event_id = Uuid::new_v4().to_string(); @@ -756,10 +759,12 @@ impl OrchestrationEventService { source_agent_id: sender_agent_id.to_string(), attempt_count: 0, detail: PendingEventDetail::Message { + sequence: 0, message_id: message_id.clone(), addresses: target_agent_ids.to_vec(), subject: subject.clone(), message_body: message_body.clone(), + occurred_at: occurred_at.clone(), }, }; self.pending_events @@ -890,6 +895,67 @@ impl OrchestrationEventService { ctx.emit(OrchestrationEventServiceEvent::EventsReady { conversation_id }); } + pub fn peek_pending_message_events( + &self, + conversation_id: AIConversationId, + ) -> Vec { + self.pending_events + .get(&conversation_id) + .into_iter() + .flatten() + .filter(|event| matches!(event.detail, PendingEventDetail::Message { .. })) + .cloned() + .collect() + } + + pub fn take_pending_events_by_id( + &mut self, + conversation_id: AIConversationId, + event_ids: &[String], + ) -> Vec { + if event_ids.is_empty() { + return Vec::new(); + } + + let event_ids = event_ids.iter().map(String::as_str).collect::>(); + let Some(queue) = self.pending_events.get_mut(&conversation_id) else { + return Vec::new(); + }; + + let mut removed = Vec::new(); + let mut retained = Vec::with_capacity(queue.len()); + for event in std::mem::take(queue) { + if event_ids.contains(event.event_id.as_str()) { + removed.push(event); + } else { + retained.push(event); + } + } + *queue = retained; + + if queue.is_empty() { + self.pending_events.remove(&conversation_id); + } + + removed + } + + pub fn prepend_pending_events( + &mut self, + conversation_id: AIConversationId, + mut events: Vec, + ctx: &mut ModelContext, + ) { + if events.is_empty() { + return; + } + + let queue = self.pending_events.entry(conversation_id).or_default(); + events.append(queue); + *queue = events; + ctx.emit(OrchestrationEventServiceEvent::EventsReady { conversation_id }); + } + /// Drain and return all pending events for a conversation. fn drain_pending_events(&mut self, conversation_id: &AIConversationId) -> Vec { self.pending_events @@ -933,10 +999,12 @@ impl OrchestrationEventService { for event in &deliverable { match &event.detail { PendingEventDetail::Message { + sequence: _, message_id, addresses, subject, message_body, + occurred_at: _, } => messages.push(ReceivedMessageInput { message_id: message_id.clone(), sender_agent_id: event.source_agent_id.clone(), diff --git a/app/src/ai/blocklist/orchestration_events_tests.rs b/app/src/ai/blocklist/orchestration_events_tests.rs index c3ba0a52e..147bc835c 100644 --- a/app/src/ai/blocklist/orchestration_events_tests.rs +++ b/app/src/ai/blocklist/orchestration_events_tests.rs @@ -71,10 +71,12 @@ fn message_pending_event(event_id: &str) -> PendingEvent { source_agent_id: "sender".to_string(), attempt_count: 0, detail: PendingEventDetail::Message { + sequence: 0, message_id: "message-1".to_string(), addresses: vec!["target".to_string()], subject: "subject".to_string(), message_body: "body".to_string(), + occurred_at: "2026-01-01T00:00:00Z".to_string(), }, } } @@ -269,10 +271,12 @@ fn test_did_event_round_trip_through_server_matches_message_event_by_message_id( source_agent_id: "sender".to_string(), attempt_count: 0, detail: PendingEventDetail::Message { + sequence: 0, message_id: "message-1".to_string(), addresses: vec!["target".to_string()], subject: "subject".to_string(), message_body: "body".to_string(), + occurred_at: "2026-01-01T00:00:00Z".to_string(), }, }; @@ -359,65 +363,53 @@ fn test_lifecycle_event_type_from_proto_includes_cancelled_and_blocked() { } #[test] -fn restored_v1_child_conversation_re_registers_lifecycle_subscription() { - use crate::ai::agent::conversation::AIConversation; - use warp_core::features::FeatureFlag; - use warpui::{App, EntityId}; - - App::test((), |mut app| async move { - // V1 path is gated on `!OrchestrationV2`. - let _v1_guard = FeatureFlag::OrchestrationV2.override_enabled(false); - - let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); - - // Build a parent conversation with a server token; under V1 the - // parent's `server_conversation_token` is the agent identifier the - // child subscribes to. - let parent_token = "parent-token-v1"; - let mut parent_conversation = AIConversation::new(false); - parent_conversation.set_server_conversation_token(parent_token.to_string()); - let parent_conversation_id = parent_conversation.id(); - - // Build a child conversation pointing at the parent. - let mut child_conversation = AIConversation::new(false); - child_conversation.set_parent_conversation_id(parent_conversation_id); - let child_conversation_id = child_conversation.id(); - - let terminal_view_id = EntityId::new(); - history_model.update(&mut app, |model, ctx| { - model.restore_conversations( - terminal_view_id, - vec![parent_conversation, child_conversation], - ctx, - ); - }); - - // Drive the OrchestrationEventService through its standard - // `handle_history_event` entry point; `restore_conversations` already - // emitted `RestoredConversations`, so we replay it explicitly through - // the service to keep this test independent of subscription wiring. - let service = app.add_singleton_model(|_| OrchestrationEventService::default()); - service.update(&mut app, |svc, ctx| { - svc.handle_history_event( - &BlocklistAIHistoryEvent::RestoredConversations { - terminal_view_id, - conversation_ids: vec![parent_conversation_id, child_conversation_id], - }, - ctx, - ); - }); - - service.read(&app, |svc, _| { - let routes = svc - .lifecycle_subscription_routes - .get(&child_conversation_id) - .expect("expected V1 lifecycle route to be registered for the child"); - assert_eq!(routes.len(), 1, "expected exactly one route"); - assert_eq!(routes[0].target_agent_id, parent_token); - assert!( - routes[0].subscribed_event_types.is_none(), - "restore re-registers with `None` (subscribe to all event types)" - ); - }); - }); +fn test_pending_message_helpers_peek_and_take_only_selected_messages() { + let conversation_id = crate::ai::agent::conversation::AIConversationId::new(); + let mut service = OrchestrationEventService::new_without_subscriptions(); + service.pending_events.insert( + conversation_id, + vec![ + lifecycle_pending_event( + "lifecycle-1", + "child-a", + api::LifecycleEventType::InProgress, + 0, + ), + message_pending_event("message-event-1"), + lifecycle_pending_event( + "lifecycle-2", + "child-b", + api::LifecycleEventType::Succeeded, + 0, + ), + message_pending_event("message-event-2"), + ], + ); + + let peeked = service.peek_pending_message_events(conversation_id); + assert_eq!(peeked.len(), 2); + assert_eq!(peeked[0].event_id, "message-event-1"); + assert_eq!(peeked[1].event_id, "message-event-2"); + + let removed = service.take_pending_events_by_id( + conversation_id, + &["message-event-2".to_string(), "message-event-1".to_string()], + ); + assert_eq!(removed.len(), 2); + assert_eq!(removed[0].event_id, "message-event-1"); + assert_eq!(removed[1].event_id, "message-event-2"); + + let remaining = service + .pending_events + .get(&conversation_id) + .expect("lifecycle events should remain queued"); + assert_eq!(remaining.len(), 2); + assert!(matches!( + remaining[0].detail, + PendingEventDetail::Lifecycle { .. } + )); + assert!(matches!( + remaining[1].detail, + PendingEventDetail::Lifecycle { .. } + )); } diff --git a/app/src/server/server_api.rs b/app/src/server/server_api.rs index 7405abb78..62b85323f 100644 --- a/app/src/server/server_api.rs +++ b/app/src/server/server_api.rs @@ -470,6 +470,16 @@ impl ServerApi { .collect()) } + async fn ambient_agent_headers_for_task( + &self, + task_id: &AmbientAgentTaskId, + ) -> Result> { + let mut headers = self.ambient_agent_headers().await?; + headers.retain(|(name, _)| *name != CLOUD_AGENT_ID_HEADER); + headers.push((CLOUD_AGENT_ID_HEADER, task_id.to_string())); + Ok(headers) + } + fn create_oauth_client() -> self::auth::OAuth2Client { let server_root = Url::parse(&ChannelState::server_root_url()).expect("Server root URL must be valid"); diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index 6e859d363..fed8a37ff 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -1904,6 +1904,7 @@ impl AIClient for ServerApi { struct UpdateBody { sequence: i64, } + self.patch_public_api_unit( &format!("agent/runs/{run_id}/event-sequence"), &UpdateBody { sequence }, @@ -2267,12 +2268,6 @@ fn convert_harness(harness: warp_graphql::ai::AgentHarness) -> AIAgentHarness { warp_graphql::ai::AgentHarness::Oz => AIAgentHarness::Oz, warp_graphql::ai::AgentHarness::ClaudeCode => AIAgentHarness::ClaudeCode, warp_graphql::ai::AgentHarness::Gemini => AIAgentHarness::Gemini, - warp_graphql::ai::AgentHarness::Other(value) => { - report_error!(anyhow!( - "Invalid AgentHarness '{value}'. Make sure to update client GraphQL types!" - )); - AIAgentHarness::Unknown - } } } diff --git a/app/src/server/server_api/harness_support.rs b/app/src/server/server_api/harness_support.rs index 7f0a982c5..9bfe27ee2 100644 --- a/app/src/server/server_api/harness_support.rs +++ b/app/src/server/server_api/harness_support.rs @@ -10,9 +10,10 @@ use mockall::automock; use super::ServerApi; use crate::ai::agent::conversation::AIConversationId; -#[cfg(not(target_family = "wasm"))] use crate::ai::agent_sdk::retry::with_bounded_retry; +use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::artifacts::Artifact; +use crate::server::server_api::auth::AuthClient; /// A presigned upload target returned by the server. #[derive(Debug, Clone, serde::Deserialize)] @@ -86,10 +87,8 @@ pub struct ResolvedHarnessPrompt { pub prompt: String, #[serde(default)] pub system_prompt: Option, - /// Optional user-turn preamble for resumed third-party harness sessions. The harness - /// decides how to surface this — Claude Code prepends it to the user-turn prompt fed - /// into the CLI so the agent treats it as immediate intent rather than background - /// system context. Empty when no resumption is in effect. + /// Optional user-turn preamble for resumed third-party harness sessions. + /// Each harness decides how to surface this. #[serde(default)] pub resumption_prompt: Option, } @@ -152,22 +151,117 @@ pub trait HarnessSupportClient: 'static + Send + Sync { request: &SnapshotUploadRequest, ) -> Result>; - /// Download the raw third-party harness transcript bytes for the current task's - /// conversation. - /// - /// Hits `GET /harness-support/transcript`, which redirects to a signed GCS URL. - /// The conversation is resolved from the task's `agent_conversation_id` server-side, - /// so callers do not pass a conversation id. Each harness deserializes the returned - /// bytes into its own envelope shape (e.g. Claude Code parses - /// `ClaudeTranscriptEnvelope`). Transient failures retry with bounded exponential - /// backoff; permanent 4xx (e.g. 404 "no transcript") fail fast so the caller can - /// surface a resume-specific error. + /// Download the raw third-party harness transcript bytes for the current + /// task's conversation. async fn fetch_transcript(&self) -> Result; /// Get an HTTP client to use with [`UploadTarget`]s for saving blobs. fn http_client(&self) -> &http_client::Client; } +impl ServerApi { + async fn get_public_api_response_for_task( + &self, + task_id: &AmbientAgentTaskId, + path: &str, + ) -> Result { + let auth_token = self + .get_or_refresh_access_token() + .await + .context("Failed to get access token for API request")?; + + let url = format!("{}/api/v1/{}", crate::ChannelState::server_root_url(), path); + + let mut request = self.client.get(&url); + if let Some(token) = auth_token.as_bearer_token() { + request = request.bearer_auth(token); + } + + for (name, value) in self.ambient_agent_headers_for_task(task_id).await? { + request = request.header(name, value); + } + + let response = request + .send() + .await + .with_context(|| format!("Failed to send API request to {url}"))?; + + if response.status().is_success() { + Ok(response) + } else { + Err(Self::error_from_response(response).await) + } + } + + async fn post_public_api_response_for_task( + &self, + task_id: &AmbientAgentTaskId, + path: &str, + body: &B, + ) -> Result + where + B: serde::Serialize, + { + let auth_token = self + .get_or_refresh_access_token() + .await + .context("Failed to get access token for API request")?; + + let url = format!("{}/api/v1/{}", crate::ChannelState::server_root_url(), path); + + let mut request = self.client.post(&url).json(body); + if let Some(token) = auth_token.as_bearer_token() { + request = request.bearer_auth(token); + } + + for (name, value) in self.ambient_agent_headers_for_task(task_id).await? { + request = request.header(name, value); + } + + let response = request + .send() + .await + .with_context(|| format!("Failed to send API request to {url}"))?; + + if response.status().is_success() { + Ok(response) + } else { + Err(Self::error_from_response(response).await) + } + } + + pub(crate) async fn resolve_prompt_for_task( + &self, + task_id: &AmbientAgentTaskId, + request: ResolvePromptRequest, + ) -> Result { + let response = self + .post_public_api_response_for_task(task_id, "harness-support/resolve-prompt", &request) + .await?; + let url = response.url().clone(); + response + .json::() + .await + .with_context(|| format!("Failed to deserialize response from {url}")) + } + + pub(crate) async fn fetch_transcript_for_task( + &self, + task_id: &AmbientAgentTaskId, + ) -> Result { + with_bounded_retry("fetch task-scoped harness-support transcript", || async { + let response = self + .get_public_api_response_for_task(task_id, "harness-support/transcript") + .await?; + response + .bytes() + .await + .context("Failed to read harness-support transcript response body") + }) + .await + } +} + #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] impl HarnessSupportClient for ServerApi { @@ -253,25 +347,16 @@ impl HarnessSupportClient for ServerApi { } async fn fetch_transcript(&self) -> Result { - #[cfg(not(target_family = "wasm"))] - { - with_bounded_retry("fetch harness-support transcript", || async { - let response = self - .get_public_api_response("harness-support/transcript") - .await?; - response - .bytes() - .await - .context("Failed to read harness-support transcript body") - }) - .await - } - #[cfg(target_family = "wasm")] - { - unreachable!( - "fetch_transcript is not supported on wasm; agent_sdk is not built on this target" - ); - } + with_bounded_retry("fetch harness-support transcript", || async { + let response = self + .get_public_api_response("harness-support/transcript") + .await?; + response + .bytes() + .await + .context("Failed to read harness-support transcript response body") + }) + .await } fn http_client(&self) -> &http_client::Client { diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index e8a47b0c0..36f816cf6 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -4612,6 +4612,9 @@ impl TerminalView { if let BlocklistAIControllerEvent::SentRequest { model_id, .. } = event { self.maybe_insert_aws_bedrock_login_banner(model_id, ctx); } + if let BlocklistAIControllerEvent::ExecuteLocalHarnessCommand { command } = event { + self.execute_command_or_set_pending(command, ctx); + } if let BlocklistAIControllerEvent::FinishedReceivingOutput { conversation_id, .. } = event diff --git a/crates/http_client/src/lib.rs b/crates/http_client/src/lib.rs index f10875d09..2944038f1 100644 --- a/crates/http_client/src/lib.rs +++ b/crates/http_client/src/lib.rs @@ -217,6 +217,13 @@ impl Client { ) } + pub fn patch(&self, url: U) -> RequestBuilder<'_> { + self.builder( + self.wrapped.patch(url.clone()), + Self::include_warp_http_headers(url), + ) + } + pub fn delete(&self, url: U) -> RequestBuilder<'_> { self.builder( self.wrapped.delete(url.clone()), diff --git a/crates/warp_cli/src/agent.rs b/crates/warp_cli/src/agent.rs index 647ba7632..503523e93 100644 --- a/crates/warp_cli/src/agent.rs +++ b/crates/warp_cli/src/agent.rs @@ -278,7 +278,11 @@ pub struct RunAgentArgs { #[command(flatten)] pub snapshot: SnapshotArgs, /// Identifier for the task that spawned this agent, used to report progress. - #[arg(long = "task-id", hide = true, conflicts_with_all = ["prompt", "saved_prompt", "file"])] + /// + /// Mutually exclusive with `--conversation`: when `--task-id` is set the conversation id + /// (if any) is read off the server-side task metadata, so the caller never needs to pass + /// both. + #[arg(long = "task-id", hide = true, conflicts_with_all = ["prompt", "saved_prompt", "file", "conversation"])] pub task_id: Option, /// Whether we are running the agent in a sandboxed environment. diff --git a/crates/warp_cli/src/lib_tests.rs b/crates/warp_cli/src/lib_tests.rs index 69cbbc6ea..56396586c 100644 --- a/crates/warp_cli/src/lib_tests.rs +++ b/crates/warp_cli/src/lib_tests.rs @@ -1401,6 +1401,24 @@ fn agent_run_cloud_accepts_snapshot_flags() { ); } +#[test] +fn agent_run_rejects_task_id_with_conversation() { + let err = Args::try_parse_from([ + "warp", + "agent", + "run", + "--task-id", + "task-123", + "--conversation", + "conv-123", + ]) + .unwrap_err() + .to_string(); + + assert!(err.contains("--task-id")); + assert!(err.contains("--conversation")); +} + #[test] fn agent_run_cloud_accepts_computer_use_flag() { let args = Args::try_parse_from([ From 91013fcae5bd94d10c354783d752165a39042a42 Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Tue, 28 Apr 2026 08:51:14 -0400 Subject: [PATCH 02/13] Fix bug with parent sending messages to children --- .../action_model/execute/send_message.rs | 20 +++-- app/src/pane_group/pane/terminal_pane.rs | 8 +- app/src/server/server_api/ai.rs | 14 ++++ app/src/server/server_api/ai_test.rs | 24 ++++++ app/src/server/server_api/harness_support.rs | 2 +- app/src/terminal/view.rs | 35 ++++++-- app/src/terminal/view_test.rs | 82 +++++++++++++++++++ 7 files changed, 167 insertions(+), 18 deletions(-) diff --git a/app/src/ai/blocklist/action_model/execute/send_message.rs b/app/src/ai/blocklist/action_model/execute/send_message.rs index 9818371ce..449f9d080 100644 --- a/app/src/ai/blocklist/action_model/execute/send_message.rs +++ b/app/src/ai/blocklist/action_model/execute/send_message.rs @@ -57,14 +57,19 @@ impl SendMessageToAgentExecutor { let message_body = message.clone(); if FeatureFlag::OrchestrationV2.is_enabled() { - let sender_run_id = BlocklistAIHistoryModel::as_ref(ctx) + let (sender_run_id, task_id) = BlocklistAIHistoryModel::as_ref(ctx) .conversation(&conversation_id) - .and_then(|c| c.run_id()) - .map(|s| s.to_string()) - .unwrap_or_default(); + .map(|conversation| { + ( + conversation.run_id().unwrap_or_default(), + conversation.task_id(), + ) + }) + .unwrap_or_else(|| (String::new(), None)); let log_addresses = addresses.clone(); let log_subject = subject.clone(); let log_sender_run_id = sender_run_id.clone(); + let server_api = ServerApiProvider::as_ref(ctx).get(); let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); let request = SendAgentMessageRequest { to: addresses, @@ -73,7 +78,12 @@ impl SendMessageToAgentExecutor { sender_run_id, }; return ActionExecution::new_async( - async move { ai_client.send_agent_message(request).await }, + async move { + match task_id { + Some(task_id) => server_api.send_agent_message_for_task(&task_id, request).await, + None => ai_client.send_agent_message(request).await, + } + }, move |result, ctx| match result { Ok(response) => { let message_id = diff --git a/app/src/pane_group/pane/terminal_pane.rs b/app/src/pane_group/pane/terminal_pane.rs index cba52d7fe..0a5935cef 100644 --- a/app/src/pane_group/pane/terminal_pane.rs +++ b/app/src/pane_group/pane/terminal_pane.rs @@ -1223,16 +1223,16 @@ fn handle_terminal_view_event( ); new_terminal_view.update(ctx, |terminal_view, ctx| { - terminal_view.execute_command_or_set_pending( - &command, - ctx, - ); terminal_view.enter_agent_view( None, Some(conversation_id), AgentViewEntryOrigin::ChildAgent, ctx, ); + terminal_view.execute_command_or_set_pending( + &command, + ctx, + ); }); } else { create_error_child_agent_conversation( diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index fed8a37ff..883b5a8ac 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -979,6 +979,20 @@ fn into_file_artifact_record( } } +impl ServerApi { + pub(crate) async fn send_agent_message_for_task( + &self, + task_id: &AmbientAgentTaskId, + request: SendAgentMessageRequest, + ) -> anyhow::Result { + let response = self + .post_public_api_response_for_task(task_id, "agent/messages", &request) + .await?; + let response = response.json::().await?; + Ok(response) + } +} + #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] impl AIClient for ServerApi { diff --git a/app/src/server/server_api/ai_test.rs b/app/src/server/server_api/ai_test.rs index ed188c61b..d8e6e84a6 100644 --- a/app/src/server/server_api/ai_test.rs +++ b/app/src/server/server_api/ai_test.rs @@ -1,6 +1,9 @@ use chrono::TimeZone; use chrono::Utc; +use futures::executor::block_on; +use super::super::auth::CLOUD_AGENT_ID_HEADER; +use super::super::ServerApi; use super::{ build_list_agent_runs_url, AgentMessageHeader, AgentRunEvent, AgentSource, AmbientAgentTaskState, Artifact, ArtifactDownloadResponse, ArtifactType, ExecutionLocation, @@ -8,6 +11,27 @@ use super::{ }; use crate::notebooks::NotebookId; +#[test] +fn ambient_agent_headers_for_task_overrides_existing_cloud_agent_header() { + let server_api = ServerApi::new_for_test(); + let ambient_task_id = "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(); + let task_scoped_id = "123e4567-e89b-12d3-a456-426614174000".parse().unwrap(); + + server_api.set_ambient_agent_task_id(Some(ambient_task_id)); + + let cloud_agent_headers: Vec<_> = + block_on(server_api.ambient_agent_headers_for_task(&task_scoped_id)) + .unwrap() + .into_iter() + .filter(|(name, _)| *name == CLOUD_AGENT_ID_HEADER) + .collect(); + + assert_eq!( + cloud_agent_headers, + vec![(CLOUD_AGENT_ID_HEADER, task_scoped_id.to_string())] + ); +} + #[test] fn test_deserialize_file_artifact_download_response() { let json = r#"{ diff --git a/app/src/server/server_api/harness_support.rs b/app/src/server/server_api/harness_support.rs index 9bfe27ee2..6d403db8e 100644 --- a/app/src/server/server_api/harness_support.rs +++ b/app/src/server/server_api/harness_support.rs @@ -193,7 +193,7 @@ impl ServerApi { } } - async fn post_public_api_response_for_task( + pub(crate) async fn post_public_api_response_for_task( &self, task_id: &AmbientAgentTaskId, path: &str, diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index 36f816cf6..d636594c0 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -11622,6 +11622,32 @@ impl TerminalView { } } + fn child_conversation_id_for_cli_status_updates( + &self, + ctx: &AppContext, + ) -> Option { + if let Some(conversation_id) = BlocklistAIHistoryModel::as_ref(ctx) + .active_conversation(self.view_id) + .and_then(|conversation| { + conversation + .is_child_agent_conversation() + .then_some(conversation.id()) + }) + { + return Some(conversation_id); + } + + let mut child_conversation_ids = BlocklistAIHistoryModel::as_ref(ctx) + .all_live_conversations_for_terminal_view(self.view_id) + .filter(|conversation| conversation.is_child_agent_conversation()) + .map(|conversation| conversation.id()); + let child_conversation_id = child_conversation_ids.next()?; + child_conversation_ids + .next() + .is_none() + .then_some(child_conversation_id) + } + /// If the startup auto-open setting is enabled, auto-opens rich input for a /// CLI agent session. Called after creating a command-detected session or /// registering a listener so rich input is shown immediately. @@ -11705,14 +11731,7 @@ impl TerminalView { return; } - let active_child_conversation_id = BlocklistAIHistoryModel::as_ref(ctx) - .active_conversation(self.view_id) - .and_then(|conversation| { - conversation - .is_child_agent_conversation() - .then_some(conversation.id()) - }); - if let Some(conversation_id) = active_child_conversation_id { + if let Some(conversation_id) = self.child_conversation_id_for_cli_status_updates(ctx) { BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { history_model.update_conversation_status( self.view_id, diff --git a/app/src/terminal/view_test.rs b/app/src/terminal/view_test.rs index 8c61cbb17..5b828706d 100644 --- a/app/src/terminal/view_test.rs +++ b/app/src/terminal/view_test.rs @@ -4158,6 +4158,88 @@ fn cli_session_status_updates_active_child_conversation() { }) } +#[test] +fn cli_session_status_updates_single_child_conversation_without_agent_view() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let _agent_view = FeatureFlag::AgentView.override_enabled(true); + + let terminal = add_window_with_terminal(&mut app, None); + + let child_conversation_id = terminal.update(&mut app, |view, ctx| { + let parent_conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.start_new_conversation(view.view_id, false, false, ctx) + }); + let child_conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.start_new_child_conversation( + view.view_id, + "Agent 2".to_string(), + parent_conversation_id, + ctx, + ) + }); + + CLIAgentSessionsModel::handle(ctx).update(ctx, |sessions, ctx| { + sessions.set_session( + view.view_id, + CLIAgentSession { + agent: CLIAgent::Claude, + status: CLIAgentSessionStatus::InProgress, + session_context: CLIAgentSessionContext::default(), + input_state: CLIAgentInputState::Closed, + should_auto_toggle_input: false, + listener: None, + remote_host: None, + plugin_version: None, + draft_text: None, + custom_command_prefix: None, + }, + ctx, + ); + }); + + child_conversation_id + }); + + terminal.read(&app, |_view, ctx| { + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&child_conversation_id) + .expect("child conversation should exist"); + assert_eq!(conversation.status(), &ConversationStatus::InProgress); + }); + + terminal.update(&mut app, |view, ctx| { + CLIAgentSessionsModel::handle(ctx).update(ctx, |sessions, ctx| { + sessions.update_from_event( + view.view_id, + &CLIAgentEvent { + v: 1, + agent: CLIAgent::Claude, + event: CLIAgentEventType::Stop, + session_id: None, + cwd: None, + project: None, + payload: CLIAgentEventPayload { + response: Some("Done".to_owned()), + ..Default::default() + }, + }, + ctx, + ); + }); + }); + + terminal.read(&app, |_view, ctx| { + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&child_conversation_id) + .expect("child conversation should exist"); + assert_eq!(conversation.status(), &ConversationStatus::Success); + }); + }) +} + #[test] fn manual_dismiss_disables_auto_toggle_for_session() { App::test((), |mut app| async move { From 932ca2761e562ea93d05fa0c50541ec077aee3e3 Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Tue, 28 Apr 2026 14:23:56 -0400 Subject: [PATCH 03/13] fix child not sending message to parent --- app/src/ai/agent_events/message_hydrator.rs | 63 ++- .../agent_sdk/driver/harness/claude_code.rs | 414 ++++-------------- .../harness/claude_code/parent_bridge.rs | 13 +- .../driver/harness/claude_code_tests.rs | 21 + app/src/ai/blocklist/action_model/execute.rs | 3 + .../action_model/execute/send_message.rs | 145 +++++- .../execute/send_message_tests.rs | 60 +++ app/src/ai/blocklist/controller.rs | 78 +++- .../blocklist/orchestration_event_poller.rs | 17 +- app/src/pane_group/child_agent.rs | 34 +- app/src/pane_group/mod.rs | 14 + app/src/pane_group/mod_tests.rs | 140 ++++++ app/src/pane_group/pane/terminal_pane.rs | 10 +- app/src/server/server_api/ai.rs | 30 ++ 14 files changed, 674 insertions(+), 368 deletions(-) create mode 100644 app/src/ai/blocklist/action_model/execute/send_message_tests.rs diff --git a/app/src/ai/agent_events/message_hydrator.rs b/app/src/ai/agent_events/message_hydrator.rs index e73d78c35..7d5711181 100644 --- a/app/src/ai/agent_events/message_hydrator.rs +++ b/app/src/ai/agent_events/message_hydrator.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use std::time::Duration; +use crate::ai::ambient_agents::AmbientAgentTaskId; use anyhow::{anyhow, Context, Result}; #[cfg(not(target_family = "wasm"))] use futures::future::Either; @@ -9,6 +10,7 @@ use warpui::r#async::Timer; use crate::ai::agent::ReceivedMessageInput; use crate::server::server_api::ai::{AIClient, AgentRunEvent, ReadAgentMessageResponse}; +use crate::server::server_api::ServerApi; pub(crate) const DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT: Duration = Duration::from_secs(5); @@ -17,6 +19,8 @@ pub(crate) const DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT: Duration = Duration::from_ #[derive(Clone)] pub(crate) struct MessageHydrator { ai_client: Arc, + task_scoped_server_api: Option>, + task_id: Option, #[cfg_attr(target_family = "wasm", allow(dead_code))] fetch_timeout: Duration, } @@ -26,16 +30,40 @@ impl MessageHydrator { Self::with_fetch_timeout(ai_client, DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT) } + pub(crate) fn for_task(server_api: Arc, task_id: AmbientAgentTaskId) -> Self { + let ai_client: Arc = server_api.clone(); + Self { + ai_client, + task_scoped_server_api: Some(server_api), + task_id: Some(task_id), + fetch_timeout: DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT, + } + } + pub(crate) fn with_fetch_timeout( ai_client: Arc, fetch_timeout: Duration, ) -> Self { Self { ai_client, + task_scoped_server_api: None, + task_id: None, fetch_timeout, } } + async fn read_message(&self, message_id: &str) -> Result { + match (self.task_scoped_server_api.as_ref(), self.task_id) { + (Some(server_api), Some(task_id)) => { + server_api + .read_agent_message_for_task(&task_id, message_id) + .await + } + _ => self.ai_client.read_agent_message(message_id).await, + } + .with_context(|| format!("Failed to read agent message {message_id}")) + } + pub(crate) async fn hydrate_event_for_recipient( &self, event: &AgentRunEvent, @@ -55,6 +83,17 @@ impl MessageHydrator { return None; } }; + if message.body.is_empty() { + log::warn!( + "Hydrated empty-body agent message: message_id={} event_sequence={} recipient_run_id={} sender_run_id={} subject={:?} task_id={:?}", + message.message_id, + event.sequence, + recipient_run_id, + message.sender_run_id, + message.subject, + self.task_id.map(|task_id| task_id.to_string()) + ); + } Some(ReceivedMessageInput { message_id: message.message_id, @@ -70,15 +109,13 @@ impl MessageHydrator { &self, message_id: &str, ) -> Result { - let read_message = self.ai_client.read_agent_message(message_id); + let read_message = self.read_message(message_id); let timeout = Timer::after(self.fetch_timeout); futures::pin_mut!(read_message); futures::pin_mut!(timeout); match futures::future::select(read_message, timeout).await { - Either::Left((result, _)) => { - result.with_context(|| format!("Failed to read agent message {message_id}")) - } + Either::Left((result, _)) => result, Either::Right(_) => Err(anyhow!("Timed out reading agent message {message_id}")), } } @@ -88,10 +125,7 @@ impl MessageHydrator { &self, message_id: &str, ) -> Result { - self.ai_client - .read_agent_message(message_id) - .await - .with_context(|| format!("Failed to read agent message {message_id}")) + self.read_message(message_id).await } pub(crate) async fn read_message_from_event_with_timeout( @@ -105,10 +139,15 @@ impl MessageHydrator { } pub(crate) async fn mark_message_delivered(&self, message_id: &str) -> Result<()> { - self.ai_client - .mark_message_delivered(message_id) - .await - .with_context(|| format!("Failed to mark agent message {message_id} as delivered")) + match (self.task_scoped_server_api.as_ref(), self.task_id) { + (Some(server_api), Some(task_id)) => { + server_api + .mark_message_delivered_for_task(&task_id, message_id) + .await + } + _ => self.ai_client.mark_message_delivered(message_id).await, + } + .with_context(|| format!("Failed to mark agent message {message_id} as delivered")) } pub(crate) async fn mark_messages_delivered_best_effort<'a, I>( diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code.rs b/app/src/ai/agent_sdk/driver/harness/claude_code.rs index c56e0029f..5e6416cf2 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; +use std::ffi::OsString; use std::fmt::Write as _; -use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -9,10 +9,10 @@ use async_trait::async_trait; use parking_lot::Mutex; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; +use shell_words::quote as shell_quote; use tempfile::NamedTempFile; use uuid::Uuid; use warp_cli::agent::Harness; -use warp_core::safe_warn; use warpui::{ModelHandle, ModelSpawner}; use crate::ai::agent::conversation::AIConversationId; @@ -28,10 +28,14 @@ use crate::terminal::CLIAgent; use super::super::terminal::{CommandHandle, TerminalDriver}; use super::super::{AgentDriver, AgentDriverError}; +use super::claude_transcript::{ + claude_config_dir, read_envelope, write_envelope, write_session_index_entry, ClaudeResumeInfo, + ClaudeTranscriptEnvelope, +}; use super::json_utils::{read_json_file_or_default, write_json_file}; use super::{ - cli_agent_session_status, write_temp_file, HarnessCleanupDisposition, HarnessRunner, - ManagedSecretValue, ResumePayload, SavePoint, ThirdPartyHarness, + cli_agent_session_status, task_env_vars, write_temp_file, HarnessCleanupDisposition, + HarnessRunner, ManagedSecretValue, ResumePayload, SavePoint, ThirdPartyHarness, }; mod parent_bridge; @@ -52,13 +56,6 @@ use parent_bridge::{ MESSAGE_BRIDGE_CONTEXT_PREAMBLE, }; -#[derive(Debug)] -pub(crate) struct ClaudeResumeInfo { - pub(crate) conversation_id: AIConversationId, - pub(crate) session_id: Uuid, - pub(crate) envelope: ClaudeTranscriptEnvelope, -} - #[derive(Debug, Clone)] pub(crate) struct ClaudeWakeMessage { pub(crate) sequence: i64, @@ -133,6 +130,8 @@ impl ClaudeHarness { pub(crate) async fn prepare_local_wake_command( server_api: Arc, + task_id: AmbientAgentTaskId, + parent_run_id: Option, working_dir: Option, mut remote: ClaudeWakeRemoteContext, pending_messages: Vec, @@ -152,7 +151,7 @@ impl ClaudeHarness { let state_dir = parent_bridge_root()?.join(remote.session_id.to_string()); ensure_parent_bridge_state_dir(&state_dir)?; - let hydrator = MessageHydrator::new(server_api); + let hydrator = MessageHydrator::for_task(server_api, task_id); acknowledge_parent_bridge_hook_output(&hydrator, &state_dir).await?; for record in pending_messages .into_iter() @@ -167,16 +166,42 @@ impl ClaudeHarness { std::fs::write(&prompt_path, remote.wake_prompt.as_bytes()) .with_context(|| format!("Failed to write {}", prompt_path.display()))?; - Ok(claude_command( + let command = claude_command( CLIAgent::Claude.command_prefix(), &remote.session_id, &prompt_path.display().to_string(), None, true, - )) + ); + let env_vars = task_env_vars(Some(&task_id), parent_run_id.as_deref(), Harness::Claude); + + Ok(prefix_command_with_env_vars(command, env_vars)) } } +fn prefix_command_with_env_vars(command: String, env_vars: HashMap) -> String { + if env_vars.is_empty() { + return command; + } + + let mut env_pairs = env_vars + .into_iter() + .map(|(key, value)| { + ( + key.to_string_lossy().into_owned(), + value.to_string_lossy().into_owned(), + ) + }) + .collect::>(); + env_pairs.sort_unstable_by(|(left, _), (right, _)| left.cmp(right)); + + let assignments = env_pairs + .into_iter() + .map(|(key, value)| format!("{key}={}", shell_quote(&value))) + .collect::>() + .join(" "); + format!("env {assignments} {command}") +} #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] impl ThirdPartyHarness for ClaudeHarness { @@ -206,6 +231,10 @@ impl ThirdPartyHarness for ClaudeHarness { }) } + /// Fetch the Claude Code transcript for the current task's conversation and wrap it + /// into a [`ResumePayload::Claude`]. Maps a server 404 to + /// [`AgentDriverError::ConversationResumeStateMissing`] tagged as the `claude` harness + /// so the user sees a resume-specific error rather than a generic load failure. async fn fetch_resume_payload( &self, conversation_id: &AIConversationId, @@ -216,6 +245,8 @@ impl ThirdPartyHarness for ClaudeHarness { .fetch_transcript() .await .map_err(|err| { + // A 404 from the server maps to "no stored transcript" so the CLI can tell + // the user the prior run never saved state. let message = format!("{err:#}").to_lowercase(); if message.contains("status 404") { AgentDriverError::ConversationResumeStateMissing { @@ -231,9 +262,10 @@ impl ThirdPartyHarness for ClaudeHarness { "Failed to deserialize Claude transcript for {conversation_id_str}: {err:#}" )) })?; + let session_id = envelope.uuid; Ok(Some(ResumePayload::Claude(ClaudeResumeInfo { conversation_id: *conversation_id, - session_id: envelope.uuid, + session_id, envelope, }))) } @@ -249,9 +281,14 @@ impl ThirdPartyHarness for ClaudeHarness { terminal_driver: ModelHandle, resume: Option, ) -> Result, AgentDriverError> { + // Extract the Claude variant; any other variant is ignored since it belongs to a + // different harness. Today there are no other variants, but this keeps the shape + // ready for future CLI-specific payloads. let claude_resume = resume.map(|payload| match payload { ResumePayload::Claude(info) => info, }); + // Claude treats the user-turn message as immediate intent, so the resumption preamble + // is most reliable when prepended directly to the prompt that gets piped into the CLI. let owned_prompt = match resumption_prompt { Some(preamble) if !preamble.is_empty() => format!("{preamble}\n\n{prompt}"), _ => prompt.to_string(), @@ -279,6 +316,11 @@ const CLAUDE_WAKE_PROMPT_FILE_NAME: &str = "wake-turn-prompt.txt"; /// Build the shell command that launches the Claude CLI for a given session and /// prompt file. +/// +/// When `resuming` is true we pass `--resume ` so Claude picks up the +/// existing on-disk session; otherwise we pass `--session-id ` to pin a +/// fresh session to that id. If `system_prompt_path` is provided, the CLI is +/// told to append its contents to the base system prompt. fn claude_command( cli_name: &str, session_id: &Uuid, @@ -322,10 +364,14 @@ struct ClaudeHarnessRunner { parent_bridge: Option, /// Lazily cached output of `claude --version`. claude_version: Mutex>, + /// When resuming an existing conversation, we pin the runner's server conversation id + /// up front instead of calling `create_external_conversation` in [`HarnessRunner::start`]. + /// Subsequent saves overwrite the same GCS objects keyed by this id. preexisting_conversation_id: Option, } impl ClaudeHarnessRunner { + #[allow(clippy::too_many_arguments)] fn new( cli_command: &str, prompt: &str, @@ -346,20 +392,25 @@ impl ClaudeHarnessRunner { session_id, mut envelope, }) => { + // Rehydrate the stored envelope under the current working directory so + // `claude --resume ` finds the jsonl under ~/.claude/projects//. + // The original envelope's cwd usually points at the cloud sandbox path, which + // doesn't exist locally. envelope.cwd = working_dir.to_path_buf(); - let config_root = claude_config_dir().map_err(|error| { + let config_root = claude_config_dir().map_err(|e| { AgentDriverError::ConfigBuildFailed( - error.context("Failed to resolve Claude config dir"), + e.context("Failed to resolve Claude config dir"), ) })?; - write_envelope(&envelope, &config_root).map_err(|error| { + write_envelope(&envelope, &config_root).map_err(|e| { AgentDriverError::ConfigBuildFailed( - error.context("Failed to rehydrate Claude transcript"), + e.context("Failed to rehydrate Claude transcript"), ) })?; - if let Err(error) = write_session_index_entry(session_id, working_dir, &config_root) - { - log::warn!("Failed to update Claude sessions-index.json: {error:#}"); + // Index write is best-effort: upstream Claude versions vary in how they use + // `sessions-index.json`, so losing the index entry shouldn't abort the run. + if let Err(e) = write_session_index_entry(session_id, working_dir, &config_root) { + log::warn!("Failed to update Claude sessions-index.json: {e:#}"); } (session_id, Some(conversation_id), true) } @@ -506,23 +557,28 @@ impl HarnessRunner for ClaudeHarnessRunner { &self, foreground: &ModelSpawner, ) -> Result { + // When resuming, we already have a server conversation id from the prior run. + // Otherwise create a fresh external conversation record for this run. + // TODO(REMOTE-1149): `create_external_conversation` currently won't work for local CLI + // runs. We should either support it or have a fallback. let conversation_id = match self.preexisting_conversation_id { - Some(conversation_id) => { - log::info!("Resuming external conversation {conversation_id}"); - conversation_id + Some(id) => { + log::info!("Resuming external conversation {id}"); + id + } + None => { + let id = self + .client + .create_external_conversation(CLAUDE_CODE_FORMAT) + .await + .map_err(|e| { + log::error!("Failed to create external conversation: {e}"); + AgentDriverError::ConfigBuildFailed(e) + })?; + log::info!("Created external conversation {id}"); + id } - None => self - .client - .create_external_conversation(CLAUDE_CODE_FORMAT) - .await - .map_err(|e| { - log::error!("Failed to create external conversation: {e}"); - AgentDriverError::ConfigBuildFailed(e) - })?, }; - if self.preexisting_conversation_id.is_none() { - log::info!("Created external conversation {conversation_id}"); - } self.start_parent_bridge(foreground) .await .map_err(AgentDriverError::ConfigBuildFailed)?; @@ -657,284 +713,6 @@ async fn upload_transcript( .with_context(|| format!("Failed to get transcript upload target for {conversation_id}"))?; upload_to_target(client.http_client(), &target, body).await } - -// ─── Transcript envelope ────────────────────────────────────────────────────── - -/// JSON envelope sent to the server representing a complete Claude Code session. -/// -/// Bundles the main session transcript, any subagent transcripts, and -/// per-agent TODO lists assembled from the Claude state directory. -#[derive(Debug, PartialEq, Serialize, Deserialize)] -pub(crate) struct ClaudeTranscriptEnvelope { - /// The directory that the Claude Code session started in. - cwd: PathBuf, - /// Unique session identifier. - uuid: Uuid, - /// Claude Code version, if available. - #[serde(default, skip_serializing_if = "Option::is_none")] - claude_version: Option, - /// List of messages in the main agent conversation. - entries: Vec, - /// Messages in each subagent conversation, keyed by the agent filename (e.g. `"agent-aac0b7f3db6bccfaf"`). - subagents: HashMap>, - /// TODO lists for each agent, keyed on the session and agent (e.g. `"-agent-"`). - todos: HashMap, -} - -/// Encode a filesystem path as a Claude config directory name, matching the -/// Claude CLI convention of replacing every `/` with `-`. -/// -/// Example: `/Users/ben/src/foo` → `-Users-ben-src-foo` -fn encode_cwd(cwd: &Path) -> String { - cwd.to_string_lossy().replace(['/', '.'], "-") -} - -/// Resolve the Claude config directory. -/// -/// Reads `$CLAUDE_CONFIG_DIR` if set, otherwise falls back to `~/.claude`. -// -/// TODO(REMOTE-1209): Use the transcript path reported by our hook. -fn claude_config_dir() -> Result { - if let Ok(dir) = std::env::var("CLAUDE_CONFIG_DIR") { - return Ok(PathBuf::from(dir)); - } - dirs::home_dir() - .map(|h| h.join(".claude")) - .ok_or_else(|| anyhow::anyhow!("could not determine home directory")) -} - -/// Assemble a [`ClaudeTranscriptEnvelope`] from the Claude config directory. -/// -/// Reads: -/// - `/projects//.jsonl` - main transcript -/// - `/projects///subagents/*.jsonl` - subagents -/// - `/todos/-agent-*.json` - per-agent todo lists -/// -/// If the main JSONL does not exist yet (e.g. during an early periodic save) -/// the envelope is returned with an empty `entries` list rather than an error. -fn read_envelope( - session_uuid: Uuid, - cwd: &Path, - config_root: &Path, -) -> Result { - let encoded = encode_cwd(cwd); - let projects_dir = config_root.join("projects").join(&encoded); - - // Main session transcript. - let session_file = projects_dir.join(format!("{session_uuid}.jsonl")); - let entries = read_jsonl(&session_file)?; - - // Subagents are stored in a directory named after the session UUID. - let mut subagents: HashMap> = HashMap::new(); - let subagents_dir = projects_dir - .join(session_uuid.to_string()) - .join("subagents"); - if subagents_dir.is_dir() { - for entry in std::fs::read_dir(&subagents_dir) - .with_context(|| format!("Failed to read subagents dir {}", subagents_dir.display()))? - { - let entry = entry?; - let path = entry.path(); - if path.extension().and_then(|e| e.to_str()) != Some("jsonl") { - continue; - } - let Some(stem) = path.file_stem().and_then(|s| s.to_str()) else { - continue; - }; - subagents.insert(stem.to_owned(), read_jsonl(&path)?); - } - } - - // Per-agent todo lists. - let mut todos: HashMap = HashMap::new(); - let todos_dir = config_root.join("todos"); - let todos_prefix = format!("{session_uuid}-agent-"); - if todos_dir.is_dir() { - for entry in std::fs::read_dir(&todos_dir) - .with_context(|| format!("Failed to read todos dir {}", todos_dir.display()))? - { - let entry = entry?; - let path = entry.path(); - if path.extension().and_then(|e| e.to_str()) != Some("json") { - continue; - } - let Some(stem) = path.file_stem().and_then(|s| s.to_str()) else { - continue; - }; - if !stem.starts_with(&todos_prefix) { - continue; - } - match std::fs::read_to_string(&path) { - Ok(content) => match serde_json::from_str(&content) { - Ok(value) => { - todos.insert(stem.to_owned(), value); - } - Err(e) => log::warn!("Failed to parse todos file {}: {e}", path.display()), - }, - Err(e) => log::warn!("Failed to read todos file {}: {e}", path.display()), - } - } - } - - Ok(ClaudeTranscriptEnvelope { - cwd: cwd.to_path_buf(), - uuid: session_uuid, - claude_version: None, - entries, - subagents, - todos, - }) -} - -/// Write a [`ClaudeTranscriptEnvelope`] back to disk using the same layout -/// that Claude Code uses. -/// -/// Creates: -/// - `/projects//.jsonl` - main transcript -/// - `/projects///subagents/.jsonl` - subagents -/// - `/todos/.json` - per-agent todo lists -fn write_envelope(envelope: &ClaudeTranscriptEnvelope, config_root: &Path) -> Result<()> { - let encoded = encode_cwd(&envelope.cwd); - let projects_dir = config_root.join("projects").join(&encoded); - std::fs::create_dir_all(&projects_dir) - .with_context(|| format!("Failed to create {}", projects_dir.display()))?; - - // Main session JSONL. - let session_file = projects_dir.join(format!("{}.jsonl", envelope.uuid)); - std::fs::write(&session_file, entries_to_jsonl(&envelope.entries)?) - .with_context(|| format!("Failed to write {}", session_file.display()))?; - - // Subagent JSONLs. - if !envelope.subagents.is_empty() { - let subagents_dir = projects_dir - .join(envelope.uuid.to_string()) - .join("subagents"); - std::fs::create_dir_all(&subagents_dir) - .with_context(|| format!("Failed to create {}", subagents_dir.display()))?; - for (stem, entries) in &envelope.subagents { - let path = subagents_dir.join(format!("{stem}.jsonl")); - std::fs::write(&path, entries_to_jsonl(entries)?) - .with_context(|| format!("Failed to write {}", path.display()))?; - } - } - - // Per-agent todo lists. - if !envelope.todos.is_empty() { - let todos_dir = config_root.join("todos"); - std::fs::create_dir_all(&todos_dir) - .with_context(|| format!("Failed to create {}", todos_dir.display()))?; - for (stem, value) in &envelope.todos { - let path = todos_dir.join(format!("{stem}.json")); - std::fs::write(&path, serde_json::to_vec(value)?) - .with_context(|| format!("Failed to write {}", path.display()))?; - } - } - - Ok(()) -} - -/// Filename of Claude's global session index. -const SESSIONS_INDEX_FILENAME: &str = "sessions-index.json"; - -/// Upsert an entry for `session_uuid` into `/sessions-index.json` so Claude's -/// `claude --resume ` lookup can find the rehydrated jsonl. -/// -/// Best-effort: callers should log a warning on failure rather than aborting the run. -fn write_session_index_entry(session_uuid: Uuid, cwd: &Path, config_root: &Path) -> Result<()> { - let index_path = config_root.join(SESSIONS_INDEX_FILENAME); - - let mut index: serde_json::Map = match std::fs::read_to_string(&index_path) { - Ok(content) => match serde_json::from_str::(&content) { - Ok(Value::Object(map)) => map, - Ok(_) => { - safe_warn!( - safe: ("sessions-index.json is not a JSON object; overwriting"), - full: ("sessions-index.json at {} is not a JSON object; overwriting", index_path.display()) - ); - serde_json::Map::new() - } - Err(error) => { - safe_warn!( - safe: ("Failed to parse sessions-index.json; overwriting"), - full: ("Failed to parse sessions-index.json at {}: {error}; overwriting", index_path.display()) - ); - serde_json::Map::new() - } - }, - Err(error) if error.kind() == std::io::ErrorKind::NotFound => serde_json::Map::new(), - Err(error) => { - return Err(anyhow::Error::from(error) - .context(format!("Failed to read {}", index_path.display()))); - } - }; - - let encoded = encode_cwd(cwd); - let transcript_path = format!("projects/{encoded}/{session_uuid}.jsonl"); - let entry = serde_json::json!({ - "sessionId": session_uuid.to_string(), - "cwd": cwd.to_string_lossy(), - "projectPath": encoded, - "transcriptPath": transcript_path, - }); - index.insert(session_uuid.to_string(), entry); - - if let Some(parent) = index_path.parent() { - std::fs::create_dir_all(parent) - .with_context(|| format!("Failed to create {}", parent.display()))?; - } - std::fs::write( - &index_path, - serde_json::to_vec_pretty(&Value::Object(index)) - .context("Failed to serialize sessions-index.json")?, - ) - .with_context(|| format!("Failed to write {}", index_path.display()))?; - Ok(()) -} -/// Serialize a slice of JSON values as a JSONL byte string (one value per line). -fn entries_to_jsonl(entries: &[Value]) -> Result> { - let mut buf = Vec::new(); - for entry in entries { - serde_json::to_writer(&mut buf, entry)?; - buf.push(b'\n'); - } - Ok(buf) -} - -/// Read a JSONL file, returning one parsed [`Value`] per non-blank line. -/// -/// Lines that fail to parse as JSON are skipped with a warning rather than -/// causing the entire read to fail. A missing file returns an empty [`Vec`]. -fn read_jsonl(path: &Path) -> Result> { - let file = match std::fs::File::open(path) { - Ok(f) => f, - Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), - Err(e) => { - return Err( - anyhow::Error::from(e).context(format!("Failed to open {}", path.display())) - ); - } - }; - let reader = BufReader::new(file); - let mut entries = Vec::new(); - for line in reader.lines() { - let line = line.with_context(|| format!("Failed to read line from {}", path.display()))?; - let trimmed = line.trim(); - if trimmed.is_empty() { - continue; - } - match serde_json::from_str(trimmed) { - Ok(value) => entries.push(value), - Err(e) => { - safe_warn!( - safe: ("Skipping malformed JSONL entry"), - full: ("Skipping malformed JSONL entry in {}: {e}", path.display()) - ); - } - } - } - Ok(entries) -} - fn prepare_claude_environment_config( working_dir: &Path, secrets: &HashMap, @@ -994,7 +772,7 @@ const CLAUDE_JSON_FILE_NAME: &str = ".claude.json"; const CLAUDE_SETTINGS_FILE_NAME: &str = "settings.json"; const ANTHROPIC_API_KEY_SUFFIX_LEN: usize = 20; -#[derive(Default, Deserialize, Serialize)] +#[derive(Default, Deserialize, Serialize, Debug)] #[serde(rename_all = "camelCase")] struct ClaudeConfig { #[serde(default)] @@ -1009,7 +787,7 @@ struct ClaudeConfig { extra: Map, } -#[derive(Default, Deserialize, Serialize)] +#[derive(Default, Deserialize, Serialize, Debug)] #[serde(rename_all = "camelCase")] struct CustomApiKeyResponses { #[serde(default)] @@ -1018,7 +796,7 @@ struct CustomApiKeyResponses { extra: Map, } -#[derive(Default, Deserialize, Serialize)] +#[derive(Default, Deserialize, Serialize, Debug)] #[serde(rename_all = "camelCase")] struct ClaudeProjectConfig { #[serde(default)] @@ -1027,7 +805,7 @@ struct ClaudeProjectConfig { extra: Map, } -#[derive(Default, Deserialize, Serialize)] +#[derive(Default, Deserialize, Serialize, Debug)] #[serde(rename_all = "camelCase")] struct ClaudeSettings { #[serde(default)] diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs b/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs index 6f20b2176..36eccc75b 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs @@ -28,6 +28,7 @@ use crate::ai::agent_events::{ AgentEventDriverConfig, MessageHydrator, ServerApiAgentEventSource, }; use crate::ai::agent_sdk::driver::{AgentDriver, OZ_MESSAGE_LISTENER_STATE_ROOT_ENV}; +use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::server::server_api::ai::AIClient; use crate::server::server_api::ai::AgentRunEvent; use crate::server::server_api::ServerApi; @@ -150,6 +151,12 @@ struct SelectedMessageBridgeMessages { } impl MessageBridge { + fn hydrator(&self, server_api: Arc) -> MessageHydrator { + match self.run_id.parse::() { + Ok(task_id) => MessageHydrator::for_task(server_api, task_id), + Err(_) => MessageHydrator::new(server_api), + } + } pub(super) fn new(run_id: String, session_id: Uuid) -> Result { Ok(Self { run_id, @@ -197,8 +204,7 @@ impl MessageBridge { if !self.state_dir.exists() { return Ok(()); } - - let hydrator = MessageHydrator::new(server_api); + let hydrator = self.hydrator(server_api); let _guard = self.state_lock.lock().await; acknowledge_parent_bridge_hook_output(&hydrator, &self.state_dir).await?; prepare_parent_bridge_hook_output( @@ -213,8 +219,7 @@ impl MessageBridge { if !self.state_dir.exists() { return Ok(()); } - - let hydrator = MessageHydrator::new(server_api); + let hydrator = self.hydrator(server_api); let _guard = self.state_lock.lock().await; acknowledge_parent_bridge_hook_output(&hydrator, &self.state_dir).await } diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs b/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs index 5775e9870..243f0ca46 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs @@ -6,9 +6,12 @@ use std::sync::Arc; use tempfile::TempDir; use uuid::Uuid; +use warp_cli::{OZ_HARNESS_ENV, OZ_PARENT_RUN_ID_ENV, OZ_RUN_ID_ENV}; use super::*; use crate::ai::agent_events::MessageHydrator; +use crate::ai::agent_sdk::driver::harness::claude_transcript::encode_cwd; +use crate::ai::agent_sdk::driver::OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV; use crate::server::server_api::ai::{MockAIClient, ReadAgentMessageResponse}; use crate::server::server_api::ServerApiProvider; @@ -675,9 +678,13 @@ fn prepare_local_wake_command_rehydrates_transcript_and_stages_messages() { body: "Inspect the failing tests first.".to_string(), occurred_at: "2026-04-17T15:46:00Z".to_string(), }; + let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440010".parse().unwrap(); + let parent_run_id = "parent-run-456".to_string(); let command = futures::executor::block_on(ClaudeHarness::prepare_local_wake_command( ServerApiProvider::new_for_test().get(), + task_id, + Some(parent_run_id.clone()), Some(working_dir.clone()), remote, vec![message.clone()], @@ -690,8 +697,22 @@ fn prepare_local_wake_command_rehydrates_transcript_and_stages_messages() { parent_bridge_surfaced_message_path(&state_dir, message.sequence, &message.message_id); assert!(command.contains("--resume")); + assert!(command.starts_with("env ")); assert!(command.contains(&session_id.to_string())); assert!(command.contains(CLAUDE_WAKE_PROMPT_FILE_NAME)); + assert!(command.contains(&format!( + "{OZ_RUN_ID_ENV}={}", + shell_quote(&task_id.to_string()) + ))); + assert!(command.contains(&format!( + "{OZ_PARENT_RUN_ID_ENV}={}", + shell_quote(&parent_run_id) + ))); + assert!(command.contains(&format!("{OZ_HARNESS_ENV}={}", shell_quote("claude")))); + assert!(command.contains(&format!( + "{OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV}={}", + shell_quote("1") + ))); assert_eq!( fs::read_to_string(&prompt_path).unwrap(), "resume prompt\n\nwake prompt" diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index 5681dd3ce..0ef183d53 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -418,6 +418,9 @@ impl BlocklistAIActionExecutor { id: Option, ctx: &mut ModelContext, ) { + self.send_message_executor.update(ctx, |executor, _| { + executor.set_ambient_agent_task_id(id); + }); self.request_computer_use_executor .update(ctx, |executor, _| { executor.set_ambient_agent_task_id(id); diff --git a/app/src/ai/blocklist/action_model/execute/send_message.rs b/app/src/ai/blocklist/action_model/execute/send_message.rs index 449f9d080..79172496c 100644 --- a/app/src/ai/blocklist/action_model/execute/send_message.rs +++ b/app/src/ai/blocklist/action_model/execute/send_message.rs @@ -1,9 +1,18 @@ +use std::time::Duration; + +use anyhow::anyhow; +#[cfg(not(target_family = "wasm"))] +use futures::future::Either; use futures::{future::BoxFuture, FutureExt}; -use warpui::{Entity, ModelContext, SingletonEntity}; +#[cfg(not(target_family = "wasm"))] +use warpui::r#async::Timer; +use warpui::{AppContext, Entity, ModelContext, SingletonEntity}; use crate::ai::agent::{ - AIAgentAction, AIAgentActionResultType, AIAgentActionType, SendMessageToAgentResult, + conversation::AIConversationId, AIAgentAction, AIAgentActionResultType, AIAgentActionType, + SendMessageToAgentResult, }; +use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::blocklist::history_model::BlocklistAIHistoryModel; use crate::ai::blocklist::orchestration_events::{OrchestrationEventService, SendMessageResult}; use crate::ai::blocklist::telemetry::{ @@ -11,18 +20,111 @@ use crate::ai::blocklist::telemetry::{ TeamAgentCommunicationFailureReason, TeamAgentCommunicationKind, TeamAgentCommunicationTransport, TeamAgentOrchestrationVersion, }; -use crate::server::server_api::ai::SendAgentMessageRequest; +use crate::server::server_api::ai::{SendAgentMessageRequest, SendAgentMessageResponse}; use crate::server::server_api::ServerApiProvider; use warp_core::features::FeatureFlag; use warp_core::send_telemetry_from_ctx; use super::{ActionExecution, AnyActionExecution, ExecuteActionInput, PreprocessActionInput}; -pub struct SendMessageToAgentExecutor; +const SEND_AGENT_MESSAGE_TIMEOUT: Duration = Duration::from_secs(15); + +pub struct SendMessageToAgentExecutor { + ambient_agent_task_id: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SendMessageTaskResolution { + ConversationTask, + AmbientTaskFallback, + NoTaskContext, +} + +fn sender_run_id_and_task_id_for_send( + conversation_id: AIConversationId, + ambient_agent_task_id: Option, + ctx: &AppContext, +) -> ( + String, + Option, + SendMessageTaskResolution, +) { + let conversation = BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id); + let conversation_task_id = conversation.and_then(|conversation| conversation.task_id()); + let (task_id, task_resolution) = match (conversation_task_id, ambient_agent_task_id) { + (Some(task_id), _) => (Some(task_id), SendMessageTaskResolution::ConversationTask), + (None, Some(task_id)) => ( + Some(task_id), + SendMessageTaskResolution::AmbientTaskFallback, + ), + (None, None) => (None, SendMessageTaskResolution::NoTaskContext), + }; + let sender_run_id = conversation + .and_then(|conversation| conversation.run_id()) + .or_else(|| task_id.map(|task_id| task_id.to_string())) + .unwrap_or_default(); + (sender_run_id, task_id, task_resolution) +} + +#[cfg(not(target_family = "wasm"))] +async fn send_agent_message_with_timeout( + server_api: std::sync::Arc, + ai_client: std::sync::Arc, + task_id: Option, + request: SendAgentMessageRequest, +) -> anyhow::Result { + let task_id_for_timeout = task_id.map(|task_id| task_id.to_string()); + let send_message = async move { + match task_id { + Some(task_id) => { + server_api + .send_agent_message_for_task(&task_id, request) + .await + } + None => ai_client.send_agent_message(request).await, + } + }; + let timeout = Timer::after(SEND_AGENT_MESSAGE_TIMEOUT); + futures::pin_mut!(send_message); + futures::pin_mut!(timeout); + + match futures::future::select(send_message, timeout).await { + Either::Left((result, _)) => result, + Either::Right(_) => Err(anyhow!( + "Timed out sending orchestration message{}", + task_id_for_timeout + .map(|task_id| format!(" for task {task_id}")) + .unwrap_or_default() + )), + } +} + +#[cfg(target_family = "wasm")] +async fn send_agent_message_with_timeout( + server_api: std::sync::Arc, + ai_client: std::sync::Arc, + task_id: Option, + request: SendAgentMessageRequest, +) -> anyhow::Result { + match task_id { + Some(task_id) => { + server_api + .send_agent_message_for_task(&task_id, request) + .await + } + None => ai_client.send_agent_message(request).await, + } +} impl SendMessageToAgentExecutor { pub fn new() -> Self { - Self + Self { + ambient_agent_task_id: None, + } + } + + pub fn set_ambient_agent_task_id(&mut self, id: Option) { + self.ambient_agent_task_id = id; } pub(super) fn should_autoexecute( @@ -57,20 +159,21 @@ impl SendMessageToAgentExecutor { let message_body = message.clone(); if FeatureFlag::OrchestrationV2.is_enabled() { - let (sender_run_id, task_id) = BlocklistAIHistoryModel::as_ref(ctx) - .conversation(&conversation_id) - .map(|conversation| { - ( - conversation.run_id().unwrap_or_default(), - conversation.task_id(), - ) - }) - .unwrap_or_else(|| (String::new(), None)); + let (sender_run_id, task_id, task_resolution) = sender_run_id_and_task_id_for_send( + conversation_id, + self.ambient_agent_task_id, + ctx, + ); let log_addresses = addresses.clone(); let log_subject = subject.clone(); let log_sender_run_id = sender_run_id.clone(); + let log_task_id = task_id.map(|task_id| task_id.to_string()); + let log_body_len = message_body.chars().count(); let server_api = ServerApiProvider::as_ref(ctx).get(); let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); + log::info!( + "Sending orchestration message: conversation_id={conversation_id:?} resolution={task_resolution:?} sender_run_id={log_sender_run_id:?} task_id={log_task_id:?} target_agent_ids={log_addresses:?} subject={log_subject:?} body_len={log_body_len}" + ); let request = SendAgentMessageRequest { to: addresses, subject, @@ -79,15 +182,15 @@ impl SendMessageToAgentExecutor { }; return ActionExecution::new_async( async move { - match task_id { - Some(task_id) => server_api.send_agent_message_for_task(&task_id, request).await, - None => ai_client.send_agent_message(request).await, - } + send_agent_message_with_timeout(server_api, ai_client, task_id, request).await }, move |result, ctx| match result { Ok(response) => { let message_id = response.message_ids.into_iter().next().unwrap_or_default(); + log::info!( + "Sent orchestration message: conversation_id={conversation_id:?} resolution={task_resolution:?} sender_run_id={log_sender_run_id:?} task_id={log_task_id:?} target_agent_ids={log_addresses:?} subject={log_subject:?} body_len={log_body_len} message_id={message_id:?}" + ); AIAgentActionResultType::SendMessageToAgent( SendMessageToAgentResult::Success { message_id }, ) @@ -113,7 +216,7 @@ impl SendMessageToAgentExecutor { ctx ); log::warn!( - "Failed to send child-agent message via server API: conversation_id={conversation_id:?} sender_run_id={log_sender_run_id:?} target_agent_ids={log_addresses:?} subject={log_subject:?} error={err:#}" + "Failed to send child-agent message via server API: conversation_id={conversation_id:?} resolution={task_resolution:?} sender_run_id={log_sender_run_id:?} task_id={log_task_id:?} target_agent_ids={log_addresses:?} subject={log_subject:?} body_len={log_body_len} error={err:#}" ); AIAgentActionResultType::SendMessageToAgent( SendMessageToAgentResult::Error(error_message), @@ -155,3 +258,7 @@ impl Default for SendMessageToAgentExecutor { impl Entity for SendMessageToAgentExecutor { type Event = (); } + +#[cfg(test)] +#[path = "send_message_tests.rs"] +mod tests; diff --git a/app/src/ai/blocklist/action_model/execute/send_message_tests.rs b/app/src/ai/blocklist/action_model/execute/send_message_tests.rs new file mode 100644 index 000000000..a0e2ac50c --- /dev/null +++ b/app/src/ai/blocklist/action_model/execute/send_message_tests.rs @@ -0,0 +1,60 @@ +use super::*; +use crate::ai::blocklist::BlocklistAIHistoryModel; +use warpui::{App, EntityId}; + +#[test] +fn sender_run_id_and_task_id_for_send_falls_back_to_ambient_task_id() { + App::test((), |mut app| async move { + let terminal_view_id = EntityId::new(); + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + let conversation_id = history_model.update(&mut app, |history_model, ctx| { + history_model.start_new_conversation(terminal_view_id, false, false, ctx) + }); + let ambient_task_id = "11111111-1111-1111-1111-111111111111" + .parse() + .expect("valid ambient task id"); + + let (sender_run_id, task_id, task_resolution) = app.read(|ctx| { + sender_run_id_and_task_id_for_send(conversation_id, Some(ambient_task_id), ctx) + }); + + assert_eq!(sender_run_id, ambient_task_id.to_string()); + assert_eq!(task_id, Some(ambient_task_id)); + assert_eq!( + task_resolution, + SendMessageTaskResolution::AmbientTaskFallback + ); + }); +} + +#[test] +fn sender_run_id_and_task_id_for_send_prefers_conversation_task_id() { + App::test((), |mut app| async move { + let terminal_view_id = EntityId::new(); + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + let conversation_id = history_model.update(&mut app, |history_model, ctx| { + history_model.start_new_conversation(terminal_view_id, false, false, ctx) + }); + let conversation_task_id = "22222222-2222-2222-2222-222222222222" + .parse() + .expect("valid conversation task id"); + let ambient_task_id = "33333333-3333-3333-3333-333333333333" + .parse() + .expect("valid ambient task id"); + + history_model.update(&mut app, |history_model, _| { + history_model + .conversation_mut(&conversation_id) + .expect("conversation exists") + .set_task_id(conversation_task_id); + }); + + let (sender_run_id, task_id, task_resolution) = app.read(|ctx| { + sender_run_id_and_task_id_for_send(conversation_id, Some(ambient_task_id), ctx) + }); + + assert_eq!(sender_run_id, conversation_task_id.to_string()); + assert_eq!(task_id, Some(conversation_task_id)); + assert_eq!(task_resolution, SendMessageTaskResolution::ConversationTask); + }); +} diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 8934869f2..78469bb08 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -79,6 +79,7 @@ use std::sync::Arc; use std::time::Duration; use warp_cli::agent::Harness; use warp_core::assertions::safe_assert; +use warp_graphql::ai::AgentTaskState; use warp_multi_agent_api::{message, Task, ToolType}; use warpui::r#async::{SpawnedFutureHandle, Timer}; @@ -1536,7 +1537,7 @@ impl BlocklistAIController { &self, conversation_id: AIConversationId, ctx: &ModelContext, - ) -> Option<(AmbientAgentTaskId, Option)> { + ) -> Option<(AmbientAgentTaskId, Option, Option)> { let conversation = BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id)?; if !conversation.is_child_agent_conversation() || conversation.is_remote_child() { return None; @@ -1549,6 +1550,18 @@ impl BlocklistAIController { .current_working_directory() .cloned() .map(PathBuf::from), + conversation + .parent_agent_id() + .map(str::to_owned) + .or_else(|| { + conversation + .parent_conversation_id() + .and_then(|parent_conversation_id| { + BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&parent_conversation_id) + }) + .and_then(|parent_conversation| parent_conversation.run_id()) + }), )) } @@ -1642,7 +1655,8 @@ impl BlocklistAIController { return true; } - let Some((task_id, working_dir)) = self.local_claude_wake_candidate(conversation_id, ctx) + let Some((task_id, working_dir, parent_run_id)) = + self.local_claude_wake_candidate(conversation_id, ctx) else { return false; }; @@ -1659,6 +1673,7 @@ impl BlocklistAIController { .iter() .map(|event| event.event_id.clone()) .collect_vec(); + let pending_message_count = pending_message_event_ids.len(); let server_api = ServerApiProvider::as_ref(ctx).get(); let handle = ctx.spawn( async move { @@ -1668,9 +1683,17 @@ impl BlocklistAIController { .as_ref() .and_then(|snapshot| snapshot.harness.as_ref()) .map(|config| config.harness_type); + log::info!( + "Evaluating dormant Claude wake: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={pending_message_count} server_task_state={:?} harness={harness:?}", + task.state + ); if task.state != AmbientAgentTaskState::Succeeded || harness != Some(Harness::Claude) { + log::info!( + "Skipping dormant Claude wake: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={pending_message_count} server_task_state={:?} harness={harness:?}", + task.state + ); return Ok::, anyhow::Error>(None); } @@ -1722,18 +1745,58 @@ impl BlocklistAIController { }); return; } + log::info!( + "Prepared dormant Claude wake messages: conversation_id={conversation_id:?} task_id={task_id} message_count={}", + wake_messages.len() + ); + for wake_message in &wake_messages { + log::info!( + "Dormant Claude wake message: conversation_id={conversation_id:?} task_id={task_id} message_id={} sequence={} subject={:?} body_len={}", + wake_message.message_id, + wake_message.sequence, + wake_message.subject, + wake_message.body.chars().count() + ); + } let removed_events_for_retry = removed_events.clone(); let server_api = ServerApiProvider::as_ref(ctx).get(); + let wake_message_count = wake_messages.len(); let handle = ctx.spawn( async move { - ClaudeHarness::prepare_local_wake_command( - server_api, + log::info!( + "Preparing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={wake_message_count}" + ); + let command = ClaudeHarness::prepare_local_wake_command( + server_api.clone(), + task_id, + parent_run_id, working_dir, remote, wake_messages, ) - .await + .await?; + log::info!( + "Reopening dormant Claude task before wake command: conversation_id={conversation_id:?} task_id={task_id}" + ); + server_api + .update_agent_task( + task_id, + Some(AgentTaskState::InProgress), + None, + None, + None, + ) + .await + .map_err(|err| { + anyhow!( + "Failed to reopen dormant Claude task {task_id} before wake: {err:#}" + ) + })?; + log::info!( + "Reopened dormant Claude task before wake command: conversation_id={conversation_id:?} task_id={task_id}" + ); + Ok::<_, anyhow::Error>(command) }, move |me, result, ctx| { me.pending_local_claude_wakes.remove(&conversation_id); @@ -1750,6 +1813,9 @@ impl BlocklistAIController { return; } + log::info!( + "Executing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id}" + ); BlocklistAIHistoryModel::handle(ctx).update( ctx, |history_model, ctx| { @@ -1767,7 +1833,7 @@ impl BlocklistAIController { } Err(err) => { log::warn!( - "Failed to finalize dormant Claude wake for {conversation_id:?}: {err:#}" + "Failed to finalize dormant Claude wake for {conversation_id:?} task_id={task_id}: {err:#}" ); OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { svc.prepend_pending_events( diff --git a/app/src/ai/blocklist/orchestration_event_poller.rs b/app/src/ai/blocklist/orchestration_event_poller.rs index 104a2d616..0a71035a6 100644 --- a/app/src/ai/blocklist/orchestration_event_poller.rs +++ b/app/src/ai/blocklist/orchestration_event_poller.rs @@ -112,6 +112,12 @@ pub enum OrchestrationEventPollerEvent { } impl OrchestrationEventPoller { + fn message_hydrator_for_run_id(&self, run_id: &str) -> MessageHydrator { + match run_id.parse::() { + Ok(task_id) => MessageHydrator::for_task(self.server_api.clone(), task_id), + Err(_) => MessageHydrator::new(self.ai_client.clone()), + } + } pub fn new(ctx: &mut ModelContext) -> Self { let provider = ServerApiProvider::as_ref(ctx); let ai_client = provider.get_ai_client(); @@ -750,7 +756,8 @@ impl OrchestrationEventPoller { } } - let hydrator = MessageHydrator::new(self.ai_client.clone()); + let hydrator = + self.message_hydrator_for_run_id(conversation.run_id().as_deref().unwrap_or_default()); ctx.spawn( async move { hydrator @@ -784,15 +791,14 @@ impl OrchestrationEventPoller { .copied() .unwrap_or(0); - let ai_client = self.ai_client.clone(); - let hydrator = MessageHydrator::new(ai_client.clone()); - // Capture own run_id to filter out self-originated lifecycle events. let self_run_id = BlocklistAIHistoryModel::as_ref(ctx) .conversation(&conversation_id) .and_then(|c| c.run_id()) .map(|s| s.to_string()) .unwrap_or_default(); + let ai_client = self.ai_client.clone(); + let hydrator = self.message_hydrator_for_run_id(&self_run_id); struct PollResult { events: Vec, @@ -1023,7 +1029,6 @@ impl OrchestrationEventPoller { .unwrap_or(0); let server_api = self.server_api.clone(); - let ai_client = self.ai_client.clone(); let self_run_id = BlocklistAIHistoryModel::as_ref(ctx) .conversation(&conversation_id) @@ -1050,7 +1055,7 @@ impl OrchestrationEventPoller { let config = AgentEventDriverConfig::retry_forever(watched.clone(), cursor); let source = ServerApiAgentEventSource::new(server_api); - let hydrator = MessageHydrator::new(ai_client); + let hydrator = self.message_hydrator_for_run_id(&self_run_id); ctx.spawn( async move { diff --git a/app/src/pane_group/child_agent.rs b/app/src/pane_group/child_agent.rs index ed633684d..6d678115c 100644 --- a/app/src/pane_group/child_agent.rs +++ b/app/src/pane_group/child_agent.rs @@ -1,5 +1,7 @@ -use std::{collections::HashMap, ffi::OsString}; +use std::{collections::HashMap, ffi::OsString, path::PathBuf}; +use crate::ai::ambient_agents::AmbientAgentTaskId; +use crate::ai::attachment_utils::attachments_download_dir; use warpui::{EntityId, SingletonEntity, ViewContext, ViewHandle}; use crate::ai::agent::conversation::{AIConversationId, ConversationStatus}; @@ -15,6 +17,31 @@ pub(crate) struct HiddenChildAgentConversation { pub terminal_view_id: EntityId, pub conversation_id: AIConversationId, } +#[derive(Clone, Debug)] +pub(crate) struct HiddenChildAgentTaskContext { + pub task_id: AmbientAgentTaskId, + pub working_dir: Option, +} + +pub(crate) fn apply_hidden_child_agent_task_context( + terminal_view: &ViewHandle, + task_context: &HiddenChildAgentTaskContext, + ctx: &mut ViewContext, +) { + let task_id = task_context.task_id; + let working_dir = task_context.working_dir.clone(); + + terminal_view.update(ctx, move |terminal_view, ctx| { + terminal_view + .ai_controller() + .update(ctx, |controller, ctx| { + controller.set_ambient_agent_task_id(Some(task_id), ctx); + if let Some(working_dir) = working_dir.as_deref() { + controller.set_attachments_download_dir(attachments_download_dir(working_dir)); + } + }); + }); +} fn propagate_parent_agent_settings( group: &PaneGroup, @@ -72,6 +99,7 @@ pub(crate) fn create_hidden_child_agent_conversation( name: String, parent_conversation_id: AIConversationId, env_vars: HashMap, + task_context: Option, ctx: &mut ViewContext, ) -> Option { let new_pane_id = @@ -84,6 +112,9 @@ pub(crate) fn create_hidden_child_agent_conversation( let terminal_view_id = new_terminal_view.id(); propagate_parent_agent_settings(group, parent_pane_id, terminal_view_id, ctx); + if let Some(task_context) = task_context.as_ref() { + apply_hidden_child_agent_task_context(&new_terminal_view, task_context, ctx); + } let conversation_id = start_new_child_conversation(terminal_view_id, name, parent_conversation_id, ctx); @@ -117,6 +148,7 @@ fn create_error_child_agent_conversation_context( name.clone(), parent_conversation_id, HashMap::new(), + None, ctx, ) { return Some((Some(terminal_view), terminal_view_id, conversation_id)); diff --git a/app/src/pane_group/mod.rs b/app/src/pane_group/mod.rs index cd6477e18..9ce7b332a 100644 --- a/app/src/pane_group/mod.rs +++ b/app/src/pane_group/mod.rs @@ -170,6 +170,7 @@ pub mod focus_state; pub mod pane; pub mod tree; pub mod working_directories; +use child_agent::{apply_hidden_child_agent_task_context, HiddenChildAgentTaskContext}; use focus_state::PaneGroupFocusState; @@ -3083,10 +3084,23 @@ impl PaneGroup { ctx: &mut ViewContext, ) { let child_id = child_conversation.id(); + let child_task_context = + child_conversation + .task_id() + .map(|task_id| HiddenChildAgentTaskContext { + task_id, + working_dir: child_conversation + .current_working_directory() + .or_else(|| child_conversation.initial_working_directory()) + .map(PathBuf::from), + }); let new_pane_id = self.insert_terminal_pane_hidden_for_child_agent(parent_pane_id, HashMap::new(), ctx); if let Some(new_terminal_view) = self.terminal_view_from_pane_id(new_pane_id, ctx) { + if let Some(task_context) = child_task_context.as_ref() { + apply_hidden_child_agent_task_context(&new_terminal_view, task_context, ctx); + } new_terminal_view.update(ctx, |terminal_view, ctx| { terminal_view.restore_conversation_after_view_creation( RestoredAIConversation::new(child_conversation), diff --git a/app/src/pane_group/mod_tests.rs b/app/src/pane_group/mod_tests.rs index b2d21601e..bc3815b79 100644 --- a/app/src/pane_group/mod_tests.rs +++ b/app/src/pane_group/mod_tests.rs @@ -2,8 +2,13 @@ use crate::terminal::cli_agent_sessions::CLIAgentSessionsModel; use crate::{ ai::{ active_agent_views_model::ActiveAgentViewsModel, + agent::{ + conversation::{AIConversation, AIConversationId}, + PassiveSuggestionTrigger, + }, agent_conversations_model::AgentConversationsModel, ambient_agents::github_auth_notifier::GitHubAuthNotifier, + ambient_agents::AmbientAgentTaskId, blocklist::BlocklistAIHistoryModel, document::ai_document_model::AIDocumentModel, execution_profiles::profiles::AIExecutionProfilesModel, @@ -18,6 +23,7 @@ use crate::{ AIRequestUsageModel, }, auth::auth_manager::AuthManager, + changelog_model::ChangelogModel, cloud_object::model::persistence::CloudModel, context_chips::prompt::Prompt, experiments, @@ -64,8 +70,11 @@ use crate::{ use repo_metadata::RepoMetadataModel; use repo_metadata::{repositories::DetectedRepositories, watcher::DirectoryWatcher}; use std::collections::HashMap; +use uuid::Uuid; +use warp_core::features::FeatureFlag; use watcher::HomeDirectoryWatcher; +use super::child_agent::{create_hidden_child_agent_conversation, HiddenChildAgentTaskContext}; use super::*; use crate::terminal::resizable_data::ResizableData; use ai::{ @@ -84,6 +93,7 @@ fn initialize_app(app: &mut App) { initialize_settings_for_tests(app); app.add_singleton_model(|_ctx| ServerApiProvider::new_for_test()); + app.add_singleton_model(|ctx| ChangelogModel::new(ServerApiProvider::as_ref(ctx).get())); app.add_singleton_model(|_| AuthStateProvider::new_for_test()); app.add_singleton_model(AppTelemetryContextProvider::new_context_provider); app.add_singleton_model(AuthManager::new_for_test); @@ -125,6 +135,7 @@ fn initialize_app(app: &mut App) { app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); app.add_singleton_model(|_| CLIAgentSessionsModel::new()); app.add_singleton_model(|_| ActiveAgentViewsModel::new()); + app.add_singleton_model(crate::ai::blocklist::BlocklistAIPermissions::new); app.add_singleton_model(AgentNotificationsModel::new); app.add_singleton_model(|ctx| { AIExecutionProfilesModel::new(&crate::LaunchMode::new_for_unit_test(), ctx) @@ -222,6 +233,50 @@ fn new_notebook(ctx: &mut ViewContext) -> ViewHandle { ctx.add_typed_action_view(NotebookView::new) } +fn new_ambient_agent_task_id() -> AmbientAgentTaskId { + Uuid::new_v4().to_string().parse().unwrap() +} + +fn start_parent_conversation( + panes: &PaneGroup, + parent_pane_id: PaneId, + ctx: &mut ViewContext, +) -> AIConversationId { + let parent_terminal_view_id = panes + .terminal_view_from_pane_id(parent_pane_id, ctx) + .expect("parent pane should have a terminal view") + .id(); + + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.start_new_conversation(parent_terminal_view_id, false, false, ctx) + }) +} + +fn request_ambient_agent_task_id_for_hidden_child( + panes: &PaneGroup, + child_conversation_id: AIConversationId, + child_pane_id: PaneId, + ctx: &mut ViewContext, +) -> Option { + let terminal_view = panes + .terminal_view_from_pane_id(child_pane_id, ctx) + .expect("child pane should have a terminal view"); + let ai_controller = terminal_view.as_ref(ctx).ai_controller().clone(); + + ai_controller.update(ctx, |controller, ctx| { + controller + .build_passive_suggestions_request_params( + Some(child_conversation_id), + PassiveSuggestionTrigger::FilesChanged, + vec![], + ctx, + ) + .expect("child pane should build passive suggestion request params") + .1 + .ambient_agent_task_id + }) +} + struct PreAttachReturnsFalsePane { pane_id: PaneId, pane_configuration: ModelHandle, @@ -376,6 +431,91 @@ fn test_insert_hidden_child_agent_pane_keeps_focus_and_active_session() { }); } +#[test] +fn test_hidden_child_creation_applies_ambient_task_id_to_controller() { + let _orchestration_v2 = FeatureFlag::OrchestrationV2.override_enabled(true); + + App::test((), |mut app| async move { + initialize_app(&mut app); + let pane_group = mock_pane_group(&mut app, Default::default()); + + pane_group.update(&mut app, |panes, ctx| { + let parent_pane_id = get_newly_created_pane_id(panes, &[]); + let parent_conversation_id = start_parent_conversation(panes, parent_pane_id, ctx); + let task_id = new_ambient_agent_task_id(); + + let child = create_hidden_child_agent_conversation( + panes, + parent_pane_id, + "Agent 1".to_string(), + parent_conversation_id, + HashMap::new(), + Some(HiddenChildAgentTaskContext { + task_id, + working_dir: None, + }), + ctx, + ) + .expect("fresh hidden child conversation should be created"); + + let child_pane_id = panes + .child_agent_panes + .get(&child.conversation_id) + .copied() + .expect("fresh hidden child pane should be tracked"); + + assert_eq!( + request_ambient_agent_task_id_for_hidden_child( + panes, + child.conversation_id, + child_pane_id, + ctx, + ), + Some(task_id) + ); + }); + }); +} + +#[test] +fn test_restored_hidden_child_pane_reapplies_ambient_task_id_to_controller() { + let _orchestration_v2 = FeatureFlag::OrchestrationV2.override_enabled(true); + + App::test((), |mut app| async move { + initialize_app(&mut app); + let pane_group = mock_pane_group(&mut app, Default::default()); + + pane_group.update(&mut app, |panes, ctx| { + let parent_pane_id = get_newly_created_pane_id(panes, &[]); + let parent_conversation_id = start_parent_conversation(panes, parent_pane_id, ctx); + let task_id = new_ambient_agent_task_id(); + + let mut child_conversation = AIConversation::new(false); + child_conversation.set_parent_conversation_id(parent_conversation_id); + child_conversation.set_task_id(task_id); + let child_conversation_id = child_conversation.id(); + + panes.create_hidden_child_agent_pane(child_conversation, parent_pane_id, ctx); + + let child_pane_id = panes + .child_agent_panes + .get(&child_conversation_id) + .copied() + .expect("restored hidden child pane should be tracked"); + + assert_eq!( + request_ambient_agent_task_id_for_hidden_child( + panes, + child_conversation_id, + child_pane_id, + ctx, + ), + Some(task_id) + ); + }); + }); +} + #[test] fn test_active_session_id_reset_on_last_pane_close() { App::test((), |mut app| async move { diff --git a/app/src/pane_group/pane/terminal_pane.rs b/app/src/pane_group/pane/terminal_pane.rs index 0a5935cef..612bef0be 100644 --- a/app/src/pane_group/pane/terminal_pane.rs +++ b/app/src/pane_group/pane/terminal_pane.rs @@ -30,7 +30,7 @@ use crate::{ app_state::{AmbientAgentPaneSnapshot, LeafContents, TerminalPaneSnapshot}, pane_group::child_agent::{ create_error_child_agent_conversation, create_hidden_child_agent_conversation, - HiddenChildAgentConversation, + HiddenChildAgentConversation, HiddenChildAgentTaskContext, }, pane_group::{self, Direction, Event::OpenConversationHistory, PaneGroup}, persistence::{BlockCompleted, ModelEvent}, @@ -1129,6 +1129,7 @@ fn handle_terminal_view_event( request.name, request.parent_conversation_id, HashMap::new(), + None, ctx, ) { register_legacy_local_lifecycle_subscription( @@ -1164,6 +1165,7 @@ fn handle_terminal_view_event( } => { let startup_directory = group.startup_path_for_new_session(Some(terminal_pane_id), ctx); + let launch_startup_directory = startup_directory.clone(); let ai_client = ServerApiProvider::handle(ctx).as_ref(ctx).get_ai_client(); let parent_pane_id = pane_id; let request_name = request.name.clone(); @@ -1183,7 +1185,7 @@ fn handle_terminal_view_event( harness_type, parent_run_id, shell_type, - startup_directory, + launch_startup_directory, ai_client, ) .await @@ -1207,6 +1209,10 @@ fn handle_terminal_view_event( request_name.clone(), parent_conversation_id, env_vars, + Some(HiddenChildAgentTaskContext { + task_id, + working_dir: startup_directory.clone(), + }), ctx, ) { BlocklistAIHistoryModel::handle(ctx).update( diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index 883b5a8ac..84183e527 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -991,6 +991,36 @@ impl ServerApi { let response = response.json::().await?; Ok(response) } + + pub(crate) async fn mark_message_delivered_for_task( + &self, + task_id: &AmbientAgentTaskId, + message_id: &str, + ) -> anyhow::Result<(), anyhow::Error> { + self.post_public_api_response_for_task( + task_id, + &format!("agent/messages/{message_id}/delivered"), + &(), + ) + .await?; + Ok(()) + } + + pub(crate) async fn read_agent_message_for_task( + &self, + task_id: &AmbientAgentTaskId, + message_id: &str, + ) -> anyhow::Result { + let response = self + .post_public_api_response_for_task( + task_id, + &format!("agent/messages/{message_id}/read"), + &(), + ) + .await?; + let response = response.json::().await?; + Ok(response) + } } #[cfg_attr(not(target_family = "wasm"), async_trait)] From 392fad8737936ebf75e523dbff24baa9bfe10686 Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Tue, 28 Apr 2026 14:53:41 -0400 Subject: [PATCH 04/13] Handle non-retryable orchestration restore failures Co-Authored-By: Oz --- .../blocklist/orchestration_event_poller.rs | 45 ++++-- .../orchestration_event_poller_tests.rs | 131 +++++++++++++++++- 2 files changed, 160 insertions(+), 16 deletions(-) diff --git a/app/src/ai/blocklist/orchestration_event_poller.rs b/app/src/ai/blocklist/orchestration_event_poller.rs index 0a71035a6..da8ed8130 100644 --- a/app/src/ai/blocklist/orchestration_event_poller.rs +++ b/app/src/ai/blocklist/orchestration_event_poller.rs @@ -10,7 +10,7 @@ use crate::ai::agent_events::{ }; use crate::ai::ambient_agents::{AmbientAgentTask, AmbientAgentTaskId}; use crate::server::server_api::ai::{AIClient, AgentRunEvent}; -use crate::server::server_api::{ServerApi, ServerApiProvider}; +use crate::server::server_api::{AIApiError, ServerApi, ServerApiProvider}; use anyhow::anyhow; use async_trait::async_trait; use futures::channel::mpsc; @@ -613,20 +613,43 @@ impl OrchestrationEventPoller { }, move |me, result, ctx| match result { Ok(task) => me.finish_restore_fetch(conversation_id, sqlite_cursor, task, ctx), - Err(err) => { - log::warn!( - "Failed to restore orchestration event state for {conversation_id:?}: {err:#}" - ); - me.start_restore_fetch_retry_timer( - conversation_id, - retry_task_id_string.clone(), - ctx, - ); - } + Err(err) => me.handle_restore_fetch_error( + conversation_id, + retry_task_id_string.clone(), + err, + ctx, + ), }, ); } + fn handle_restore_fetch_error( + &mut self, + conversation_id: AIConversationId, + task_id: String, + err: anyhow::Error, + ctx: &mut ModelContext, + ) { + let should_retry = err + .downcast_ref::() + .map(AIApiError::is_retryable) + .unwrap_or(true); + + if should_retry { + log::warn!( + "Failed to restore orchestration event state for {conversation_id:?} task_id={task_id}; retrying: {err:#}" + ); + self.start_restore_fetch_retry_timer(conversation_id, task_id, ctx); + return; + } + + log::warn!( + "Failed to restore orchestration event state for {conversation_id:?} task_id={task_id}; using persisted run_ids/cursor without server task metadata: {err:#}" + ); + self.restore_fetch_failures.remove(&conversation_id); + self.maybe_start_delivery_after_restore(conversation_id, ctx); + } + fn finish_restore_fetch( &mut self, conversation_id: AIConversationId, diff --git a/app/src/ai/blocklist/orchestration_event_poller_tests.rs b/app/src/ai/blocklist/orchestration_event_poller_tests.rs index c9775dacb..9150227b9 100644 --- a/app/src/ai/blocklist/orchestration_event_poller_tests.rs +++ b/app/src/ai/blocklist/orchestration_event_poller_tests.rs @@ -7,10 +7,11 @@ use crate::ai::agent_events::{ use crate::ai::ambient_agents::{AmbientAgentTask, AmbientAgentTaskState}; use crate::persistence::{model::AgentConversationData, ModelEvent}; use crate::server::server_api::ai::MockAIClient; -use crate::server::server_api::ServerApiProvider; +use crate::server::server_api::{AIApiError, ServerApiProvider}; use crate::test_util::settings::initialize_settings_for_tests; use crate::{GlobalResourceHandles, GlobalResourceHandlesProvider}; use chrono::Utc; +use http::StatusCode; use std::collections::HashSet; use std::sync::Arc; use warp_multi_agent_api as api; @@ -200,12 +201,14 @@ fn finish_restore_fetch_merges_server_cursor_and_child_runs() { ); }); + let mut ai_client = MockAIClient::new(); + ai_client + .expect_poll_agent_events() + .returning(|_, _, _| Ok(vec![])); + let server_api = ServerApiProvider::new_for_test().get(); let poller = app.add_singleton_model(move |_| { - OrchestrationEventPoller::new_with_clients_for_test( - Arc::new(MockAIClient::new()), - server_api, - ) + OrchestrationEventPoller::new_with_clients_for_test(Arc::new(ai_client), server_api) }); let child_run_id = uuid::Uuid::new_v4().to_string(); @@ -232,6 +235,124 @@ fn finish_restore_fetch_merges_server_cursor_and_child_runs() { }); } +#[test] +fn non_retryable_restore_fetch_failure_falls_back_to_persisted_delivery_state() { + App::test((), |mut app| async move { + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); + let terminal_view_id = EntityId::new(); + let conversation_id = AIConversationId::new(); + let run_id = uuid::Uuid::new_v4().to_string(); + + history_model.update(&mut app, |history_model, ctx| { + history_model.restore_conversations( + terminal_view_id, + vec![restored_conversation( + conversation_id, + run_id.clone(), + Some(17), + )], + ctx, + ); + history_model.update_conversation_status( + terminal_view_id, + conversation_id, + ConversationStatus::Success, + ctx, + ); + }); + + let mut ai_client = MockAIClient::new(); + ai_client + .expect_poll_agent_events() + .returning(|_, _, _| Ok(vec![])); + + let server_api = ServerApiProvider::new_for_test().get(); + let poller = app.add_singleton_model(move |_| { + OrchestrationEventPoller::new_with_clients_for_test(Arc::new(ai_client), server_api) + }); + + poller.update(&mut app, |poller, ctx| { + poller + .watched_run_ids + .insert(conversation_id, HashSet::from([run_id.clone()])); + poller.restore_fetch_failures.insert(conversation_id, 1); + + poller.handle_restore_fetch_error( + conversation_id, + run_id.clone(), + anyhow::Error::new(AIApiError::ErrorStatus( + StatusCode::NOT_FOUND, + "missing task".to_string(), + )), + ctx, + ); + + assert!(!poller.restore_fetch_failures.contains_key(&conversation_id)); + assert!( + poller.poll_in_flight.contains(&conversation_id) + || poller.sse_connections.contains_key(&conversation_id) + ); + }); + }); +} + +#[test] +fn retryable_restore_fetch_failure_keeps_retry_state() { + App::test((), |mut app| async move { + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); + let terminal_view_id = EntityId::new(); + let conversation_id = AIConversationId::new(); + let run_id = uuid::Uuid::new_v4().to_string(); + + history_model.update(&mut app, |history_model, ctx| { + history_model.restore_conversations( + terminal_view_id, + vec![restored_conversation( + conversation_id, + run_id.clone(), + Some(17), + )], + ctx, + ); + history_model.update_conversation_status( + terminal_view_id, + conversation_id, + ConversationStatus::Success, + ctx, + ); + }); + + let server_api = ServerApiProvider::new_for_test().get(); + let poller = app.add_singleton_model(move |_| { + OrchestrationEventPoller::new_with_clients_for_test( + Arc::new(MockAIClient::new()), + server_api, + ) + }); + + poller.update(&mut app, |poller, ctx| { + poller + .watched_run_ids + .insert(conversation_id, HashSet::from([run_id.clone()])); + poller.restore_fetch_failures.insert(conversation_id, 1); + + poller.handle_restore_fetch_error( + conversation_id, + run_id.clone(), + anyhow::Error::new(AIApiError::ErrorStatus( + StatusCode::INTERNAL_SERVER_ERROR, + "server error".to_string(), + )), + ctx, + ); + + assert!(poller.restore_fetch_failures.contains_key(&conversation_id)); + assert!(!poller.poll_in_flight.contains(&conversation_id)); + assert!(!poller.sse_connections.contains_key(&conversation_id)); + }); + }); +} + #[test] fn build_pending_events_preserves_message_sequence_and_timestamp() { let pending = build_pending_events( From f7f554ddfae245df3f3376a84fa06f8f965f50bc Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Tue, 28 Apr 2026 15:24:04 -0400 Subject: [PATCH 05/13] Align QUALITY-482 port with source branch Co-Authored-By: Oz --- app/src/ai/agent_sdk/common.rs | 52 ++-- app/src/ai/agent_sdk/driver.rs | 41 ++- .../agent_sdk/driver/error_classification.rs | 21 +- app/src/ai/agent_sdk/driver/harness/gemini.rs | 3 + app/src/ai/agent_sdk/driver/harness/mod.rs | 41 ++- app/src/ai/agent_sdk/driver/snapshot_tests.rs | 4 - app/src/ai/agent_sdk/mod.rs | 173 ++++++----- app/src/ai/blocklist/controller.rs | 42 +-- .../blocklist/orchestration_event_poller.rs | 294 ------------------ app/src/ai/blocklist/orchestration_events.rs | 36 --- app/src/server/server_api/ai.rs | 6 + app/src/server/server_api/harness_support.rs | 78 +++-- crates/http_client/src/lib.rs | 7 - 13 files changed, 274 insertions(+), 524 deletions(-) diff --git a/app/src/ai/agent_sdk/common.rs b/app/src/ai/agent_sdk/common.rs index 5683d8f16..22ac02192 100644 --- a/app/src/ai/agent_sdk/common.rs +++ b/app/src/ai/agent_sdk/common.rs @@ -5,7 +5,14 @@ use std::future::Future; use std::sync::Arc; use std::time::Duration; -use crate::ai::agent::conversation::{AIAgentHarness, ServerAIConversationMetadata}; +use futures::TryFutureExt; +use inquire::{InquireError, Select}; +use warp_cli::agent::Harness; +use warp_cli::environment::{EnvironmentCreateArgs, EnvironmentUpdateArgs}; +use warpui::r#async::FutureExt; +use warpui::{AppContext, GetSingletonModelHandle, SingletonEntity as _, UpdateModel}; + +use crate::ai::agent::conversation::ServerAIConversationMetadata; use crate::ai::agent_sdk::driver::{AgentDriverError, WARP_DRIVE_SYNC_TIMEOUT}; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::cloud_environments::CloudAmbientAgentEnvironment; @@ -18,12 +25,6 @@ use crate::server::server_api::ai::AIClient; use crate::server::server_api::ServerApiProvider; use crate::workspaces::update_manager::TeamUpdateManager; use crate::workspaces::user_workspaces::UserWorkspaces; -use futures::TryFutureExt; -use inquire::{InquireError, Select}; -use warp_cli::agent::Harness; -use warp_cli::environment::{EnvironmentCreateArgs, EnvironmentUpdateArgs}; -use warpui::r#async::FutureExt; -use warpui::{AppContext, GetSingletonModelHandle, SingletonEntity as _, UpdateModel}; /// How long to wait for workspace metadata to refresh. pub const WORKSPACE_METADATA_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); @@ -143,8 +144,14 @@ pub fn refresh_warp_drive( .map_err(|_| anyhow::anyhow!("Timed out waiting for Warp Drive to sync")) } -/// Fetch the conversation's server metadata and validate that its harness -/// matches the caller's `--harness` choice. +/// Fetch the conversation's server metadata and validate that its harness matches the caller's +/// `--harness` choice. Returns the metadata on success so the caller can reuse it (e.g. for the +/// server conversation token). +/// +/// Called up-front before any task/config-build logic consumes `args.harness`, so a mismatch +/// error surfaces before side effects like task creation. We deliberately do NOT auto-upgrade +/// the harness: `Harness::Oz` default with a Claude conversation id is treated as a mismatch +/// and errors out. pub(super) async fn fetch_and_validate_conversation_harness( ai_client: Arc, conversation_id: &str, @@ -153,7 +160,7 @@ pub(super) async fn fetch_and_validate_conversation_harness( let metadata = ai_client .list_ai_conversation_metadata(Some(vec![conversation_id.to_string()])) .await - .map_err(|error| AgentDriverError::ConversationLoadFailed(format!("{error:#}")))? + .map_err(|e| AgentDriverError::ConversationLoadFailed(format!("{e:#}")))? .into_iter() .next() .ok_or_else(|| { @@ -162,36 +169,17 @@ pub(super) async fn fetch_and_validate_conversation_harness( )) })?; - let expected = harness_label(metadata.harness); - let got = harness_label_from_cli(args_harness); - if expected != got { + if metadata.harness != args_harness { return Err(AgentDriverError::ConversationHarnessMismatch { conversation_id: conversation_id.to_string(), - expected: expected.to_string(), - got: got.to_string(), + expected: Harness::from(metadata.harness).to_string(), + got: args_harness.to_string(), }); } Ok(metadata) } -fn harness_label(harness: AIAgentHarness) -> &'static str { - match harness { - AIAgentHarness::Oz => "oz", - AIAgentHarness::ClaudeCode => "claude", - AIAgentHarness::Gemini => "gemini", - } -} - -fn harness_label_from_cli(harness: Harness) -> &'static str { - match harness { - Harness::Oz => "oz", - Harness::Claude => "claude", - Harness::OpenCode => "opencode", - Harness::Gemini => "gemini", - } -} - /// Format an object owner for display in the CLI. pub fn format_owner(owner: &Owner) -> &'static str { // TODO: For potentially-shared objects, consider looking up the particular user/team name. diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index f792fcc0d..25268f858 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -197,6 +197,16 @@ impl IdleTimeoutSender { } } +/// How to resume an existing conversation when starting an agent run. +/// +/// The Oz harness restores the full conversation transcript into the terminal pane and treats +/// any new prompt as a follow-up; third-party harnesses round-trip a harness-specific payload +/// (see [`ResumePayload`]) instead. +pub enum ResumeOptions { + Oz(Box), + ThirdParty(Box), +} + /// Options for initializing the agent driver. pub struct AgentDriverOptions { /// Initial working directory for the agent's terminal session. @@ -211,12 +221,10 @@ pub struct AgentDriverOptions { pub should_share: bool, /// How long to keep the session alive after the agent run completes, if at all. pub idle_on_complete: Option, - /// Conversation to restore when creating the terminal. - /// Any ambient agent prompt will be sent as a follow-up to this conversation. - pub conversation_restoration: Option, - /// If set, resume an existing third-party-harness conversation instead of - /// starting fresh. - pub resume_payload: Option, + /// If set, resume an existing conversation instead of starting fresh. The variant + /// determines which harness-specific path is taken (Oz transcript restore vs. + /// third-party-harness payload rehydration). + pub resume: Option, /// Cloud providers to configure within the agent's session. pub cloud_providers: Vec>, /// Resolved environment configuration, if any. @@ -404,6 +412,15 @@ pub enum AgentDriverError { expected: String, got: String, }, + #[error( + "Task {task_id} was created with the {expected} harness, but --harness {got} was requested. \ + Re-run with --harness {expected} (or omit --harness to match) to continue this task." + )] + TaskHarnessMismatch { + task_id: String, + expected: String, + got: String, + }, #[error( "Conversation {conversation_id} has no stored transcript for the {harness} harness. \ The prior run may have crashed before saving any state." @@ -452,8 +469,7 @@ impl AgentDriver { should_share, idle_on_complete, secrets, - conversation_restoration, - resume_payload, + resume, cloud_providers, environment, selected_harness, @@ -462,6 +478,15 @@ impl AgentDriver { snapshot_script_timeout, } = options; + // Split the unified resume option into the two internal slots that the rest of + // the driver consumes: terminal-driven Oz transcript restoration vs. third-party + // harness payload rehydration. + let (conversation_restoration, resume_payload) = match resume { + Some(ResumeOptions::Oz(restoration)) => (Some(*restoration), None), + Some(ResumeOptions::ThirdParty(payload)) => (None, Some(*payload)), + None => (None, None), + }; + safe_info!( safe: ("Initializing agent driver: share={should_share}, idle_on_complete={idle_on_complete:?}"), full: ( diff --git a/app/src/ai/agent_sdk/driver/error_classification.rs b/app/src/ai/agent_sdk/driver/error_classification.rs index 3353936d7..47d0f23d1 100644 --- a/app/src/ai/agent_sdk/driver/error_classification.rs +++ b/app/src/ai/agent_sdk/driver/error_classification.rs @@ -248,11 +248,7 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS PlatformErrorCode::InternalError, ), ), - AgentDriverError::ConversationHarnessMismatch { - conversation_id, - expected, - got, - } => ( + AgentDriverError::ConversationHarnessMismatch { conversation_id, expected, got } => ( AgentTaskState::Failed, TaskStatusUpdate::with_error_code( format!( @@ -262,10 +258,17 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS PlatformErrorCode::EnvironmentSetupFailed, ), ), - AgentDriverError::ConversationResumeStateMissing { - harness, - conversation_id, - } => ( + AgentDriverError::TaskHarnessMismatch { task_id, expected, got } => ( + AgentTaskState::Failed, + TaskStatusUpdate::with_error_code( + format!( + "Task {task_id} was created with the {expected} harness, but --harness {got} was requested. \ + Re-run with --harness {expected} (or omit --harness) to continue this task." + ), + PlatformErrorCode::EnvironmentSetupFailed, + ), + ), + AgentDriverError::ConversationResumeStateMissing { harness, conversation_id } => ( AgentTaskState::Failed, TaskStatusUpdate::with_error_code( format!( diff --git a/app/src/ai/agent_sdk/driver/harness/gemini.rs b/app/src/ai/agent_sdk/driver/harness/gemini.rs index 88a1a0021..02aa12348 100644 --- a/app/src/ai/agent_sdk/driver/harness/gemini.rs +++ b/app/src/ai/agent_sdk/driver/harness/gemini.rs @@ -74,6 +74,9 @@ impl ThirdPartyHarness for GeminiHarness { terminal_driver: ModelHandle, _resume: Option, ) -> Result, AgentDriverError> { + // Gemini does not support conversation resume yet. When it does, it will add its + // own `ResumePayload::Gemini(..)` variant and override `fetch_resume_payload`, + // and decide how to surface the user-turn resumption preamble. let client: Arc = server_api; Ok(Box::new(GeminiHarnessRunner::new( self.cli_agent().command_prefix(), diff --git a/app/src/ai/agent_sdk/driver/harness/mod.rs b/app/src/ai/agent_sdk/driver/harness/mod.rs index 787f9e3e4..b5afd4be7 100644 --- a/app/src/ai/agent_sdk/driver/harness/mod.rs +++ b/app/src/ai/agent_sdk/driver/harness/mod.rs @@ -35,15 +35,20 @@ use super::{ }; mod claude_code; +pub(crate) mod claude_transcript; mod gemini; mod json_utils; - -use claude_code::ClaudeResumeInfo; pub(crate) use claude_code::{ClaudeHarness, ClaudeWakeMessage, ClaudeWakeRemoteContext}; +use claude_transcript::ClaudeResumeInfo; use gemini::GeminiHarness; /// Harness-agnostic payload describing how to resume an existing conversation. +/// +/// Each variant carries the data a specific harness needs to rehydrate state before its CLI +/// launches. Harnesses match on the variant they produce and ignore others; new CLIs that +/// want resume support add a new variant and override [`ThirdPartyHarness::fetch_resume_payload`]. pub(crate) enum ResumePayload { + /// Claude Code session state fetched from the server's transcript endpoint. Claude(ClaudeResumeInfo), } @@ -82,6 +87,15 @@ pub(crate) trait ThirdPartyHarness: Send + Sync { } /// Fetch the harness-specific resume payload for an existing conversation. + /// + /// The driver calls this when the user passes `--conversation ` and the harness + /// matches the stored conversation's harness. Harnesses that don't support resume + /// use the default impl, which returns `Ok(None)` and causes the run to start fresh. + /// + /// Implementations download the raw transcript via [`HarnessSupportClient::fetch_transcript`] + /// (which derives the conversation from the current task's `agent_conversation_id`) and + /// own all harness-specific deserialization and error mapping (e.g. a 404 maps to + /// [`AgentDriverError::ConversationResumeStateMissing`] tagged with the harness label). async fn fetch_resume_payload( &self, _conversation_id: &AIConversationId, @@ -91,6 +105,15 @@ pub(crate) trait ThirdPartyHarness: Send + Sync { } /// Build a runner for executing this harness with the given prompt. + /// + /// If `resume` is `Some`, the harness matches on its own [`ResumePayload`] variant and + /// reuses the stored session/conversation ids instead of minting fresh ones. Variants + /// belonging to other harnesses are ignored. + /// + /// `resumption_prompt`, when non-empty, is a short user-turn preamble the server emits + /// during a resumed session. Each harness decides exactly how to surface it (e.g. Claude + /// prepends it to the user-turn prompt that gets piped into the CLI). Harnesses that + /// don't yet support resumption can ignore it. #[allow(clippy::too_many_arguments)] fn build_runner( &self, @@ -134,12 +157,16 @@ impl fmt::Debug for HarnessKind { } /// Build a [`HarnessKind`] for the given [`Harness`]. -pub(crate) fn harness_kind(harness: Harness) -> HarnessKind { +/// +/// We shouldn't ever get a `--harness unknown` here because clap should handle +/// it. +pub(crate) fn harness_kind(harness: Harness) -> Result { match harness { - Harness::Oz => HarnessKind::Oz, - Harness::Claude => HarnessKind::ThirdParty(Box::new(ClaudeHarness)), - Harness::OpenCode => HarnessKind::Unsupported(Harness::OpenCode), - Harness::Gemini => HarnessKind::ThirdParty(Box::new(GeminiHarness)), + Harness::Oz => Ok(HarnessKind::Oz), + Harness::Claude => Ok(HarnessKind::ThirdParty(Box::new(ClaudeHarness))), + Harness::OpenCode => Ok(HarnessKind::Unsupported(Harness::OpenCode)), + Harness::Gemini => Ok(HarnessKind::ThirdParty(Box::new(GeminiHarness))), + Harness::Unknown => Err(AgentDriverError::InvalidRuntimeState), } } diff --git a/app/src/ai/agent_sdk/driver/snapshot_tests.rs b/app/src/ai/agent_sdk/driver/snapshot_tests.rs index 31c0cb4d5..7c4b42542 100644 --- a/app/src/ai/agent_sdk/driver/snapshot_tests.rs +++ b/app/src/ai/agent_sdk/driver/snapshot_tests.rs @@ -142,10 +142,6 @@ impl HarnessSupportClient for TestClient { Ok(targets) } - async fn fetch_transcript(&self) -> Result { - unimplemented!("not used by upload_snapshot_from_declarations_file") - } - fn http_client(&self) -> &http_client::Client { &self.http } diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index 6b9c34711..2d374c6b6 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -401,7 +401,7 @@ fn build_merged_config_and_task( model: model_override, profile: args.profile.clone(), mcp_specs: runtime_mcp_specs, - harness: harness_kind(args.harness), + harness: harness_kind(args.harness)?, }; Ok((merged_config, task)) @@ -461,7 +461,7 @@ fn build_server_side_task( model: model_override, profile, mcp_specs: runtime_mcp_specs, - harness: harness_kind(args.harness), + harness: harness_kind(args.harness)?, }; Ok((config, task)) @@ -568,47 +568,42 @@ impl AgentDriverRunner { // Set up and run the driver, reporting any errors back to the server. let result: Result<(), AgentDriverError> = async { - // Pull relevant variables out of args before moving it into the - // build_driver_options_and_task future. + // Pull relevant variables out of args before moving it into the closure. let share_requests = args.share.share.clone(); let bedrock_inference_role = args.bedrock_inference_role.clone(); let args_harness = args.harness; - - let resume_context = if let Some(conversation_id) = args.conversation.clone() { - let metadata = common::fetch_and_validate_conversation_harness( + // `--conversation` path (user-invoked local resume): validate before any task side + // effects so mismatches fail fast. The `--task-id` path derives its conversation id + // from the server-side task metadata inside `build_driver_options_and_task`. Both + // can currently be passed together (the worker server-side appends `--conversation` + // alongside `--task-id` for Slack/Linear followups); when both are set, the explicit + // `--conversation` value wins via the merge below. + if let Some(conversation_id) = args.conversation.as_deref() { + common::fetch_and_validate_conversation_harness( server_api.clone(), - &conversation_id, + conversation_id, args_harness, ) .await?; - Some((conversation_id, metadata)) - } else { - None - }; + } + let resume_conversation_id = args.conversation.clone(); - let (mut driver_options, task, environment_id, task_conversation_id) = + // Build driver options and task, handling task creation or existing task setup. + // For the `--task-id` path, `task_conversation_id` is the `conversation_id` read off + // the fetched `AmbientAgentTask` (set by the server when linking the task to an + // existing conversation, e.g. via `run-cloud --conversation`). + let (mut driver_options, task, task_conversation_id) = Self::build_driver_options_and_task(&foreground, args, &server_api).await?; // Update the effective task ID so errors are reported correctly. // This only matters if we created a task ID locally. task_id = driver_options.task_id.or(task_id); - let resume_context = match (resume_context, task_conversation_id) { - (Some(context), _) => Some(context), - (None, Some(conversation_id)) => { - let metadata = common::fetch_and_validate_conversation_harness( - server_api.clone(), - &conversation_id, - args_harness, - ) - .await?; - Some((conversation_id, metadata)) - } - (None, None) => None, - }; - // Resolve the environment. We make sure this happens after resolving the task ID - // so that errors are reported. - Self::resolve_environment(&foreground, environment_id, &mut driver_options).await?; + // The `--task-id` branch already validated `args_harness` against the task's harness + // setting inside `build_driver_options_and_task`; the conversation that the task spawned + // necessarily uses the same harness, so no extra conversation-metadata roundtrip is + // needed here. Just merge the task's linked conversation id into the resume target. + let resume_conversation_id = resume_conversation_id.or(task_conversation_id); let bedrock_task_id = driver_options.task_id.map(|id| id.to_string()); @@ -658,13 +653,11 @@ impl AgentDriverRunner { } // Pull conversation information, if we have it - if let Some((conversation_id, resume_metadata)) = resume_context { - Self::load_conversation_information( + if let Some(conversation_id) = resume_conversation_id { + driver_options.resume = Self::load_conversation_information( &foreground, conversation_id, - resume_metadata, &task.harness, - &mut driver_options, ) .await?; } @@ -755,13 +748,15 @@ impl AgentDriverRunner { /// Build the AgentDriverOptions and Task, handling task creation or existing task setup. /// - /// Returns the driver options, the task, the unresolved environment ID (if any), and any - /// conversation id read off server-side task metadata. + /// The third tuple element is the conversation id read off the server-side task metadata + /// on the `--task-id` branch. It's `None` when no task id was passed or when the task is + /// not linked to a conversation; callers use it to drive `--task-id`-implied resume + /// without requiring the caller to also pass `--conversation`. async fn build_driver_options_and_task( foreground: &ModelSpawner, args: RunAgentArgs, server_api: &Arc, - ) -> Result<(AgentDriverOptions, Task, Option, Option), AgentDriverError> { + ) -> Result<(AgentDriverOptions, Task, Option), AgentDriverError> { // Get the working directory let working_dir = match args.cwd.as_ref() { Some(dir) => dunce::canonicalize(dir) @@ -796,8 +791,7 @@ impl AgentDriverRunner { should_share, idle_on_complete: args.idle_on_complete.map(|d| d.into()), secrets: Default::default(), - conversation_restoration: None, - resume_payload: None, + resume: None, cloud_providers: Vec::new(), environment: None, selected_harness: args.harness, @@ -819,7 +813,9 @@ impl AgentDriverRunner { let environment_id = merged_config.environment_id.clone(); - // Handle secrets/attachments fetch (existing task) or task creation (new run) + // Handle secrets/attachments fetch (existing task) or task creation (new run). + // The existing-task branch also surfaces the task's `conversation_id` (if any) so + // the caller can wire up resume without a separate `--conversation` arg. let task_conversation_id = if let Some(task_id_str) = task_id_str { Self::fetch_secrets_and_attachments( foreground, @@ -851,8 +847,10 @@ impl AgentDriverRunner { .await?; None }; + // Resolve environment and cloud providers. + Self::resolve_environment(foreground, environment_id, &mut driver_options).await?; - Ok((driver_options, task, environment_id, task_conversation_id)) + Ok((driver_options, task, task_conversation_id)) } /// Creates a new task on the server for this agent run, sets the task ID on the driver @@ -906,6 +904,10 @@ impl AgentDriverRunner { /// When starting an agent run from an existing task_id, fetch secrets, task metadata, /// and task attachments (images and files) from the server and update the driver options. + /// + /// Returns the task's `conversation_id` when the server has linked the task to an existing + /// AI conversation (e.g. a `run-cloud --conversation` spawn). The caller uses this to drive + /// transcript rehydration without a separate `--conversation` CLI arg. async fn fetch_secrets_and_attachments( foreground: &ModelSpawner, task_id_str: String, @@ -1025,15 +1027,42 @@ impl AgentDriverRunner { } } }; - let (parent_run_id, task_conversation_id) = match task_metadata_result { - Ok(Some(task_metadata)) => (task_metadata.parent_run_id, task_metadata.conversation_id), - Ok(None) => (None, None), + let (parent_run_id, task_conversation_id, task_harness) = match task_metadata_result { + Ok(Some(task_metadata)) => { + // The task's harness is stored on the snapshot; if absent, it's the default Oz. + let task_harness = task_metadata + .agent_config_snapshot + .as_ref() + .and_then(|c| c.harness.as_ref()) + .map(|h| h.harness_type) + .unwrap_or(Harness::Oz); + ( + task_metadata.parent_run_id, + task_metadata.conversation_id, + Some(task_harness), + ) + } + Ok(None) => (None, None, None), Err(err) => { log::warn!("Failed to fetch task metadata: {err:#}"); - (None, None) + (None, None, None) } }; + // Validate the requested `--harness` against the task's harness setting. This avoids the + // extra conversation-metadata roundtrip that would otherwise be needed downstream when the + // task is linked to an existing conversation, since task harness and conversation harness + // always match (the task spawned the conversation). + if let Some(task_harness) = task_harness { + if task_harness != driver_options.selected_harness { + return Err(AgentDriverError::TaskHarnessMismatch { + task_id: task_id_str, + expected: task_harness.to_string(), + got: driver_options.selected_harness.to_string(), + }); + } + } + // Set the task ID on the ServerApi so it's sent with all subsequent requests. foreground .spawn(move |_, ctx| { @@ -1061,17 +1090,24 @@ impl AgentDriverRunner { } /// If we are starting this agent run from an existing conversation, load the conversation - /// data from the server and set the relevant driver options. + /// data from the server and return the harness-specific [`ResumeOptions`] payload that the + /// caller plugs onto [`AgentDriverOptions::resume`]. + /// + /// `harness` is the resolved harness from the task config (already validated against the + /// conversation's metadata up-front by [`common::fetch_and_validate_conversation_harness`]). + /// + /// For the Oz harness, fetches the full conversation and returns a [`driver::ResumeOptions::Oz`]. + /// For third-party harnesses, delegates to [`ThirdPartyHarness::fetch_resume_payload`] and + /// wraps the returned payload (if any) in [`driver::ResumeOptions::ThirdParty`]; each harness + /// owns its server call and error mapping. Returns `None` if a third-party harness has no + /// resume payload to surface. async fn load_conversation_information( foreground: &ModelSpawner, conversation_id: String, - resume_metadata: crate::ai::agent::conversation::ServerAIConversationMetadata, harness: &HarnessKind, - driver_options: &mut AgentDriverOptions, - ) -> Result<(), AgentDriverError> { + ) -> Result, AgentDriverError> { match harness { HarnessKind::Oz => { - let _ = resume_metadata; let server_api = foreground .spawn(|_, ctx| { ServerApiProvider::handle(ctx) @@ -1096,36 +1132,33 @@ impl AgentDriverRunner { "Failed to convert conversation data to AIConversation".into(), ) })?; - - driver_options.conversation_restoration = - Some(ConversationRestorationInNewPaneType::Historical { + Ok(Some(driver::ResumeOptions::Oz(Box::new( + ConversationRestorationInNewPaneType::Historical { conversation, should_use_live_appearance: false, ambient_agent_task_id: None, - }); + }, + )))) } - HarnessKind::ThirdParty(harness) => { - let _ = resume_metadata; + HarnessKind::ThirdParty(h) => { let harness_support_client = foreground .spawn(|_, ctx| ServerApiProvider::as_ref(ctx).get_harness_support_client()) .await?; - let resume_conversation_id = AIConversationId::try_from(conversation_id) + let resume_conversation_id = AIConversationId::try_from(conversation_id.clone()) .map_err(|err| AgentDriverError::ConversationLoadFailed(format!("{err:#}")))?; - driver_options.resume_payload = harness - .fetch_resume_payload(&resume_conversation_id, harness_support_client) - .await?; - } - HarnessKind::Unsupported(harness) => { - return Err(AgentDriverError::HarnessSetupFailed { - harness: harness.to_string(), - reason: format!( - "The {harness} harness is only supported for local child agent launches." - ), - }); + Ok( + h.fetch_resume_payload(&resume_conversation_id, harness_support_client) + .await? + .map(|payload| driver::ResumeOptions::ThirdParty(Box::new(payload))), + ) } + HarnessKind::Unsupported(harness) => Err(AgentDriverError::HarnessSetupFailed { + harness: harness.to_string(), + reason: format!( + "The {harness} harness is only supported for local child agent launches." + ), + }), } - - Ok(()) } /// Resolve the environment and store into `driver_options`. @@ -1354,7 +1387,7 @@ fn resolve_orchestration_harness_label() -> &'static str { Some(Harness::Claude) => "claude", Some(Harness::OpenCode) => "opencode", Some(Harness::Gemini) => "gemini", - _ => "unknown", + Some(Harness::Unknown) | None => "unknown", } } diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 78469bb08..e4bc1f56f 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -613,8 +613,15 @@ impl BlocklistAIController { .remove(&conversation_id) .unwrap_or_default(); - let cancellation_reason = - self.cancel_active_conversation_for_follow_up(conversation_id, ctx); + let ai_history_model = BlocklistAIHistoryModel::as_ref(ctx); + let active_conversation_id = ai_history_model.active_conversation_id(self.terminal_view_id); + let cancellation_reason = CancellationReason::FollowUpSubmitted { + is_for_same_conversation: active_conversation_id + .is_some_and(|id| id == conversation_id), + }; + if let Some(active_conversation_id) = active_conversation_id { + self.cancel_conversation_progress(active_conversation_id, cancellation_reason, ctx); + } if let Some(slash_command_request) = SlashCommandRequest::from_query(query.as_str()) { slash_command_request.send_request(self, is_queued_prompt, ctx); @@ -1216,40 +1223,9 @@ impl BlocklistAIController { slash_command: SlashCommandRequest, ctx: &mut ModelContext, ) { - // Slash commands are a fresh user turn; mirror `send_query`'s - // cancel-and-resend so we don't trip `send_request_input`'s in-flight - // invariant. - if let Some(conversation_id) = slash_command.conversation_id(self, ctx) { - self.cancel_active_conversation_for_follow_up(conversation_id, ctx); - } slash_command.send_request(self, /*is_queued_prompt*/ false, ctx); } - /// Cancel any in-flight progress on the active conversation in preparation - /// for sending a follow-up turn that will land on `target_conversation_id`. - /// Without this pre-cancel, [`Self::send_request_input`] would trip its - /// in-flight invariant when the new turn re-uses an existing conversation. - /// - /// Returns the [`CancellationReason::FollowUpSubmitted`] reason used so - /// callers can reuse it for downstream side effects (e.g. cancelling - /// pending actions on the target conversation). - fn cancel_active_conversation_for_follow_up( - &mut self, - target_conversation_id: AIConversationId, - ctx: &mut ModelContext, - ) -> CancellationReason { - let active_conversation_id = - BlocklistAIHistoryModel::as_ref(ctx).active_conversation_id(self.terminal_view_id); - let reason = CancellationReason::FollowUpSubmitted { - is_for_same_conversation: active_conversation_id - .is_some_and(|id| id == target_conversation_id), - }; - if let Some(active_conversation_id) = active_conversation_id { - self.cancel_conversation_progress(active_conversation_id, reason, ctx); - } - reason - } - /// Same as [`Self::send_slash_command_request`] but marks the emitted `SentRequest` /// event as a queued prompt submission so UI subscribers (e.g. the input editor) /// don't clear the input buffer on the auto-send. diff --git a/app/src/ai/blocklist/orchestration_event_poller.rs b/app/src/ai/blocklist/orchestration_event_poller.rs index da8ed8130..941f9cb61 100644 --- a/app/src/ai/blocklist/orchestration_event_poller.rs +++ b/app/src/ai/blocklist/orchestration_event_poller.rs @@ -102,9 +102,6 @@ pub struct OrchestrationEventPoller { /// Monotonic counter for SSE connection generations. Ensures stale /// callbacks from replaced connections are discarded. next_sse_generation: u64, - /// Consecutive failure count for the post-restore `get_ambient_agent_task` - /// fetch (resets on success). Drives exponential backoff for retries. - restore_fetch_failures: HashMap, } pub enum OrchestrationEventPollerEvent { @@ -155,35 +152,6 @@ impl OrchestrationEventPoller { poll_in_flight: HashSet::new(), sse_connections: HashMap::new(), next_sse_generation: 0, - restore_fetch_failures: HashMap::new(), - } - } - - /// Constructs a poller wired to the supplied (mock) clients instead of - /// looking them up via `ServerApiProvider`. Lets unit tests inject a - /// `MockAIClient` while still subscribing to `BlocklistAIHistoryModel`. - #[cfg(test)] - pub(super) fn new_with_clients_for_test( - ai_client: Arc, - server_api: Arc, - ctx: &mut ModelContext, - ) -> Self { - let history_model = BlocklistAIHistoryModel::handle(ctx); - ctx.subscribe_to_model(&history_model, |me, event, ctx| { - me.handle_history_event(event, ctx); - }); - Self { - ai_client, - server_api, - watched_run_ids: HashMap::new(), - event_cursor: HashMap::new(), - poll_backoff_index: HashMap::new(), - pending_delivery: HashMap::new(), - conversation_statuses: HashMap::new(), - poll_in_flight: HashSet::new(), - sse_connections: HashMap::new(), - next_sse_generation: 0, - restore_fetch_failures: HashMap::new(), } } @@ -249,7 +217,6 @@ impl OrchestrationEventPoller { self.pending_delivery.remove(conversation_id); self.conversation_statuses.remove(conversation_id); self.poll_in_flight.remove(conversation_id); - self.restore_fetch_failures.remove(conversation_id); // SSE cleanup // task's next send to fail, which terminates the task. self.sse_connections.remove(conversation_id); @@ -267,233 +234,6 @@ impl OrchestrationEventPoller { | BlocklistAIHistoryEvent::SplitConversation { .. } | BlocklistAIHistoryEvent::UpdatedConversationMetadata { .. } | BlocklistAIHistoryEvent::UpdatedConversationArtifacts { .. } => {} - BlocklistAIHistoryEvent::RestoredConversations { - conversation_ids, .. - } => { - self.on_restored_conversations(conversation_ids.clone(), ctx); - } - } - } - - /// Handles restoration of conversations on startup (or driver re-attach). - /// - /// Re-establishes orchestration event delivery state that is not persisted - /// directly in memory: watched run_ids, the per-conversation event cursor, - /// and — for `Success` parents with watched children — the poll/SSE loop. - fn on_restored_conversations( - &mut self, - conversation_ids: Vec, - ctx: &mut ModelContext, - ) { - // Orchestration v2 owns the events endpoints and the cursor model. - // V1 conversations may carry a run_id but the v2-only event APIs - // would return spurious 4xx responses, so skip restore entirely - // when V2 is disabled. - if !FeatureFlag::OrchestrationV2.is_enabled() { - return; - } - - for conv_id in conversation_ids { - let (run_id, cursor, status, is_viewer) = { - let history_model = BlocklistAIHistoryModel::as_ref(ctx); - let Some(conversation) = history_model.conversation(&conv_id) else { - continue; - }; - let is_viewer = conversation.is_viewing_shared_session(); - let run_id = conversation.run_id(); - let cursor = conversation.last_event_sequence().unwrap_or(0); - let status = conversation.status().clone(); - (run_id, cursor, status, is_viewer) - }; - - // Shared-session viewers receive updates through session sharing; - // polling here would re-inject events the session has already - // processed. - if is_viewer { - continue; - } - - // Initialize the in-memory cursor from the persisted SQLite value. - // A later server `GET /agent/runs/{run_id}` response may advance - // it to `max(SQLite, server)` before delivery starts. - // - // Note: a status transition arriving in the window before - // finish_restore_fetch completes may trigger - // start_event_delivery with only the SQLite cursor. This is - // acceptable — worst case is one extra batch of duplicate - // events. - self.event_cursor.insert(conv_id, cursor); - self.conversation_statuses.insert(conv_id, status.clone()); - - // Register the conversation's own run_id so lifecycle events for - // self are correctly filtered and the SSE/poll loop has a set - // of run_ids to open against. - if let Some(ref own) = run_id { - self.watched_run_ids - .entry(conv_id) - .or_default() - .insert(own.clone()); - } - - // No run_id means we can't query the server for children or for - // the canonical cursor. There's nothing more to do here; if a - // run_id gets assigned later the standard self-registration path - // will pick it up. - let Some(run_id) = run_id else { - self.maybe_start_delivery_after_restore(conv_id, &status, ctx); - continue; - }; - - let Ok(task_id) = run_id.parse::() - else { - log::warn!("could not parse run_id {run_id:?} for {conv_id:?}"); - self.maybe_start_delivery_after_restore(conv_id, &status, ctx); - continue; - }; - - self.spawn_restore_fetch(conv_id, task_id, cursor, ctx); - } - } - - /// Issues `GET /agent/runs/{task_id}` and routes the result through - /// `finish_restore_fetch`. Used both for the initial post-restore fetch - /// and for backoff-driven retries. - fn spawn_restore_fetch( - &mut self, - conv_id: AIConversationId, - task_id: crate::ai::ambient_agents::AmbientAgentTaskId, - sqlite_cursor: i64, - ctx: &mut ModelContext, - ) { - let ai_client = self.ai_client.clone(); - ctx.spawn( - async move { ai_client.get_ambient_agent_task(&task_id).await }, - move |me, run_result, ctx| { - me.finish_restore_fetch(conv_id, task_id, sqlite_cursor, run_result, ctx); - }, - ); - } - - /// Completes the post-restore async fetch by merging the server cursor, - /// installing the server-reported child run_ids, and — if the parent is - /// `Success` — starting event delivery. On a server-fetch failure, - /// schedules a retry with exponential backoff: V2 children always have a - /// server-side `ai_tasks` row, so the server is the authoritative source - /// for the watched run_id set, and any local fallback would be incomplete - /// anyway. Without network connectivity event delivery wouldn't function, - /// so retrying is the right behavior. - fn finish_restore_fetch( - &mut self, - conv_id: AIConversationId, - task_id: crate::ai::ambient_agents::AmbientAgentTaskId, - sqlite_cursor: i64, - run_result: anyhow::Result, - ctx: &mut ModelContext, - ) { - match run_result { - Ok(task) => { - // If the conversation was removed while the fetch was in-flight, - // the removal handler already cleaned up all poller state. Return - // early to avoid recreating watched_run_ids for a deleted conversation. - if !self.event_cursor.contains_key(&conv_id) { - self.restore_fetch_failures.remove(&conv_id); - return; - } - - // Reset the retry counter on success. - self.restore_fetch_failures.remove(&conv_id); - - // Merge the server cursor: use the max of SQLite and server - // values so we don't re-deliver events the client already - // acknowledged locally. - let server_seq = task.last_event_sequence.unwrap_or(0); - let merged = sqlite_cursor.max(server_seq); - self.event_cursor.insert(conv_id, merged); - - // The server response includes `children` inline on - // `AmbientAgentTask`; this is the authoritative set of - // direct child run_ids for the parent. - // - // Insert children and reconnect SSE once if any new run_ids - // were added and a connection is already open (e.g. because a - // status transition raced with this fetch and opened SSE with - // only the parent's own run_id). - let had_sse = self.sse_connections.contains_key(&conv_id); - let watched = self.watched_run_ids.entry(conv_id).or_default(); - let mut any_new_children = false; - for child in task.children { - if watched.insert(child) { - any_new_children = true; - } - } - if any_new_children && had_sse { - self.reconnect_sse(conv_id, ctx); - } - - let status = BlocklistAIHistoryModel::as_ref(ctx) - .conversation(&conv_id) - .map(|c| c.status().clone()) - .unwrap_or(ConversationStatus::Success); - self.maybe_start_delivery_after_restore(conv_id, &status, ctx); - } - Err(err) => { - log::warn!("Restore: get_agent_run failed for {conv_id:?}: {err:#}; will retry"); - self.start_restore_fetch_retry_timer(conv_id, task_id, sqlite_cursor, ctx); - } - } - } - - /// Schedules a retry of the post-restore `get_ambient_agent_task` fetch - /// after an exponential backoff. The backoff schedule reuses - /// `POLL_BACKOFF_STEPS` (1s, 2s, 5s, 10s capped) keyed on a per-conversation - /// failure counter. The counter resets on success. - fn start_restore_fetch_retry_timer( - &mut self, - conv_id: AIConversationId, - task_id: crate::ai::ambient_agents::AmbientAgentTaskId, - sqlite_cursor: i64, - ctx: &mut ModelContext, - ) { - let failures = self - .restore_fetch_failures - .entry(conv_id) - .and_modify(|c| *c += 1) - .or_insert(1); - let step_index = failures.saturating_sub(1).min(POLL_BACKOFF_STEPS.len() - 1); - let backoff = Duration::from_secs(POLL_BACKOFF_STEPS[step_index]); - ctx.spawn( - async move { Timer::after(backoff).await }, - move |me, _, ctx| { - // The conversation may have been removed in the meantime; - // if so, drop the retry. Otherwise re-issue the fetch. - if !me.event_cursor.contains_key(&conv_id) { - me.restore_fetch_failures.remove(&conv_id); - return; - } - me.spawn_restore_fetch(conv_id, task_id, sqlite_cursor, ctx); - }, - ); - } - - /// Starts event delivery for a restored conversation if the parent is - /// currently `Success` and has at least one watched run_id. `InProgress` - /// parents are deferred to `on_conversation_status_updated` once they - /// next transition to `Success`. - fn maybe_start_delivery_after_restore( - &mut self, - conv_id: AIConversationId, - status: &ConversationStatus, - ctx: &mut ModelContext, - ) { - let has_watched = self - .watched_run_ids - .get(&conv_id) - .is_some_and(|w| !w.is_empty()); - if !has_watched { - return; - } - if matches!(status, ConversationStatus::Success) { - self.start_event_delivery(conv_id, ctx); } } @@ -916,40 +656,6 @@ impl OrchestrationEventPoller { ); } - // Persist the cursor to SQLite so that after a restart we can resume - // event delivery from this sequence number without re-delivering - // events the parent has already acted on. - BlocklistAIHistoryModel::handle(ctx).update(ctx, |model, ctx| { - model.update_event_sequence(conversation_id, max_seq, ctx); - }); - - // Also persist the cursor to the server so driver / cloud restarts - // can resume without local SQLite state. Fire-and-forget: log on - // failure, don't block event delivery. The server persists the - // cursor on `ai_tasks.last_event_sequence`. - let own_run_id = BlocklistAIHistoryModel::as_ref(ctx) - .conversation(&conversation_id) - .and_then(|c| c.run_id()); - if let Some(run_id) = own_run_id { - // TODO: consider debouncing this server write (see - // specs/replay-agent-events-on-restore/TECH.md Risks). - let ai_client = self.ai_client.clone(); - ctx.spawn( - async move { - ai_client - .update_event_sequence_on_server(&run_id, max_seq) - .await - }, - move |_, result, _| { - if let Err(err) = result { - log::warn!( - "Failed to persist event cursor to server for {conversation_id:?}: {err:#}" - ); - } - }, - ); - } - // Track message IDs for server-side mark_delivered calls. let message_ids: Vec = events .iter() diff --git a/app/src/ai/blocklist/orchestration_events.rs b/app/src/ai/blocklist/orchestration_events.rs index febb47c63..6a6a6b418 100644 --- a/app/src/ai/blocklist/orchestration_events.rs +++ b/app/src/ai/blocklist/orchestration_events.rs @@ -359,42 +359,6 @@ impl OrchestrationEventService { } => { for conversation_id in conversation_ids { self.sync_conversation_status(*conversation_id, ctx); - // Under V1 local lifecycle dispatch, child status - // transitions are forwarded to the parent via - // `lifecycle_subscription_routes`. That map is not - // persisted, so re-register subscriptions for each - // restored child whose parent is loaded locally so that - // child status transitions continue to propagate after - // a restart. V2 uses the server event log and does not - // need this. - if !FeatureFlag::OrchestrationV2.is_enabled() { - let parent_agent_id = { - let history_model = BlocklistAIHistoryModel::as_ref(ctx); - let Some(child_conv) = history_model.conversation(conversation_id) - else { - continue; - }; - if !child_conv.is_child_agent_conversation() { - continue; - } - child_conv - .parent_conversation_id() - .and_then(|pid| history_model.conversation(&pid)) - .and_then(|p| p.server_conversation_token()) - .map(|t| t.as_str().to_string()) - }; - if let Some(parent_agent_id) = parent_agent_id { - // `None` event-type filter = subscribe to all - // lifecycle types. The original filter (if any) - // is not persisted; subscribing broader than the - // original is acceptable per the tech spec. - self.register_lifecycle_subscription( - *conversation_id, - parent_agent_id, - None, - ); - } - } } } BlocklistAIHistoryEvent::RemoveConversation { diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index 84183e527..a5a536006 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -2312,6 +2312,12 @@ fn convert_harness(harness: warp_graphql::ai::AgentHarness) -> AIAgentHarness { warp_graphql::ai::AgentHarness::Oz => AIAgentHarness::Oz, warp_graphql::ai::AgentHarness::ClaudeCode => AIAgentHarness::ClaudeCode, warp_graphql::ai::AgentHarness::Gemini => AIAgentHarness::Gemini, + warp_graphql::ai::AgentHarness::Other(value) => { + report_error!(anyhow!( + "Invalid AgentHarness '{value}'. Make sure to update client GraphQL types!" + )); + AIAgentHarness::Unknown + } } } diff --git a/app/src/server/server_api/harness_support.rs b/app/src/server/server_api/harness_support.rs index 6d403db8e..e4b9947ac 100644 --- a/app/src/server/server_api/harness_support.rs +++ b/app/src/server/server_api/harness_support.rs @@ -10,6 +10,7 @@ use mockall::automock; use super::ServerApi; use crate::ai::agent::conversation::AIConversationId; +#[cfg(not(target_family = "wasm"))] use crate::ai::agent_sdk::retry::with_bounded_retry; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::artifacts::Artifact; @@ -87,8 +88,10 @@ pub struct ResolvedHarnessPrompt { pub prompt: String, #[serde(default)] pub system_prompt: Option, - /// Optional user-turn preamble for resumed third-party harness sessions. - /// Each harness decides how to surface this. + /// Optional user-turn preamble for resumed third-party harness sessions. The harness + /// decides how to surface this — Claude Code prepends it to the user-turn prompt fed + /// into the CLI so the agent treats it as immediate intent rather than background + /// system context. Empty when no resumption is in effect. #[serde(default)] pub resumption_prompt: Option, } @@ -151,8 +154,16 @@ pub trait HarnessSupportClient: 'static + Send + Sync { request: &SnapshotUploadRequest, ) -> Result>; - /// Download the raw third-party harness transcript bytes for the current - /// task's conversation. + /// Download the raw third-party harness transcript bytes for the current task's + /// conversation. + /// + /// Hits `GET /harness-support/transcript`, which redirects to a signed GCS URL. + /// The conversation is resolved from the task's `agent_conversation_id` server-side, + /// so callers do not pass a conversation id. Each harness deserializes the returned + /// bytes into its own envelope shape (e.g. Claude Code parses + /// `ClaudeTranscriptEnvelope`). Transient failures retry with bounded exponential + /// backoff; permanent 4xx (e.g. 404 "no transcript") fail fast so the caller can + /// surface a resume-specific error. async fn fetch_transcript(&self) -> Result; /// Get an HTTP client to use with [`UploadTarget`]s for saving blobs. @@ -249,16 +260,26 @@ impl ServerApi { &self, task_id: &AmbientAgentTaskId, ) -> Result { - with_bounded_retry("fetch task-scoped harness-support transcript", || async { - let response = self - .get_public_api_response_for_task(task_id, "harness-support/transcript") - .await?; - response - .bytes() - .await - .context("Failed to read harness-support transcript response body") - }) - .await + #[cfg(not(target_family = "wasm"))] + { + with_bounded_retry("fetch task-scoped harness-support transcript", || async { + let response = self + .get_public_api_response_for_task(task_id, "harness-support/transcript") + .await?; + response + .bytes() + .await + .context("Failed to read task-scoped harness-support transcript body") + }) + .await + } + #[cfg(target_family = "wasm")] + { + let _ = task_id; + unreachable!( + "fetch_transcript_for_task is not supported on wasm; agent_sdk is not built on this target" + ); + } } } @@ -347,16 +368,25 @@ impl HarnessSupportClient for ServerApi { } async fn fetch_transcript(&self) -> Result { - with_bounded_retry("fetch harness-support transcript", || async { - let response = self - .get_public_api_response("harness-support/transcript") - .await?; - response - .bytes() - .await - .context("Failed to read harness-support transcript response body") - }) - .await + #[cfg(not(target_family = "wasm"))] + { + with_bounded_retry("fetch harness-support transcript", || async { + let response = self + .get_public_api_response("harness-support/transcript") + .await?; + response + .bytes() + .await + .context("Failed to read harness-support transcript body") + }) + .await + } + #[cfg(target_family = "wasm")] + { + unreachable!( + "fetch_transcript is not supported on wasm; agent_sdk is not built on this target" + ); + } } fn http_client(&self) -> &http_client::Client { diff --git a/crates/http_client/src/lib.rs b/crates/http_client/src/lib.rs index 2944038f1..936925b8b 100644 --- a/crates/http_client/src/lib.rs +++ b/crates/http_client/src/lib.rs @@ -203,13 +203,6 @@ impl Client { ) } - pub fn patch(&self, url: U) -> RequestBuilder<'_> { - self.builder( - self.wrapped.patch(url.clone()), - Self::include_warp_http_headers(url), - ) - } - pub fn put(&self, url: U) -> RequestBuilder<'_> { self.builder( self.wrapped.put(url.clone()), From 43804ab249e47aca2251dd6078544b3fda1c1fef Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Wed, 29 Apr 2026 07:29:47 -0400 Subject: [PATCH 06/13] Wake remote child agents on new events Co-Authored-By: Oz --- app/src/ai/agent/api/convert_conversation.rs | 2 + app/src/ai/agent/conversation.rs | 9 +- app/src/ai/agent/conversation_tests.rs | 24 ++ app/src/ai/agent_conversations_model_tests.rs | 7 + app/src/ai/agent_sdk/ambient.rs | 244 +++++++++++++++--- app/src/ai/agent_sdk/ambient_tests.rs | 106 +++++++- app/src/ai/agent_sdk/mod.rs | 47 +++- app/src/ai/agent_sdk/mod_tests.rs | 36 ++- app/src/ai/blocklist/controller.rs | 195 ++++++++++++-- app/src/ai/blocklist/history_model.rs | 2 + app/src/ai/blocklist/history_model_test.rs | 1 + .../orchestration_event_poller_tests.rs | 1 + .../ai/conversation_details_panel_tests.rs | 2 + app/src/pane_group/mod.rs | 39 +++ app/src/pane_group/mod_tests.rs | 59 +++++ app/src/server/server_api.rs | 34 +++ app/src/server/server_api/ai.rs | 42 +++ app/src/server/server_api/harness_support.rs | 2 +- app/src/terminal/view/ambient_agent/model.rs | 4 +- app/src/terminal/view/load_ai_conversation.rs | 1 + crates/persistence/src/model.rs | 32 +++ 21 files changed, 818 insertions(+), 71 deletions(-) diff --git a/app/src/ai/agent/api/convert_conversation.rs b/app/src/ai/agent/api/convert_conversation.rs index b77e430cd..943e9e456 100644 --- a/app/src/ai/agent/api/convert_conversation.rs +++ b/app/src/ai/agent/api/convert_conversation.rs @@ -82,6 +82,7 @@ pub fn convert_conversation_data_to_ai_conversation( parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, @@ -97,6 +98,7 @@ pub fn convert_conversation_data_to_ai_conversation( parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, // TODO: Populate run_id from server metadata once it is exposed // in ServerAIConversationMetadata. For cloud conversations that // were spawned via the server API, the run_id is created at task diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index 05a2fbbe8..13bd3306c 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -355,6 +355,7 @@ impl AIConversation { parent_agent_id, agent_name, parent_conversation_id, + is_remote_child, run_id, autoexecute_override, last_event_sequence, @@ -380,6 +381,7 @@ impl AIConversation { let parent_conversation_id = data .parent_conversation_id .and_then(|id| AIConversationId::try_from(id).ok()); + let is_remote_child = data.is_remote_child; let run_id = data.run_id; let autoexecute_override = if FeatureFlag::RememberFastForwardState.is_enabled() { data.autoexecute_override @@ -399,6 +401,7 @@ impl AIConversation { parent_agent_id, agent_name, parent_conversation_id, + is_remote_child, run_id, autoexecute_override, last_event_sequence, @@ -413,6 +416,7 @@ impl AIConversation { None, None, None, + false, None, AIConversationAutoexecuteMode::default(), None, @@ -457,7 +461,7 @@ impl AIConversation { parent_agent_id, agent_name, parent_conversation_id, - is_remote_child: false, + is_remote_child, last_event_sequence, }) } @@ -809,7 +813,7 @@ impl AIConversation { /// Returns true if this conversation was spawned by a parent orchestrator agent. pub fn is_child_agent_conversation(&self) -> bool { - self.parent_conversation_id.is_some() + self.parent_conversation_id.is_some() || self.parent_agent_id.is_some() } /// Returns true if this is a placeholder for a child agent executing on a @@ -2830,6 +2834,7 @@ impl AIConversation { parent_agent_id: self.parent_agent_id.clone(), agent_name: self.agent_name.clone(), parent_conversation_id: self.parent_conversation_id.map(|id| id.to_string()), + is_remote_child: self.is_remote_child, run_id: self.task_id.map(|id| id.to_string()), autoexecute_override: Some(self.autoexecute_override.into()), last_event_sequence: self.last_event_sequence, diff --git a/app/src/ai/agent/conversation_tests.rs b/app/src/ai/agent/conversation_tests.rs index 535286b3f..3ba96b4f1 100644 --- a/app/src/ai/agent/conversation_tests.rs +++ b/app/src/ai/agent/conversation_tests.rs @@ -132,6 +132,30 @@ fn restored_conversation_uses_persisted_last_event_sequence() { assert_eq!(conversation.last_event_sequence(), Some(42)); } +#[test] +fn restored_conversation_uses_persisted_remote_child_marker() { + let conversation_data: AgentConversationData = + serde_json::from_str(r#"{"server_conversation_token":null,"is_remote_child":true}"#) + .unwrap(); + + let conversation = restored_conversation(Some(conversation_data)); + + assert!(conversation.is_remote_child()); +} + +#[test] +fn child_conversation_detection_uses_parent_agent_id() { + let conversation_data: AgentConversationData = serde_json::from_str( + r#"{"server_conversation_token":null,"parent_agent_id":"parent-run-id"}"#, + ) + .unwrap(); + + let conversation = restored_conversation(Some(conversation_data)); + + assert!(conversation.is_child_agent_conversation()); + assert_eq!(conversation.parent_conversation_id(), None); +} + #[test] fn restored_conversation_defaults_unknown_persisted_autoexecute_override() { let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true); diff --git a/app/src/ai/agent_conversations_model_tests.rs b/app/src/ai/agent_conversations_model_tests.rs index 3042520a3..10c0f8eff 100644 --- a/app/src/ai/agent_conversations_model_tests.rs +++ b/app/src/ai/agent_conversations_model_tests.rs @@ -152,6 +152,7 @@ fn test_display_status_uses_matching_conversation_for_in_progress_task() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.clone()), autoexecute_override: None, last_event_sequence: None, @@ -204,6 +205,7 @@ fn test_display_status_updates_when_blocked_conversation_resumes() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.clone()), autoexecute_override: None, last_event_sequence: None, @@ -280,6 +282,7 @@ fn test_display_status_terminal_task_state_overrides_matching_conversation() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.clone()), autoexecute_override: None, last_event_sequence: None, @@ -332,6 +335,7 @@ fn test_status_filter_uses_display_status_for_task_backed_conversations() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.clone()), autoexecute_override: None, last_event_sequence: None, @@ -767,6 +771,7 @@ fn test_get_tasks_and_conversations_prefers_task_when_task_id_matches_conversati parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.clone()), autoexecute_override: None, last_event_sequence: None, @@ -824,6 +829,7 @@ fn test_get_tasks_and_conversations_prefers_task_when_server_token_matches() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, @@ -880,6 +886,7 @@ fn test_get_tasks_and_conversations_keeps_unrelated_tasks_and_conversations() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, diff --git a/app/src/ai/agent_sdk/ambient.rs b/app/src/ai/agent_sdk/ambient.rs index 2f696004c..1cc32c1a7 100644 --- a/app/src/ai/agent_sdk/ambient.rs +++ b/app/src/ai/agent_sdk/ambient.rs @@ -1,4 +1,5 @@ //! Commands to interact with ambient agents on Warp's platform. +use std::future::Future; use std::io::Write as _; use std::sync::Arc; use std::time::Duration; @@ -8,7 +9,7 @@ use crate::ai::ambient_agents::spawn::{ }; use crate::ai::ambient_agents::task::HarnessConfig; use crate::ai::ambient_agents::AmbientAgentTaskState; -use crate::ai::ambient_agents::{AgentConfigSnapshot, AmbientAgentTask}; +use crate::ai::ambient_agents::{AgentConfigSnapshot, AmbientAgentTask, AmbientAgentTaskId}; use crate::ai::artifacts::Artifact; use crate::auth::AuthStateProvider; use crate::server::server_api::ai::{ @@ -39,7 +40,7 @@ use warp_cli::{ }; use warp_core::channel::ChannelState; use warp_core::features::FeatureFlag; -use warpui::r#async::Timer; +use warpui::r#async::{FutureExt as _, Timer}; use warpui::{ platform::TerminationMode, r#async::Spawnable, AppContext, ModelContext, SingletonEntity, }; @@ -50,10 +51,11 @@ use crate::ai::agent_sdk::driver::attachments::{ use crate::cloud_object::model::persistence::CloudModel; use crate::server::ids::{ServerId, SyncId}; -use super::common::{EnvironmentChoice, ResolveConfigurationError}; +use super::common::{parse_ambient_task_id, EnvironmentChoice, ResolveConfigurationError}; const MAX_LINE_WIDTH: usize = 90; const STREAM_RETRY_BACKOFF_STEPS: &[u64] = &[1, 2, 5, 10]; +const SEND_AGENT_MESSAGE_TIMEOUT: Duration = Duration::from_secs(15); /// Singleton model that runs async work for ambient agent CLI commands. struct AmbientAgentRunner; @@ -640,17 +642,35 @@ impl AmbientAgentRunner { output_format: OutputFormat, ctx: &mut ModelContext, ) -> anyhow::Result<()> { - let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); + let provider = ServerApiProvider::as_ref(ctx); + let ai_client = provider.get_ai_client(); + let server_api = provider.get(); + let scoped_task_id = task_id_for_message_send(&args.sender_run_id)?; let future = async move { - let response = ai_client - .send_agent_message(SendAgentMessageRequest { - to: args.to, - subject: args.subject, - body: args.body, - sender_run_id: args.sender_run_id, - }) - .await?; + let request = SendAgentMessageRequest { + to: args.to, + subject: args.subject, + body: args.body, + sender_run_id: args.sender_run_id, + }; + let log_context = SendAgentMessageLogContext::new(&request, scoped_task_id.as_ref()); + let send_message = async move { + match scoped_task_id { + Some(task_id) => { + server_api + .send_agent_message_for_task(&task_id, request) + .await + } + None => ai_client.send_agent_message(request).await, + } + }; + let response = send_agent_message_result_with_timeout( + send_message, + &log_context, + SEND_AGENT_MESSAGE_TIMEOUT, + ) + .await?; print_send_message_response(&response, output_format)?; Ok(()) }; @@ -658,25 +678,31 @@ impl AmbientAgentRunner { Ok(()) } + fn list_messages( &self, args: MessageListArgs, output_format: OutputFormat, ctx: &mut ModelContext, ) -> anyhow::Result<()> { - let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); + let provider = ServerApiProvider::as_ref(ctx); + let ai_client = provider.get_ai_client(); + let server_api = provider.get(); let future = async move { - let messages = ai_client - .list_agent_messages( - &args.run_id, - ListAgentMessagesRequest { - unread_only: args.unread, - since: args.since, - limit: args.limit, - }, - ) - .await?; + let request = ListAgentMessagesRequest { + unread_only: args.unread, + since: args.since, + limit: args.limit, + }; + let messages = match task_id_from_run_id(&args.run_id) { + Some(task_id) => { + server_api + .list_agent_messages_for_task(&task_id, &args.run_id, request) + .await? + } + None => ai_client.list_agent_messages(&args.run_id, request).await?, + }; super::output::print_list(messages, output_format); Ok(()) }; @@ -708,10 +734,20 @@ impl AmbientAgentRunner { output_format: OutputFormat, ctx: &mut ModelContext, ) -> anyhow::Result<()> { - let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); + let provider = ServerApiProvider::as_ref(ctx); + let ai_client = provider.get_ai_client(); + let server_api = provider.get(); + let scoped_task_id = task_id_from_oz_run_id_env()?; let future = async move { - let message = ai_client.read_agent_message(&args.message_id).await?; + let message = match scoped_task_id { + Some(task_id) => { + server_api + .read_agent_message_for_task(&task_id, &args.message_id) + .await? + } + None => ai_client.read_agent_message(&args.message_id).await?, + }; print_read_message_response(&message, output_format)?; Ok(()) }; @@ -726,10 +762,20 @@ impl AmbientAgentRunner { output_format: OutputFormat, ctx: &mut ModelContext, ) -> anyhow::Result<()> { - let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); + let provider = ServerApiProvider::as_ref(ctx); + let ai_client = provider.get_ai_client(); + let server_api = provider.get(); + let scoped_task_id = task_id_from_oz_run_id_env()?; let future = async move { - ai_client.mark_message_delivered(&args.message_id).await?; + match scoped_task_id { + Some(task_id) => { + server_api + .mark_message_delivered_for_task(&task_id, &args.message_id) + .await? + } + None => ai_client.mark_message_delivered(&args.message_id).await?, + } print_mark_message_delivered_result(&args.message_id, output_format)?; Ok(()) }; @@ -933,6 +979,121 @@ fn write_stream_record(record: &T) -> anyhow::Result<()> { stdout.flush().context("unable to flush stdout")?; Ok(()) } + +fn task_id_from_run_id(run_id: &str) -> Option { + run_id.parse().ok() +} + +fn task_id_from_oz_run_id_env() -> anyhow::Result> { + match std::env::var(warp_cli::OZ_RUN_ID_ENV) { + Ok(run_id) => parse_ambient_task_id(&run_id, "Invalid OZ_RUN_ID").map(Some), + Err(std::env::VarError::NotPresent) => Ok(None), + Err(std::env::VarError::NotUnicode(_)) => Err(anyhow!( + "{} is set but is not valid Unicode", + warp_cli::OZ_RUN_ID_ENV + )), + } +} + +fn task_id_for_message_send(sender_run_id: &str) -> anyhow::Result> { + match task_id_from_run_id(sender_run_id) { + Some(task_id) => Ok(Some(task_id)), + None => task_id_from_oz_run_id_env(), + } +} + +#[derive(Debug, Clone)] +struct SendAgentMessageLogContext { + sender_run_id: String, + task_id: Option, + target_agent_ids: Vec, + subject: String, + body_len: usize, +} + +impl SendAgentMessageLogContext { + fn new(request: &SendAgentMessageRequest, task_id: Option<&AmbientAgentTaskId>) -> Self { + Self { + sender_run_id: request.sender_run_id.clone(), + task_id: task_id.map(|task_id| task_id.to_string()), + target_agent_ids: request.to.clone(), + subject: request.subject.clone(), + body_len: request.body.chars().count(), + } + } + + fn error_context(&self) -> String { + format!( + "Failed to send agent message (sender_run_id={:?}, task_id={:?}, target_agent_ids={:?})", + self.sender_run_id, self.task_id, self.target_agent_ids + ) + } +} + +async fn send_agent_message_result_with_timeout( + send_message: F, + log_context: &SendAgentMessageLogContext, + timeout: Duration, +) -> anyhow::Result +where + F: Future>, +{ + log::info!( + "Sending ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={}", + log_context.sender_run_id, + log_context.task_id, + log_context.target_agent_ids, + log_context.subject, + log_context.body_len + ); + + match send_message.with_timeout(timeout).await { + Ok(Ok(response)) => { + log::info!( + "Sent ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} message_ids={:?}", + log_context.sender_run_id, + log_context.task_id, + log_context.target_agent_ids, + log_context.subject, + log_context.body_len, + response.message_ids + ); + Ok(response) + } + Ok(Err(err)) => { + let err = err.context(log_context.error_context()); + log::warn!( + "Failed to send ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} error={err:#}", + log_context.sender_run_id, + log_context.task_id, + log_context.target_agent_ids, + log_context.subject, + log_context.body_len + ); + eprintln!("{err:#}"); + Err(err) + } + Err(_) => { + let err = anyhow!( + "Timed out sending agent message after {timeout:?} (sender_run_id={:?}, task_id={:?}, target_agent_ids={:?})", + log_context.sender_run_id, + log_context.task_id, + log_context.target_agent_ids + ); + log::warn!( + "Timed out sending ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} timeout={timeout:?}", + log_context.sender_run_id, + log_context.task_id, + log_context.target_agent_ids, + log_context.subject, + log_context.body_len + ); + eprintln!("{err:#}"); + Err(err) + } + } +} + async fn watch_messages_forever( server_api: Arc, ai_client: Arc, @@ -940,15 +1101,25 @@ async fn watch_messages_forever( ) -> anyhow::Result<()> { let run_id = args.run_id; let watched_run_ids = vec![run_id.clone()]; + let scoped_task_id = task_id_from_run_id(&run_id); let mut last_seen_sequence = args.since_sequence; let mut initial_connect = true; let mut failures = 0usize; loop { - let mut stream = match server_api - .stream_agent_events(&watched_run_ids, last_seen_sequence) - .await - { + let stream_result = match scoped_task_id.as_ref() { + Some(task_id) => { + server_api + .stream_agent_events_for_task(task_id, &watched_run_ids, last_seen_sequence) + .await + } + None => { + server_api + .stream_agent_events(&watched_run_ids, last_seen_sequence) + .await + } + }; + let mut stream = match stream_result { Ok(stream) => { if !initial_connect { eprintln!( @@ -1004,8 +1175,15 @@ async fn watch_messages_forever( last_seen_sequence = event.sequence; continue; }; - - let message = match ai_client.read_agent_message(&message_id).await { + let message_result = match scoped_task_id.as_ref() { + Some(task_id) => { + server_api + .read_agent_message_for_task(task_id, &message_id) + .await + } + None => ai_client.read_agent_message(&message_id).await, + }; + let message = match message_result { Ok(message) => message, Err(err) => { failures += 1; diff --git a/app/src/ai/agent_sdk/ambient_tests.rs b/app/src/ai/agent_sdk/ambient_tests.rs index 4522cad6c..f0d3488ef 100644 --- a/app/src/ai/agent_sdk/ambient_tests.rs +++ b/app/src/ai/agent_sdk/ambient_tests.rs @@ -1,7 +1,8 @@ -//! Unit tests for `filter_from_args`. Verifies the clap enums are faithfully translated into -//! `TaskListFilter` without dropping any fields. +//! Unit tests for ambient agent CLI argument mapping and message helpers. +use anyhow::anyhow; use chrono::{TimeZone, Utc}; +use futures::executor::block_on; use warp_cli::json_filter::JsonOutput; use warp_cli::task::{ @@ -12,6 +13,19 @@ use warp_cli::task::{ use super::*; use crate::server::server_api::ai::{ArtifactType, ExecutionLocation, RunSortBy, RunSortOrder}; +const TASK_ID: &str = "00000000-0000-0000-0000-000000000001"; +const OTHER_TASK_ID: &str = "00000000-0000-0000-0000-000000000002"; + +fn send_log_context() -> SendAgentMessageLogContext { + SendAgentMessageLogContext { + sender_run_id: TASK_ID.to_string(), + task_id: Some(TASK_ID.to_string()), + target_agent_ids: vec!["parent-agent".to_string()], + subject: "status".to_string(), + body_len: 12, + } +} + /// A `ListTasksArgs` whose fields are all at their defaults. fn empty_args() -> ListTasksArgs { ListTasksArgs { @@ -167,3 +181,91 @@ fn every_field_maps_through() { assert_eq!(filter.sort_order, Some(RunSortOrder::Asc)); assert_eq!(filter.cursor.as_deref(), Some("abcd==")); } + +#[test] +fn task_id_from_run_id_accepts_task_uuid() { + let task_id = task_id_from_run_id(TASK_ID).expect("valid task id"); + + assert_eq!(task_id.to_string(), TASK_ID); +} + +#[test] +fn task_id_from_run_id_ignores_non_task_ids() { + assert!(task_id_from_run_id("local-child-run").is_none()); +} + +#[test] +#[serial_test::serial] +fn task_id_for_message_send_prefers_sender_run_id() { + std::env::set_var(warp_cli::OZ_RUN_ID_ENV, OTHER_TASK_ID); + let task_id = task_id_for_message_send(TASK_ID) + .expect("valid task id") + .expect("task id"); + std::env::remove_var(warp_cli::OZ_RUN_ID_ENV); + + assert_eq!(task_id.to_string(), TASK_ID); +} + +#[test] +#[serial_test::serial] +fn task_id_for_message_send_falls_back_to_oz_run_id() { + std::env::set_var(warp_cli::OZ_RUN_ID_ENV, TASK_ID); + let task_id = task_id_for_message_send("local-child-run") + .expect("valid env task id") + .expect("task id"); + std::env::remove_var(warp_cli::OZ_RUN_ID_ENV); + + assert_eq!(task_id.to_string(), TASK_ID); +} + +#[test] +#[serial_test::serial] +fn task_id_from_oz_run_id_env_rejects_invalid_value() { + std::env::set_var(warp_cli::OZ_RUN_ID_ENV, "not-a-task-id"); + let err = task_id_from_oz_run_id_env().expect_err("invalid task id"); + std::env::remove_var(warp_cli::OZ_RUN_ID_ENV); + + assert!(err.to_string().contains("Invalid OZ_RUN_ID")); +} + +#[test] +fn send_agent_message_result_returns_success_before_timeout() { + let response = block_on(send_agent_message_result_with_timeout( + async { + Ok(SendAgentMessageResponse { + message_ids: vec!["message-1".to_string()], + }) + }, + &send_log_context(), + Duration::from_secs(1), + )) + .expect("send should succeed"); + + assert_eq!(response.message_ids, ["message-1"]); +} + +#[test] +fn send_agent_message_result_surfaces_request_errors() { + let err = block_on(send_agent_message_result_with_timeout( + async { Err::(anyhow!("server rejected request")) }, + &send_log_context(), + Duration::from_secs(1), + )) + .expect_err("send should fail"); + let rendered = format!("{err:#}"); + + assert!(rendered.contains("Failed to send agent message")); + assert!(rendered.contains("server rejected request")); +} + +#[test] +fn send_agent_message_result_times_out() { + let err = block_on(send_agent_message_result_with_timeout( + futures::future::pending::>(), + &send_log_context(), + Duration::from_millis(1), + )) + .expect_err("send should time out"); + + assert!(err.to_string().contains("Timed out sending agent message")); +} diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index 2d374c6b6..aae46e1d6 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -467,6 +467,24 @@ fn build_server_side_task( Ok((config, task)) } +fn reconcile_task_harness( + task_id: &str, + selected_harness: &mut Harness, + task_harness: Harness, +) -> Result { + if *selected_harness == Harness::Oz { + *selected_harness = task_harness; + } else if task_harness != *selected_harness { + return Err(AgentDriverError::TaskHarnessMismatch { + task_id: task_id.to_string(), + expected: task_harness.to_string(), + got: selected_harness.to_string(), + }); + } + + harness_kind(*selected_harness) +} + /// Resolve a `Prompt` to a plain string. fn resolve_prompt(prompt: &Prompt, ctx: &AppContext) -> Result { match prompt { @@ -571,6 +589,7 @@ impl AgentDriverRunner { // Pull relevant variables out of args before moving it into the closure. let share_requests = args.share.share.clone(); let bedrock_inference_role = args.bedrock_inference_role.clone(); + let has_task_id = args.task_id.is_some(); let args_harness = args.harness; // `--conversation` path (user-invoked local resume): validate before any task side // effects so mismatches fail fast. The `--task-id` path derives its conversation id @@ -578,13 +597,15 @@ impl AgentDriverRunner { // can currently be passed together (the worker server-side appends `--conversation` // alongside `--task-id` for Slack/Linear followups); when both are set, the explicit // `--conversation` value wins via the merge below. - if let Some(conversation_id) = args.conversation.as_deref() { - common::fetch_and_validate_conversation_harness( - server_api.clone(), - conversation_id, - args_harness, - ) - .await?; + if !has_task_id { + if let Some(conversation_id) = args.conversation.as_deref() { + common::fetch_and_validate_conversation_harness( + server_api.clone(), + conversation_id, + args_harness, + ) + .await?; + } } let resume_conversation_id = args.conversation.clone(); @@ -1054,13 +1075,11 @@ impl AgentDriverRunner { // task is linked to an existing conversation, since task harness and conversation harness // always match (the task spawned the conversation). if let Some(task_harness) = task_harness { - if task_harness != driver_options.selected_harness { - return Err(AgentDriverError::TaskHarnessMismatch { - task_id: task_id_str, - expected: task_harness.to_string(), - got: driver_options.selected_harness.to_string(), - }); - } + task.harness = reconcile_task_harness( + &task_id_str, + &mut driver_options.selected_harness, + task_harness, + )?; } // Set the task ID on the ServerApi so it's sent with all subsequent requests. diff --git a/app/src/ai/agent_sdk/mod_tests.rs b/app/src/ai/agent_sdk/mod_tests.rs index 4091d9343..4870d9968 100644 --- a/app/src/ai/agent_sdk/mod_tests.rs +++ b/app/src/ai/agent_sdk/mod_tests.rs @@ -1,12 +1,14 @@ +use super::{command_requires_auth, command_to_telemetry_event, reconcile_task_harness}; use serde_json::json; use warp_cli::{ + agent::Harness, artifact::{ArtifactCommand, DownloadArtifactArgs, GetArtifactArgs, UploadArtifactArgs}, task::{MessageCommand, MessageSendArgs, MessageWatchArgs, TaskCommand}, CliCommand, }; use warp_core::telemetry::TelemetryEvent; -use super::{command_requires_auth, command_to_telemetry_event}; +const TASK_ID: &str = "00000000-0000-0000-0000-000000000001"; #[test] fn logout_does_not_require_auth() { @@ -128,6 +130,38 @@ fn run_message_send_telemetry_defaults_to_unknown_harness() { assert_eq!(event.payload(), Some(json!({ "harness": "unknown" }))); } +#[test] +fn reconcile_task_harness_adopts_task_harness_when_cli_uses_default() { + let mut selected_harness = Harness::Oz; + let harness = reconcile_task_harness(TASK_ID, &mut selected_harness, Harness::Claude) + .expect("default harness should adopt task harness"); + + assert_eq!(selected_harness, Harness::Claude); + assert_eq!(harness.harness(), Harness::Claude); +} + +#[test] +fn reconcile_task_harness_allows_matching_explicit_harness() { + let mut selected_harness = Harness::Claude; + let harness = reconcile_task_harness(TASK_ID, &mut selected_harness, Harness::Claude) + .expect("matching harness should succeed"); + + assert_eq!(selected_harness, Harness::Claude); + assert_eq!(harness.harness(), Harness::Claude); +} + +#[test] +fn reconcile_task_harness_rejects_explicit_mismatch() { + let mut selected_harness = Harness::Gemini; + let err = reconcile_task_harness(TASK_ID, &mut selected_harness, Harness::Claude) + .expect_err("mismatched harness should fail"); + + assert_eq!(selected_harness, Harness::Gemini); + assert!(err.to_string().contains("Task")); + assert!(err.to_string().contains("--harness gemini")); + assert!(err.to_string().contains("claude")); +} + #[test] #[serial_test::serial] fn run_message_watch_telemetry_defaults_to_unknown_harness() { diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index e4bc1f56f..72fca84af 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -53,7 +53,7 @@ use crate::network::NetworkStatus; use crate::notebooks::editor::model::FileLinkResolutionContext; use crate::persistence::ModelEvent; use crate::search::slash_command_menu::static_commands::commands; -use crate::server::server_api::ai::AIClient; +use crate::server::server_api::ai::{AIClient, RunFollowupRequest}; use crate::server::server_api::{AIApiError, ServerApiProvider}; use crate::terminal::model::block::{ formatted_terminal_contents_for_input, BlockId, CURSOR_MARKER, @@ -326,6 +326,8 @@ pub struct BlocklistAIController { pending_auto_resume_handles: HashMap, /// Pending dormant Claude wake preparations for success-idle child conversations. pending_local_claude_wakes: HashMap, + /// Pending remote child wake follow-up submissions. + pending_remote_child_wakes: HashMap, /// Passive conversations explicitly requested to follow up after actions complete. pending_passive_follow_ups: HashSet, /// Passive suggestion results that should be included with the next request @@ -565,6 +567,7 @@ impl BlocklistAIController { attachments_download_dir: None, pending_auto_resume_handles: HashMap::new(), pending_local_claude_wakes: HashMap::new(), + pending_remote_child_wakes: HashMap::new(), pending_passive_follow_ups: HashSet::new(), pending_passive_suggestion_results: HashMap::new(), } @@ -1491,22 +1494,27 @@ impl BlocklistAIController { let owns = BlocklistAIHistoryModel::as_ref(ctx) .all_live_conversations_for_terminal_view(self.terminal_view_id) .any(|conversation| conversation.id() == conversation_id); - if !owns { - return false; - } - - if self + let has_active_stream = self .in_flight_response_streams - .has_active_stream_for_conversation(conversation_id, ctx) - { + .has_active_stream_for_conversation(conversation_id, ctx); + let Some(conversation) = + BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) + else { + log::info!( + "Pending events are not ready: conversation_id={conversation_id:?} reason=conversation_missing owns_conversation={owns} has_active_stream={has_active_stream}" + ); + return false; + }; + let is_success = matches!(conversation.status(), ConversationStatus::Success); + if !owns || has_active_stream || !is_success { + log::info!( + "Pending events are not ready: conversation_id={conversation_id:?} owns_conversation={owns} has_active_stream={has_active_stream} status={:?}", + conversation.status() + ); return false; } - BlocklistAIHistoryModel::as_ref(ctx) - .conversation(&conversation_id) - .is_some_and(|conversation| { - matches!(conversation.status(), ConversationStatus::Success) - }) + true } fn local_claude_wake_candidate( @@ -1514,13 +1522,31 @@ impl BlocklistAIController { conversation_id: AIConversationId, ctx: &ModelContext, ) -> Option<(AmbientAgentTaskId, Option, Option)> { - let conversation = BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id)?; + let Some(conversation) = + BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) + else { + log::info!( + "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=conversation_missing" + ); + return None; + }; if !conversation.is_child_agent_conversation() || conversation.is_remote_child() { + log::info!( + "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=not_local_child is_child_agent_conversation={} is_remote_child={}", + conversation.is_child_agent_conversation(), + conversation.is_remote_child() + ); return None; } + let Some(task_id) = conversation.task_id() else { + log::info!( + "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=missing_task_id" + ); + return None; + }; Some(( - conversation.task_id()?, + task_id, self.active_session .as_ref(ctx) .current_working_directory() @@ -1541,6 +1567,121 @@ impl BlocklistAIController { )) } + fn remote_child_wake_candidate( + &self, + conversation_id: AIConversationId, + ctx: &ModelContext, + ) -> Option { + let Some(conversation) = + BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) + else { + log::info!( + "Skipping remote child wake candidate: conversation_id={conversation_id:?} reason=conversation_missing" + ); + return None; + }; + if !conversation.is_remote_child() { + return None; + } + let Some(task_id) = conversation.task_id() else { + log::info!( + "Skipping remote child wake candidate: conversation_id={conversation_id:?} reason=missing_task_id" + ); + return None; + }; + + Some(task_id) + } + + fn remote_child_wake_followup_message(pending_message_count: usize) -> String { + format!( + "You have received {pending_message_count} new parent-agent message(s). \ + Read all unread agent messages, treat the latest parent instructions as authoritative, \ + continue from the current task state, and reply when appropriate." + ) + } + + fn maybe_submit_remote_child_wake( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) -> bool { + if self + .pending_remote_child_wakes + .contains_key(&conversation_id) + { + log::info!("Remote child wake already pending: conversation_id={conversation_id:?}"); + return true; + } + + let Some(task_id) = self.remote_child_wake_candidate(conversation_id, ctx) else { + return false; + }; + + let pending_message_events = OrchestrationEventService::handle(ctx) + .update(ctx, |svc, _| { + svc.peek_pending_message_events(conversation_id) + }); + if pending_message_events.is_empty() { + log::info!( + "Skipping generic pending-event injection for remote child with no pending message events: conversation_id={conversation_id:?} task_id={task_id}" + ); + return true; + } + + let pending_message_event_ids = pending_message_events + .iter() + .map(|event| event.event_id.clone()) + .collect_vec(); + let pending_message_count = pending_message_event_ids.len(); + let message = Self::remote_child_wake_followup_message(pending_message_count); + let server_api = ServerApiProvider::as_ref(ctx).get(); + let handle = ctx.spawn( + async move { + log::info!( + "Submitting remote child wake follow-up: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={pending_message_count}" + ); + server_api + .submit_run_followup(&task_id, RunFollowupRequest { message }) + .await + }, + move |me, result, ctx| { + me.pending_remote_child_wakes.remove(&conversation_id); + match result { + Ok(()) => { + let removed_events = + OrchestrationEventService::handle(ctx).update(ctx, |svc, _| { + svc.take_pending_events_by_id( + conversation_id, + &pending_message_event_ids, + ) + }); + if removed_events.len() != pending_message_event_ids.len() { + log::info!( + "Remote child wake follow-up submitted but pending message set changed: conversation_id={conversation_id:?} task_id={task_id} expected_message_event_count={} actual_removed_event_count={}", + pending_message_event_ids.len(), + removed_events.len() + ); + } + log::info!( + "Submitted remote child wake follow-up: conversation_id={conversation_id:?} task_id={task_id} message_count={pending_message_count}" + ); + me.handle_pending_events_ready(conversation_id, ctx); + } + Err(err) => { + log::warn!( + "Failed to submit remote child wake follow-up for {conversation_id:?} task_id={task_id}: {err:#}" + ); + me.schedule_pending_events_ready_retry(conversation_id, ctx); + } + } + }, + ); + self.pending_remote_child_wakes + .insert(conversation_id, handle); + true + } + fn pending_event_to_claude_wake_message(event: &PendingEvent) -> Option { let PendingEventDetail::Message { sequence, @@ -1628,6 +1769,7 @@ impl BlocklistAIController { .pending_local_claude_wakes .contains_key(&conversation_id) { + log::info!("Dormant Claude wake already pending: conversation_id={conversation_id:?}"); return true; } @@ -1642,6 +1784,9 @@ impl BlocklistAIController { svc.peek_pending_message_events(conversation_id) }); if pending_message_events.is_empty() { + log::info!( + "Skipping dormant Claude wake preparation: conversation_id={conversation_id:?} reason=no_pending_message_events" + ); return false; } @@ -1683,6 +1828,9 @@ impl BlocklistAIController { let remote = match result { Ok(Some(remote)) => remote, Ok(None) => { + log::info!( + "Falling back to generic pending-event injection after dormant Claude wake eligibility check: conversation_id={conversation_id:?} task_id={task_id}" + ); me.inject_pending_events_for_request(conversation_id, ctx); return; } @@ -1703,6 +1851,11 @@ impl BlocklistAIController { svc.take_pending_events_by_id(conversation_id, &pending_message_event_ids) }); if removed_events.len() != pending_message_event_ids.len() { + log::info!( + "Aborting dormant Claude wake because the pending message set changed: conversation_id={conversation_id:?} task_id={task_id} expected_message_event_count={} actual_removed_event_count={}", + pending_message_event_ids.len(), + removed_events.len() + ); if !removed_events.is_empty() { OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { svc.prepend_pending_events(conversation_id, removed_events, ctx); @@ -1716,6 +1869,11 @@ impl BlocklistAIController { .filter_map(Self::pending_event_to_claude_wake_message) .collect_vec(); if wake_messages.len() != removed_events.len() { + log::info!( + "Aborting dormant Claude wake because some removed events were not message events: conversation_id={conversation_id:?} task_id={task_id} removed_event_count={} wake_message_count={}", + removed_events.len(), + wake_messages.len() + ); OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { svc.prepend_pending_events(conversation_id, removed_events, ctx); }); @@ -1779,6 +1937,9 @@ impl BlocklistAIController { match result { Ok(command) => { if !me.conversation_ready_for_pending_events(conversation_id, ctx) { + log::info!( + "Requeueing dormant Claude wake messages because the conversation stopped being ready before execution: conversation_id={conversation_id:?} task_id={task_id}" + ); OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { svc.prepend_pending_events( conversation_id, @@ -1841,6 +2002,10 @@ impl BlocklistAIController { return; } + if self.maybe_submit_remote_child_wake(conversation_id, ctx) { + return; + } + if self.maybe_prepare_local_claude_wake(conversation_id, ctx) { return; } diff --git a/app/src/ai/blocklist/history_model.rs b/app/src/ai/blocklist/history_model.rs index e5e5149a6..668963cbb 100644 --- a/app/src/ai/blocklist/history_model.rs +++ b/app/src/ai/blocklist/history_model.rs @@ -1090,6 +1090,7 @@ impl BlocklistAIHistoryModel { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: Some(source_conversation.autoexecute_override().into()), last_event_sequence: None, @@ -1243,6 +1244,7 @@ impl BlocklistAIHistoryModel { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: Some(conversation.autoexecute_override().into()), last_event_sequence: None, diff --git a/app/src/ai/blocklist/history_model_test.rs b/app/src/ai/blocklist/history_model_test.rs index 6dd4d2d06..9c7e9df65 100644 --- a/app/src/ai/blocklist/history_model_test.rs +++ b/app/src/ai/blocklist/history_model_test.rs @@ -1174,6 +1174,7 @@ fn test_find_by_token_after_insert_forked_conversation_from_tasks() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, diff --git a/app/src/ai/blocklist/orchestration_event_poller_tests.rs b/app/src/ai/blocklist/orchestration_event_poller_tests.rs index 9150227b9..4b79b5e8e 100644 --- a/app/src/ai/blocklist/orchestration_event_poller_tests.rs +++ b/app/src/ai/blocklist/orchestration_event_poller_tests.rs @@ -142,6 +142,7 @@ fn restored_conversation( parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(run_id), autoexecute_override: None, last_event_sequence, diff --git a/app/src/ai/conversation_details_panel_tests.rs b/app/src/ai/conversation_details_panel_tests.rs index 001898a40..aecfd76a8 100644 --- a/app/src/ai/conversation_details_panel_tests.rs +++ b/app/src/ai/conversation_details_panel_tests.rs @@ -130,6 +130,7 @@ fn test_from_task_includes_linked_directory_when_run_id_matches() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: Some(task_id.to_string()), autoexecute_override: None, last_event_sequence: None, @@ -245,6 +246,7 @@ fn test_from_task_includes_linked_directory_when_server_token_matches() { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, diff --git a/app/src/pane_group/mod.rs b/app/src/pane_group/mod.rs index 9ce7b332a..8c117c75e 100644 --- a/app/src/pane_group/mod.rs +++ b/app/src/pane_group/mod.rs @@ -3084,6 +3084,45 @@ impl PaneGroup { ctx: &mut ViewContext, ) { let child_id = child_conversation.id(); + if child_conversation.is_remote_child() { + let Some(task_id) = child_conversation.task_id() else { + log::warn!( + "Cannot restore remote child conversation {child_id:?} without a task ID" + ); + return; + }; + + let new_pane_id = + self.insert_ambient_agent_pane_hidden_for_child_agent(parent_pane_id, ctx); + + if let Some(new_terminal_view) = self.terminal_view_from_pane_id(new_pane_id, ctx) { + new_terminal_view.update(ctx, |terminal_view, ctx| { + terminal_view.restore_conversation_after_view_creation( + RestoredAIConversation::new(child_conversation), + true, + ctx, + ); + terminal_view.enter_agent_view( + None, + Some(child_id), + AgentViewEntryOrigin::CloudAgent, + ctx, + ); + terminal_view + .ambient_agent_view_model() + .update(ctx, |model, ctx| { + model.set_conversation_id(Some(child_id)); + model.enter_viewing_existing_session(task_id, ctx); + }); + }); + + self.child_agent_panes.insert(child_id, new_pane_id.into()); + } else { + log::error!("Failed to get terminal view for remote child agent pane {child_id:?}"); + self.discard_pane(new_pane_id.into(), ctx); + } + return; + } let child_task_context = child_conversation .task_id() diff --git a/app/src/pane_group/mod_tests.rs b/app/src/pane_group/mod_tests.rs index bc3815b79..42d46bc33 100644 --- a/app/src/pane_group/mod_tests.rs +++ b/app/src/pane_group/mod_tests.rs @@ -277,6 +277,25 @@ fn request_ambient_agent_task_id_for_hidden_child( }) } +fn ambient_child_session_state( + panes: &PaneGroup, + child_pane_id: PaneId, + ctx: &mut ViewContext, +) -> (Option, bool, Option) { + let terminal_view = panes + .terminal_view_from_pane_id(child_pane_id, ctx) + .expect("child pane should have a terminal view"); + let terminal_view_ref = terminal_view.as_ref(ctx); + let active_conversation_id = terminal_view_ref.active_conversation_id(ctx); + let ambient_model = terminal_view_ref.ambient_agent_view_model().as_ref(ctx); + + ( + ambient_model.task_id(), + ambient_model.is_agent_running(), + active_conversation_id, + ) +} + struct PreAttachReturnsFalsePane { pane_id: PaneId, pane_configuration: ModelHandle, @@ -516,6 +535,46 @@ fn test_restored_hidden_child_pane_reapplies_ambient_task_id_to_controller() { }); } +#[test] +fn test_restored_remote_hidden_child_pane_enters_existing_ambient_session() { + let _orchestration_v2 = FeatureFlag::OrchestrationV2.override_enabled(true); + + App::test((), |mut app| async move { + initialize_app(&mut app); + let pane_group = mock_pane_group(&mut app, Default::default()); + + pane_group.update(&mut app, |panes, ctx| { + let parent_pane_id = get_newly_created_pane_id(panes, &[]); + let parent_conversation_id = start_parent_conversation(panes, parent_pane_id, ctx); + let task_id = new_ambient_agent_task_id(); + + let mut child_conversation = AIConversation::new(false); + child_conversation.set_parent_conversation_id(parent_conversation_id); + child_conversation.set_task_id(task_id); + child_conversation.mark_as_remote_child(); + let child_conversation_id = child_conversation.id(); + + panes.create_hidden_child_agent_pane(child_conversation, parent_pane_id, ctx); + + let child_pane_id = panes + .child_agent_panes + .get(&child_conversation_id) + .copied() + .expect("restored remote hidden child pane should be tracked"); + + let (ambient_task_id, is_agent_running, active_conversation_id) = + ambient_child_session_state(panes, child_pane_id, ctx); + + assert_eq!(ambient_task_id, Some(task_id)); + assert!( + is_agent_running, + "remote child restore should view the existing ambient session" + ); + assert_eq!(active_conversation_id, Some(child_conversation_id)); + }); + }); +} + #[test] fn test_active_session_id_reset_on_last_pane_close() { App::test((), |mut app| async move { diff --git a/app/src/server/server_api.rs b/app/src/server/server_api.rs index 62b85323f..264038bd1 100644 --- a/app/src/server/server_api.rs +++ b/app/src/server/server_api.rs @@ -701,6 +701,40 @@ impl ServerApi { Ok(request.eventsource()) } + pub async fn stream_agent_events_for_task( + &self, + task_id: &AmbientAgentTaskId, + run_ids: &[String], + since_sequence: i64, + ) -> Result { + debug_assert!(!run_ids.is_empty(), "run_ids must not be empty"); + let auth_token = self + .get_or_refresh_access_token() + .await + .context("Failed to get access token for SSE stream")?; + + let run_ids_param: String = run_ids + .iter() + .map(|id| format!("run_ids[]={}", urlencoding::encode(id))) + .collect::>() + .join("&"); + let url = format!( + "{}/api/v1/agent/events/stream?{run_ids_param}&since={since_sequence}", + ChannelState::rtc_http_url() + ); + + let mut request = self.client.get(&url); + if let Some(token) = auth_token.as_bearer_token() { + request = request.bearer_auth(token); + } + + for (name, value) in self.ambient_agent_headers_for_task(task_id).await? { + request = request.header(name, value); + } + + Ok(request.eventsource()) + } + /// Sends a POST request to a public API endpoint and returns the raw response on success. async fn post_public_api_response( &self, diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index a5a536006..452b9ad32 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -209,6 +209,11 @@ pub struct SpawnAgentRequest { // --- Orchestrations V2 messaging types --- +#[derive(Debug, Clone, serde::Serialize)] +pub struct RunFollowupRequest { + pub message: String, +} + #[derive(Debug, Clone, serde::Serialize)] pub struct SendAgentMessageRequest { pub to: Vec, @@ -905,6 +910,12 @@ pub trait AIClient: 'static + Send + Sync { // --- Orchestrations V2 messaging --- + async fn submit_run_followup( + &self, + run_id: &AmbientAgentTaskId, + request: RunFollowupRequest, + ) -> anyhow::Result<(), anyhow::Error>; + async fn send_agent_message( &self, request: SendAgentMessageRequest, @@ -992,6 +1003,28 @@ impl ServerApi { Ok(response) } + pub(crate) async fn list_agent_messages_for_task( + &self, + task_id: &AmbientAgentTaskId, + run_id: &str, + request: ListAgentMessagesRequest, + ) -> anyhow::Result, anyhow::Error> { + let mut params = vec![format!("limit={}", request.limit)]; + if request.unread_only { + params.push("unread=true".to_string()); + } + if let Some(since) = request.since { + params.push(format!("since={}", urlencoding::encode(&since))); + } + + let path = format!("agent/messages/{run_id}?{}", params.join("&")); + let response = self + .get_public_api_response_for_task(task_id, &path) + .await?; + let response = response.json::>().await?; + Ok(response) + } + pub(crate) async fn mark_message_delivered_for_task( &self, task_id: &AmbientAgentTaskId, @@ -1896,6 +1929,15 @@ impl AIClient for ServerApi { // --- Orchestrations V2 messaging --- + async fn submit_run_followup( + &self, + run_id: &AmbientAgentTaskId, + request: RunFollowupRequest, + ) -> anyhow::Result<(), anyhow::Error> { + self.post_public_api_unit(&format!("agent/runs/{run_id}/followups"), &request) + .await + } + async fn send_agent_message( &self, request: SendAgentMessageRequest, diff --git a/app/src/server/server_api/harness_support.rs b/app/src/server/server_api/harness_support.rs index e4b9947ac..301f6fd4a 100644 --- a/app/src/server/server_api/harness_support.rs +++ b/app/src/server/server_api/harness_support.rs @@ -171,7 +171,7 @@ pub trait HarnessSupportClient: 'static + Send + Sync { } impl ServerApi { - async fn get_public_api_response_for_task( + pub(crate) async fn get_public_api_response_for_task( &self, task_id: &AmbientAgentTaskId, path: &str, diff --git a/app/src/terminal/view/ambient_agent/model.rs b/app/src/terminal/view/ambient_agent/model.rs index 4e6d20203..df0829e3e 100644 --- a/app/src/terminal/view/ambient_agent/model.rs +++ b/app/src/terminal/view/ambient_agent/model.rs @@ -420,9 +420,7 @@ impl AmbientAgentViewModel { // Store the task ID for later use self.task_id = Some(task_id); - if matches!(self.status, Status::NotAmbientAgent) { - self.status = Status::AgentRunning; - } + self.status = Status::AgentRunning; // Fetch the task so we can set the correct environment (instead of defaulting to the most // recently-used one) and the correct harness (so non-oz viewers know to use the diff --git a/app/src/terminal/view/load_ai_conversation.rs b/app/src/terminal/view/load_ai_conversation.rs index b83b80f4b..1573d06e7 100644 --- a/app/src/terminal/view/load_ai_conversation.rs +++ b/app/src/terminal/view/load_ai_conversation.rs @@ -986,6 +986,7 @@ impl TerminalView { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, diff --git a/crates/persistence/src/model.rs b/crates/persistence/src/model.rs index f99220c60..e5069c619 100644 --- a/crates/persistence/src/model.rs +++ b/crates/persistence/src/model.rs @@ -1005,6 +1005,10 @@ impl<'de> Deserialize<'de> for PersistedAutoexecuteMode { }) } } +fn is_false(value: &bool) -> bool { + !*value +} + // Serializes to `conversation_data` column in `agent_conversations`. #[derive(Debug, Serialize, Deserialize, Clone)] pub struct AgentConversationData { @@ -1029,6 +1033,10 @@ pub struct AgentConversationData { /// The local conversation ID of the parent conversation. #[serde(default, skip_serializing_if = "Option::is_none")] pub parent_conversation_id: Option, + /// True when this conversation is a parent-side placeholder for a child + /// agent executing on a remote worker. + #[serde(default, skip_serializing_if = "is_false")] + pub is_remote_child: bool, /// The server-assigned run identifier (`ai_tasks.id`) for v2 orchestration. /// For local agents this arrives via StreamInit; for cloud agents it will /// come from SpawnAgentResponse once the local→cloud spawn path is wired. @@ -1347,6 +1355,7 @@ mod tests { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: Some(42), @@ -1356,6 +1365,27 @@ mod tests { assert_eq!(roundtripped.last_event_sequence, Some(42)); } + #[test] + fn agent_conversation_data_roundtrips_remote_child_marker() { + let data = AgentConversationData { + server_conversation_token: None, + conversation_usage_metadata: None, + reverted_action_ids: None, + forked_from_server_conversation_token: None, + artifacts_json: None, + parent_agent_id: None, + agent_name: None, + parent_conversation_id: None, + is_remote_child: true, + run_id: None, + autoexecute_override: None, + last_event_sequence: None, + }; + let json = serde_json::to_string(&data).expect("serialize"); + let roundtripped: AgentConversationData = serde_json::from_str(&json).expect("deserialize"); + assert!(roundtripped.is_remote_child); + } + #[test] fn agent_conversation_data_deserializes_legacy_payload_without_last_event_sequence() { // Legacy rows persisted before this feature landed omit the field @@ -1364,6 +1394,7 @@ mod tests { let data: AgentConversationData = serde_json::from_str(legacy_json).expect("legacy rows must deserialize"); assert_eq!(data.last_event_sequence, None); + assert!(!data.is_remote_child); } #[test] @@ -1377,6 +1408,7 @@ mod tests { parent_agent_id: None, agent_name: None, parent_conversation_id: None, + is_remote_child: false, run_id: None, autoexecute_override: None, last_event_sequence: None, From e20b9027fdd92a19236a511abf41f09de37317c8 Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Wed, 29 Apr 2026 08:33:14 -0400 Subject: [PATCH 07/13] Address remote agent wake review comments Co-Authored-By: Oz --- app/src/ai/blocklist/controller.rs | 113 ++++++++++++------ .../blocklist/orchestration_event_poller.rs | 3 + .../orchestration_event_poller_tests.rs | 38 ++++++ app/src/ai/blocklist/orchestration_events.rs | 31 +++++ .../blocklist/orchestration_events_tests.rs | 52 ++++++++ crates/warp_cli/src/agent.rs | 8 +- crates/warp_cli/src/lib_tests.rs | 18 ++- 7 files changed, 216 insertions(+), 47 deletions(-) diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 72fca84af..dbf39fcb4 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -69,6 +69,7 @@ use crate::workspaces::user_workspaces::UserWorkspaces; use crate::{send_telemetry_from_ctx, server::telemetry::TelemetryEvent}; use anyhow::anyhow; use chrono::{DateTime, Local}; +use futures::future::Either; use itertools::Itertools; use parking_lot::FairMutex; use pending_response_streams::PendingResponseStreams; @@ -88,6 +89,8 @@ use super::orchestration_events::{ }; use warpui::{AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; +const REMOTE_CHILD_WAKE_FOLLOWUP_TIMEOUT: Duration = Duration::from_secs(15); + #[derive(Debug, Clone)] pub struct SessionContext { session_type: Option, @@ -1641,9 +1644,16 @@ impl BlocklistAIController { log::info!( "Submitting remote child wake follow-up: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={pending_message_count}" ); - server_api - .submit_run_followup(&task_id, RunFollowupRequest { message }) - .await + let submit = Box::pin( + server_api.submit_run_followup(&task_id, RunFollowupRequest { message }), + ); + let timeout = Box::pin(Timer::after(REMOTE_CHILD_WAKE_FOLLOWUP_TIMEOUT)); + match futures::future::select(submit, timeout).await { + Either::Left((result, _)) => result, + Either::Right((_, _)) => Err(anyhow!( + "Timed out submitting remote child wake follow-up for task {task_id}" + )), + } }, move |me, result, ctx| { me.pending_remote_child_wakes.remove(&conversation_id); @@ -1910,26 +1920,6 @@ impl BlocklistAIController { wake_messages, ) .await?; - log::info!( - "Reopening dormant Claude task before wake command: conversation_id={conversation_id:?} task_id={task_id}" - ); - server_api - .update_agent_task( - task_id, - Some(AgentTaskState::InProgress), - None, - None, - None, - ) - .await - .map_err(|err| { - anyhow!( - "Failed to reopen dormant Claude task {task_id} before wake: {err:#}" - ) - })?; - log::info!( - "Reopened dormant Claude task before wake command: conversation_id={conversation_id:?} task_id={task_id}" - ); Ok::<_, anyhow::Error>(command) }, move |me, result, ctx| { @@ -1950,23 +1940,72 @@ impl BlocklistAIController { return; } - log::info!( - "Executing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id}" - ); - BlocklistAIHistoryModel::handle(ctx).update( - ctx, - |history_model, ctx| { - history_model.update_conversation_status( - me.terminal_view_id, - conversation_id, - ConversationStatus::InProgress, - ctx, + let server_api = ServerApiProvider::as_ref(ctx).get(); + let removed_events_for_retry = removed_events_for_retry.clone(); + let handle = ctx.spawn( + async move { + log::info!( + "Reopening dormant Claude task before wake command: conversation_id={conversation_id:?} task_id={task_id}" + ); + server_api + .update_agent_task( + task_id, + Some(AgentTaskState::InProgress), + None, + None, + None, + ) + .await + .map_err(|err| { + anyhow!( + "Failed to reopen dormant Claude task {task_id} before wake: {err:#}" + ) + })?; + log::info!( + "Reopened dormant Claude task before wake command: conversation_id={conversation_id:?} task_id={task_id}" ); + Ok::<_, anyhow::Error>(command) + }, + move |me, result, ctx| { + me.pending_local_claude_wakes.remove(&conversation_id); + match result { + Ok(command) => { + log::info!( + "Executing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id}" + ); + BlocklistAIHistoryModel::handle(ctx).update( + ctx, + |history_model, ctx| { + history_model.update_conversation_status( + me.terminal_view_id, + conversation_id, + ConversationStatus::InProgress, + ctx, + ); + }, + ); + ctx.emit( + BlocklistAIControllerEvent::ExecuteLocalHarnessCommand { + command, + }, + ); + } + Err(err) => { + log::warn!( + "Failed to reopen dormant Claude task for {conversation_id:?} task_id={task_id}: {err:#}" + ); + OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { + svc.prepend_pending_events( + conversation_id, + removed_events_for_retry.clone(), + ctx, + ); + }); + } + } }, ); - ctx.emit(BlocklistAIControllerEvent::ExecuteLocalHarnessCommand { - command, - }); + me.pending_local_claude_wakes.insert(conversation_id, handle); } Err(err) => { log::warn!( diff --git a/app/src/ai/blocklist/orchestration_event_poller.rs b/app/src/ai/blocklist/orchestration_event_poller.rs index 941f9cb61..36b835c7e 100644 --- a/app/src/ai/blocklist/orchestration_event_poller.rs +++ b/app/src/ai/blocklist/orchestration_event_poller.rs @@ -274,6 +274,9 @@ impl OrchestrationEventPoller { conversation_ids: &[AIConversationId], ctx: &mut ModelContext, ) { + if !FeatureFlag::OrchestrationV2.is_enabled() { + return; + } for conversation_id in conversation_ids { let Some(conversation) = BlocklistAIHistoryModel::as_ref(ctx).conversation(conversation_id) diff --git a/app/src/ai/blocklist/orchestration_event_poller_tests.rs b/app/src/ai/blocklist/orchestration_event_poller_tests.rs index 4b79b5e8e..75e63d7da 100644 --- a/app/src/ai/blocklist/orchestration_event_poller_tests.rs +++ b/app/src/ai/blocklist/orchestration_event_poller_tests.rs @@ -14,6 +14,7 @@ use chrono::Utc; use http::StatusCode; use std::collections::HashSet; use std::sync::Arc; +use warp_core::features::FeatureFlag; use warp_multi_agent_api as api; use warpui::{App, EntityId}; @@ -180,6 +181,43 @@ fn ambient_agent_task( } } +#[test] +fn restored_conversations_skip_v2_polling_when_orchestration_v2_disabled() { + App::test((), |mut app| async move { + let _orchestration_v2 = FeatureFlag::OrchestrationV2.override_enabled(false); + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[])); + let terminal_view_id = EntityId::new(); + let conversation_id = AIConversationId::new(); + let run_id = uuid::Uuid::new_v4().to_string(); + + history_model.update(&mut app, |history_model, ctx| { + history_model.restore_conversations( + terminal_view_id, + vec![restored_conversation( + conversation_id, + run_id.clone(), + Some(17), + )], + ctx, + ); + }); + + let server_api = ServerApiProvider::new_for_test().get(); + let poller = app.add_singleton_model(move |_| { + OrchestrationEventPoller::new_with_clients_for_test( + Arc::new(MockAIClient::new()), + server_api, + ) + }); + + poller.update(&mut app, |poller, ctx| { + poller.on_restored_conversations(&[conversation_id], ctx); + assert!(poller.watched_run_ids.is_empty()); + assert!(poller.event_cursor.is_empty()); + assert!(poller.restore_fetch_failures.is_empty()); + }); + }); +} #[test] fn finish_restore_fetch_merges_server_cursor_and_child_runs() { App::test((), |mut app| async move { diff --git a/app/src/ai/blocklist/orchestration_events.rs b/app/src/ai/blocklist/orchestration_events.rs index 6a6a6b418..1a62acbd8 100644 --- a/app/src/ai/blocklist/orchestration_events.rs +++ b/app/src/ai/blocklist/orchestration_events.rs @@ -359,6 +359,7 @@ impl OrchestrationEventService { } => { for conversation_id in conversation_ids { self.sync_conversation_status(*conversation_id, ctx); + self.restore_v1_lifecycle_subscription(*conversation_id, ctx); } } BlocklistAIHistoryEvent::RemoveConversation { @@ -417,6 +418,36 @@ impl OrchestrationEventService { .insert(conversation_id, conversation.status().clone()); } + fn restore_v1_lifecycle_subscription( + &mut self, + source_conversation_id: AIConversationId, + ctx: &ModelContext, + ) { + if FeatureFlag::OrchestrationV2.is_enabled() { + return; + } + + let target_agent_id = { + let history_model = BlocklistAIHistoryModel::as_ref(ctx); + let Some(conversation) = history_model.conversation(&source_conversation_id) else { + return; + }; + if !conversation.is_child_agent_conversation() { + return; + } + + conversation + .parent_conversation_id() + .and_then(|parent_id| history_model.conversation(&parent_id)) + .and_then(|parent| parent.orchestration_agent_id()) + .or_else(|| conversation.parent_agent_id().map(str::to_string)) + }; + + if let Some(target_agent_id) = target_agent_id { + self.register_lifecycle_subscription(source_conversation_id, target_agent_id, None); + } + } + fn on_conversation_status_updated( &mut self, conversation_id: AIConversationId, diff --git a/app/src/ai/blocklist/orchestration_events_tests.rs b/app/src/ai/blocklist/orchestration_events_tests.rs index 147bc835c..18a87c037 100644 --- a/app/src/ai/blocklist/orchestration_events_tests.rs +++ b/app/src/ai/blocklist/orchestration_events_tests.rs @@ -1,7 +1,10 @@ #![allow(deprecated)] use super::*; +use crate::ai::blocklist::BlocklistAIHistoryModel; use std::collections::HashSet; +use warp_core::features::FeatureFlag; use warp_multi_agent_api as api; +use warpui::{App, EntityId}; // Helper for constructing lifecycle pending events with minimal boilerplate. // Tests use this to focus on queue/coalescing behavior rather than payload setup. @@ -413,3 +416,52 @@ fn test_pending_message_helpers_peek_and_take_only_selected_messages() { PendingEventDetail::Lifecycle { .. } )); } + +#[test] +fn test_restored_v1_child_reregisters_lifecycle_subscription() { + App::test((), |mut app| async move { + let _orchestration_v2 = FeatureFlag::OrchestrationV2.override_enabled(false); + let terminal_view_id = EntityId::new(); + let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + let service = app.add_model(|_| OrchestrationEventService::new_without_subscriptions()); + + let (_parent_conversation_id, child_conversation_id) = + history_model.update(&mut app, |history_model, ctx| { + let parent_conversation_id = + history_model.start_new_conversation(terminal_view_id, false, false, ctx); + history_model.set_server_conversation_token_for_conversation( + parent_conversation_id, + "parent-token".to_string(), + ); + let child_conversation_id = history_model.start_new_child_conversation( + terminal_view_id, + "child".to_string(), + parent_conversation_id, + ctx, + ); + history_model.set_server_conversation_token_for_conversation( + child_conversation_id, + "child-token".to_string(), + ); + (parent_conversation_id, child_conversation_id) + }); + + service.update(&mut app, |service, ctx| { + service.handle_history_event( + &BlocklistAIHistoryEvent::RestoredConversations { + terminal_view_id, + conversation_ids: vec![child_conversation_id], + }, + ctx, + ); + + let routes = service + .lifecycle_subscription_routes + .get(&child_conversation_id) + .expect("restored V1 child should have a lifecycle subscription"); + assert_eq!(routes.len(), 1); + assert_eq!(routes[0].target_agent_id, "parent-token"); + assert_eq!(routes[0].subscribed_event_types, None); + }); + }); +} diff --git a/crates/warp_cli/src/agent.rs b/crates/warp_cli/src/agent.rs index 503523e93..3478d59d0 100644 --- a/crates/warp_cli/src/agent.rs +++ b/crates/warp_cli/src/agent.rs @@ -279,10 +279,10 @@ pub struct RunAgentArgs { pub snapshot: SnapshotArgs, /// Identifier for the task that spawned this agent, used to report progress. /// - /// Mutually exclusive with `--conversation`: when `--task-id` is set the conversation id - /// (if any) is read off the server-side task metadata, so the caller never needs to pass - /// both. - #[arg(long = "task-id", hide = true, conflicts_with_all = ["prompt", "saved_prompt", "file", "conversation"])] + /// When `--conversation` is omitted, the conversation id is read off the server-side + /// task metadata. Some worker follow-up call sites still pass both flags, so keep + /// accepting the compatibility shape until all producers have been updated. + #[arg(long = "task-id", hide = true, conflicts_with_all = ["prompt", "saved_prompt", "file"])] pub task_id: Option, /// Whether we are running the agent in a sandboxed environment. diff --git a/crates/warp_cli/src/lib_tests.rs b/crates/warp_cli/src/lib_tests.rs index 56396586c..bc5add087 100644 --- a/crates/warp_cli/src/lib_tests.rs +++ b/crates/warp_cli/src/lib_tests.rs @@ -1402,8 +1402,8 @@ fn agent_run_cloud_accepts_snapshot_flags() { } #[test] -fn agent_run_rejects_task_id_with_conversation() { - let err = Args::try_parse_from([ +fn agent_run_accepts_task_id_with_conversation_for_worker_followups() { + let args = Args::try_parse_from([ "warp", "agent", "run", @@ -1412,11 +1412,17 @@ fn agent_run_rejects_task_id_with_conversation() { "--conversation", "conv-123", ]) - .unwrap_err() - .to_string(); + .unwrap(); + + let Some(Command::CommandLine(boxed_cmd)) = args.command else { + panic!("Expected `warp agent run` command"); + }; + let CliCommand::Agent(AgentCommand::Run(run_args)) = boxed_cmd.as_ref() else { + panic!("Expected `warp agent run` command"); + }; - assert!(err.contains("--task-id")); - assert!(err.contains("--conversation")); + assert_eq!(run_args.task_id.as_deref(), Some("task-123")); + assert_eq!(run_args.conversation.as_deref(), Some("conv-123")); } #[test] From 943b7b54d9c0a285e686a169f49caf1d6fac0da9 Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Wed, 29 Apr 2026 09:11:51 -0400 Subject: [PATCH 08/13] Fix WASM build for dormant Claude wake path Co-Authored-By: Oz --- app/src/ai/blocklist/controller.rs | 30 ++++++++++++++++++++++++------ 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index dbf39fcb4..3a6bd40e6 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -30,7 +30,11 @@ use crate::ai::agent::{ PassiveSuggestionTriggerType, RunningCommand, }; use crate::ai::agent::{DocumentContentAttachmentSource, FileContext}; -use crate::ai::ambient_agents::{AmbientAgentTaskId, AmbientAgentTaskState}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::{ClaudeHarness, ClaudeWakeMessage, ClaudeWakeRemoteContext}; +use crate::ai::ambient_agents::AmbientAgentTaskId; +#[cfg(not(target_family = "wasm"))] +use crate::ai::ambient_agents::AmbientAgentTaskState; use crate::ai::document::ai_document_model::{ AIDocumentId, AIDocumentModel, AIDocumentUserEditStatus, }; @@ -42,7 +46,6 @@ use crate::ai::{ FinishedAIAgentOutput, RenderableAIError, RequestCost, RequestMetadata, StaticQueryType, UserQueryMode, }, - agent_sdk::{ClaudeHarness, ClaudeWakeMessage, ClaudeWakeRemoteContext}, llms::LLMPreferences, AIRequestUsageModel, }; @@ -75,18 +78,21 @@ use parking_lot::FairMutex; use pending_response_streams::PendingResponseStreams; use session_sharing_protocol::common::ParticipantId; use std::collections::{HashMap, HashSet}; +#[cfg(not(target_family = "wasm"))] use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; +#[cfg(not(target_family = "wasm"))] use warp_cli::agent::Harness; use warp_core::assertions::safe_assert; +#[cfg(not(target_family = "wasm"))] use warp_graphql::ai::AgentTaskState; use warp_multi_agent_api::{message, Task, ToolType}; use warpui::r#async::{SpawnedFutureHandle, Timer}; -use super::orchestration_events::{ - OrchestrationEventService, OrchestrationEventServiceEvent, PendingEvent, PendingEventDetail, -}; +use super::orchestration_events::{OrchestrationEventService, OrchestrationEventServiceEvent}; +#[cfg(not(target_family = "wasm"))] +use super::orchestration_events::{PendingEvent, PendingEventDetail}; use warpui::{AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; const REMOTE_CHILD_WAKE_FOLLOWUP_TIMEOUT: Duration = Duration::from_secs(15); @@ -1520,6 +1526,16 @@ impl BlocklistAIController { true } + #[cfg(target_family = "wasm")] + fn maybe_prepare_local_claude_wake( + &mut self, + _conversation_id: AIConversationId, + _ctx: &mut ModelContext, + ) -> bool { + false + } + + #[cfg(not(target_family = "wasm"))] fn local_claude_wake_candidate( &self, conversation_id: AIConversationId, @@ -1692,6 +1708,7 @@ impl BlocklistAIController { true } + #[cfg(not(target_family = "wasm"))] fn pending_event_to_claude_wake_message(event: &PendingEvent) -> Option { let PendingEventDetail::Message { sequence, @@ -1770,6 +1787,7 @@ impl BlocklistAIController { } } + #[cfg(not(target_family = "wasm"))] fn maybe_prepare_local_claude_wake( &mut self, conversation_id: AIConversationId, @@ -1966,7 +1984,7 @@ impl BlocklistAIController { ); Ok::<_, anyhow::Error>(command) }, - move |me, result, ctx| { + move |me, result, ctx: &mut ModelContext| { me.pending_local_claude_wakes.remove(&conversation_id); match result { Ok(command) => { From 1f88ed012771a36ae73a7ebd38f8523ea403d2e0 Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Wed, 29 Apr 2026 10:02:05 -0400 Subject: [PATCH 09/13] Fix WASM clippy for native wake paths Co-Authored-By: Oz --- app/src/ai/blocklist/action_model/execute/send_message.rs | 6 ++++-- app/src/ai/blocklist/controller.rs | 1 + app/src/ai/blocklist/orchestration_events.rs | 3 +++ app/src/pane_group/pane/terminal_pane.rs | 4 +++- app/src/server/server_api/ai.rs | 1 + 5 files changed, 12 insertions(+), 3 deletions(-) diff --git a/app/src/ai/blocklist/action_model/execute/send_message.rs b/app/src/ai/blocklist/action_model/execute/send_message.rs index 79172496c..6e42c09a0 100644 --- a/app/src/ai/blocklist/action_model/execute/send_message.rs +++ b/app/src/ai/blocklist/action_model/execute/send_message.rs @@ -1,10 +1,11 @@ -use std::time::Duration; - +#[cfg(not(target_family = "wasm"))] use anyhow::anyhow; #[cfg(not(target_family = "wasm"))] use futures::future::Either; use futures::{future::BoxFuture, FutureExt}; #[cfg(not(target_family = "wasm"))] +use std::time::Duration; +#[cfg(not(target_family = "wasm"))] use warpui::r#async::Timer; use warpui::{AppContext, Entity, ModelContext, SingletonEntity}; @@ -27,6 +28,7 @@ use warp_core::send_telemetry_from_ctx; use super::{ActionExecution, AnyActionExecution, ExecuteActionInput, PreprocessActionInput}; +#[cfg(not(target_family = "wasm"))] const SEND_AGENT_MESSAGE_TIMEOUT: Duration = Duration::from_secs(15); pub struct SendMessageToAgentExecutor { diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 3a6bd40e6..e139ddb1f 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -334,6 +334,7 @@ pub struct BlocklistAIController { /// 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, /// Pending remote child wake follow-up submissions. pending_remote_child_wakes: HashMap, diff --git a/app/src/ai/blocklist/orchestration_events.rs b/app/src/ai/blocklist/orchestration_events.rs index 1a62acbd8..c37a8ac05 100644 --- a/app/src/ai/blocklist/orchestration_events.rs +++ b/app/src/ai/blocklist/orchestration_events.rs @@ -56,11 +56,13 @@ impl LifecycleEventDetailStage { #[derive(Debug, Clone)] pub enum PendingEventDetail { Message { + #[cfg_attr(target_family = "wasm", allow(dead_code))] sequence: i64, message_id: String, addresses: Vec, subject: String, message_body: String, + #[cfg_attr(target_family = "wasm", allow(dead_code))] occurred_at: String, }, Lifecycle { @@ -935,6 +937,7 @@ impl OrchestrationEventService { removed } + #[cfg_attr(target_family = "wasm", allow(dead_code))] pub fn prepend_pending_events( &mut self, conversation_id: AIConversationId, diff --git a/app/src/pane_group/pane/terminal_pane.rs b/app/src/pane_group/pane/terminal_pane.rs index 612bef0be..75b2477f2 100644 --- a/app/src/pane_group/pane/terminal_pane.rs +++ b/app/src/pane_group/pane/terminal_pane.rs @@ -30,7 +30,7 @@ use crate::{ app_state::{AmbientAgentPaneSnapshot, LeafContents, TerminalPaneSnapshot}, pane_group::child_agent::{ create_error_child_agent_conversation, create_hidden_child_agent_conversation, - HiddenChildAgentConversation, HiddenChildAgentTaskContext, + HiddenChildAgentConversation, }, pane_group::{self, Direction, Event::OpenConversationHistory, PaneGroup}, persistence::{BlockCompleted, ModelEvent}, @@ -56,6 +56,8 @@ use crate::{ #[cfg(feature = "local_fs")] use crate::ai::blocklist::BlocklistAIHistoryEvent; #[cfg(not(target_family = "wasm"))] +use crate::pane_group::child_agent::HiddenChildAgentTaskContext; +#[cfg(not(target_family = "wasm"))] use crate::server::server_api::ServerApiProvider; use warp_core::execution_mode::AppExecutionMode; diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index 452b9ad32..039722cfa 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -1003,6 +1003,7 @@ impl ServerApi { Ok(response) } + #[cfg_attr(target_family = "wasm", allow(dead_code))] pub(crate) async fn list_agent_messages_for_task( &self, task_id: &AmbientAgentTaskId, From b0ba7fd632c1778451702a5ed0f77e0503274594 Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Thu, 30 Apr 2026 09:20:03 -0400 Subject: [PATCH 10/13] Prepare client for server-driven child wakes Co-Authored-By: Oz --- .../agent_sdk/driver/harness/claude_code.rs | 182 +------- .../driver/harness/claude_code/wake_driver.rs | 230 +++++++++ app/src/ai/agent_sdk/driver/harness/mod.rs | 2 +- app/src/ai/agent_sdk/mod.rs | 3 +- app/src/ai/blocklist/controller.rs | 438 +++++------------- .../blocklist/orchestration_event_streamer.rs | 21 +- .../orchestration_event_streamer_tests.rs | 24 + 7 files changed, 414 insertions(+), 486 deletions(-) create mode 100644 app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code.rs b/app/src/ai/agent_sdk/driver/harness/claude_code.rs index 5e6416cf2..78fb88533 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code.rs @@ -1,5 +1,4 @@ use std::collections::HashMap; -use std::ffi::OsString; use std::fmt::Write as _; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -9,18 +8,14 @@ use async_trait::async_trait; use parking_lot::Mutex; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -use shell_words::quote as shell_quote; use tempfile::NamedTempFile; use uuid::Uuid; use warp_cli::agent::Harness; use warpui::{ModelHandle, ModelSpawner}; use crate::ai::agent::conversation::AIConversationId; -use crate::ai::agent_events::MessageHydrator; use crate::ai::ambient_agents::AmbientAgentTaskId; -use crate::server::server_api::harness_support::{ - upload_to_target, HarnessSupportClient, ResolvePromptRequest, -}; +use crate::server::server_api::harness_support::{upload_to_target, HarnessSupportClient}; use crate::server::server_api::ServerApi; use crate::terminal::model::block::BlockId; use crate::terminal::model::session::ExecuteCommandOptions; @@ -34,174 +29,32 @@ use super::claude_transcript::{ }; use super::json_utils::{read_json_file_or_default, write_json_file}; use super::{ - cli_agent_session_status, task_env_vars, write_temp_file, HarnessCleanupDisposition, - HarnessRunner, ManagedSecretValue, ResumePayload, SavePoint, ThirdPartyHarness, + cli_agent_session_status, write_temp_file, HarnessCleanupDisposition, HarnessRunner, + ManagedSecretValue, ResumePayload, SavePoint, ThirdPartyHarness, }; mod parent_bridge; +mod wake_driver; #[cfg(test)] use super::super::OZ_MESSAGE_LISTENER_STATE_ROOT_ENV; -use parent_bridge::{ - acknowledge_parent_bridge_hook_output, ensure_parent_bridge_state_dir, - parent_bridge_max_context_chars, parent_bridge_root, prepare_parent_bridge_hook_output, - stage_parent_bridge_message, MessageBridge, MessageBridgeCleanupDisposition, - MessageBridgeMessageRecord, -}; #[cfg(test)] use parent_bridge::{ + acknowledge_parent_bridge_hook_output, ensure_parent_bridge_state_dir, parent_bridge_char_count, parent_bridge_event_cursor_file, parent_bridge_hook_output_ack_file, - parent_bridge_hook_output_file, parent_bridge_staged_message_path, - parent_bridge_surfaced_message_path, read_parent_bridge_event_cursor, - render_parent_bridge_message_block, write_parent_bridge_event_cursor, MessageBridgeHookOutput, - MESSAGE_BRIDGE_CONTEXT_PREAMBLE, + parent_bridge_hook_output_file, parent_bridge_root, parent_bridge_staged_message_path, + parent_bridge_surfaced_message_path, prepare_parent_bridge_hook_output, + read_parent_bridge_event_cursor, render_parent_bridge_message_block, + stage_parent_bridge_message, write_parent_bridge_event_cursor, MessageBridgeHookOutput, + MessageBridgeMessageRecord, MESSAGE_BRIDGE_CONTEXT_PREAMBLE, }; - -#[derive(Debug, Clone)] -pub(crate) struct ClaudeWakeMessage { - pub(crate) sequence: i64, - pub(crate) message_id: String, - pub(crate) sender_run_id: String, - pub(crate) subject: String, - pub(crate) body: String, - pub(crate) occurred_at: String, -} - -impl From for MessageBridgeMessageRecord { - fn from(value: ClaudeWakeMessage) -> Self { - Self { - sequence: value.sequence, - message_id: value.message_id, - sender_run_id: value.sender_run_id, - subject: value.subject, - body: value.body, - occurred_at: value.occurred_at, - } - } -} - -#[derive(Debug)] -pub(crate) struct ClaudeWakeRemoteContext { - session_id: Uuid, - envelope: ClaudeTranscriptEnvelope, - wake_prompt: String, -} +use parent_bridge::{MessageBridge, MessageBridgeCleanupDisposition}; +#[cfg(test)] +use shell_words::quote as shell_quote; +pub(crate) use wake_driver::ClaudeWakeMessage; +#[cfg(test)] +use wake_driver::{ClaudeWakeRemoteContext, CLAUDE_WAKE_PROMPT_FILE_NAME}; pub(crate) struct ClaudeHarness; - -impl ClaudeHarness { - pub(crate) async fn fetch_local_wake_remote_context( - task_id: AmbientAgentTaskId, - server_api: Arc, - ) -> Result { - let resolved = server_api - .resolve_prompt_for_task( - &task_id, - ResolvePromptRequest { - skill: None, - attachments_dir: None, - }, - ) - .await - .with_context(|| format!("Failed to resolve Claude wake prompt for task {task_id}"))?; - let bytes = server_api - .fetch_transcript_for_task(&task_id) - .await - .with_context(|| format!("Failed to fetch Claude transcript for task {task_id}"))?; - let envelope: ClaudeTranscriptEnvelope = - serde_json::from_slice(&bytes).with_context(|| { - format!("Failed to deserialize Claude transcript for wake task {task_id}") - })?; - let wake_prompt = match resolved.resumption_prompt { - Some(resumption_prompt) if !resumption_prompt.is_empty() => { - format!( - "{resumption_prompt} - -{CLAUDE_WAKE_PROMPT}" - ) - } - _ => CLAUDE_WAKE_PROMPT.to_string(), - }; - Ok(ClaudeWakeRemoteContext { - session_id: envelope.uuid, - envelope, - wake_prompt, - }) - } - - pub(crate) async fn prepare_local_wake_command( - server_api: Arc, - task_id: AmbientAgentTaskId, - parent_run_id: Option, - working_dir: Option, - mut remote: ClaudeWakeRemoteContext, - pending_messages: Vec, - ) -> Result { - let working_dir = working_dir.unwrap_or_else(|| remote.envelope.cwd.clone()); - prepare_claude_environment_config(&working_dir, &HashMap::new()) - .context("Failed to prepare Claude environment for wake")?; - - remote.envelope.cwd = working_dir.clone(); - let config_root = claude_config_dir().context("Failed to resolve Claude config dir")?; - write_envelope(&remote.envelope, &config_root) - .context("Failed to rehydrate Claude transcript for wake")?; - if let Err(error) = write_session_index_entry(remote.session_id, &working_dir, &config_root) - { - log::warn!("Failed to update Claude sessions-index.json for wake: {error:#}"); - } - - let state_dir = parent_bridge_root()?.join(remote.session_id.to_string()); - ensure_parent_bridge_state_dir(&state_dir)?; - let hydrator = MessageHydrator::for_task(server_api, task_id); - acknowledge_parent_bridge_hook_output(&hydrator, &state_dir).await?; - for record in pending_messages - .into_iter() - .map(MessageBridgeMessageRecord::from) - { - stage_parent_bridge_message(&state_dir, &record)?; - } - prepare_parent_bridge_hook_output(&hydrator, &state_dir, parent_bridge_max_context_chars()) - .await?; - - let prompt_path = state_dir.join(CLAUDE_WAKE_PROMPT_FILE_NAME); - std::fs::write(&prompt_path, remote.wake_prompt.as_bytes()) - .with_context(|| format!("Failed to write {}", prompt_path.display()))?; - - let command = claude_command( - CLIAgent::Claude.command_prefix(), - &remote.session_id, - &prompt_path.display().to_string(), - None, - true, - ); - let env_vars = task_env_vars(Some(&task_id), parent_run_id.as_deref(), Harness::Claude); - - Ok(prefix_command_with_env_vars(command, env_vars)) - } -} -fn prefix_command_with_env_vars(command: String, env_vars: HashMap) -> String { - if env_vars.is_empty() { - return command; - } - - let mut env_pairs = env_vars - .into_iter() - .map(|(key, value)| { - ( - key.to_string_lossy().into_owned(), - value.to_string_lossy().into_owned(), - ) - }) - .collect::>(); - env_pairs.sort_unstable_by(|(left, _), (right, _)| left.cmp(right)); - - let assignments = env_pairs - .into_iter() - .map(|(key, value)| format!("{key}={}", shell_quote(&value))) - .collect::>() - .join(" "); - - format!("env {assignments} {command}") -} #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] impl ThirdPartyHarness for ClaudeHarness { @@ -310,9 +163,6 @@ impl ThirdPartyHarness for ClaudeHarness { const CLAUDE_CODE_FORMAT: &str = "claude_code_cli"; /// Command used to exit claude. const CLAUDE_EXIT_COMMAND: &str = "/exit"; -const CLAUDE_WAKE_PROMPT: &str = - "New lead-agent messages are available. Read the latest lead-agent updates and continue the task accordingly."; -const CLAUDE_WAKE_PROMPT_FILE_NAME: &str = "wake-turn-prompt.txt"; /// Build the shell command that launches the Claude CLI for a given session and /// prompt file. diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs b/app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs new file mode 100644 index 000000000..955a0c861 --- /dev/null +++ b/app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs @@ -0,0 +1,230 @@ +use std::collections::HashMap; +use std::ffi::OsString; +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::{Context, Result}; +use shell_words::quote as shell_quote; +use uuid::Uuid; +use warp_cli::agent::Harness; +use warp_graphql::ai::AgentTaskState; + +use crate::ai::agent_events::MessageHydrator; +use crate::ai::ambient_agents::{AmbientAgentTaskId, AmbientAgentTaskState}; +use crate::server::server_api::ai::AIClient; +use crate::server::server_api::harness_support::ResolvePromptRequest; +use crate::server::server_api::ServerApi; +use crate::terminal::CLIAgent; + +use super::super::claude_transcript::{ + claude_config_dir, write_envelope, write_session_index_entry, ClaudeTranscriptEnvelope, +}; +use super::super::task_env_vars; +use super::parent_bridge::{ + acknowledge_parent_bridge_hook_output, ensure_parent_bridge_state_dir, + parent_bridge_max_context_chars, parent_bridge_root, prepare_parent_bridge_hook_output, + stage_parent_bridge_message, MessageBridgeMessageRecord, +}; +use super::{claude_command, prepare_claude_environment_config, ClaudeHarness}; + +const CLAUDE_WAKE_PROMPT: &str = + "New lead-agent messages are available. Read the latest lead-agent updates and continue the task accordingly."; +pub(super) const CLAUDE_WAKE_PROMPT_FILE_NAME: &str = "wake-turn-prompt.txt"; + +#[derive(Debug, Clone)] +pub(crate) struct ClaudeWakeMessage { + pub(crate) sequence: i64, + pub(crate) message_id: String, + pub(crate) sender_run_id: String, + pub(crate) subject: String, + pub(crate) body: String, + pub(crate) occurred_at: String, +} + +impl From for MessageBridgeMessageRecord { + fn from(value: ClaudeWakeMessage) -> Self { + Self { + sequence: value.sequence, + message_id: value.message_id, + sender_run_id: value.sender_run_id, + subject: value.subject, + body: value.body, + occurred_at: value.occurred_at, + } + } +} + +#[derive(Debug)] +pub(super) struct ClaudeWakeRemoteContext { + pub(super) session_id: Uuid, + pub(super) envelope: ClaudeTranscriptEnvelope, + pub(super) wake_prompt: String, +} + +impl ClaudeHarness { + pub(crate) async fn wake_dormant_session( + server_api: Arc, + task_id: AmbientAgentTaskId, + parent_run_id: Option, + working_dir: Option, + pending_messages: Vec, + ) -> Result> { + let task = server_api.get_ambient_agent_task(&task_id).await?; + let harness = task + .agent_config_snapshot + .as_ref() + .and_then(|snapshot| snapshot.harness.as_ref()) + .map(|config| config.harness_type); + log::info!( + "Evaluating dormant Claude wake: task_id={task_id} pending_message_count={} server_task_state={:?} harness={harness:?}", + pending_messages.len(), + task.state + ); + if task.state != AmbientAgentTaskState::Succeeded || harness != Some(Harness::Claude) { + log::info!( + "Skipping dormant Claude wake: task_id={task_id} pending_message_count={} server_task_state={:?} harness={harness:?}", + pending_messages.len(), + task.state + ); + return Ok(None); + } + + let remote = Self::fetch_local_wake_remote_context(task_id, server_api.clone()).await?; + let command = Self::prepare_local_wake_command( + server_api.clone(), + task_id, + parent_run_id, + working_dir, + remote, + pending_messages, + ) + .await?; + + log::info!("Reopening dormant Claude task before wake command: task_id={task_id}"); + server_api + .update_agent_task(task_id, Some(AgentTaskState::InProgress), None, None, None) + .await + .map_err(|err| { + anyhow::anyhow!( + "Failed to reopen dormant Claude task {task_id} before wake: {err:#}" + ) + })?; + log::info!("Reopened dormant Claude task before wake command: task_id={task_id}"); + + Ok(Some(command)) + } + + async fn fetch_local_wake_remote_context( + task_id: AmbientAgentTaskId, + server_api: Arc, + ) -> Result { + let resolved = server_api + .resolve_prompt_for_task( + &task_id, + ResolvePromptRequest { + skill: None, + attachments_dir: None, + }, + ) + .await + .with_context(|| format!("Failed to resolve Claude wake prompt for task {task_id}"))?; + let bytes = server_api + .fetch_transcript_for_task(&task_id) + .await + .with_context(|| format!("Failed to fetch Claude transcript for task {task_id}"))?; + let envelope: ClaudeTranscriptEnvelope = + serde_json::from_slice(&bytes).with_context(|| { + format!("Failed to deserialize Claude transcript for wake task {task_id}") + })?; + let wake_prompt = match resolved.resumption_prompt { + Some(resumption_prompt) if !resumption_prompt.is_empty() => { + format!( + "{resumption_prompt} + +{CLAUDE_WAKE_PROMPT}" + ) + } + _ => CLAUDE_WAKE_PROMPT.to_string(), + }; + Ok(ClaudeWakeRemoteContext { + session_id: envelope.uuid, + envelope, + wake_prompt, + }) + } + + pub(super) async fn prepare_local_wake_command( + server_api: Arc, + task_id: AmbientAgentTaskId, + parent_run_id: Option, + working_dir: Option, + mut remote: ClaudeWakeRemoteContext, + pending_messages: Vec, + ) -> Result { + let working_dir = working_dir.unwrap_or_else(|| remote.envelope.cwd.clone()); + prepare_claude_environment_config(&working_dir, &HashMap::new()) + .context("Failed to prepare Claude environment for wake")?; + + remote.envelope.cwd = working_dir.clone(); + let config_root = claude_config_dir().context("Failed to resolve Claude config dir")?; + write_envelope(&remote.envelope, &config_root) + .context("Failed to rehydrate Claude transcript for wake")?; + if let Err(error) = write_session_index_entry(remote.session_id, &working_dir, &config_root) + { + log::warn!("Failed to update Claude sessions-index.json for wake: {error:#}"); + } + + let state_dir = parent_bridge_root()?.join(remote.session_id.to_string()); + ensure_parent_bridge_state_dir(&state_dir)?; + let hydrator = MessageHydrator::for_task(server_api, task_id); + acknowledge_parent_bridge_hook_output(&hydrator, &state_dir).await?; + for record in pending_messages + .into_iter() + .map(MessageBridgeMessageRecord::from) + { + stage_parent_bridge_message(&state_dir, &record)?; + } + prepare_parent_bridge_hook_output(&hydrator, &state_dir, parent_bridge_max_context_chars()) + .await?; + + let prompt_path = state_dir.join(CLAUDE_WAKE_PROMPT_FILE_NAME); + std::fs::write(&prompt_path, remote.wake_prompt.as_bytes()) + .with_context(|| format!("Failed to write {}", prompt_path.display()))?; + + let command = claude_command( + CLIAgent::Claude.command_prefix(), + &remote.session_id, + &prompt_path.display().to_string(), + None, + true, + ); + let env_vars = task_env_vars(Some(&task_id), parent_run_id.as_deref(), Harness::Claude); + + Ok(prefix_command_with_env_vars(command, env_vars)) + } +} + +fn prefix_command_with_env_vars(command: String, env_vars: HashMap) -> String { + if env_vars.is_empty() { + return command; + } + + let mut env_pairs = env_vars + .into_iter() + .map(|(key, value)| { + ( + key.to_string_lossy().into_owned(), + value.to_string_lossy().into_owned(), + ) + }) + .collect::>(); + env_pairs.sort_unstable_by(|(left, _), (right, _)| left.cmp(right)); + + let assignments = env_pairs + .into_iter() + .map(|(key, value)| format!("{key}={}", shell_quote(&value))) + .collect::>() + .join(" "); + + format!("env {assignments} {command}") +} diff --git a/app/src/ai/agent_sdk/driver/harness/mod.rs b/app/src/ai/agent_sdk/driver/harness/mod.rs index b5afd4be7..84af99652 100644 --- a/app/src/ai/agent_sdk/driver/harness/mod.rs +++ b/app/src/ai/agent_sdk/driver/harness/mod.rs @@ -38,7 +38,7 @@ mod claude_code; pub(crate) mod claude_transcript; mod gemini; mod json_utils; -pub(crate) use claude_code::{ClaudeHarness, ClaudeWakeMessage, ClaudeWakeRemoteContext}; +pub(crate) use claude_code::{ClaudeHarness, ClaudeWakeMessage}; use claude_transcript::ClaudeResumeInfo; use gemini::GeminiHarness; diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index aae46e1d6..febea0471 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -68,8 +68,7 @@ use crate::ai::skills::{ }; pub(crate) use driver::harness::{ - task_env_vars, validate_cli_installed, ClaudeHarness, ClaudeWakeMessage, - ClaudeWakeRemoteContext, ThirdPartyHarness, + task_env_vars, validate_cli_installed, ClaudeHarness, ClaudeWakeMessage, ThirdPartyHarness, }; pub use driver::AgentDriver; use telemetry::CliTelemetryEvent; diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index e139ddb1f..653f82ebc 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -31,10 +31,8 @@ use crate::ai::agent::{ }; use crate::ai::agent::{DocumentContentAttachmentSource, FileContext}; #[cfg(not(target_family = "wasm"))] -use crate::ai::agent_sdk::{ClaudeHarness, ClaudeWakeMessage, ClaudeWakeRemoteContext}; +use crate::ai::agent_sdk::{ClaudeHarness, ClaudeWakeMessage}; use crate::ai::ambient_agents::AmbientAgentTaskId; -#[cfg(not(target_family = "wasm"))] -use crate::ai::ambient_agents::AmbientAgentTaskState; use crate::ai::document::ai_document_model::{ AIDocumentId, AIDocumentModel, AIDocumentUserEditStatus, }; @@ -56,7 +54,6 @@ use crate::network::NetworkStatus; use crate::notebooks::editor::model::FileLinkResolutionContext; use crate::persistence::ModelEvent; use crate::search::slash_command_menu::static_commands::commands; -use crate::server::server_api::ai::{AIClient, RunFollowupRequest}; use crate::server::server_api::{AIApiError, ServerApiProvider}; use crate::terminal::model::block::{ formatted_terminal_contents_for_input, BlockId, CURSOR_MARKER, @@ -72,7 +69,6 @@ use crate::workspaces::user_workspaces::UserWorkspaces; use crate::{send_telemetry_from_ctx, server::telemetry::TelemetryEvent}; use anyhow::anyhow; use chrono::{DateTime, Local}; -use futures::future::Either; use itertools::Itertools; use parking_lot::FairMutex; use pending_response_streams::PendingResponseStreams; @@ -82,11 +78,7 @@ use std::collections::{HashMap, HashSet}; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; -#[cfg(not(target_family = "wasm"))] -use warp_cli::agent::Harness; use warp_core::assertions::safe_assert; -#[cfg(not(target_family = "wasm"))] -use warp_graphql::ai::AgentTaskState; use warp_multi_agent_api::{message, Task, ToolType}; use warpui::r#async::{SpawnedFutureHandle, Timer}; @@ -95,8 +87,6 @@ use super::orchestration_events::{OrchestrationEventService, OrchestrationEventS use super::orchestration_events::{PendingEvent, PendingEventDetail}; use warpui::{AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; -const REMOTE_CHILD_WAKE_FOLLOWUP_TIMEOUT: Duration = Duration::from_secs(15); - #[derive(Debug, Clone)] pub struct SessionContext { session_type: Option, @@ -336,8 +326,6 @@ pub struct BlocklistAIController { /// Pending dormant Claude wake preparations for success-idle child conversations. #[cfg_attr(target_family = "wasm", allow(dead_code))] pending_local_claude_wakes: HashMap, - /// Pending remote child wake follow-up submissions. - pending_remote_child_wakes: HashMap, /// Passive conversations explicitly requested to follow up after actions complete. pending_passive_follow_ups: HashSet, /// Passive suggestion results that should be included with the next request @@ -577,7 +565,6 @@ impl BlocklistAIController { attachments_download_dir: None, pending_auto_resume_handles: HashMap::new(), pending_local_claude_wakes: HashMap::new(), - pending_remote_child_wakes: HashMap::new(), pending_passive_follow_ups: HashSet::new(), pending_passive_suggestion_results: HashMap::new(), } @@ -1587,128 +1574,66 @@ impl BlocklistAIController { )) } - fn remote_child_wake_candidate( - &self, + #[cfg(not(target_family = "wasm"))] + fn drain_remote_child_message_events( + &mut self, conversation_id: AIConversationId, - ctx: &ModelContext, - ) -> Option { + ctx: &mut ModelContext, + ) -> bool { let Some(conversation) = BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) else { log::info!( "Skipping remote child wake candidate: conversation_id={conversation_id:?} reason=conversation_missing" ); - return None; - }; - if !conversation.is_remote_child() { - return None; - } - let Some(task_id) = conversation.task_id() else { - log::info!( - "Skipping remote child wake candidate: conversation_id={conversation_id:?} reason=missing_task_id" - ); - return None; - }; - - Some(task_id) - } - - fn remote_child_wake_followup_message(pending_message_count: usize) -> String { - format!( - "You have received {pending_message_count} new parent-agent message(s). \ - Read all unread agent messages, treat the latest parent instructions as authoritative, \ - continue from the current task state, and reply when appropriate." - ) - } - - fn maybe_submit_remote_child_wake( - &mut self, - conversation_id: AIConversationId, - ctx: &mut ModelContext, - ) -> bool { - if self - .pending_remote_child_wakes - .contains_key(&conversation_id) - { - log::info!("Remote child wake already pending: conversation_id={conversation_id:?}"); - return true; - } - - let Some(task_id) = self.remote_child_wake_candidate(conversation_id, ctx) else { return false; }; + let is_remote_child = conversation.is_remote_child(); + let task_id = conversation.task_id(); + if !is_remote_child { + return false; + } let pending_message_events = OrchestrationEventService::handle(ctx) .update(ctx, |svc, _| { svc.peek_pending_message_events(conversation_id) }); if pending_message_events.is_empty() { - log::info!( - "Skipping generic pending-event injection for remote child with no pending message events: conversation_id={conversation_id:?} task_id={task_id}" - ); - return true; + return false; } let pending_message_event_ids = pending_message_events .iter() .map(|event| event.event_id.clone()) .collect_vec(); + let removed_events = OrchestrationEventService::handle(ctx).update(ctx, |svc, _| { + svc.take_pending_events_by_id(conversation_id, &pending_message_event_ids) + }); let pending_message_count = pending_message_event_ids.len(); - let message = Self::remote_child_wake_followup_message(pending_message_count); - let server_api = ServerApiProvider::as_ref(ctx).get(); - let handle = ctx.spawn( - async move { - log::info!( - "Submitting remote child wake follow-up: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={pending_message_count}" - ); - let submit = Box::pin( - server_api.submit_run_followup(&task_id, RunFollowupRequest { message }), - ); - let timeout = Box::pin(Timer::after(REMOTE_CHILD_WAKE_FOLLOWUP_TIMEOUT)); - match futures::future::select(submit, timeout).await { - Either::Left((result, _)) => result, - Either::Right((_, _)) => Err(anyhow!( - "Timed out submitting remote child wake follow-up for task {task_id}" - )), - } - }, - move |me, result, ctx| { - me.pending_remote_child_wakes.remove(&conversation_id); - match result { - Ok(()) => { - let removed_events = - OrchestrationEventService::handle(ctx).update(ctx, |svc, _| { - svc.take_pending_events_by_id( - conversation_id, - &pending_message_event_ids, - ) - }); - if removed_events.len() != pending_message_event_ids.len() { - log::info!( - "Remote child wake follow-up submitted but pending message set changed: conversation_id={conversation_id:?} task_id={task_id} expected_message_event_count={} actual_removed_event_count={}", - pending_message_event_ids.len(), - removed_events.len() - ); - } - log::info!( - "Submitted remote child wake follow-up: conversation_id={conversation_id:?} task_id={task_id} message_count={pending_message_count}" - ); - me.handle_pending_events_ready(conversation_id, ctx); - } - Err(err) => { - log::warn!( - "Failed to submit remote child wake follow-up for {conversation_id:?} task_id={task_id}: {err:#}" - ); - me.schedule_pending_events_ready_retry(conversation_id, ctx); - } - } - }, + if removed_events.len() != pending_message_event_ids.len() { + log::info!( + "Remote child pending message set changed while draining local wake events: conversation_id={conversation_id:?} task_id={:?} expected_message_event_count={} actual_removed_event_count={}", + task_id, + pending_message_event_ids.len(), + removed_events.len() + ); + } + log::info!( + "Drained remote child message events locally because server owns wakeups: conversation_id={conversation_id:?} task_id={:?} message_count={pending_message_count}", + task_id ); - self.pending_remote_child_wakes - .insert(conversation_id, handle); true } + #[cfg(target_family = "wasm")] + fn drain_remote_child_message_events( + &mut self, + _conversation_id: AIConversationId, + _ctx: &mut ModelContext, + ) -> bool { + false + } + #[cfg(not(target_family = "wasm"))] fn pending_event_to_claude_wake_message(event: &PendingEvent) -> Option { let PendingEventDetail::Message { @@ -1823,225 +1748,118 @@ impl BlocklistAIController { .iter() .map(|event| event.event_id.clone()) .collect_vec(); - let pending_message_count = pending_message_event_ids.len(); + let removed_events = OrchestrationEventService::handle(ctx).update(ctx, |svc, _| { + svc.take_pending_events_by_id(conversation_id, &pending_message_event_ids) + }); + if removed_events.len() != pending_message_event_ids.len() { + log::info!( + "Aborting dormant Claude wake because the pending message set changed: conversation_id={conversation_id:?} task_id={task_id} expected_message_event_count={} actual_removed_event_count={}", + pending_message_event_ids.len(), + removed_events.len() + ); + if !removed_events.is_empty() { + OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { + svc.prepend_pending_events(conversation_id, removed_events, ctx); + }); + } + return true; + } + + let wake_messages = removed_events + .iter() + .filter_map(Self::pending_event_to_claude_wake_message) + .collect_vec(); + if wake_messages.len() != removed_events.len() { + log::info!( + "Aborting dormant Claude wake because some removed events were not message events: conversation_id={conversation_id:?} task_id={task_id} removed_event_count={} wake_message_count={}", + removed_events.len(), + wake_messages.len() + ); + OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { + svc.prepend_pending_events(conversation_id, removed_events, ctx); + }); + return true; + } + log::info!( + "Prepared dormant Claude wake messages: conversation_id={conversation_id:?} task_id={task_id} message_count={}", + wake_messages.len() + ); + for wake_message in &wake_messages { + log::info!( + "Dormant Claude wake message: conversation_id={conversation_id:?} task_id={task_id} message_id={} sequence={} subject={:?} body_len={}", + wake_message.message_id, + wake_message.sequence, + wake_message.subject, + wake_message.body.chars().count() + ); + } + + let removed_events_for_retry = removed_events.clone(); let server_api = ServerApiProvider::as_ref(ctx).get(); + let wake_message_count = wake_messages.len(); let handle = ctx.spawn( async move { - let task = server_api.get_ambient_agent_task(&task_id).await?; - let harness = task - .agent_config_snapshot - .as_ref() - .and_then(|snapshot| snapshot.harness.as_ref()) - .map(|config| config.harness_type); log::info!( - "Evaluating dormant Claude wake: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={pending_message_count} server_task_state={:?} harness={harness:?}", - task.state + "Preparing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={wake_message_count}" ); - if task.state != AmbientAgentTaskState::Succeeded - || harness != Some(Harness::Claude) - { - log::info!( - "Skipping dormant Claude wake: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={pending_message_count} server_task_state={:?} harness={harness:?}", - task.state - ); - return Ok::, anyhow::Error>(None); - } - - ClaudeHarness::fetch_local_wake_remote_context(task_id, server_api.clone()) - .await - .map(Some) + ClaudeHarness::wake_dormant_session( + server_api.clone(), + task_id, + parent_run_id, + working_dir, + wake_messages, + ) + .await }, move |me, result, ctx| { me.pending_local_claude_wakes.remove(&conversation_id); - - let remote = match result { - Ok(Some(remote)) => remote, + match result { + Ok(Some(command)) => { + log::info!( + "Executing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id}" + ); + BlocklistAIHistoryModel::handle(ctx).update( + ctx, + |history_model, ctx| { + history_model.update_conversation_status( + me.terminal_view_id, + conversation_id, + ConversationStatus::InProgress, + ctx, + ); + }, + ); + ctx.emit(BlocklistAIControllerEvent::ExecuteLocalHarnessCommand { + command, + }); + } Ok(None) => { log::info!( "Falling back to generic pending-event injection after dormant Claude wake eligibility check: conversation_id={conversation_id:?} task_id={task_id}" ); + OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { + svc.prepend_pending_events( + conversation_id, + removed_events_for_retry.clone(), + ctx, + ); + }); me.inject_pending_events_for_request(conversation_id, ctx); - return; } Err(err) => { log::warn!( - "Failed to prepare dormant Claude wake context for {conversation_id:?}: {err:#}" + "Failed to prepare dormant Claude wake command for {conversation_id:?} task_id={task_id}: {err:#}" ); - me.schedule_pending_events_ready_retry(conversation_id, ctx); - return; - } - }; - - if !me.conversation_ready_for_pending_events(conversation_id, ctx) { - return; - } - - let removed_events = OrchestrationEventService::handle(ctx).update(ctx, |svc, _| { - svc.take_pending_events_by_id(conversation_id, &pending_message_event_ids) - }); - if removed_events.len() != pending_message_event_ids.len() { - log::info!( - "Aborting dormant Claude wake because the pending message set changed: conversation_id={conversation_id:?} task_id={task_id} expected_message_event_count={} actual_removed_event_count={}", - pending_message_event_ids.len(), - removed_events.len() - ); - if !removed_events.is_empty() { OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { - svc.prepend_pending_events(conversation_id, removed_events, ctx); + svc.prepend_pending_events( + conversation_id, + removed_events_for_retry.clone(), + ctx, + ); }); + me.schedule_pending_events_ready_retry(conversation_id, ctx); } - return; } - - let wake_messages = removed_events - .iter() - .filter_map(Self::pending_event_to_claude_wake_message) - .collect_vec(); - if wake_messages.len() != removed_events.len() { - log::info!( - "Aborting dormant Claude wake because some removed events were not message events: conversation_id={conversation_id:?} task_id={task_id} removed_event_count={} wake_message_count={}", - removed_events.len(), - wake_messages.len() - ); - OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { - svc.prepend_pending_events(conversation_id, removed_events, ctx); - }); - return; - } - log::info!( - "Prepared dormant Claude wake messages: conversation_id={conversation_id:?} task_id={task_id} message_count={}", - wake_messages.len() - ); - for wake_message in &wake_messages { - log::info!( - "Dormant Claude wake message: conversation_id={conversation_id:?} task_id={task_id} message_id={} sequence={} subject={:?} body_len={}", - wake_message.message_id, - wake_message.sequence, - wake_message.subject, - wake_message.body.chars().count() - ); - } - - let removed_events_for_retry = removed_events.clone(); - let server_api = ServerApiProvider::as_ref(ctx).get(); - let wake_message_count = wake_messages.len(); - let handle = ctx.spawn( - async move { - log::info!( - "Preparing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={wake_message_count}" - ); - let command = ClaudeHarness::prepare_local_wake_command( - server_api.clone(), - task_id, - parent_run_id, - working_dir, - remote, - wake_messages, - ) - .await?; - Ok::<_, anyhow::Error>(command) - }, - move |me, result, ctx| { - me.pending_local_claude_wakes.remove(&conversation_id); - match result { - Ok(command) => { - if !me.conversation_ready_for_pending_events(conversation_id, ctx) { - log::info!( - "Requeueing dormant Claude wake messages because the conversation stopped being ready before execution: conversation_id={conversation_id:?} task_id={task_id}" - ); - OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { - svc.prepend_pending_events( - conversation_id, - removed_events_for_retry.clone(), - ctx, - ); - }); - return; - } - - let server_api = ServerApiProvider::as_ref(ctx).get(); - let removed_events_for_retry = removed_events_for_retry.clone(); - let handle = ctx.spawn( - async move { - log::info!( - "Reopening dormant Claude task before wake command: conversation_id={conversation_id:?} task_id={task_id}" - ); - server_api - .update_agent_task( - task_id, - Some(AgentTaskState::InProgress), - None, - None, - None, - ) - .await - .map_err(|err| { - anyhow!( - "Failed to reopen dormant Claude task {task_id} before wake: {err:#}" - ) - })?; - log::info!( - "Reopened dormant Claude task before wake command: conversation_id={conversation_id:?} task_id={task_id}" - ); - Ok::<_, anyhow::Error>(command) - }, - move |me, result, ctx: &mut ModelContext| { - me.pending_local_claude_wakes.remove(&conversation_id); - match result { - Ok(command) => { - log::info!( - "Executing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id}" - ); - BlocklistAIHistoryModel::handle(ctx).update( - ctx, - |history_model, ctx| { - history_model.update_conversation_status( - me.terminal_view_id, - conversation_id, - ConversationStatus::InProgress, - ctx, - ); - }, - ); - ctx.emit( - BlocklistAIControllerEvent::ExecuteLocalHarnessCommand { - command, - }, - ); - } - Err(err) => { - log::warn!( - "Failed to reopen dormant Claude task for {conversation_id:?} task_id={task_id}: {err:#}" - ); - OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { - svc.prepend_pending_events( - conversation_id, - removed_events_for_retry.clone(), - ctx, - ); - }); - } - } - }, - ); - me.pending_local_claude_wakes.insert(conversation_id, handle); - } - Err(err) => { - log::warn!( - "Failed to finalize dormant Claude wake for {conversation_id:?} task_id={task_id}: {err:#}" - ); - OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { - svc.prepend_pending_events( - conversation_id, - removed_events_for_retry.clone(), - ctx, - ); - }); - } - } - }, - ); - me.pending_local_claude_wakes.insert(conversation_id, handle); }, ); self.pending_local_claude_wakes @@ -2059,8 +1877,8 @@ impl BlocklistAIController { if !self.conversation_ready_for_pending_events(conversation_id, ctx) { return; } - - if self.maybe_submit_remote_child_wake(conversation_id, ctx) { + if self.drain_remote_child_message_events(conversation_id, ctx) { + self.handle_pending_events_ready(conversation_id, ctx); return; } diff --git a/app/src/ai/blocklist/orchestration_event_streamer.rs b/app/src/ai/blocklist/orchestration_event_streamer.rs index 08191c54f..af238bf55 100644 --- a/app/src/ai/blocklist/orchestration_event_streamer.rs +++ b/app/src/ai/blocklist/orchestration_event_streamer.rs @@ -52,6 +52,7 @@ struct SseForwardingConsumer { tx: mpsc::UnboundedSender, self_run_id: String, hydrator: MessageHydrator, + hydrate_new_messages: bool, } #[cfg_attr(target_family = "wasm", async_trait(?Send))] @@ -61,10 +62,13 @@ impl AgentEventConsumer for SseForwardingConsumer { &mut self, event: AgentRunEvent, ) -> anyhow::Result { - let fetched_message = self - .hydrator - .hydrate_event_for_recipient(&event, &self.self_run_id) - .await; + let fetched_message = if self.hydrate_new_messages { + self.hydrator + .hydrate_event_for_recipient(&event, &self.self_run_id) + .await + } else { + None + }; self.tx .unbounded_send(SseStreamItem { @@ -636,10 +640,9 @@ impl OrchestrationEventStreamer { } // Track message IDs for server-side mark_delivered calls. - let message_ids: Vec = events + let message_ids: Vec = messages .iter() - .filter(|e| e.event_type == "new_message" && e.run_id == self_run_id) - .filter_map(|e| e.ref_id.clone()) + .map(|message| message.message_id.clone()) .collect(); if !message_ids.is_empty() { self.pending_delivery @@ -701,6 +704,9 @@ impl OrchestrationEventStreamer { .and_then(|c| c.run_id()) .map(|s| s.to_string()) .unwrap_or_default(); + let hydrate_new_messages = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .is_some_and(|conversation| !conversation.is_remote_child()); let (tx, rx) = mpsc::unbounded(); let generation = self.next_sse_generation; @@ -729,6 +735,7 @@ impl OrchestrationEventStreamer { tx, self_run_id, hydrator, + hydrate_new_messages, }; run_agent_event_driver(source, config, &mut consumer).await }, diff --git a/app/src/ai/blocklist/orchestration_event_streamer_tests.rs b/app/src/ai/blocklist/orchestration_event_streamer_tests.rs index 63ea5755e..d197060d4 100644 --- a/app/src/ai/blocklist/orchestration_event_streamer_tests.rs +++ b/app/src/ai/blocklist/orchestration_event_streamer_tests.rs @@ -278,6 +278,30 @@ fn build_pending_events_preserves_message_sequence_and_timestamp() { assert_eq!(message_id, "message-123"); assert_eq!(event_occurred_at, occurred_at); } + +#[tokio::test] +async fn sse_forwarding_consumer_skips_message_hydration_when_disabled() { + use futures::StreamExt; + + let (tx, mut rx) = futures::channel::mpsc::unbounded(); + let mut ai_client = crate::server::server_api::ai::MockAIClient::new(); + ai_client.expect_read_agent_message().times(0); + let ai_client: Arc = Arc::new(ai_client); + let hydrator = MessageHydrator::new(ai_client); + let mut consumer = SseForwardingConsumer { + tx, + self_run_id: "child-run".to_string(), + hydrator, + hydrate_new_messages: false, + }; + let event = make_run_event("new_message", "child-run", Some("message-123")); + + consumer.on_event(event).await.unwrap(); + + let item = rx.next().await.expect("expected forwarded event"); + assert_eq!(item.event.ref_id.as_deref(), Some("message-123")); + assert!(item.fetched_message.is_none()); +} #[test] fn finish_restore_fetch_uses_server_cursor_when_sqlite_is_absent() { use crate::ai::agent::conversation::AIConversation; From 4b00a4052b702f3103db9728e0518a793cb46a42 Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Thu, 30 Apr 2026 11:32:26 -0400 Subject: [PATCH 11/13] fix failing presubmit checks --- Cargo.lock | 1 + app/Cargo.toml | 1 + app/src/ai/agent/conversation.rs | 1 + app/src/ai/agent_sdk/driver/harness/codex.rs | 447 ++++++++++++++++++ app/src/ai/agent_sdk/driver/harness/mod.rs | 3 + app/src/ai/agent_sdk/mod.rs | 1 + app/src/ai/ambient_agents/task.rs | 1 + .../history_model/conversation_loader.rs | 2 +- app/src/ai/harness_display.rs | 6 +- app/src/pane_group/mod.rs | 57 ++- .../pane_group/pane/local_harness_launch.rs | 5 +- app/src/server/server_api/ai.rs | 41 +- app/src/terminal/cli_agent.rs | 2 +- .../terminal/view/ambient_agent/view_impl.rs | 1 + crates/warp_cli/src/agent.rs | 8 +- 15 files changed, 543 insertions(+), 34 deletions(-) create mode 100644 app/src/ai/agent_sdk/driver/harness/codex.rs diff --git a/Cargo.lock b/Cargo.lock index 5487866fb..e8f7c82a3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14358,6 +14358,7 @@ dependencies = [ "tokio", "tokio-util", "toml 0.8.23", + "toml_edit 0.25.6+spec-1.1.0", "tracing", "typed-path 0.10.0", "ui_components", diff --git a/app/Cargo.toml b/app/Cargo.toml index 6cd6009fd..7fba4b339 100644 --- a/app/Cargo.toml +++ b/app/Cargo.toml @@ -199,6 +199,7 @@ tikv-jemallocator = { version = "0.6", optional = true, features = [ "override_allocator_on_supported_platforms", ] } toml = "0.8.13" +toml_edit.workspace = true tracing.workspace = true ui_components.workspace = true unicase = "2.7.0" diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index 13bd3306c..905d46eca 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -3521,6 +3521,7 @@ pub enum AIAgentHarness { Oz, ClaudeCode, Gemini, + Codex, Unknown, } diff --git a/app/src/ai/agent_sdk/driver/harness/codex.rs b/app/src/ai/agent_sdk/driver/harness/codex.rs new file mode 100644 index 000000000..28a99a67d --- /dev/null +++ b/app/src/ai/agent_sdk/driver/harness/codex.rs @@ -0,0 +1,447 @@ +use std::collections::HashMap; +use std::fs; +use std::path::Path; +use std::sync::Arc; + +use anyhow::{Context, Result}; +use async_trait::async_trait; +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; +use warp_cli::agent::Harness; +use warp_managed_secrets::ManagedSecretValue; +use warpui::{ModelHandle, ModelSpawner}; + +use crate::ai::agent::conversation::AIConversationId; +use crate::ai::ambient_agents::AmbientAgentTaskId; +use crate::server::server_api::harness_support::HarnessSupportClient; +use crate::server::server_api::ServerApi; +use crate::terminal::model::block::BlockId; +use crate::terminal::CLIAgent; + +use super::super::terminal::{CommandHandle, TerminalDriver}; +use super::super::{AgentDriver, AgentDriverError}; +use super::json_utils::read_json_file_or_default; +use super::{write_temp_file, HarnessRunner, ResumePayload, SavePoint, ThirdPartyHarness}; + +pub(crate) struct CodexHarness; + +/// Format slug sent to the server when creating a Codex conversation. +const CODEX_CLI_FORMAT: &str = "codex_cli"; +/// Slash command Codex's TUI recognises as a graceful shutdown. +const CODEX_EXIT_COMMAND: &str = "/exit"; + +#[cfg_attr(not(target_family = "wasm"), async_trait)] +#[cfg_attr(target_family = "wasm", async_trait(?Send))] +impl ThirdPartyHarness for CodexHarness { + fn harness(&self) -> Harness { + Harness::Codex + } + + fn cli_agent(&self) -> CLIAgent { + CLIAgent::Codex + } + + fn install_docs_url(&self) -> Option<&'static str> { + Some("https://developers.openai.com/codex/cli") + } + + fn prepare_environment_config( + &self, + working_dir: &Path, + system_prompt: Option<&str>, + secrets: &HashMap, + ) -> Result<(), AgentDriverError> { + prepare_codex_environment_config(working_dir, system_prompt, secrets).map_err(|error| { + AgentDriverError::HarnessConfigSetupFailed { + harness: self.cli_agent().command_prefix().to_owned(), + error, + } + }) + } + + fn build_runner( + &self, + prompt: &str, + system_prompt: Option<&str>, + _resumption_prompt: Option<&str>, + working_dir: &Path, + _task_id: Option, + server_api: Arc, + terminal_driver: ModelHandle, + _resume: Option, + ) -> Result, AgentDriverError> { + // TODO(REMOTE-1503): support resume for Codex. + let client: Arc = server_api; + Ok(Box::new(CodexHarnessRunner::new( + self.cli_agent().command_prefix(), + prompt, + system_prompt, + working_dir, + client, + terminal_driver, + )?)) + } +} + +/// Build the shell command that launches the Codex TUI. +/// +/// `--dangerously-bypass-approvals-and-sandbox` disables both the sandbox and approval +/// prompts so the agent can run autonomously. +fn codex_command(cli_name: &str, prompt_path: &str) -> String { + format!("{cli_name} --dangerously-bypass-approvals-and-sandbox \"$(cat '{prompt_path}')\"") +} + +enum CodexRunnerState { + Preexec, + Running { + conversation_id: AIConversationId, + block_id: BlockId, + }, +} + +struct CodexHarnessRunner { + command: String, + /// Held so the temp file is cleaned up when the runner is dropped. + _temp_prompt_file: tempfile::NamedTempFile, + client: Arc, + terminal_driver: ModelHandle, + state: Mutex, +} + +impl CodexHarnessRunner { + fn new( + cli_command: &str, + prompt: &str, + _system_prompt: Option<&str>, + _working_dir: &Path, + client: Arc, + terminal_driver: ModelHandle, + ) -> Result { + let temp_file = write_temp_file("oz_prompt_", prompt)?; + let prompt_path = temp_file.path().display().to_string(); + + Ok(Self { + command: codex_command(cli_command, &prompt_path), + _temp_prompt_file: temp_file, + client, + terminal_driver, + state: Mutex::new(CodexRunnerState::Preexec), + }) + } +} + +#[cfg_attr(not(target_family = "wasm"), async_trait)] +#[cfg_attr(target_family = "wasm", async_trait(?Send))] +impl HarnessRunner for CodexHarnessRunner { + async fn start( + &self, + foreground: &ModelSpawner, + ) -> Result { + let conversation_id = self + .client + .create_external_conversation(CODEX_CLI_FORMAT) + .await + .map_err(|e| { + log::error!("Failed to create external conversation: {e}"); + AgentDriverError::ConfigBuildFailed(e) + })?; + log::info!("Created external conversation {conversation_id}"); + + let command = self.command.clone(); + let terminal_driver = self.terminal_driver.clone(); + let command_handle = foreground + .spawn(move |_, ctx| { + terminal_driver.update(ctx, |driver, ctx| driver.execute_command(&command, ctx)) + }) + .await?? + .await?; + + *self.state.lock() = CodexRunnerState::Running { + conversation_id, + block_id: command_handle.block_id().clone(), + }; + + Ok(command_handle) + } + + async fn exit(&self, foreground: &ModelSpawner) -> Result<()> { + log::info!("Sending /exit to Codex CLI"); + let terminal_driver = self.terminal_driver.clone(); + foreground + .spawn(move |_, ctx| { + terminal_driver.update(ctx, |driver, ctx| { + driver.send_text_to_cli(CODEX_EXIT_COMMAND.to_string(), ctx); + }); + }) + .await + .map_err(|_| anyhow::anyhow!("Agent driver dropped while sending /exit")) + } + + async fn save_conversation( + &self, + save_point: SavePoint, + foreground: &ModelSpawner, + ) -> Result<()> { + if matches!(save_point, SavePoint::Periodic) + && !super::has_running_cli_agent(&self.terminal_driver, foreground).await + { + log::debug!("Will not save conversation, Codex not in progress"); + return Ok(()); + } + + let (conversation_id, block_id) = match &*self.state.lock() { + CodexRunnerState::Preexec => { + log::warn!("save_conversation called before start"); + return Ok(()); + } + CodexRunnerState::Running { + conversation_id, + block_id, + } => (*conversation_id, block_id.clone()), + }; + + // TODO(REMOTE-1504) Also save the conversation transcript. + super::upload_current_block_snapshot( + foreground, + &self.terminal_driver, + self.client.as_ref(), + conversation_id, + block_id, + ) + .await + } +} + +const CODEX_CONFIG_DIR: &str = ".codex"; +const CODEX_AGENTS_OVERRIDE_FILE_NAME: &str = "AGENTS.override.md"; +const CODEX_AUTH_FILE_NAME: &str = "auth.json"; +const CODEX_CONFIG_TOML_FILE_NAME: &str = "config.toml"; +const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; +const CODEX_AUTH_MODE_API_KEY: &str = "apikey"; +/// Lowercase string Codex's `TrustLevel` enum serializes to (codex +/// `protocol/src/config_types.rs::TrustLevel`). +const CODEX_TRUST_LEVEL_TRUSTED: &str = "trusted"; +/// Top-level config key codex reads to override the built-in `openai` provider's base URL +/// (codex `core/src/config/mod.rs`). +const CODEX_OPENAI_BASE_URL_KEY: &str = "openai_base_url"; +/// US data-residency endpoint. Our OpenAI keys are issued under a US-residency project, +/// which rejects requests to the global host with `401 incorrect_hostname`. +/// TODO(REMOTE-1509): plumb a region-tagged auth secret instead of hardcoding the URL. +const CODEX_OPENAI_BASE_URL: &str = "https://us.api.openai.com/v1"; + +fn prepare_codex_environment_config( + working_dir: &Path, + system_prompt: Option<&str>, + secrets: &HashMap, +) -> Result<()> { + let home_dir = + dirs::home_dir().ok_or_else(|| anyhow::anyhow!("could not determine home directory"))?; + let codex_dir = home_dir.join(CODEX_CONFIG_DIR); + + if let Some(prompt) = system_prompt { + write_codex_agents_override(&codex_dir, prompt)?; + } + + match resolve_openai_api_key(secrets) { + Some(api_key) => prepare_codex_auth(&codex_dir.join(CODEX_AUTH_FILE_NAME), &api_key)?, + None => log::info!("No OPENAI_API_KEY available; skipping Codex auth.json seed"), + } + + prepare_codex_config_toml(&codex_dir.join(CODEX_CONFIG_TOML_FILE_NAME), working_dir)?; + Ok(()) +} + +fn write_codex_agents_override(codex_dir: &Path, system_prompt: &str) -> Result<()> { + fs::create_dir_all(codex_dir).with_context(|| { + format!( + "Failed to create Codex config dir at {}", + codex_dir.display() + ) + })?; + + // Note: this currently works because we are only doing this for cloud agents; if we enable + // this for local runs we'll want to make sure we don't clobber any existing file overrides. + let prompt_path = codex_dir.join(CODEX_AGENTS_OVERRIDE_FILE_NAME); + fs::write(&prompt_path, system_prompt).with_context(|| { + format!( + "Failed to write Codex system prompt to {}", + prompt_path.display() + ) + }) +} + +/// Mirrors the subset of Codex's `AuthDotJson` (codex `login/src/auth/storage.rs`) that we +/// need to seed. Unknown fields (`tokens`, `last_refresh`, `agent_identity`, ...) are +/// preserved via `extra` so we don't clobber an existing login. +#[derive(Default, Deserialize, Serialize, Debug)] +struct CodexAuthDotJson { + #[serde(default, skip_serializing_if = "Option::is_none")] + auth_mode: Option, + #[serde( + rename = "OPENAI_API_KEY", + default, + skip_serializing_if = "Option::is_none" + )] + openai_api_key: Option, + #[serde(flatten)] + extra: Map, +} + +fn prepare_codex_auth(auth_path: &Path, api_key: &str) -> Result<()> { + let mut auth: CodexAuthDotJson = read_json_file_or_default(auth_path)?; + auth.openai_api_key = Some(api_key.to_owned()); + if auth.auth_mode.is_none() { + auth.auth_mode = Some(CODEX_AUTH_MODE_API_KEY.to_owned()); + } + write_codex_auth_json(auth_path, &auth) +} + +/// Write Codex's `auth.json` with restrictive (0o600) permissions, mirroring how +/// codex sets up this file itself. +fn write_codex_auth_json(path: &Path, auth: &CodexAuthDotJson) -> Result<()> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("Failed to create {}", parent.display()))?; + } + let bytes = serde_json::to_vec_pretty(auth).context("Failed to serialize Codex auth.json")?; + + #[cfg(unix)] + { + use std::io::Write as _; + use std::os::unix::fs::OpenOptionsExt; + let mut file = fs::OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .mode(0o600) + .open(path) + .with_context(|| format!("Failed to open {} for writing", path.display()))?; + file.write_all(&bytes) + .with_context(|| format!("Failed to write {}", path.display()))?; + } + #[cfg(not(unix))] + fs::write(path, &bytes).with_context(|| format!("Failed to write {}", path.display()))?; + + Ok(()) +} + +/// Returns the OpenAI API key for Codex auth, preferring the `OPENAI_API_KEY` env +/// var so the seeded `auth.json` matches the credential the launched Codex process +/// will see. [`AgentDriver::new`] skips a managed `OPENAI_API_KEY` secret when the +/// env var is already set, so we mirror that precedence here. +fn resolve_openai_api_key(secrets: &HashMap) -> Option { + if let Ok(value) = std::env::var(OPENAI_API_KEY_ENV) { + let trimmed = value.trim(); + if !trimmed.is_empty() { + return Some(trimmed.to_owned()); + } + } + if let Some(ManagedSecretValue::RawValue { value }) = secrets.get(OPENAI_API_KEY_ENV) { + let trimmed = value.trim(); + if !trimmed.is_empty() { + return Some(trimmed.to_owned()); + } + } + None +} + +/// Edit `~/.codex/config.toml` via `toml_edit` to seed the harness defaults +/// while preserving anything that might already exist there. We handle: +/// - project trust: for a working dir and all of its git repo subdirectories, +/// set the projects to `trusted`. +/// - base URL: set `openai_base_url = ""` so we +/// hit the regional host our API keys require. +fn prepare_codex_config_toml(config_toml_path: &Path, working_dir: &Path) -> Result<()> { + let existing = match fs::read_to_string(config_toml_path) { + Ok(content) => content, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => String::new(), + Err(e) => { + return Err(anyhow::Error::from(e).context(format!( + "Failed to read Codex config.toml at {}", + config_toml_path.display() + ))); + } + }; + let mut doc: toml_edit::DocumentMut = existing.parse().with_context(|| { + format!( + "Failed to parse Codex config.toml at {}", + config_toml_path.display() + ) + })?; + + set_codex_openai_base_url(&mut doc, CODEX_OPENAI_BASE_URL); + + let canonical = working_dir.canonicalize().with_context(|| { + format!( + "Failed to canonicalize Codex working dir at {}", + working_dir.display() + ) + })?; + let project_key = canonical.to_string_lossy().into_owned(); + set_codex_project_trust_level(&mut doc, &project_key, CODEX_TRUST_LEVEL_TRUSTED); + + // Codex's trust check is not recursive (see openai/codex#19426) -- since we + // clone the git repos into workspace/ for cloud agents, we usually have git + // repo children that we also want to trust. + for child_repo in find_child_git_repos(&canonical) { + let key = child_repo.to_string_lossy().into_owned(); + set_codex_project_trust_level(&mut doc, &key, CODEX_TRUST_LEVEL_TRUSTED); + } + + if let Some(parent) = config_toml_path.parent() { + fs::create_dir_all(parent).with_context(|| { + format!("Failed to create Codex config dir at {}", parent.display()) + })?; + } + fs::write(config_toml_path, doc.to_string()).with_context(|| { + format!( + "Failed to write Codex config.toml at {}", + config_toml_path.display() + ) + }) +} + +/// Set the top-level `openai_base_url` key, overwriting any existing value. +fn set_codex_openai_base_url(doc: &mut toml_edit::DocumentMut, base_url: &str) { + doc[CODEX_OPENAI_BASE_URL_KEY] = toml_edit::value(base_url); +} + +/// Return immediate subdirectories of `dir` that contain a `.git`. +fn find_child_git_repos(dir: &Path) -> Vec { + let Ok(entries) = fs::read_dir(dir) else { + return Vec::new(); + }; + entries + .flatten() + .filter_map(|entry| { + let path = entry.path(); + (path.is_dir() && path.join(".git").exists()).then_some(path) + }) + .collect() +} + +/// Insert/update `[projects.""] trust_level = `. +/// +/// Codex itself always writes `projects` as an explicit table, so we don't +/// handle the inline-table form here. +fn set_codex_project_trust_level( + doc: &mut toml_edit::DocumentMut, + project_key: &str, + trust_level: &str, +) { + if !doc.contains_table("projects") { + let mut projects_tbl = toml_edit::Table::new(); + projects_tbl.set_implicit(true); + doc.insert("projects", toml_edit::Item::Table(projects_tbl)); + } + let proj_tbl = doc["projects"] + .as_table_mut() + .expect("projects table inserted above") + .entry(project_key) + .or_insert_with(toml_edit::table) + .as_table_mut() + .expect("project entry is a table"); + proj_tbl.set_implicit(false); + proj_tbl["trust_level"] = toml_edit::value(trust_level); +} diff --git a/app/src/ai/agent_sdk/driver/harness/mod.rs b/app/src/ai/agent_sdk/driver/harness/mod.rs index 84af99652..376340003 100644 --- a/app/src/ai/agent_sdk/driver/harness/mod.rs +++ b/app/src/ai/agent_sdk/driver/harness/mod.rs @@ -36,10 +36,12 @@ use super::{ mod claude_code; pub(crate) mod claude_transcript; +mod codex; mod gemini; mod json_utils; pub(crate) use claude_code::{ClaudeHarness, ClaudeWakeMessage}; use claude_transcript::ClaudeResumeInfo; +use codex::CodexHarness; use gemini::GeminiHarness; /// Harness-agnostic payload describing how to resume an existing conversation. @@ -164,6 +166,7 @@ pub(crate) fn harness_kind(harness: Harness) -> Result Ok(HarnessKind::Oz), Harness::Claude => Ok(HarnessKind::ThirdParty(Box::new(ClaudeHarness))), + Harness::Codex => Ok(HarnessKind::ThirdParty(Box::new(CodexHarness))), Harness::OpenCode => Ok(HarnessKind::Unsupported(Harness::OpenCode)), Harness::Gemini => Ok(HarnessKind::ThirdParty(Box::new(GeminiHarness))), Harness::Unknown => Err(AgentDriverError::InvalidRuntimeState), diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index febea0471..55ed6a6b5 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -1405,6 +1405,7 @@ fn resolve_orchestration_harness_label() -> &'static str { Some(Harness::Claude) => "claude", Some(Harness::OpenCode) => "opencode", Some(Harness::Gemini) => "gemini", + Some(Harness::Codex) => "codex", Some(Harness::Unknown) | None => "unknown", } } diff --git a/app/src/ai/ambient_agents/task.rs b/app/src/ai/ambient_agents/task.rs index 375b1d48a..926add972 100644 --- a/app/src/ai/ambient_agents/task.rs +++ b/app/src/ai/ambient_agents/task.rs @@ -92,6 +92,7 @@ pub(crate) fn harness_from_name(name: &str) -> Harness { "claude" => Harness::Claude, "opencode" => Harness::OpenCode, "gemini" => Harness::Gemini, + "codex" => Harness::Codex, "oz" => Harness::Oz, other => { log::warn!("Unknown harness config name: {other:?}; treating as Unknown"); diff --git a/app/src/ai/blocklist/history_model/conversation_loader.rs b/app/src/ai/blocklist/history_model/conversation_loader.rs index 8a9221453..b029c88c9 100644 --- a/app/src/ai/blocklist/history_model/conversation_loader.rs +++ b/app/src/ai/blocklist/history_model/conversation_loader.rs @@ -135,7 +135,7 @@ pub async fn load_conversation_from_server( } } } - AIAgentHarness::ClaudeCode | AIAgentHarness::Gemini => { + AIAgentHarness::ClaudeCode | AIAgentHarness::Gemini | AIAgentHarness::Codex => { if !FeatureFlag::AgentHarness.is_enabled() { log::warn!("Ignoring non-Oz conversation {conversation_id}: AgentHarness flag is disabled"); return None; diff --git a/app/src/ai/harness_display.rs b/app/src/ai/harness_display.rs index db287b80e..8bc99de10 100644 --- a/app/src/ai/harness_display.rs +++ b/app/src/ai/harness_display.rs @@ -9,7 +9,7 @@ use warp_cli::agent::Harness; use crate::ai::agent::conversation::AIAgentHarness; use crate::ai::blocklist::CLAUDE_ORANGE; -use crate::terminal::cli_agent::GEMINI_BLUE; +use crate::terminal::cli_agent::{GEMINI_BLUE, OPENAI_COLOR}; use crate::ui_components::icons::Icon; /// User-visible display name for a [`Harness`]. @@ -19,6 +19,7 @@ pub fn display_name(harness: Harness) -> &'static str { Harness::Claude => "Claude Code", Harness::OpenCode => "OpenCode", Harness::Gemini => "Gemini CLI", + Harness::Codex => "Codex", Harness::Unknown => "Unknown", } } @@ -30,6 +31,7 @@ pub fn icon_for(harness: Harness) -> Icon { Harness::Claude => Icon::ClaudeLogo, Harness::OpenCode => Icon::OpenCodeLogo, Harness::Gemini => Icon::GeminiLogo, + Harness::Codex => Icon::OpenAILogo, Harness::Unknown => Icon::HelpCircle, } } @@ -42,6 +44,7 @@ pub fn brand_color(harness: Harness) -> Option { Harness::Claude => Some(CLAUDE_ORANGE), Harness::OpenCode => None, Harness::Gemini => Some(GEMINI_BLUE), + Harness::Codex => Some(OPENAI_COLOR), Harness::Unknown => None, } } @@ -54,6 +57,7 @@ impl From for Harness { AIAgentHarness::Oz => Harness::Oz, AIAgentHarness::ClaudeCode => Harness::Claude, AIAgentHarness::Gemini => Harness::Gemini, + AIAgentHarness::Codex => Harness::Codex, AIAgentHarness::Unknown => Harness::Unknown, } } diff --git a/app/src/pane_group/mod.rs b/app/src/pane_group/mod.rs index 8c117c75e..3131bea79 100644 --- a/app/src/pane_group/mod.rs +++ b/app/src/pane_group/mod.rs @@ -213,6 +213,25 @@ const MINIMUM_PANE_SIZE: f32 = 50.; const MINIMUM_PANE_SIZE_UDI: f32 = 190.; const KEYBOARD_RESIZE_DELTA: f32 = 10.; +type AmbientAgentViewModelHandle = + ModelHandle; + +trait AmbientAgentViewModelHandleExt<'a> { + fn into_optional_handle(self) -> Option<&'a AmbientAgentViewModelHandle>; +} + +impl<'a> AmbientAgentViewModelHandleExt<'a> for &'a AmbientAgentViewModelHandle { + fn into_optional_handle(self) -> Option<&'a AmbientAgentViewModelHandle> { + Some(self) + } +} + +impl<'a> AmbientAgentViewModelHandleExt<'a> for Option<&'a AmbientAgentViewModelHandle> { + fn into_optional_handle(self) -> Option<&'a AmbientAgentViewModelHandle> { + self + } +} + fn get_minimum_pane_size(app: &AppContext) -> f32 { use crate::settings::InputSettings; if InputSettings::as_ref(app).is_universal_developer_input_enabled(app) { @@ -3096,6 +3115,7 @@ impl PaneGroup { self.insert_ambient_agent_pane_hidden_for_child_agent(parent_pane_id, ctx); if let Some(new_terminal_view) = self.terminal_view_from_pane_id(new_pane_id, ctx) { + let mut restored = false; new_terminal_view.update(ctx, |terminal_view, ctx| { terminal_view.restore_conversation_after_view_creation( RestoredAIConversation::new(child_conversation), @@ -3108,15 +3128,27 @@ impl PaneGroup { AgentViewEntryOrigin::CloudAgent, ctx, ); - terminal_view + let Some(ambient_agent_view_model) = terminal_view .ambient_agent_view_model() - .update(ctx, |model, ctx| { - model.set_conversation_id(Some(child_id)); - model.enter_viewing_existing_session(task_id, ctx); - }); + .into_optional_handle() + .cloned() + else { + return; + }; + ambient_agent_view_model.update(ctx, |model, ctx| { + model.set_conversation_id(Some(child_id)); + model.enter_viewing_existing_session(task_id, ctx); + }); + restored = true; }); - - self.child_agent_panes.insert(child_id, new_pane_id.into()); + if restored { + self.child_agent_panes.insert(child_id, new_pane_id.into()); + } else { + log::error!( + "Failed to restore remote child agent pane {child_id:?}: missing ambient agent view model" + ); + self.discard_pane(new_pane_id.into(), ctx); + } } else { log::error!("Failed to get terminal view for remote child agent pane {child_id:?}"); self.discard_pane(new_pane_id.into(), ctx); @@ -3832,6 +3864,7 @@ impl PaneGroup { let harness = match cli_conversation.metadata.harness { AIAgentHarness::ClaudeCode => Some(Harness::Claude), AIAgentHarness::Gemini => Some(Harness::Gemini), + AIAgentHarness::Codex => Some(Harness::Codex), AIAgentHarness::Oz => None, AIAgentHarness::Unknown => Some(Harness::Unknown), }; @@ -3844,11 +3877,15 @@ impl PaneGroup { ); // Keep the viewer's AmbientAgentViewModel harness in sync with the loaded run. if let Some(harness) = harness { - view.ambient_agent_view_model() - .clone() - .update(ctx, |model, ctx| { + if let Some(ambient_agent_view_model) = view + .ambient_agent_view_model() + .into_optional_handle() + .cloned() + { + ambient_agent_view_model.update(ctx, |model, ctx| { model.set_harness(harness, ctx); }); + } } // 3p runs have no materialized AIConversation, so enter agent view with a // fresh vehicle conversation and retag the restored snapshot block onto it so diff --git a/app/src/pane_group/pane/local_harness_launch.rs b/app/src/pane_group/pane/local_harness_launch.rs index b65844aa1..abd6f4f9e 100644 --- a/app/src/pane_group/pane/local_harness_launch.rs +++ b/app/src/pane_group/pane/local_harness_launch.rs @@ -59,7 +59,9 @@ pub(super) fn build_local_opencode_child_command(prompt: &str) -> String { fn local_child_task_config(harness: Harness) -> Option { match harness { - Harness::Oz | Harness::OpenCode | Harness::Gemini | Harness::Unknown => None, + Harness::Oz | Harness::OpenCode | Harness::Gemini | Harness::Codex | Harness::Unknown => { + None + } Harness::Claude => Some(AgentConfigSnapshot { harness: Some(HarnessConfig::from_harness_type(harness)), ..Default::default() @@ -123,6 +125,7 @@ pub(super) async fn prepare_local_harness_child_launch( .map_err(|error: AgentDriverError| error.to_string())?; build_local_opencode_child_command(&prompt) } + Harness::Codex => unreachable!("normalize_local_child_harness filters out Codex"), Harness::Gemini => unreachable!("normalize_local_child_harness filters out Gemini"), }; diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index b47993852..5b68666c0 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -207,12 +207,11 @@ pub struct SpawnAgentRequest { pub referenced_attachments: Vec, } -// --- Orchestrations V2 messaging types --- - #[derive(Debug, Clone, serde::Serialize)] pub struct RunFollowupRequest { pub message: String, } +// --- Orchestrations V2 messaging types --- #[derive(Debug, Clone, serde::Serialize)] pub struct SendAgentMessageRequest { @@ -661,6 +660,9 @@ pub(crate) fn build_list_agent_runs_url(limit: i32, filter: &TaskListFilter) -> url } +pub(crate) fn build_run_followup_url(run_id: &AmbientAgentTaskId) -> String { + format!("agent/runs/{run_id}/followups") +} struct ListRunsResponse { runs: Vec, @@ -829,6 +831,11 @@ pub trait AIClient: 'static + Send + Sync { &self, task_id: &AmbientAgentTaskId, ) -> anyhow::Result; + async fn submit_run_followup( + &self, + run_id: &AmbientAgentTaskId, + request: RunFollowupRequest, + ) -> anyhow::Result<(), anyhow::Error>; async fn get_scheduled_agent_history( &self, @@ -910,12 +917,6 @@ pub trait AIClient: 'static + Send + Sync { // --- Orchestrations V2 messaging --- - async fn submit_run_followup( - &self, - run_id: &AmbientAgentTaskId, - request: RunFollowupRequest, - ) -> anyhow::Result<(), anyhow::Error>; - async fn send_agent_message( &self, request: SendAgentMessageRequest, @@ -1531,6 +1532,14 @@ impl AIClient for ServerApi { .await?; Ok(response) } + async fn submit_run_followup( + &self, + run_id: &AmbientAgentTaskId, + request: RunFollowupRequest, + ) -> anyhow::Result<(), anyhow::Error> { + self.post_public_api_unit(&build_run_followup_url(run_id), &request) + .await + } async fn get_scheduled_agent_history( &self, @@ -1923,15 +1932,6 @@ impl AIClient for ServerApi { // --- Orchestrations V2 messaging --- - async fn submit_run_followup( - &self, - run_id: &AmbientAgentTaskId, - request: RunFollowupRequest, - ) -> anyhow::Result<(), anyhow::Error> { - self.post_public_api_unit(&format!("agent/runs/{run_id}/followups"), &request) - .await - } - async fn send_agent_message( &self, request: SendAgentMessageRequest, @@ -2332,9 +2332,12 @@ fn convert_harness(harness: warp_graphql::ai::AgentHarness) -> AIAgentHarness { warp_graphql::ai::AgentHarness::Oz => AIAgentHarness::Oz, warp_graphql::ai::AgentHarness::ClaudeCode => AIAgentHarness::ClaudeCode, warp_graphql::ai::AgentHarness::Gemini => AIAgentHarness::Gemini, - warp_graphql::ai::AgentHarness::Other(value) => { + other => { + if format!("{other:?}") == "Codex" { + return AIAgentHarness::Codex; + } report_error!(anyhow!( - "Invalid AgentHarness '{value}'. Make sure to update client GraphQL types!" + "Invalid AgentHarness '{other:?}'. Make sure to update client GraphQL types!" )); AIAgentHarness::Unknown } diff --git a/app/src/terminal/cli_agent.rs b/app/src/terminal/cli_agent.rs index fd909f3d5..91d890c49 100644 --- a/app/src/terminal/cli_agent.rs +++ b/app/src/terminal/cli_agent.rs @@ -40,7 +40,7 @@ pub(crate) const GEMINI_BLUE: ColorU = ColorU { }; /// OpenAI brand color (dark gray/black) -const OPENAI_COLOR: ColorU = ColorU { +pub(crate) const OPENAI_COLOR: ColorU = ColorU { r: 0, g: 0, b: 0, diff --git a/app/src/terminal/view/ambient_agent/view_impl.rs b/app/src/terminal/view/ambient_agent/view_impl.rs index f8817c198..00b116302 100644 --- a/app/src/terminal/view/ambient_agent/view_impl.rs +++ b/app/src/terminal/view/ambient_agent/view_impl.rs @@ -488,6 +488,7 @@ impl TerminalView { Harness::Claude => matches!(cli_agent, CLIAgent::Claude), Harness::OpenCode => matches!(cli_agent, CLIAgent::OpenCode), Harness::Gemini => matches!(cli_agent, CLIAgent::Gemini), + Harness::Codex => matches!(cli_agent, CLIAgent::Codex), Harness::Unknown => false, } } diff --git a/crates/warp_cli/src/agent.rs b/crates/warp_cli/src/agent.rs index 3478d59d0..155dccfa6 100644 --- a/crates/warp_cli/src/agent.rs +++ b/crates/warp_cli/src/agent.rs @@ -134,6 +134,9 @@ pub enum Harness { /// Delegate to the `gemini` CLI. #[value(name = "gemini")] Gemini, + /// Delegate to the `codex` CLI. + #[value(name = "codex")] + Codex, /// A harness produced by a newer client/server that this client doesn't /// recognize. Surfaced via deserialization fallbacks (e.g. unknown GraphQL /// enum values, unknown `harness_type` strings); never selectable from the @@ -151,7 +154,8 @@ impl Harness { pub fn parse_local_child_harness(value: &str) -> Option { match Self::parse_orchestration_harness(value) { Some(harness @ (Self::Claude | Self::OpenCode)) => Some(harness), - Some(Self::Oz) | Some(Self::Gemini) | Some(Self::Unknown) | None => None, + Some(Self::Oz) | Some(Self::Gemini) | Some(Self::Codex) | Some(Self::Unknown) + | None => None, } } @@ -161,6 +165,7 @@ impl Harness { Self::Claude => "Claude Code", Self::OpenCode => "OpenCode", Self::Gemini => "Gemini CLI", + Self::Codex => "Codex", Self::Unknown => "Unknown", } } @@ -173,6 +178,7 @@ impl fmt::Display for Harness { Harness::Claude => "claude", Harness::OpenCode => "opencode", Harness::Gemini => "gemini", + Harness::Codex => "codex", Harness::Unknown => "unknown", }; f.write_str(name) From 9729b180bcee1d5af45b1060a529a943c18c42be Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Thu, 30 Apr 2026 18:54:08 -0400 Subject: [PATCH 12/13] Address child wake review feedback Co-Authored-By: Oz --- app/src/ai/agent_sdk/ambient.rs | 117 +++--- app/src/ai/agent_sdk/ambient_tests.rs | 55 --- .../agent_sdk/driver/harness/claude_code.rs | 1 - .../harness/claude_code/parent_bridge.rs | 5 + .../driver/harness/claude_code/wake_driver.rs | 106 +++--- .../driver/harness/claude_code_tests.rs | 21 +- app/src/ai/agent_sdk/driver/harness/codex.rs | 3 + app/src/ai/agent_sdk/driver/harness/mod.rs | 2 +- app/src/ai/agent_sdk/mod.rs | 2 +- app/src/ai/blocklist/controller.rs | 360 ++++-------------- .../ai/blocklist/controller/slash_command.rs | 18 +- .../blocklist/orchestration_event_streamer.rs | 28 +- app/src/ai/blocklist/orchestration_events.rs | 64 +--- .../blocklist/orchestration_events_tests.rs | 33 +- app/src/server/server_api/ai.rs | 28 +- app/src/terminal/view/ambient_agent/model.rs | 4 +- crates/graphql/src/api/ai.rs | 1 + crates/warp_graphql_schema/api/schema.graphql | 1 + 18 files changed, 259 insertions(+), 590 deletions(-) diff --git a/app/src/ai/agent_sdk/ambient.rs b/app/src/ai/agent_sdk/ambient.rs index 1cc32c1a7..d0d00821d 100644 --- a/app/src/ai/agent_sdk/ambient.rs +++ b/app/src/ai/agent_sdk/ambient.rs @@ -1,5 +1,4 @@ //! Commands to interact with ambient agents on Warp's platform. -use std::future::Future; use std::io::Write as _; use std::sync::Arc; use std::time::Duration; @@ -40,7 +39,7 @@ use warp_cli::{ }; use warp_core::channel::ChannelState; use warp_core::features::FeatureFlag; -use warpui::r#async::{FutureExt as _, Timer}; +use warpui::r#async::Timer; use warpui::{ platform::TerminationMode, r#async::Spawnable, AppContext, ModelContext, SingletonEntity, }; @@ -55,7 +54,6 @@ use super::common::{parse_ambient_task_id, EnvironmentChoice, ResolveConfigurati const MAX_LINE_WIDTH: usize = 90; const STREAM_RETRY_BACKOFF_STEPS: &[u64] = &[1, 2, 5, 10]; -const SEND_AGENT_MESSAGE_TIMEOUT: Duration = Duration::from_secs(15); /// Singleton model that runs async work for ambient agent CLI commands. struct AmbientAgentRunner; @@ -655,6 +653,7 @@ impl AmbientAgentRunner { sender_run_id: args.sender_run_id, }; let log_context = SendAgentMessageLogContext::new(&request, scoped_task_id.as_ref()); + log_context.log_start(); let send_message = async move { match scoped_task_id { Some(task_id) => { @@ -665,12 +664,18 @@ impl AmbientAgentRunner { None => ai_client.send_agent_message(request).await, } }; - let response = send_agent_message_result_with_timeout( - send_message, - &log_context, - SEND_AGENT_MESSAGE_TIMEOUT, - ) - .await?; + let response = match send_message.await { + Ok(response) => { + log_context.log_success(&response); + response + } + Err(err) => { + let err = err.context(log_context.error_context()); + log_context.log_error(&err); + eprintln!("{err:#}"); + return Err(err); + } + }; print_send_message_response(&response, output_format)?; Ok(()) }; @@ -1028,69 +1033,39 @@ impl SendAgentMessageLogContext { self.sender_run_id, self.task_id, self.target_agent_ids ) } -} -async fn send_agent_message_result_with_timeout( - send_message: F, - log_context: &SendAgentMessageLogContext, - timeout: Duration, -) -> anyhow::Result -where - F: Future>, -{ - log::info!( - "Sending ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={}", - log_context.sender_run_id, - log_context.task_id, - log_context.target_agent_ids, - log_context.subject, - log_context.body_len - ); - - match send_message.with_timeout(timeout).await { - Ok(Ok(response)) => { - log::info!( - "Sent ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} message_ids={:?}", - log_context.sender_run_id, - log_context.task_id, - log_context.target_agent_ids, - log_context.subject, - log_context.body_len, - response.message_ids - ); - Ok(response) - } - Ok(Err(err)) => { - let err = err.context(log_context.error_context()); - log::warn!( - "Failed to send ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} error={err:#}", - log_context.sender_run_id, - log_context.task_id, - log_context.target_agent_ids, - log_context.subject, - log_context.body_len - ); - eprintln!("{err:#}"); - Err(err) - } - Err(_) => { - let err = anyhow!( - "Timed out sending agent message after {timeout:?} (sender_run_id={:?}, task_id={:?}, target_agent_ids={:?})", - log_context.sender_run_id, - log_context.task_id, - log_context.target_agent_ids - ); - log::warn!( - "Timed out sending ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} timeout={timeout:?}", - log_context.sender_run_id, - log_context.task_id, - log_context.target_agent_ids, - log_context.subject, - log_context.body_len - ); - eprintln!("{err:#}"); - Err(err) - } + fn log_start(&self) { + log::info!( + "Sending ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={}", + self.sender_run_id, + self.task_id, + self.target_agent_ids, + self.subject, + self.body_len + ); + } + + fn log_success(&self, response: &SendAgentMessageResponse) { + log::info!( + "Sent ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} message_ids={:?}", + self.sender_run_id, + self.task_id, + self.target_agent_ids, + self.subject, + self.body_len, + response.message_ids + ); + } + + fn log_error(&self, err: &anyhow::Error) { + log::warn!( + "Failed to send ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} error={err:#}", + self.sender_run_id, + self.task_id, + self.target_agent_ids, + self.subject, + self.body_len + ); } } diff --git a/app/src/ai/agent_sdk/ambient_tests.rs b/app/src/ai/agent_sdk/ambient_tests.rs index f0d3488ef..aa298fa89 100644 --- a/app/src/ai/agent_sdk/ambient_tests.rs +++ b/app/src/ai/agent_sdk/ambient_tests.rs @@ -1,8 +1,5 @@ //! Unit tests for ambient agent CLI argument mapping and message helpers. -use anyhow::anyhow; - use chrono::{TimeZone, Utc}; -use futures::executor::block_on; use warp_cli::json_filter::JsonOutput; use warp_cli::task::{ @@ -16,16 +13,6 @@ use crate::server::server_api::ai::{ArtifactType, ExecutionLocation, RunSortBy, const TASK_ID: &str = "00000000-0000-0000-0000-000000000001"; const OTHER_TASK_ID: &str = "00000000-0000-0000-0000-000000000002"; -fn send_log_context() -> SendAgentMessageLogContext { - SendAgentMessageLogContext { - sender_run_id: TASK_ID.to_string(), - task_id: Some(TASK_ID.to_string()), - target_agent_ids: vec!["parent-agent".to_string()], - subject: "status".to_string(), - body_len: 12, - } -} - /// A `ListTasksArgs` whose fields are all at their defaults. fn empty_args() -> ListTasksArgs { ListTasksArgs { @@ -227,45 +214,3 @@ fn task_id_from_oz_run_id_env_rejects_invalid_value() { assert!(err.to_string().contains("Invalid OZ_RUN_ID")); } - -#[test] -fn send_agent_message_result_returns_success_before_timeout() { - let response = block_on(send_agent_message_result_with_timeout( - async { - Ok(SendAgentMessageResponse { - message_ids: vec!["message-1".to_string()], - }) - }, - &send_log_context(), - Duration::from_secs(1), - )) - .expect("send should succeed"); - - assert_eq!(response.message_ids, ["message-1"]); -} - -#[test] -fn send_agent_message_result_surfaces_request_errors() { - let err = block_on(send_agent_message_result_with_timeout( - async { Err::(anyhow!("server rejected request")) }, - &send_log_context(), - Duration::from_secs(1), - )) - .expect_err("send should fail"); - let rendered = format!("{err:#}"); - - assert!(rendered.contains("Failed to send agent message")); - assert!(rendered.contains("server rejected request")); -} - -#[test] -fn send_agent_message_result_times_out() { - let err = block_on(send_agent_message_result_with_timeout( - futures::future::pending::>(), - &send_log_context(), - Duration::from_millis(1), - )) - .expect_err("send should time out"); - - assert!(err.to_string().contains("Timed out sending agent message")); -} diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code.rs b/app/src/ai/agent_sdk/driver/harness/claude_code.rs index 78fb88533..d1e546c1d 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code.rs @@ -50,7 +50,6 @@ use parent_bridge::{ use parent_bridge::{MessageBridge, MessageBridgeCleanupDisposition}; #[cfg(test)] use shell_words::quote as shell_quote; -pub(crate) use wake_driver::ClaudeWakeMessage; #[cfg(test)] use wake_driver::{ClaudeWakeRemoteContext, CLAUDE_WAKE_PROMPT_FILE_NAME}; diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs b/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs index 36eccc75b..cba5002ef 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs @@ -653,6 +653,11 @@ async fn read_parent_bridge_resume_cursor( run_id: &str, state_dir: &Path, ) -> Result { + // The server cursor is the durable cross-client source of truth, but the + // bridge also keeps a local cursor for same-machine recovery. If Warp or + // Claude restarts after the bridge has staged events locally but before the + // server cursor update is visible, the local cursor prevents replaying + // messages already handed to this Claude session. let local_sequence = read_parent_bridge_event_cursor(state_dir)?; let Ok(task_id) = run_id.parse() else { return Ok(local_sequence); diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs b/app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs index 955a0c861..d29540a91 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs @@ -3,6 +3,7 @@ use std::ffi::OsString; use std::path::PathBuf; use std::sync::Arc; +use crate::ai::agent::conversation::{AIConversation, ConversationStatus}; use anyhow::{Context, Result}; use shell_words::quote as shell_quote; use uuid::Uuid; @@ -21,9 +22,7 @@ use super::super::claude_transcript::{ }; use super::super::task_env_vars; use super::parent_bridge::{ - acknowledge_parent_bridge_hook_output, ensure_parent_bridge_state_dir, - parent_bridge_max_context_chars, parent_bridge_root, prepare_parent_bridge_hook_output, - stage_parent_bridge_message, MessageBridgeMessageRecord, + acknowledge_parent_bridge_hook_output, ensure_parent_bridge_state_dir, parent_bridge_root, }; use super::{claude_command, prepare_claude_environment_config, ClaudeHarness}; @@ -31,29 +30,6 @@ const CLAUDE_WAKE_PROMPT: &str = "New lead-agent messages are available. Read the latest lead-agent updates and continue the task accordingly."; pub(super) const CLAUDE_WAKE_PROMPT_FILE_NAME: &str = "wake-turn-prompt.txt"; -#[derive(Debug, Clone)] -pub(crate) struct ClaudeWakeMessage { - pub(crate) sequence: i64, - pub(crate) message_id: String, - pub(crate) sender_run_id: String, - pub(crate) subject: String, - pub(crate) body: String, - pub(crate) occurred_at: String, -} - -impl From for MessageBridgeMessageRecord { - fn from(value: ClaudeWakeMessage) -> Self { - Self { - sequence: value.sequence, - message_id: value.message_id, - sender_run_id: value.sender_run_id, - subject: value.subject, - body: value.body, - occurred_at: value.occurred_at, - } - } -} - #[derive(Debug)] pub(super) struct ClaudeWakeRemoteContext { pub(super) session_id: Uuid, @@ -61,14 +37,30 @@ pub(super) struct ClaudeWakeRemoteContext { pub(super) wake_prompt: String, } +struct ClaudeWakeCandidate { + task_id: AmbientAgentTaskId, + parent_run_id: Option, + working_dir: Option, +} + impl ClaudeHarness { pub(crate) async fn wake_dormant_session( server_api: Arc, - task_id: AmbientAgentTaskId, - parent_run_id: Option, + conversation: AIConversation, + parent_conversation: Option, working_dir: Option, - pending_messages: Vec, ) -> Result> { + let Some(candidate) = + Self::local_wake_candidate(&conversation, parent_conversation.as_ref(), working_dir) + else { + return Ok(None); + }; + let ClaudeWakeCandidate { + task_id, + parent_run_id, + working_dir, + } = candidate; + let task = server_api.get_ambient_agent_task(&task_id).await?; let harness = task .agent_config_snapshot @@ -76,14 +68,12 @@ impl ClaudeHarness { .and_then(|snapshot| snapshot.harness.as_ref()) .map(|config| config.harness_type); log::info!( - "Evaluating dormant Claude wake: task_id={task_id} pending_message_count={} server_task_state={:?} harness={harness:?}", - pending_messages.len(), + "Evaluating dormant Claude wake: task_id={task_id} server_task_state={:?} harness={harness:?}", task.state ); if task.state != AmbientAgentTaskState::Succeeded || harness != Some(Harness::Claude) { log::info!( - "Skipping dormant Claude wake: task_id={task_id} pending_message_count={} server_task_state={:?} harness={harness:?}", - pending_messages.len(), + "Skipping dormant Claude wake: task_id={task_id} server_task_state={:?} harness={harness:?}", task.state ); return Ok(None); @@ -96,7 +86,6 @@ impl ClaudeHarness { parent_run_id, working_dir, remote, - pending_messages, ) .await?; @@ -114,6 +103,45 @@ impl ClaudeHarness { Ok(Some(command)) } + fn local_wake_candidate( + conversation: &AIConversation, + parent_conversation: Option<&AIConversation>, + working_dir: Option, + ) -> Option { + let conversation_id = conversation.id(); + if !matches!(conversation.status(), ConversationStatus::Success) { + log::info!( + "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=not_success status={:?}", + conversation.status() + ); + return None; + } + if !conversation.is_child_agent_conversation() || conversation.is_remote_child() { + log::info!( + "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=not_local_child is_child_agent_conversation={} is_remote_child={}", + conversation.is_child_agent_conversation(), + conversation.is_remote_child() + ); + return None; + } + let Some(task_id) = conversation.task_id() else { + log::info!( + "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=missing_task_id" + ); + return None; + }; + let parent_run_id = conversation + .parent_agent_id() + .map(str::to_owned) + .or_else(|| parent_conversation.and_then(AIConversation::run_id)); + + Some(ClaudeWakeCandidate { + task_id, + parent_run_id, + working_dir, + }) + } + async fn fetch_local_wake_remote_context( task_id: AmbientAgentTaskId, server_api: Arc, @@ -159,7 +187,6 @@ impl ClaudeHarness { parent_run_id: Option, working_dir: Option, mut remote: ClaudeWakeRemoteContext, - pending_messages: Vec, ) -> Result { let working_dir = working_dir.unwrap_or_else(|| remote.envelope.cwd.clone()); prepare_claude_environment_config(&working_dir, &HashMap::new()) @@ -178,15 +205,6 @@ impl ClaudeHarness { ensure_parent_bridge_state_dir(&state_dir)?; let hydrator = MessageHydrator::for_task(server_api, task_id); acknowledge_parent_bridge_hook_output(&hydrator, &state_dir).await?; - for record in pending_messages - .into_iter() - .map(MessageBridgeMessageRecord::from) - { - stage_parent_bridge_message(&state_dir, &record)?; - } - prepare_parent_bridge_hook_output(&hydrator, &state_dir, parent_bridge_max_context_chars()) - .await?; - let prompt_path = state_dir.join(CLAUDE_WAKE_PROMPT_FILE_NAME); std::fs::write(&prompt_path, remote.wake_prompt.as_bytes()) .with_context(|| format!("Failed to write {}", prompt_path.display()))?; diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs b/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs index 243f0ca46..d2639b5cb 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs @@ -646,7 +646,7 @@ fn resolve_suffix_returns_none_for_short_key() { #[test] #[serial_test::serial] -fn prepare_local_wake_command_rehydrates_transcript_and_stages_messages() { +fn prepare_local_wake_command_rehydrates_transcript_without_staging_messages() { let home_dir = TempDir::new().unwrap(); let claude_config_dir = TempDir::new().unwrap(); let bridge_state_root = TempDir::new().unwrap(); @@ -670,14 +670,6 @@ fn prepare_local_wake_command_rehydrates_transcript_and_stages_messages() { }, wake_prompt: "resume prompt\n\nwake prompt".to_string(), }; - let message = ClaudeWakeMessage { - sequence: 42, - message_id: "message-1".to_string(), - sender_run_id: "parent-run-123".to_string(), - subject: "Please pivot".to_string(), - body: "Inspect the failing tests first.".to_string(), - occurred_at: "2026-04-17T15:46:00Z".to_string(), - }; let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440010".parse().unwrap(); let parent_run_id = "parent-run-456".to_string(); @@ -687,14 +679,11 @@ fn prepare_local_wake_command_rehydrates_transcript_and_stages_messages() { Some(parent_run_id.clone()), Some(working_dir.clone()), remote, - vec![message.clone()], )) .unwrap(); let state_dir = bridge_state_root.path().join(session_id.to_string()); let prompt_path = state_dir.join(CLAUDE_WAKE_PROMPT_FILE_NAME); - let surfaced_path = - parent_bridge_surfaced_message_path(&state_dir, message.sequence, &message.message_id); assert!(command.contains("--resume")); assert!(command.starts_with("env ")); @@ -717,13 +706,7 @@ fn prepare_local_wake_command_rehydrates_transcript_and_stages_messages() { fs::read_to_string(&prompt_path).unwrap(), "resume prompt\n\nwake prompt" ); - assert!(surfaced_path.exists()); - assert!(parent_bridge_hook_output_file(&state_dir).exists()); - let surfaced_record: MessageBridgeMessageRecord = - serde_json::from_slice(&fs::read(&surfaced_path).unwrap()).unwrap(); - assert_eq!(surfaced_record.sender_run_id, message.sender_run_id); - assert_eq!(surfaced_record.subject, message.subject); - assert_eq!(surfaced_record.body, message.body); + assert!(!parent_bridge_hook_output_file(&state_dir).exists()); let restored_envelope = read_envelope(session_id, &working_dir, claude_config_dir.path()).unwrap(); diff --git a/app/src/ai/agent_sdk/driver/harness/codex.rs b/app/src/ai/agent_sdk/driver/harness/codex.rs index 28a99a67d..f9c8a251c 100644 --- a/app/src/ai/agent_sdk/driver/harness/codex.rs +++ b/app/src/ai/agent_sdk/driver/harness/codex.rs @@ -310,6 +310,7 @@ fn write_codex_auth_json(path: &Path, auth: &CodexAuthDotJson) -> Result<()> { { use std::io::Write as _; use std::os::unix::fs::OpenOptionsExt; + use std::os::unix::fs::PermissionsExt; let mut file = fs::OpenOptions::new() .write(true) .create(true) @@ -317,6 +318,8 @@ fn write_codex_auth_json(path: &Path, auth: &CodexAuthDotJson) -> Result<()> { .mode(0o600) .open(path) .with_context(|| format!("Failed to open {} for writing", path.display()))?; + file.set_permissions(fs::Permissions::from_mode(0o600)) + .with_context(|| format!("Failed to set permissions on {}", path.display()))?; file.write_all(&bytes) .with_context(|| format!("Failed to write {}", path.display()))?; } diff --git a/app/src/ai/agent_sdk/driver/harness/mod.rs b/app/src/ai/agent_sdk/driver/harness/mod.rs index 376340003..fdbcd2e3a 100644 --- a/app/src/ai/agent_sdk/driver/harness/mod.rs +++ b/app/src/ai/agent_sdk/driver/harness/mod.rs @@ -39,7 +39,7 @@ pub(crate) mod claude_transcript; mod codex; mod gemini; mod json_utils; -pub(crate) use claude_code::{ClaudeHarness, ClaudeWakeMessage}; +pub(crate) use claude_code::ClaudeHarness; use claude_transcript::ClaudeResumeInfo; use codex::CodexHarness; use gemini::GeminiHarness; diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index 55ed6a6b5..397cd304c 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -68,7 +68,7 @@ use crate::ai::skills::{ }; pub(crate) use driver::harness::{ - task_env_vars, validate_cli_installed, ClaudeHarness, ClaudeWakeMessage, ThirdPartyHarness, + task_env_vars, validate_cli_installed, ClaudeHarness, ThirdPartyHarness, }; pub use driver::AgentDriver; use telemetry::CliTelemetryEvent; diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 653f82ebc..91681a58d 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -31,7 +31,7 @@ use crate::ai::agent::{ }; use crate::ai::agent::{DocumentContentAttachmentSource, FileContext}; #[cfg(not(target_family = "wasm"))] -use crate::ai::agent_sdk::{ClaudeHarness, ClaudeWakeMessage}; +use crate::ai::agent_sdk::ClaudeHarness; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::document::ai_document_model::{ AIDocumentId, AIDocumentModel, AIDocumentUserEditStatus, @@ -83,8 +83,6 @@ use warp_multi_agent_api::{message, Task, ToolType}; use warpui::r#async::{SpawnedFutureHandle, Timer}; use super::orchestration_events::{OrchestrationEventService, OrchestrationEventServiceEvent}; -#[cfg(not(target_family = "wasm"))] -use super::orchestration_events::{PendingEvent, PendingEventDetail}; use warpui::{AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; #[derive(Debug, Clone)] @@ -1524,138 +1522,98 @@ impl BlocklistAIController { } #[cfg(not(target_family = "wasm"))] - fn local_claude_wake_candidate( - &self, - conversation_id: AIConversationId, - ctx: &ModelContext, - ) -> Option<(AmbientAgentTaskId, Option, Option)> { - let Some(conversation) = - BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) - else { - log::info!( - "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=conversation_missing" - ); - return None; - }; - if !conversation.is_child_agent_conversation() || conversation.is_remote_child() { - log::info!( - "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=not_local_child is_child_agent_conversation={} is_remote_child={}", - conversation.is_child_agent_conversation(), - conversation.is_remote_child() - ); - return None; - } - let Some(task_id) = conversation.task_id() else { - log::info!( - "Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=missing_task_id" - ); - return None; - }; - - Some(( - task_id, - self.active_session - .as_ref(ctx) - .current_working_directory() - .cloned() - .map(PathBuf::from), - conversation - .parent_agent_id() - .map(str::to_owned) - .or_else(|| { - conversation - .parent_conversation_id() - .and_then(|parent_conversation_id| { - BlocklistAIHistoryModel::as_ref(ctx) - .conversation(&parent_conversation_id) - }) - .and_then(|parent_conversation| parent_conversation.run_id()) - }), - )) - } - - #[cfg(not(target_family = "wasm"))] - fn drain_remote_child_message_events( + fn maybe_prepare_local_claude_wake( &mut self, conversation_id: AIConversationId, ctx: &mut ModelContext, ) -> bool { - let Some(conversation) = - BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) - else { - log::info!( - "Skipping remote child wake candidate: conversation_id={conversation_id:?} reason=conversation_missing" - ); - return false; - }; - let is_remote_child = conversation.is_remote_child(); - let task_id = conversation.task_id(); - if !is_remote_child { - return false; + if self + .pending_local_claude_wakes + .contains_key(&conversation_id) + { + log::info!("Dormant Claude wake already pending: conversation_id={conversation_id:?}"); + return true; } - let pending_message_events = OrchestrationEventService::handle(ctx) - .update(ctx, |svc, _| { - svc.peek_pending_message_events(conversation_id) - }); - if pending_message_events.is_empty() { + let has_pending_events = OrchestrationEventService::handle(ctx) + .update(ctx, |svc, _| svc.has_pending_events(conversation_id)); + if !has_pending_events { return false; } - let pending_message_event_ids = pending_message_events - .iter() - .map(|event| event.event_id.clone()) - .collect_vec(); - let removed_events = OrchestrationEventService::handle(ctx).update(ctx, |svc, _| { - svc.take_pending_events_by_id(conversation_id, &pending_message_event_ids) - }); - let pending_message_count = pending_message_event_ids.len(); - if removed_events.len() != pending_message_event_ids.len() { + let history_model = BlocklistAIHistoryModel::as_ref(ctx); + let Some(conversation) = history_model.conversation(&conversation_id).cloned() else { log::info!( - "Remote child pending message set changed while draining local wake events: conversation_id={conversation_id:?} task_id={:?} expected_message_event_count={} actual_removed_event_count={}", - task_id, - pending_message_event_ids.len(), - removed_events.len() + "Skipping dormant Claude wake preparation: conversation_id={conversation_id:?} reason=conversation_missing" ); - } - log::info!( - "Drained remote child message events locally because server owns wakeups: conversation_id={conversation_id:?} task_id={:?} message_count={pending_message_count}", - task_id - ); - true - } - - #[cfg(target_family = "wasm")] - fn drain_remote_child_message_events( - &mut self, - _conversation_id: AIConversationId, - _ctx: &mut ModelContext, - ) -> bool { - false - } - - #[cfg(not(target_family = "wasm"))] - fn pending_event_to_claude_wake_message(event: &PendingEvent) -> Option { - let PendingEventDetail::Message { - sequence, - message_id, - subject, - message_body, - occurred_at, - .. - } = &event.detail - else { - return None; + return false; }; + let parent_conversation = conversation + .parent_conversation_id() + .and_then(|parent_conversation_id| history_model.conversation(&parent_conversation_id)) + .cloned(); + let working_dir = self + .active_session + .as_ref(ctx) + .current_working_directory() + .cloned() + .map(PathBuf::from); + let task_id = conversation.task_id(); - Some(ClaudeWakeMessage { - sequence: *sequence, - message_id: message_id.clone(), - sender_run_id: event.source_agent_id.clone(), - subject: subject.clone(), - body: message_body.clone(), - occurred_at: occurred_at.clone(), - }) + let server_api = ServerApiProvider::as_ref(ctx).get(); + let handle = ctx.spawn( + async move { + log::info!( + "Preparing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id:?}" + ); + ClaudeHarness::wake_dormant_session( + server_api.clone(), + conversation, + parent_conversation, + working_dir, + ) + .await + }, + move |me, result, ctx| { + me.pending_local_claude_wakes.remove(&conversation_id); + match result { + Ok(Some(command)) => { + log::info!( + "Executing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id:?}" + ); + BlocklistAIHistoryModel::handle(ctx).update( + ctx, + |history_model, ctx| { + history_model.update_conversation_status( + me.terminal_view_id, + conversation_id, + ConversationStatus::InProgress, + ctx, + ); + }, + ); + ctx.emit(BlocklistAIControllerEvent::ExecuteLocalHarnessCommand { + command, + }); + } + Ok(None) => { + log::info!( + "Falling back to generic pending-event injection after dormant Claude wake eligibility check: conversation_id={conversation_id:?} task_id={task_id:?}" + ); + me.inject_pending_events_for_request(conversation_id, ctx); + } + Err(err) => { + log::warn!( + "Failed to prepare dormant Claude wake command for {conversation_id:?} task_id={task_id:?}: {err:#}" + ); + me.schedule_pending_events_ready_retry(conversation_id, ctx); + } + } + }, + ); + self.pending_local_claude_wakes + .insert(conversation_id, handle); + true } fn schedule_pending_events_ready_retry( @@ -1713,160 +1671,6 @@ impl BlocklistAIController { } } - #[cfg(not(target_family = "wasm"))] - fn maybe_prepare_local_claude_wake( - &mut self, - conversation_id: AIConversationId, - ctx: &mut ModelContext, - ) -> bool { - if self - .pending_local_claude_wakes - .contains_key(&conversation_id) - { - log::info!("Dormant Claude wake already pending: conversation_id={conversation_id:?}"); - return true; - } - - let Some((task_id, working_dir, parent_run_id)) = - self.local_claude_wake_candidate(conversation_id, ctx) - else { - return false; - }; - - let pending_message_events = OrchestrationEventService::handle(ctx) - .update(ctx, |svc, _| { - svc.peek_pending_message_events(conversation_id) - }); - if pending_message_events.is_empty() { - log::info!( - "Skipping dormant Claude wake preparation: conversation_id={conversation_id:?} reason=no_pending_message_events" - ); - return false; - } - - let pending_message_event_ids = pending_message_events - .iter() - .map(|event| event.event_id.clone()) - .collect_vec(); - let removed_events = OrchestrationEventService::handle(ctx).update(ctx, |svc, _| { - svc.take_pending_events_by_id(conversation_id, &pending_message_event_ids) - }); - if removed_events.len() != pending_message_event_ids.len() { - log::info!( - "Aborting dormant Claude wake because the pending message set changed: conversation_id={conversation_id:?} task_id={task_id} expected_message_event_count={} actual_removed_event_count={}", - pending_message_event_ids.len(), - removed_events.len() - ); - if !removed_events.is_empty() { - OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { - svc.prepend_pending_events(conversation_id, removed_events, ctx); - }); - } - return true; - } - - let wake_messages = removed_events - .iter() - .filter_map(Self::pending_event_to_claude_wake_message) - .collect_vec(); - if wake_messages.len() != removed_events.len() { - log::info!( - "Aborting dormant Claude wake because some removed events were not message events: conversation_id={conversation_id:?} task_id={task_id} removed_event_count={} wake_message_count={}", - removed_events.len(), - wake_messages.len() - ); - OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { - svc.prepend_pending_events(conversation_id, removed_events, ctx); - }); - return true; - } - log::info!( - "Prepared dormant Claude wake messages: conversation_id={conversation_id:?} task_id={task_id} message_count={}", - wake_messages.len() - ); - for wake_message in &wake_messages { - log::info!( - "Dormant Claude wake message: conversation_id={conversation_id:?} task_id={task_id} message_id={} sequence={} subject={:?} body_len={}", - wake_message.message_id, - wake_message.sequence, - wake_message.subject, - wake_message.body.chars().count() - ); - } - - let removed_events_for_retry = removed_events.clone(); - let server_api = ServerApiProvider::as_ref(ctx).get(); - let wake_message_count = wake_messages.len(); - let handle = ctx.spawn( - async move { - log::info!( - "Preparing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id} pending_message_count={wake_message_count}" - ); - ClaudeHarness::wake_dormant_session( - server_api.clone(), - task_id, - parent_run_id, - working_dir, - wake_messages, - ) - .await - }, - move |me, result, ctx| { - me.pending_local_claude_wakes.remove(&conversation_id); - match result { - Ok(Some(command)) => { - log::info!( - "Executing dormant Claude wake command: conversation_id={conversation_id:?} task_id={task_id}" - ); - BlocklistAIHistoryModel::handle(ctx).update( - ctx, - |history_model, ctx| { - history_model.update_conversation_status( - me.terminal_view_id, - conversation_id, - ConversationStatus::InProgress, - ctx, - ); - }, - ); - ctx.emit(BlocklistAIControllerEvent::ExecuteLocalHarnessCommand { - command, - }); - } - Ok(None) => { - log::info!( - "Falling back to generic pending-event injection after dormant Claude wake eligibility check: conversation_id={conversation_id:?} task_id={task_id}" - ); - OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { - svc.prepend_pending_events( - conversation_id, - removed_events_for_retry.clone(), - ctx, - ); - }); - me.inject_pending_events_for_request(conversation_id, ctx); - } - Err(err) => { - log::warn!( - "Failed to prepare dormant Claude wake command for {conversation_id:?} task_id={task_id}: {err:#}" - ); - OrchestrationEventService::handle(ctx).update(ctx, |svc, ctx| { - svc.prepend_pending_events( - conversation_id, - removed_events_for_retry.clone(), - ctx, - ); - }); - me.schedule_pending_events_ready_retry(conversation_id, ctx); - } - } - }, - ); - self.pending_local_claude_wakes - .insert(conversation_id, handle); - true - } - /// Handles the EventsReady signal. Checks readiness, drains /// pending events from the service, and injects them into the conversation. fn handle_pending_events_ready( @@ -1877,10 +1681,6 @@ impl BlocklistAIController { if !self.conversation_ready_for_pending_events(conversation_id, ctx) { return; } - if self.drain_remote_child_message_events(conversation_id, ctx) { - self.handle_pending_events_ready(conversation_id, ctx); - return; - } if self.maybe_prepare_local_claude_wake(conversation_id, ctx) { return; diff --git a/app/src/ai/blocklist/controller/slash_command.rs b/app/src/ai/blocklist/controller/slash_command.rs index 15f6ddc8e..60aa0348f 100644 --- a/app/src/ai/blocklist/controller/slash_command.rs +++ b/app/src/ai/blocklist/controller/slash_command.rs @@ -6,8 +6,8 @@ use warpui::{AppContext, ModelContext, SingletonEntity}; use crate::{ ai::{ agent::{ - conversation::AIConversationId, AIAgentContext, AIAgentInput, CloneRepositoryURL, - EntrypointType, RequestMetadata, + conversation::AIConversationId, AIAgentContext, AIAgentInput, CancellationReason, + CloneRepositoryURL, EntrypointType, RequestMetadata, }, blocklist::agent_view::AgentViewEntryOrigin, }, @@ -90,6 +90,8 @@ impl SlashCommandRequest { if inputs.is_empty() { return; } + let active_conversation_id = BlocklistAIHistoryModel::as_ref(ctx) + .active_conversation_id(controller.terminal_view_id); // If no existing conversation, create a new one. // When AgentView is enabled, enter agent view which creates the conversation @@ -114,6 +116,18 @@ impl SlashCommandRequest { return; }; + let cancellation_reason = CancellationReason::FollowUpSubmitted { + is_for_same_conversation: active_conversation_id + .is_some_and(|id| id == conversation_id), + }; + if let Some(active_conversation_id) = active_conversation_id { + controller.cancel_conversation_progress( + active_conversation_id, + cancellation_reason, + ctx, + ); + } + let Some(conversation) = BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) else { diff --git a/app/src/ai/blocklist/orchestration_event_streamer.rs b/app/src/ai/blocklist/orchestration_event_streamer.rs index af238bf55..96b475f1d 100644 --- a/app/src/ai/blocklist/orchestration_event_streamer.rs +++ b/app/src/ai/blocklist/orchestration_event_streamer.rs @@ -1,7 +1,7 @@ use super::history_model::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel}; use super::orchestration_events::{OrchestrationEventService, PendingEvent, PendingEventDetail}; use crate::ai::agent::{ - conversation::{AIConversationId, ConversationStatus}, + conversation::{AIAgentHarness, AIConversationId, ConversationStatus}, ReceivedMessageInput, }; use crate::ai::agent_events::{ @@ -671,11 +671,35 @@ impl OrchestrationEventStreamer { conversation_id: AIConversationId, ctx: &mut ModelContext, ) { + if self.should_skip_sse_for_dormant_local_claude_child(conversation_id, ctx) { + log::info!( + "Skipping generic SSE delivery for dormant local Claude child {conversation_id:?}; parent bridge will deliver wake events" + ); + return; + } if !self.sse_connections.contains_key(&conversation_id) { self.start_sse_connection(conversation_id, ctx); } } + fn should_skip_sse_for_dormant_local_claude_child( + &self, + conversation_id: AIConversationId, + ctx: &ModelContext, + ) -> bool { + let Some(conversation) = + BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id) + else { + return false; + }; + conversation.is_child_agent_conversation() + && !conversation.is_remote_child() + && matches!(conversation.status(), ConversationStatus::Success) + && conversation + .server_metadata() + .is_some_and(|metadata| metadata.harness == AIAgentHarness::ClaudeCode) + } + /// Opens a long-lived SSE connection for `conversation_id`. Events are /// sent through an mpsc channel and drained by a periodic timer. fn start_sse_connection( @@ -840,7 +864,7 @@ impl OrchestrationEventStreamer { self.sse_connections.remove(&conversation_id); if self.watched_run_ids.contains_key(&conversation_id) { - self.start_sse_connection(conversation_id, ctx); + self.start_event_delivery(conversation_id, ctx); } } } diff --git a/app/src/ai/blocklist/orchestration_events.rs b/app/src/ai/blocklist/orchestration_events.rs index 785f6fa64..6bb8c0803 100644 --- a/app/src/ai/blocklist/orchestration_events.rs +++ b/app/src/ai/blocklist/orchestration_events.rs @@ -56,13 +56,13 @@ impl LifecycleEventDetailStage { #[derive(Debug, Clone)] pub enum PendingEventDetail { Message { - #[cfg_attr(target_family = "wasm", allow(dead_code))] + #[cfg_attr(not(test), allow(dead_code))] sequence: i64, message_id: String, addresses: Vec, subject: String, message_body: String, - #[cfg_attr(target_family = "wasm", allow(dead_code))] + #[cfg_attr(not(test), allow(dead_code))] occurred_at: String, }, Lifecycle { @@ -892,66 +892,10 @@ impl OrchestrationEventService { ctx.emit(OrchestrationEventServiceEvent::EventsReady { conversation_id }); } - pub fn peek_pending_message_events( - &self, - conversation_id: AIConversationId, - ) -> Vec { + pub fn has_pending_events(&self, conversation_id: AIConversationId) -> bool { self.pending_events .get(&conversation_id) - .into_iter() - .flatten() - .filter(|event| matches!(event.detail, PendingEventDetail::Message { .. })) - .cloned() - .collect() - } - - pub fn take_pending_events_by_id( - &mut self, - conversation_id: AIConversationId, - event_ids: &[String], - ) -> Vec { - if event_ids.is_empty() { - return Vec::new(); - } - - let event_ids = event_ids.iter().map(String::as_str).collect::>(); - let Some(queue) = self.pending_events.get_mut(&conversation_id) else { - return Vec::new(); - }; - - let mut removed = Vec::new(); - let mut retained = Vec::with_capacity(queue.len()); - for event in std::mem::take(queue) { - if event_ids.contains(event.event_id.as_str()) { - removed.push(event); - } else { - retained.push(event); - } - } - *queue = retained; - - if queue.is_empty() { - self.pending_events.remove(&conversation_id); - } - - removed - } - - #[cfg_attr(target_family = "wasm", allow(dead_code))] - pub fn prepend_pending_events( - &mut self, - conversation_id: AIConversationId, - mut events: Vec, - ctx: &mut ModelContext, - ) { - if events.is_empty() { - return; - } - - let queue = self.pending_events.entry(conversation_id).or_default(); - events.append(queue); - *queue = events; - ctx.emit(OrchestrationEventServiceEvent::EventsReady { conversation_id }); + .is_some_and(|events| !events.is_empty()) } /// Drain and return all pending events for a conversation. diff --git a/app/src/ai/blocklist/orchestration_events_tests.rs b/app/src/ai/blocklist/orchestration_events_tests.rs index 18a87c037..30055c0db 100644 --- a/app/src/ai/blocklist/orchestration_events_tests.rs +++ b/app/src/ai/blocklist/orchestration_events_tests.rs @@ -366,9 +366,10 @@ fn test_lifecycle_event_type_from_proto_includes_cancelled_and_blocked() { } #[test] -fn test_pending_message_helpers_peek_and_take_only_selected_messages() { +fn test_has_pending_events_tracks_any_event_kind() { let conversation_id = crate::ai::agent::conversation::AIConversationId::new(); let mut service = OrchestrationEventService::new_without_subscriptions(); + assert!(!service.has_pending_events(conversation_id)); service.pending_events.insert( conversation_id, vec![ @@ -388,33 +389,9 @@ fn test_pending_message_helpers_peek_and_take_only_selected_messages() { message_pending_event("message-event-2"), ], ); - - let peeked = service.peek_pending_message_events(conversation_id); - assert_eq!(peeked.len(), 2); - assert_eq!(peeked[0].event_id, "message-event-1"); - assert_eq!(peeked[1].event_id, "message-event-2"); - - let removed = service.take_pending_events_by_id( - conversation_id, - &["message-event-2".to_string(), "message-event-1".to_string()], - ); - assert_eq!(removed.len(), 2); - assert_eq!(removed[0].event_id, "message-event-1"); - assert_eq!(removed[1].event_id, "message-event-2"); - - let remaining = service - .pending_events - .get(&conversation_id) - .expect("lifecycle events should remain queued"); - assert_eq!(remaining.len(), 2); - assert!(matches!( - remaining[0].detail, - PendingEventDetail::Lifecycle { .. } - )); - assert!(matches!( - remaining[1].detail, - PendingEventDetail::Lifecycle { .. } - )); + assert!(service.has_pending_events(conversation_id)); + service.pending_events.remove(&conversation_id); + assert!(!service.has_pending_events(conversation_id)); } #[test] diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index 5b68666c0..8a8e173cc 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -207,10 +207,6 @@ pub struct SpawnAgentRequest { pub referenced_attachments: Vec, } -#[derive(Debug, Clone, serde::Serialize)] -pub struct RunFollowupRequest { - pub message: String, -} // --- Orchestrations V2 messaging types --- #[derive(Debug, Clone, serde::Serialize)] @@ -660,9 +656,6 @@ pub(crate) fn build_list_agent_runs_url(limit: i32, filter: &TaskListFilter) -> url } -pub(crate) fn build_run_followup_url(run_id: &AmbientAgentTaskId) -> String { - format!("agent/runs/{run_id}/followups") -} struct ListRunsResponse { runs: Vec, @@ -831,11 +824,6 @@ pub trait AIClient: 'static + Send + Sync { &self, task_id: &AmbientAgentTaskId, ) -> anyhow::Result; - async fn submit_run_followup( - &self, - run_id: &AmbientAgentTaskId, - request: RunFollowupRequest, - ) -> anyhow::Result<(), anyhow::Error>; async fn get_scheduled_agent_history( &self, @@ -1532,14 +1520,6 @@ impl AIClient for ServerApi { .await?; Ok(response) } - async fn submit_run_followup( - &self, - run_id: &AmbientAgentTaskId, - request: RunFollowupRequest, - ) -> anyhow::Result<(), anyhow::Error> { - self.post_public_api_unit(&build_run_followup_url(run_id), &request) - .await - } async fn get_scheduled_agent_history( &self, @@ -2332,12 +2312,10 @@ fn convert_harness(harness: warp_graphql::ai::AgentHarness) -> AIAgentHarness { warp_graphql::ai::AgentHarness::Oz => AIAgentHarness::Oz, warp_graphql::ai::AgentHarness::ClaudeCode => AIAgentHarness::ClaudeCode, warp_graphql::ai::AgentHarness::Gemini => AIAgentHarness::Gemini, - other => { - if format!("{other:?}") == "Codex" { - return AIAgentHarness::Codex; - } + warp_graphql::ai::AgentHarness::Codex => AIAgentHarness::Codex, + warp_graphql::ai::AgentHarness::Other(value) => { report_error!(anyhow!( - "Invalid AgentHarness '{other:?}'. Make sure to update client GraphQL types!" + "Invalid AgentHarness '{value}'. Make sure to update client GraphQL types!" )); AIAgentHarness::Unknown } diff --git a/app/src/terminal/view/ambient_agent/model.rs b/app/src/terminal/view/ambient_agent/model.rs index df0829e3e..4e6d20203 100644 --- a/app/src/terminal/view/ambient_agent/model.rs +++ b/app/src/terminal/view/ambient_agent/model.rs @@ -420,7 +420,9 @@ impl AmbientAgentViewModel { // Store the task ID for later use self.task_id = Some(task_id); - self.status = Status::AgentRunning; + if matches!(self.status, Status::NotAmbientAgent) { + self.status = Status::AgentRunning; + } // Fetch the task so we can set the correct environment (instead of defaulting to the most // recently-used one) and the correct harness (so non-oz viewers know to use the diff --git a/crates/graphql/src/api/ai.rs b/crates/graphql/src/api/ai.rs index de88d3608..edf06bc71 100644 --- a/crates/graphql/src/api/ai.rs +++ b/crates/graphql/src/api/ai.rs @@ -123,6 +123,7 @@ pub enum AgentHarness { Oz, ClaudeCode, Gemini, + Codex, #[cynic(fallback)] Other(String), } diff --git a/crates/warp_graphql_schema/api/schema.graphql b/crates/warp_graphql_schema/api/schema.graphql index 8cbdf50c9..0d38ffe52 100644 --- a/crates/warp_graphql_schema/api/schema.graphql +++ b/crates/warp_graphql_schema/api/schema.graphql @@ -215,6 +215,7 @@ enum AdminEnablementSetting { """The harness that produced an agent conversation.""" enum AgentHarness { CLAUDE_CODE + CODEX GEMINI OZ } From d837a14cb221b4e258eeb5269a15842abe7552c8 Mon Sep 17 00:00:00 2001 From: Katarina Jankov Date: Fri, 1 May 2026 17:09:17 -0400 Subject: [PATCH 13/13] fix lint --- app/src/ai/agent_sdk/driver/harness/mod.rs | 1 - app/src/ai/blocklist/controller.rs | 1 - .../blocklist/orchestration_event_streamer_tests.rs | 12 ++---------- app/src/pane_group/mod_tests.rs | 5 ++++- app/src/pane_group/pane/local_harness_launch.rs | 1 - 5 files changed, 6 insertions(+), 14 deletions(-) diff --git a/app/src/ai/agent_sdk/driver/harness/mod.rs b/app/src/ai/agent_sdk/driver/harness/mod.rs index 3caf678c9..fdbcd2e3a 100644 --- a/app/src/ai/agent_sdk/driver/harness/mod.rs +++ b/app/src/ai/agent_sdk/driver/harness/mod.rs @@ -169,7 +169,6 @@ pub(crate) fn harness_kind(harness: Harness) -> Result Ok(HarnessKind::ThirdParty(Box::new(CodexHarness))), Harness::OpenCode => Ok(HarnessKind::Unsupported(Harness::OpenCode)), Harness::Gemini => Ok(HarnessKind::ThirdParty(Box::new(GeminiHarness))), - Harness::Codex => Ok(HarnessKind::ThirdParty(Box::new(CodexHarness))), Harness::Unknown => Err(AgentDriverError::InvalidRuntimeState), } } diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index a94640527..f37c6c429 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -53,7 +53,6 @@ use crate::global_resource_handles::GlobalResourceHandlesProvider; use crate::network::NetworkStatus; use crate::notebooks::editor::model::FileLinkResolutionContext; use crate::persistence::ModelEvent; -use crate::search::slash_command_menu::static_commands::commands; use crate::server::server_api::{AIApiError, ServerApiProvider}; use crate::terminal::model::block::{ formatted_terminal_contents_for_input, BlockId, CURSOR_MARKER, diff --git a/app/src/ai/blocklist/orchestration_event_streamer_tests.rs b/app/src/ai/blocklist/orchestration_event_streamer_tests.rs index 7ed144be7..452f1e18d 100644 --- a/app/src/ai/blocklist/orchestration_event_streamer_tests.rs +++ b/app/src/ai/blocklist/orchestration_event_streamer_tests.rs @@ -226,16 +226,8 @@ fn restored_conversations_skip_v2_streaming_when_orchestration_v2_disabled() { streamer.read(&app, |me, _| { assert!( - !me.event_cursor.contains_key(&conversation_id), - "V2-disabled restore must not initialize event cursors" - ); - assert!( - !me.watched_run_ids.contains_key(&conversation_id), - "V2-disabled restore must not register watched run ids" - ); - assert!( - me.sse_connections.is_empty(), - "V2-disabled restore must not open SSE connections" + me.streams.is_empty(), + "V2-disabled restore must not initialize stream state" ); }); }); diff --git a/app/src/pane_group/mod_tests.rs b/app/src/pane_group/mod_tests.rs index d3a90c2e4..b3d0dbefe 100644 --- a/app/src/pane_group/mod_tests.rs +++ b/app/src/pane_group/mod_tests.rs @@ -287,7 +287,10 @@ fn ambient_child_session_state( .expect("child pane should have a terminal view"); let terminal_view_ref = terminal_view.as_ref(ctx); let active_conversation_id = terminal_view_ref.active_conversation_id(ctx); - let ambient_model = terminal_view_ref.ambient_agent_view_model().as_ref(ctx); + let ambient_model = terminal_view_ref + .ambient_agent_view_model() + .expect("child pane should have an ambient agent model") + .as_ref(ctx); ( ambient_model.task_id(), diff --git a/app/src/pane_group/pane/local_harness_launch.rs b/app/src/pane_group/pane/local_harness_launch.rs index 0412af1e2..abd6f4f9e 100644 --- a/app/src/pane_group/pane/local_harness_launch.rs +++ b/app/src/pane_group/pane/local_harness_launch.rs @@ -89,7 +89,6 @@ pub(super) async fn prepare_local_harness_child_launch( let command = match harness { Harness::Oz => unreachable!("normalize_local_child_harness filters out Oz"), Harness::Unknown => unreachable!("normalize_local_child_harness filters out Unknown"), - Harness::Codex => unreachable!("normalize_local_child_harness filters out Codex"), Harness::Claude => { let working_dir = startup_directory .or_else(|| std::env::current_dir().ok())