Preserve selected agent when starting thread in new worktree (#51628)

Richard Feldman created

When the user selected a custom agent (e.g. Codex) and chose "Start
Thread in New Worktree", the new workspace's panel ignored the selection
and defaulted to Zed Agent (NativeAgent).

The fix captures `self.selected_agent()` in
`handle_worktree_creation_requested` before the async task, passes it
through to `setup_new_workspace`, and forwards it to `external_thread`
instead of `None`.

Closes AI-84

(No release notes because it's still feature-flagged.)

Release Notes:

- N/A

Change summary

crates/agent_ui/src/agent_panel.rs | 160 +++++++++++++++++++++++++++++++
1 file changed, 159 insertions(+), 1 deletion(-)

Detailed changes

crates/agent_ui/src/agent_panel.rs 🔗

@@ -2909,6 +2909,8 @@ impl AgentPanel {
             .window_handle()
             .downcast::<workspace::MultiWorkspace>();
 
+        let selected_agent = self.selected_agent();
+
         let task = cx.spawn_in(window, async move |this, cx| {
             // Await the branch listings we kicked off earlier.
             let mut existing_branches = Vec::new();
@@ -3008,6 +3010,7 @@ impl AgentPanel {
                 non_git_paths,
                 has_non_git,
                 content,
+                selected_agent,
                 cx,
             )
             .await
@@ -3041,6 +3044,7 @@ impl AgentPanel {
         non_git_paths: Vec<PathBuf>,
         has_non_git: bool,
         content: Vec<acp::ContentBlock>,
+        selected_agent: Option<Agent>,
         cx: &mut AsyncWindowContext,
     ) -> Result<()> {
         let init: Option<
@@ -3126,7 +3130,7 @@ impl AgentPanel {
                 if let Some(panel) = workspace.panel::<AgentPanel>(cx) {
                     panel.update(cx, |panel, cx| {
                         panel.external_thread(
-                            None,
+                            selected_agent,
                             None,
                             None,
                             None,
@@ -6299,6 +6303,160 @@ mod tests {
         );
     }
 
+    #[gpui::test]
+    async fn test_worktree_creation_preserves_selected_agent(cx: &mut TestAppContext) {
+        init_test(cx);
+
+        let app_state = cx.update(|cx| {
+            cx.update_flags(true, vec!["agent-v2".to_string()]);
+            agent::ThreadStore::init_global(cx);
+            language_model::LanguageModelRegistry::test(cx);
+
+            let app_state = workspace::AppState::test(cx);
+            workspace::init(app_state.clone(), cx);
+            app_state
+        });
+
+        let fs = app_state.fs.as_fake();
+        fs.insert_tree(
+            "/project",
+            json!({
+                ".git": {},
+                "src": {
+                    "main.rs": "fn main() {}"
+                }
+            }),
+        )
+        .await;
+        fs.set_branch_name(Path::new("/project/.git"), Some("main"));
+
+        let project = Project::test(app_state.fs.clone(), [Path::new("/project")], cx).await;
+
+        let multi_workspace =
+            cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
+
+        let workspace = multi_workspace
+            .read_with(cx, |multi_workspace, _cx| {
+                multi_workspace.workspace().clone()
+            })
+            .unwrap();
+
+        workspace.update(cx, |workspace, _cx| {
+            workspace.set_random_database_id();
+        });
+
+        // Register a callback so new workspaces also get an AgentPanel.
+        cx.update(|cx| {
+            cx.observe_new(
+                |workspace: &mut Workspace,
+                 window: Option<&mut Window>,
+                 cx: &mut Context<Workspace>| {
+                    if let Some(window) = window {
+                        let project = workspace.project().clone();
+                        let text_thread_store =
+                            cx.new(|cx| TextThreadStore::fake(project.clone(), cx));
+                        let panel = cx.new(|cx| {
+                            AgentPanel::new(workspace, text_thread_store, None, window, cx)
+                        });
+                        workspace.add_panel(panel, window, cx);
+                    }
+                },
+            )
+            .detach();
+        });
+
+        let cx = &mut VisualTestContext::from_window(multi_workspace.into(), cx);
+
+        // Wait for the project to discover the git repository.
+        cx.run_until_parked();
+
+        let panel = workspace.update_in(cx, |workspace, window, cx| {
+            let text_thread_store = cx.new(|cx| TextThreadStore::fake(project.clone(), cx));
+            let panel =
+                cx.new(|cx| AgentPanel::new(workspace, text_thread_store, None, window, cx));
+            workspace.add_panel(panel.clone(), window, cx);
+            panel
+        });
+
+        cx.run_until_parked();
+
+        // Open a thread (needed so there's an active thread view).
+        panel.update_in(cx, |panel, window, cx| {
+            panel.open_external_thread_with_server(
+                Rc::new(StubAgentServer::default_response()),
+                window,
+                cx,
+            );
+        });
+
+        cx.run_until_parked();
+
+        // Set the selected agent to Codex (a custom agent) and start_thread_in
+        // to NewWorktree. We do this AFTER opening the thread because
+        // open_external_thread_with_server overrides selected_agent_type.
+        panel.update(cx, |panel, cx| {
+            panel.selected_agent_type = AgentType::Custom {
+                name: CODEX_NAME.into(),
+            };
+            panel.set_start_thread_in(&StartThreadIn::NewWorktree, cx);
+        });
+
+        // Verify the panel has the Codex agent selected.
+        panel.read_with(cx, |panel, _cx| {
+            assert_eq!(
+                panel.selected_agent_type,
+                AgentType::Custom {
+                    name: CODEX_NAME.into()
+                },
+            );
+        });
+
+        // Directly call handle_worktree_creation_requested, which is what
+        // handle_first_send_requested does when start_thread_in == NewWorktree.
+        let content = vec![acp::ContentBlock::Text(acp::TextContent::new(
+            "Hello from test",
+        ))];
+        panel.update_in(cx, |panel, window, cx| {
+            panel.handle_worktree_creation_requested(content, window, cx);
+        });
+
+        // Let the async worktree creation + workspace setup complete.
+        cx.run_until_parked();
+
+        // Find the new workspace's AgentPanel and verify it used the Codex agent.
+        let found_codex = multi_workspace
+            .read_with(cx, |multi_workspace, cx| {
+                // There should be more than one workspace now (the original + the new worktree).
+                assert!(
+                    multi_workspace.workspaces().len() > 1,
+                    "expected a new workspace to have been created, found {}",
+                    multi_workspace.workspaces().len(),
+                );
+
+                // Check the newest workspace's panel for the correct agent.
+                let new_workspace = multi_workspace
+                    .workspaces()
+                    .iter()
+                    .find(|ws| ws.entity_id() != workspace.entity_id())
+                    .expect("should find the new workspace");
+                let new_panel = new_workspace
+                    .read(cx)
+                    .panel::<AgentPanel>(cx)
+                    .expect("new workspace should have an AgentPanel");
+
+                new_panel.read(cx).selected_agent_type.clone()
+            })
+            .unwrap();
+
+        assert_eq!(
+            found_codex,
+            AgentType::Custom {
+                name: CODEX_NAME.into()
+            },
+            "the new worktree workspace should use the same agent (Codex) that was selected in the original panel",
+        );
+    }
+
     #[test]
     fn test_deserialize_legacy_serialized_panel() {
         let json = serde_json::json!({