fix: restore shared cloud sessions as ambient panes

Persist shared ambient agent session panes as ambient agent snapshots even when they were opened as plain shared-session viewers without an ambient view model. This preserves the cloud task ID for app session restoration instead of restoring as a local terminal.

Co-Authored-By: Oz <oz-agent@warp.dev>

Author: 2 people

app/src/pane_group/mod_tests.rs

@@ -74,13 +74,15 @@ use crate::{
 #[cfg(feature = "local_fs")]
 use repo_metadata::RepoMetadataModel;
 use repo_metadata::{repositories::DetectedRepositories, watcher::DirectoryWatcher};
+use session_sharing_protocol::sharer::SessionSourceType;
 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::app_state::{LeafContents, PaneNodeSnapshot};
 use crate::terminal::resizable_data::ResizableData;
 use ai::{
     index::full_source_code_embedding::manager::CodebaseIndexManager,
@@ -589,6 +591,42 @@ fn test_restored_remote_hidden_child_pane_enters_existing_ambient_session() {
     });
 }
 
+#[test]
+fn test_shared_ambient_session_without_ambient_view_model_snapshots_as_ambient_agent() {
+    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 pane_id = get_newly_created_pane_id(panes, &[]);
+            let task_id = new_ambient_agent_task_id();
+            let terminal_view = panes
+                .terminal_view_from_pane_id(pane_id, ctx)
+                .expect("initial pane should have a terminal view");
+
+            terminal_view.update(ctx, |view, _ctx| {
+                assert!(
+                    view.ambient_agent_view_model().is_none(),
+                    "plain shared-session viewers do not have cloud-mode view models"
+                );
+                let mut model = view.model.lock();
+                model.set_shared_session_status(SharedSessionStatus::reader());
+                model.set_shared_session_source_type(SessionSourceType::AmbientAgent {
+                    task_id: Some(task_id.to_string()),
+                });
+            });
+
+            let PaneNodeSnapshot::Leaf(leaf) = panes.snapshot(ctx) else {
+                panic!("expected single-pane snapshot");
+            };
+            let LeafContents::AmbientAgent(snapshot) = leaf.contents else {
+                panic!("expected ambient agent snapshot");
+            };
+            assert_eq!(snapshot.task_id, Some(task_id));
+        });
+    });
+}
+
 #[test]
 fn test_active_session_id_reset_on_last_pane_close() {
     App::test((), |mut app| async move {

app/src/pane_group/pane/terminal_pane.rs

@@ -427,7 +427,16 @@ impl PaneContent for TerminalPane {
         // Capture the current input_config from the AI input model
         let current_input_config = view.input_config(app.as_ref());
 
-        if view.model.lock().shared_session_status().is_viewer() {
+        let (is_shared_session_viewer, is_conversation_transcript_viewer, ambient_agent_task_id) = {
+            let model = view.model.lock();
+            (
+                model.shared_session_status().is_viewer(),
+                model.is_conversation_transcript_viewer(),
+                model.ambient_agent_task_id(),
+            )
+        };
+
+        if is_shared_session_viewer {
             // We save and restore ambient agent sessions
             // (restoring the shared session if it's still open and the conversation transcript otherwise).
             if let Some(ambient_model) = view.ambient_agent_view_model() {
@@ -440,6 +449,14 @@ impl PaneContent for TerminalPane {
                 });
             }
 
+            let task_id = ambient_agent_task_id;
+            if task_id.is_some() {
+                return LeafContents::AmbientAgent(AmbientAgentPaneSnapshot {
+                    uuid: self.uuid.clone(),
+                    task_id,
+                });
+            }
+
             LeafContents::Terminal(TerminalPaneSnapshot {
                 uuid: self.uuid.clone(),
                 cwd: None,
@@ -452,10 +469,10 @@ impl PaneContent for TerminalPane {
                 conversation_ids_to_restore: vec![],
                 active_conversation_id: None,
             })
-        } else if view.model.lock().is_conversation_transcript_viewer() {
+        } else if is_conversation_transcript_viewer {
             // Conversation transcript viewers (opened from the conversation list)
             // can be restored via the ambient agent task if one exists.
-            let task_id = view.model.lock().ambient_agent_task_id();
+            let task_id = ambient_agent_task_id;
             if task_id.is_some() {
                 LeafContents::AmbientAgent(AmbientAgentPaneSnapshot {
                     uuid: self.uuid.clone(),