[REMOTE-2208] Add --idle-on-fail so failed cloud sessions stay joinable

app/src/ai/agent_sdk/driver.rs

@@ -254,12 +254,15 @@ impl<T: Send + 'static> IdleTimeoutSender<T> {
         }
     }
 
-    /// End the run with `value`, deferring by `idle_on_complete` when set so the
-    /// driver stays alive long enough to accept a follow-up. Use for "graceful"
-    /// terminal statuses (Success / Blocked / Cancelled). Falls back to
-    /// immediate completion when `idle_on_complete` is `None`.
-    fn complete_with_optional_idle(&self, idle_on_complete: Option<Duration>, value: T) {
-        if let Some(idle_timeout) = idle_on_complete {
+    /// End the run with `value`, deferring by `idle_timeout` when set so the driver
+    /// stays alive long enough to accept a follow-up. Falls back to immediate
+    /// completion when `idle_timeout` is `None`.
+    ///
+    /// Callers pick the window that matches the outcome: `idle_on_complete` for
+    /// "graceful" terminal statuses (Success / Blocked / Cancelled), `idle_on_fail`
+    /// for terminal failures.
+    fn complete_with_optional_idle(&self, idle_timeout: Option<Duration>, value: T) {
+        if let Some(idle_timeout) = idle_timeout {
             self.end_run_after(idle_timeout, value);
         } else {
             self.end_run_now(value);
@@ -291,6 +294,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<Duration>,
+    /// How long to keep the session alive after the agent run ends in a terminal
+    /// failure, if at all. Sourced from `--idle-on-fail`. Deliberately separate from
+    /// `idle_on_complete`: the success window is not a fallback for the failure window.
+    pub idle_on_fail: Option<Duration>,
     /// 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).
@@ -357,6 +364,11 @@ pub struct AgentDriver {
     // and exit after this period of inactivity.
     idle_on_complete: Option<Duration>,
 
+    // Optional idle timeout after a terminal failure. If set, the process stays alive —
+    // still hosting its shared session — so the failed run can be inspected and followed
+    // up on, then exits after this period of inactivity.
+    idle_on_fail: Option<Duration>,
+
     // The conversation ID to continue (if provided).
     restored_conversation_id: Option<AIConversationId>,
 
@@ -640,6 +652,7 @@ impl AgentDriver {
             parent_run_id,
             should_share,
             idle_on_complete,
+            idle_on_fail,
             secrets,
             resume,
             cloud_providers,
@@ -665,9 +678,9 @@ impl AgentDriver {
         };
 
         safe_info!(
-            safe: ("Initializing agent driver: share={should_share}, idle_on_complete={idle_on_complete:?}"),
+            safe: ("Initializing agent driver: share={should_share}, idle_on_complete={idle_on_complete:?}, idle_on_fail={idle_on_fail:?}"),
             full: (
-                "Initializing agent driver: share={should_share}, idle_on_complete={idle_on_complete:?}, working_dir={}",
+                "Initializing agent driver: share={should_share}, idle_on_complete={idle_on_complete:?}, idle_on_fail={idle_on_fail:?}, working_dir={}",
                 working_dir.display()
             )
         );
@@ -767,6 +780,7 @@ impl AgentDriver {
             task_id,
             harness: None,
             idle_on_complete,
+            idle_on_fail,
             restored_conversation_id,
             resume_payload,
             cloud_providers,
@@ -812,6 +826,7 @@ impl AgentDriver {
             task_id: None,
             harness: None,
             idle_on_complete: None,
+            idle_on_fail: None,
             restored_conversation_id: None,
             resume_payload: None,
             cloud_providers: Vec::new(),
@@ -3392,13 +3407,29 @@ impl AgentDriver {
                                 );
                             }
                             // Errors here are terminal: in-flight recoveries surface as
-                            // TransientError (handled above).
+                            // TransientError (handled above). The failure is already
+                            // reported to the server by LocalAgentTaskSyncModel; the only
+                            // question here is whether this process sticks around.
+                            //
+                            // With `--idle-on-fail` we keep it alive so it keeps hosting
+                            // the shared session: a session is only joinable while its
+                            // host process lives, so exiting here is what makes a failed
+                            // run's session unreachable. Note this deliberately does NOT
+                            // fall back to `idle_on_complete` — the success window must
+                            // not silently extend failed runs.
                             SDKConversationOutputStatus::Error { .. } => {
-                                log::info!(
-                                    "Ambient agent idle lifecycle: event=run_completion_immediate task_id={:?} terminal_view_id={terminal_id:?} outcome=error",
-                                    me.task_id
-                                );
-                                run_exit.end_run_now(output_status);
+                                if let Some(idle_timeout) = me.idle_on_fail {
+                                    log::info!(
+                                        "Ambient agent idle lifecycle: event=idle_timeout_scheduled task_id={:?} terminal_view_id={terminal_id:?} timeout={idle_timeout:?} outcome=error",
+                                        me.task_id
+                                    );
+                                } else {
+                                    log::info!(
+                                        "Ambient agent idle lifecycle: event=run_completion_immediate task_id={:?} terminal_view_id={terminal_id:?} outcome=error",
+                                        me.task_id
+                                    );
+                                }
+                                run_exit.complete_with_optional_idle(me.idle_on_fail, output_status);
                             }
                         }
                     }
@@ -3597,11 +3628,10 @@ impl AgentDriver {
                         return;
                     }
 
-                    // Drive idle-on-complete timer for the harness exit signal.
+                    // Drive the idle timer for the harness exit signal, picking the
+                    // window that matches the outcome.
                     match status {
-                        CLIAgentSessionStatus::Success
-                        | CLIAgentSessionStatus::Failed { .. }
-                        | CLIAgentSessionStatus::Blocked { .. } => {
+                        CLIAgentSessionStatus::Success | CLIAgentSessionStatus::Blocked { .. } => {
                             if let Some(idle_timeout) = me.idle_on_complete {
                                 log::info!(
                                     "Ambient agent CLI lifecycle: event=idle_timeout_scheduled task_id={:?} terminal_view_id={terminal_view_id:?} timeout={idle_timeout:?}",
@@ -3615,6 +3645,24 @@ impl AgentDriver {
                             }
                             harness_exit.complete_with_optional_idle(me.idle_on_complete, ());
                         }
+                        CLIAgentSessionStatus::Failed { .. } => {
+                            let idle_on_fail = Self::third_party_failure_idle_window(
+                                me.idle_on_fail,
+                                me.idle_on_complete,
+                            );
+                            if let Some(idle_timeout) = idle_on_fail {
+                                log::info!(
+                                    "Ambient agent CLI lifecycle: event=idle_timeout_scheduled task_id={:?} terminal_view_id={terminal_view_id:?} timeout={idle_timeout:?} outcome=error",
+                                    me.task_id
+                                );
+                            } else {
+                                log::info!(
+                                    "Ambient agent CLI lifecycle: event=run_completion_immediate task_id={:?} terminal_view_id={terminal_view_id:?} outcome=error",
+                                    me.task_id
+                                );
+                            }
+                            harness_exit.complete_with_optional_idle(idle_on_fail, ());
+                        }
                         CLIAgentSessionStatus::InProgress => {
                             log::info!(
                                 "Ambient agent CLI lifecycle: event=idle_timeout_cancel_requested task_id={:?} terminal_view_id={terminal_view_id:?} trigger=session_in_progress",
@@ -3660,6 +3708,20 @@ impl AgentDriver {
             });
     }
 
+    /// The idle window to apply when a third-party harness session ends in failure.
+    ///
+    /// Unlike the Oz path — where a terminal error exits immediately unless
+    /// `--idle-on-fail` asks otherwise — a failed third-party harness session already
+    /// idled for `idle_on_complete` before `--idle-on-fail` existed. Falling back to the
+    /// success window preserves that behavior instead of regressing those runs to an
+    /// immediate exit, while still letting the failure-specific window take precedence.
+    fn third_party_failure_idle_window(
+        idle_on_fail: Option<Duration>,
+        idle_on_complete: Option<Duration>,
+    ) -> Option<Duration> {
+        idle_on_fail.or(idle_on_complete)
+    }
+
     /// Removes the task mapping registered for CLI agent session status updates.
     fn unregister_cli_agent_task_sync(&self, ctx: &mut ModelContext<Self>) {
         let terminal_view_id = self.terminal_driver.as_ref(ctx).terminal_view().id();

app/src/ai/agent_sdk/driver_tests.rs

@@ -577,6 +577,37 @@ fn idle_timeout_sender_complete_with_optional_idle_some_then_cancel_invalidates_
     assert_eq!(rx.try_recv().unwrap(), None);
 }
 
+#[test]
+fn third_party_failure_idle_window_prefers_failure_window() {
+    // `--idle-on-fail` wins when both windows are configured.
+    assert_eq!(
+        AgentDriver::third_party_failure_idle_window(
+            Some(Duration::from_secs(15 * 60)),
+            Some(Duration::from_secs(45 * 60)),
+        ),
+        Some(Duration::from_secs(15 * 60))
+    );
+}
+
+#[test]
+fn third_party_failure_idle_window_falls_back_to_success_window() {
+    // Back-compat: a failed third-party harness session idled for `idle_on_complete`
+    // before `--idle-on-fail` existed, so an unset failure window must not regress
+    // those runs into an immediate exit.
+    assert_eq!(
+        AgentDriver::third_party_failure_idle_window(None, Some(Duration::from_secs(45 * 60))),
+        Some(Duration::from_secs(45 * 60))
+    );
+}
+
+#[test]
+fn third_party_failure_idle_window_is_none_when_neither_is_set() {
+    assert_eq!(
+        AgentDriver::third_party_failure_idle_window(None, None),
+        None
+    );
+}
+
 #[test]
 fn task_env_vars_include_parent_run_id_when_present() {
     let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440000".parse().unwrap();

app/src/ai/agent_sdk/mod.rs

@@ -1053,6 +1053,7 @@ impl AgentDriverRunner {
                     parent_run_id: None,
                     should_share,
                     idle_on_complete: args.idle_on_complete.map(|d| d.into()),
+                    idle_on_fail: args.idle_on_fail.map(|d| d.into()),
                     secrets: Default::default(),
                     resume: None,
                     cloud_providers: Vec::new(),

crates/warp_cli/src/agent.rs

@@ -402,6 +402,26 @@ pub struct RunAgentArgs {
     )]
     pub idle_on_complete: Option<humantime::Duration>,
 
+    /// Keep the agent's session open after the conversation ends in a terminal failure.
+    ///
+    /// The failure counterpart to `--idle-on-complete`: the run is still reported as
+    /// FAILED/ERROR immediately, but the process keeps hosting its shared session so a
+    /// human can inspect the failed environment and send a follow-up. A shared session is
+    /// only reachable while the process hosting it is alive, so this flag is what makes
+    /// server-side post-failure session retention actually joinable.
+    ///
+    /// Independent of `--idle-on-complete`: neither implies the other for Oz runs.
+    ///
+    /// You can optionally provide a duration (e.g. `--idle-on-fail 15m`).
+    #[arg(
+        long = "idle-on-fail",
+        value_name = "DURATION",
+        num_args = 0..=1,
+        default_missing_value = "15m",
+        hide = true
+    )]
+    pub idle_on_fail: Option<humantime::Duration>,
+
     #[command(flatten)]
     pub snapshot: SnapshotArgs,
     /// Identifier for the task that spawned this agent, used to report progress.

crates/warp_cli/src/lib_tests.rs

@@ -691,6 +691,115 @@ fn agent_run_accepts_idle_on_complete_duration() {
     );
 }
 
+#[test]
+fn agent_run_accepts_idle_on_fail_flag() {
+    let args = Args::try_parse_from([
+        "warp",
+        "agent",
+        "run",
+        "--prompt",
+        "hello",
+        "--idle-on-fail",
+    ])
+    .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_eq!(
+        run_args.idle_on_fail,
+        Some(humantime::Duration::from(std::time::Duration::from_secs(
+            15 * 60
+        )))
+    );
+}
+
+#[test]
+fn agent_run_accepts_idle_on_fail_duration() {
+    let args = Args::try_parse_from([
+        "warp",
+        "agent",
+        "run",
+        "--prompt",
+        "hello",
+        "--idle-on-fail",
+        "30m",
+    ])
+    .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_eq!(
+        run_args.idle_on_fail,
+        Some(humantime::Duration::from(std::time::Duration::from_secs(
+            30 * 60
+        )))
+    );
+}
+
+#[test]
+fn agent_run_idle_on_fail_is_opt_in() {
+    // Post-failure session retention must stay off unless explicitly requested, so an
+    // omitted flag leaves the failure window unset (immediate exit, today's behavior).
+    let args = Args::try_parse_from(["warp", "agent", "run", "--prompt", "hello"]).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_eq!(run_args.idle_on_fail, None);
+}
+
+#[test]
+fn agent_run_accepts_idle_on_fail_alongside_idle_on_complete() {
+    // The success and failure windows are independent: the worker emits both for a run
+    // whose environment opted into post-failure retention.
+    let args = Args::try_parse_from([
+        "warp",
+        "agent",
+        "run",
+        "--prompt",
+        "hello",
+        "--idle-on-complete",
+        "10m",
+        "--idle-on-fail",
+        "15m",
+    ])
+    .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_eq!(
+        run_args.idle_on_complete,
+        Some(humantime::Duration::from(std::time::Duration::from_secs(
+            10 * 60
+        )))
+    );
+    assert_eq!(
+        run_args.idle_on_fail,
+        Some(humantime::Duration::from(std::time::Duration::from_secs(
+            15 * 60
+        )))
+    );
+}
+
 #[test]
 fn agent_run_accepts_skip_initial_turn_with_task_id_and_idle_on_complete() {
     let args = Args::try_parse_from([