fix: restore autofocus when creating worktrees (#3787)

* fix: restore autofocus when creating worktrees

refs #3766

* fix: let the client focus newly created worktrees

refs #3766

2 files changed, 151 insertions, 2 deletions

diff --git a/src/client/shell/tests/agents_worktrees_notifications.rs b/src/client/shell/tests/agents_worktrees_notifications.rs
index 10f97118..2291318e 100644
--- a/src/client/shell/tests/agents_worktrees_notifications.rs
+++ b/src/client/shell/tests/agents_worktrees_notifications.rs
@@ -978,7 +978,7 @@ fn worktree_create_previews_the_endpoint_owned_checkout_path() {
             if params.workspace_id.as_deref() == Some("ws_1")
                 && params.branch.as_deref() == Some("feature/client-shell")
                 && params.path.is_none()
-                && params.focus
+                && !params.focus
     ));
 }
 
diff --git a/src/client/shell/tests/endpoint_requests.rs b/src/client/shell/tests/endpoint_requests.rs
index f05b8462..b8805596 100644
--- a/src/client/shell/tests/endpoint_requests.rs
+++ b/src/client/shell/tests/endpoint_requests.rs
@@ -34,6 +34,11 @@ fn pending_worktree() -> (ClientShellState, Vec<ClientShellAction>) {
     let mut state = ClientShellState::new(ClientShellConfig::from_config(&Config::default()));
     state.set_snapshot(Box::new(snapshot()));
     state.set_pane_surface(surface());
+    submit_worktree(state)
+}
+
+fn submit_worktree(mut state: ClientShellState) -> (ClientShellState, Vec<ClientShellAction>) {
+    let boot_id = state.snapshot.as_ref().unwrap().boot_id.clone();
     let mut outcome = ClientShellInput::default();
     state.record_binding(
         crate::input::KeybindMatch::Action(crate::input::KeybindAction::NewWorktree),
@@ -42,7 +47,7 @@ fn pending_worktree() -> (ClientShellState, Vec<ClientShellAction>) {
     let [ClientShellAction::Endpoint { request, .. }] = &outcome.actions[..] else {
         panic!("expected worktree preparation");
     };
-    state.handle_endpoint_result("boot-1", &request.id, Ok(worktree_list_result(None)));
+    state.handle_endpoint_result(&boot_id, &request.id, Ok(worktree_list_result(None)));
     state.handle_input_bytes(b"feature/reconnect");
     let outcome = state.handle_input_bytes(b"\r");
     assert!(matches!(
@@ -59,6 +64,150 @@ fn request_id(actions: &[ClientShellAction]) -> &str {
     &request.id
 }
 
+fn worktree_created_result() -> crate::api::schema::ResponseResult {
+    serde_json::from_value(serde_json::json!({
+        "type": "worktree_created",
+        "workspace": {
+            "workspace_id": "ws_2", "number": 2, "label": "worktree",
+            "focused": false, "pane_count": 1, "tab_count": 1,
+            "active_tab_id": "tab_2", "agent_status": "idle"
+        },
+        "tab": {
+            "tab_id": "tab_2", "workspace_id": "ws_2", "number": 1,
+            "label": "worktree", "focused": false, "pane_count": 1,
+            "agent_status": "idle"
+        },
+        "root_pane": {
+            "pane_id": "pane_2", "terminal_id": "term_2", "workspace_id": "ws_2",
+            "tab_id": "tab_2", "focused": false, "agent_status": "idle", "revision": 0
+        },
+        "worktree": {
+            "path": "/repo-feature", "branch": "feature/reconnect", "is_bare": false,
+            "is_detached": false, "is_prunable": false, "is_linked_worktree": true,
+            "open_workspace_id": "ws_2", "label": "worktree"
+        }
+    }))
+    .unwrap()
+}
+
+fn add_remote(state: &mut ClientShellState) -> ClientEndpointId {
+    let profile = crate::client::endpoint::SavedSshEndpoint {
+        id: crate::client::endpoint::ProfileId::parse("0123456789abcdef0123456789abcdef").unwrap(),
+        label: "Build".into(),
+        target: "dev@build.example".into(),
+        session: "agents".into(),
+        enabled: true,
+    };
+    let remote = ClientEndpointId::Ssh(profile.id.clone());
+    state.set_endpoint_catalog(&[profile]);
+    state.set_endpoint_status(&remote, ClientEndpointStatus::Online);
+    let mut projection = snapshot();
+    projection.boot_id = "remote-boot".into();
+    state.set_endpoint_snapshot(&remote, Box::new(projection));
+    remote
+}
+
+#[test]
+fn worktree_create_leaves_server_focus_unchanged() {
+    let (_, actions) = pending_worktree();
+    let [ClientShellAction::Endpoint { request, .. }] = &actions[..] else {
+        panic!("expected worktree creation");
+    };
+    assert!(matches!(
+        &request.method,
+        crate::api::schema::Method::WorktreeCreate(params) if !params.focus
+    ));
+}
+
+#[test]
+fn worktree_create_success_focuses_returned_tab_on_its_endpoint_after_snapshot_update() {
+    for use_remote in [false, true] {
+        let mut state = ClientShellState::new(ClientShellConfig::from_config(&Config::default()));
+        state.set_snapshot(Box::new(snapshot()));
+        if use_remote {
+            let remote = add_remote(&mut state);
+            assert!(state.activate_endpoint_projection(&remote));
+        }
+        let (mut state, actions) = submit_worktree(state);
+        let endpoint_id = state.active_endpoint_id.clone();
+        let mut updated = state.snapshot.clone().unwrap();
+        let boot_id = updated.boot_id.clone();
+        updated.revision += 1;
+        updated.workspaces[0].label = "updated while creating".into();
+        state.set_endpoint_snapshot(&endpoint_id, updated);
+        let (repaint, focus) = state.handle_endpoint_result(
+            &boot_id,
+            request_id(&actions),
+            Ok(worktree_created_result()),
+        );
+        assert!(repaint);
+        assert!(state.overlay.is_none());
+        let [ClientShellAction::Endpoint {
+            endpoint_id: target,
+            boot_id: target_boot,
+            request,
+        }] = &focus[..]
+        else {
+            panic!("creation should request focus through normal client navigation");
+        };
+        assert_eq!(target, &endpoint_id);
+        assert_eq!(target_boot, &boot_id);
+        assert!(matches!(
+            &request.method,
+            crate::api::schema::Method::TabFocus(target) if target.tab_id == "tab_2"
+        ));
+        assert!(state
+            .handle_endpoint_result(
+                &boot_id,
+                request_id(&actions),
+                Ok(worktree_created_result())
+            )
+            .1
+            .is_empty());
+    }
+}
+
+#[test]
+fn worktree_create_late_success_does_not_focus_after_switching_away_and_back() {
+    let (mut state, actions) = pending_worktree();
+    let remote = add_remote(&mut state);
+    assert!(state.activate_endpoint_projection(&remote));
+    assert!(state.activate_endpoint_projection(&ClientEndpointId::Local));
+    assert!(state
+        .handle_endpoint_result(
+            "boot-1",
+            request_id(&actions),
+            Ok(worktree_created_result())
+        )
+        .1
+        .is_empty());
+}
+
+#[test]
+fn worktree_create_cancelled_or_failed_request_never_focuses() {
+    for cancel in [false, true] {
+        let (mut state, actions) = pending_worktree();
+        let id = request_id(&actions);
+        if cancel {
+            assert!(state.cancel_endpoint_request(id));
+        } else {
+            let (_, follow_up) = state.handle_endpoint_result(
+                "boot-1",
+                id,
+                Err(ClientShellEndpointError {
+                    code: Some("worktree_failed".into()),
+                    message: "creation failed".into(),
+                }),
+            );
+            assert!(follow_up.is_empty());
+        }
+        assert!(state
+            .handle_endpoint_result("boot-1", id, Ok(worktree_created_result()))
+            .1
+            .is_empty());
+    }
+}
+
 #[test]
 fn cancelling_popup_request_unblocks_input_and_ignores_late_success() {
     let (mut state, actions) = pending_popup();