fix: distinguish agent completion from startup and session changes (#4457)

refs #3338
refs #3990
refs #3916

4 files changed, 346 insertions, 82 deletions

diff --git a/src/server/headless/tests/mod.rs b/src/server/headless/tests/mod.rs
index cf6d389b..6accdb09 100644
--- a/src/server/headless/tests/mod.rs
+++ b/src/server/headless/tests/mod.rs
@@ -9,12 +9,36 @@ mod surface_delta_tests;
 #[path = "surface_interest.rs"]
 mod surface_interest_tests;
 
-fn client_shell_snapshot(message: ServerMessage) -> Box<crate::protocol::ClientShellSnapshot> {
-    let ServerMessage::EndpointControl { kind, data } = message else {
-        panic!("expected client shell snapshot");
+fn client_shell_projection(
+    receiver: &std::sync::mpsc::Receiver<Vec<u8>>,
+) -> (
+    Box<protocol::ClientShellSnapshot>,
+    protocol::endpoint::EndpointAgentCompletions,
+) {
+    let read_control = |expected| {
+        let ServerMessage::EndpointControl { kind, data } = read_server_message(
+            receiver
+                .recv_timeout(Duration::from_secs(1))
+                .expect("endpoint projection"),
+        ) else {
+            panic!("expected endpoint control {expected}");
+        };
+        assert_eq!(kind, expected);
+        data
     };
-    assert_eq!(kind, crate::protocol::endpoint::ENDPOINT_SNAPSHOT_KIND);
-    Box::new(serde_json::from_str(&data).expect("decode client shell snapshot"))
+    let completions: protocol::endpoint::EndpointAgentCompletions =
+        serde_json::from_str(&read_control(protocol::endpoint::AGENT_COMPLETIONS_KIND)).unwrap();
+    let snapshot: Box<protocol::ClientShellSnapshot> =
+        serde_json::from_str(&read_control(protocol::endpoint::ENDPOINT_SNAPSHOT_KIND)).unwrap();
+    assert_eq!(completions.boot_id, snapshot.boot_id);
+    assert_eq!(completions.revision, snapshot.revision);
+    (snapshot, completions)
+}
+
+fn client_shell_snapshot(
+    receiver: &std::sync::mpsc::Receiver<Vec<u8>>,
+) -> Box<protocol::ClientShellSnapshot> {
+    client_shell_projection(receiver).0
 }
 
 fn client_agent_view_projection(
@@ -715,6 +739,64 @@ async fn client_shell_attach_seeds_workspace() {
 }
 
 #[tokio::test]
+async fn completion_guard_endpoint_pairs_runtime_completions_with_snapshots() {
+    let mut server = test_headless_server();
+    let workspace = crate::workspace::Workspace::test_new("endpoint");
+    let pane_id = workspace.tabs[0].root_pane;
+    server.app.state.workspaces = vec![workspace];
+    server.app.state.ensure_test_terminals();
+    server.app.state.active = Some(0);
+    let (writer, control_rx, _render_rx) = test_client_writer();
+    server.handle_server_event(ServerEvent::ClientShellConnected {
+        client_id: 78,
+        surface_cols: 80,
+        surface_rows: 24,
+        cell_width_px: 0,
+        cell_height_px: 0,
+        pixel_mouse: false,
+        direct_graphics: false,
+        endpoint_keybindings: false,
+        mouse_capture: false,
+        surface_active: false,
+        surface_reuse: false,
+        surface_delta: false,
+        writer,
+    });
+    let (_, initial) = client_shell_projection(&control_rx);
+    assert!(initial.completions.is_empty());
+    for first_state in [
+        crate::detect::AgentState::Working,
+        crate::detect::AgentState::Unknown,
+    ] {
+        for state in [first_state, crate::detect::AgentState::Idle] {
+            server.app.state.handle_app_event(AppEvent::StateChanged {
+                pane_id,
+                agent: Some(crate::detect::Agent::Pi),
+                state,
+                visible_blocker: false,
+                visible_working: state == crate::detect::AgentState::Working,
+                process_exited: false,
+                observed_at: Instant::now(),
+            });
+        }
+        server.render_and_stream();
+        let (snapshot, completions) = client_shell_projection(&control_rx);
+        assert_eq!(snapshot.agents.len(), 1);
+        let agent = &snapshot.agents[0];
+        let expected =
+            (first_state == crate::detect::AgentState::Working).then_some(agent.state_change_seq);
+        assert_eq!(
+            completions.completions.get(&agent.pane_id).copied(),
+            expected
+        );
+        assert_eq!(
+            server.app.session_snapshot().agents[0].completion_seq,
+            expected
+        );
+    }
+}
+
+#[tokio::test]
 async fn client_shell_endpoint_request_uses_the_selected_connection() {
     let mut server = test_headless_server();
     server.app.state.workspaces = vec![crate::workspace::Workspace::test_new("endpoint")];
@@ -739,7 +821,7 @@ async fn client_shell_endpoint_request_uses_the_selected_connection() {
             writer,
         })
     );
-    let _initial_snapshot = control_rx.recv().expect("initial shell snapshot");
+    let _initial_snapshot = client_shell_snapshot(&control_rx);
     let boot_id = server.client_shell_boot_id.clone();
 
     assert!(
@@ -857,9 +939,7 @@ async fn client_shell_pairs_agent_view_set_replacement_and_clear_with_snapshots(
             writer,
         })
     );
-    let initial = client_shell_snapshot(read_server_message(
-        control_rx.recv().expect("initial snapshot"),
-    ));
+    let initial = client_shell_snapshot(&control_rx);
 
     let mut view = AgentViewSetParams {
         source: "example.views".into(),
@@ -875,9 +955,7 @@ async fn client_shell_pairs_agent_view_set_replacement_and_clear_with_snapshots(
     let set = client_agent_view_projection(read_server_message(
         control_rx.recv().expect("set projection"),
     ));
-    let set_snapshot = client_shell_snapshot(read_server_message(
-        control_rx.recv().expect("set snapshot"),
-    ));
+    let set_snapshot = client_shell_snapshot(&control_rx);
     assert!(set.revision > initial.revision);
     assert_eq!(set.revision, set_snapshot.revision);
     assert_eq!(
@@ -894,9 +972,7 @@ async fn client_shell_pairs_agent_view_set_replacement_and_clear_with_snapshots(
     let replacement = client_agent_view_projection(read_server_message(
         control_rx.recv().expect("replacement projection"),
     ));
-    let replacement_snapshot = client_shell_snapshot(read_server_message(
-        control_rx.recv().expect("replacement snapshot"),
-    ));
+    let replacement_snapshot = client_shell_snapshot(&control_rx);
     assert!(replacement.revision > set.revision);
     assert_eq!(replacement.revision, replacement_snapshot.revision);
     assert_eq!(
@@ -913,9 +989,7 @@ async fn client_shell_pairs_agent_view_set_replacement_and_clear_with_snapshots(
     let cleared = client_agent_view_projection(read_server_message(
         control_rx.recv().expect("clear projection"),
     ));
-    let cleared_snapshot = client_shell_snapshot(read_server_message(
-        control_rx.recv().expect("clear snapshot"),
-    ));
+    let cleared_snapshot = client_shell_snapshot(&control_rx);
     assert!(cleared.revision > replacement.revision);
     assert_eq!(cleared.revision, cleared_snapshot.revision);
     assert!(cleared.view.is_none());
@@ -967,9 +1041,7 @@ async fn client_shell_receives_metadata_then_shell_free_pane_surface() {
             writer,
         })
     );
-    let snapshot = client_shell_snapshot(read_server_message(
-        control_rx.recv().expect("shell snapshot"),
-    ));
+    let snapshot = client_shell_snapshot(&control_rx);
     assert_eq!(snapshot.workspaces.len(), 1);
     assert_eq!(snapshot.workspaces[0].label, "shell-only-label");
     assert_eq!(
@@ -1599,9 +1671,7 @@ async fn client_shell_config_diagnostics_follow_keybinding_ownership() {
             writer: local_writer,
         })
     );
-    let local_snapshot = client_shell_snapshot(read_server_message(
-        local_control.recv().expect("local shell snapshot"),
-    ));
+    let local_snapshot = client_shell_snapshot(&local_control);
     assert_eq!(
         local_snapshot.config_diagnostic.as_deref(),
         Some("theme warning")
@@ -1625,9 +1695,7 @@ async fn client_shell_config_diagnostics_follow_keybinding_ownership() {
             writer: endpoint_writer,
         })
     );
-    let endpoint_snapshot = client_shell_snapshot(read_server_message(
-        endpoint_control.recv().expect("endpoint shell snapshot"),
-    ));
+    let endpoint_snapshot = client_shell_snapshot(&endpoint_control);
     assert_eq!(
         endpoint_snapshot.config_diagnostic.as_deref(),
         Some("server keybinding warning\ntheme warning")
@@ -1658,12 +1726,8 @@ async fn client_shell_tab_focus_changes_only_the_source_connection() {
 
     let (first_control, _first_render) = connect_test_shell(&mut server, 7, 100, 30);
     let (second_control, _second_render) = connect_test_shell(&mut server, 8, 80, 24);
-    let first_initial = client_shell_snapshot(read_server_message(
-        first_control.recv().expect("first snapshot"),
-    ));
-    let second_initial = client_shell_snapshot(read_server_message(
-        second_control.recv().expect("second snapshot"),
-    ));
+    let first_initial = client_shell_snapshot(&first_control);
+    let second_initial = client_shell_snapshot(&second_control);
     assert_eq!(
         first_initial.focused_tab_id.as_deref(),
         Some(first_tab_id.as_str())
@@ -1699,9 +1763,7 @@ async fn client_shell_tab_focus_changes_only_the_source_connection() {
         first_control.try_recv().is_err(),
         "another shell must not receive a navigation replacement"
     );
-    let second_replacement = client_shell_snapshot(read_server_message(
-        second_control.recv().expect("second replacement snapshot"),
-    ));
+    let second_replacement = client_shell_snapshot(&second_control);
     assert_eq!(
         second_replacement.focused_tab_id.as_deref(),
         Some(second_tab_id.as_str())
@@ -2433,13 +2495,11 @@ async fn public_agent_focus_replaces_a_diverged_client_shell_projection() {
     let second_tab_id = server.app.public_tab_id(1, 0).unwrap();
 
     let (control_rx, render_rx) = connect_test_shell(&mut server, 9, 80, 23);
-    let _ = control_rx.recv().expect("initial snapshot");
+    let _ = client_shell_snapshot(&control_rx);
     assert!(server.focus_shell_client_on_tab(9, &second_tab_id));
     assert!(server.claim_shell_tab_geometry(9, false));
     server.render_and_stream();
-    let diverged = client_shell_snapshot(read_server_message(
-        control_rx.recv().expect("diverged snapshot"),
-    ));
+    let diverged = client_shell_snapshot(&control_rx);
     assert_eq!(
         diverged.focused_workspace_id.as_deref(),
         Some(server.app.public_workspace_id(1).as_str())
@@ -2484,9 +2544,7 @@ async fn public_agent_focus_replaces_a_diverged_client_shell_projection() {
     assert_eq!(location.focused_tab_id(), Some(first_tab_id.as_str()));
 
     server.render_and_stream();
-    let replacement = client_shell_snapshot(read_server_message(
-        control_rx.recv().expect("agent focus replacement snapshot"),
-    ));
+    let replacement = client_shell_snapshot(&control_rx);
     assert_eq!(
         replacement.focused_workspace_id.as_deref(),
         Some(first_workspace_id.as_str())
@@ -2529,10 +2587,7 @@ async fn public_api_focus_replaces_every_client_shell_projection() {
             writer,
         })
     );
-    let initial_revision = client_shell_snapshot(read_server_message(
-        control_rx.recv().expect("initial snapshot"),
-    ))
-    .revision;
+    let initial_revision = client_shell_snapshot(&control_rx).revision;
 
     let (respond_to, _response_rx) = std::sync::mpsc::channel();
     server.handle_api_request_with_shutdown_check(crate::api::ApiRequestMessage {
@@ -2551,9 +2606,7 @@ async fn public_api_focus_replaces_every_client_shell_projection() {
     assert_eq!(server.app.state.active, Some(1));
     server.render_and_stream();
 
-    let replacement = client_shell_snapshot(read_server_message(
-        control_rx.recv().expect("replacement snapshot"),
-    ));
+    let replacement = client_shell_snapshot(&control_rx);
     assert!(replacement.revision > initial_revision);
     assert_eq!(
         replacement.focused_workspace_id.as_deref(),
@@ -2781,9 +2834,7 @@ async fn client_shell_streams_and_targets_popup_terminal_content() {
             writer,
         })
     );
-    let _snapshot = client_shell_snapshot(read_server_message(
-        control_rx.recv().expect("shell snapshot"),
-    ));
+    let _snapshot = client_shell_snapshot(&control_rx);
 
     server.render_and_stream();
     let ServerMessage::PaneSurface(surface) =
@@ -6726,6 +6777,225 @@ fn notification_show_api_includes_sound_in_semantic_event() {
     }
 }
 
+fn completion_guard_server(writer: ClientWriter) -> (HeadlessServer, crate::layout::PaneId) {
+    let mut server = test_headless_server();
+    server.app.state.workspaces = ["background", "active"]
+        .map(crate::workspace::Workspace::test_new)
+        .into();
+    server.app.state.ensure_test_terminals();
+    server.app.state.active = Some(1);
+    server.app.state.toast_config.delay_seconds = 0;
+    server.clients.insert(
+        1,
+        ClientConnection::new(
+            (80, 24),
+            Default::default(),
+            1,
+            RenderEncoding::SemanticFrame,
+            Some(writer),
+        ),
+    );
+    let pane_id = server.app.state.workspaces[0].tabs[0].root_pane;
+    (server, pane_id)
+}
+
+fn completion_guard_api_report(server: &mut HeadlessServer, method: api::schema::Method) {
+    let (respond_to, response_rx) = std::sync::mpsc::channel();
+    server.handle_api_request_with_shutdown_check(api::ApiRequestMessage {
+        request: api::schema::Request {
+            id: "completion-probe".into(),
+            method,
+        },
+        respond_to,
+        response_write_complete: None,
+        stream_active: None,
+    });
+    let response = response_rx
+        .recv_timeout(Duration::from_millis(100))
+        .unwrap();
+    serde_json::from_str::<api::schema::SuccessResponse>(&response).expect("successful report");
+}
+
+fn completion_guard_notifications(
+    server: &mut HeadlessServer,
+    receiver: &std::sync::mpsc::Receiver<Vec<u8>>,
+) -> Vec<protocol::SemanticNotificationKind> {
+    server.send_to_client(
+        1,
+        ServerMessage::EndpointControl {
+            kind: "test.completion-barrier".into(),
+            data: String::new(),
+        },
+    );
+    let mut messages = Vec::new();
+    loop {
+        let message = read_server_message(
+            receiver
+                .recv_timeout(Duration::from_secs(1))
+                .expect("notification barrier"),
+        );
+        if matches!(&message, ServerMessage::EndpointControl { kind, .. } if kind == "test.completion-barrier")
+        {
+            return messages;
+        }
+        if let ServerMessage::SemanticNotification(notification) = message {
+            messages.push(notification.kind);
+        }
+    }
+}
+
+#[test]
+fn completion_guard_api_startup_blocker_respects_suppression() {
+    use protocol::SemanticNotificationKind::{Finished, NeedsAttention};
+
+    let (writer, control_rx, _render_rx) = test_client_writer();
+    let (mut server, pane_id) = completion_guard_server(writer);
+    server.handle_internal_event_with_forwarding(AppEvent::AgentProcessDetected {
+        pane_id,
+        agent: crate::detect::Agent::Pi,
+        observed_at: Instant::now(),
+    });
+    let public_pane_id = server.app.public_pane_id(0, pane_id).unwrap();
+    for (seq, state) in [
+        api::schema::PaneAgentState::Blocked,
+        api::schema::PaneAgentState::Idle,
+    ]
+    .into_iter()
+    .enumerate()
+    {
+        completion_guard_api_report(
+            &mut server,
+            api::schema::Method::PaneReportAgent(api::schema::PaneReportAgentParams {
+                pane_id: public_pane_id.clone(),
+                source: "custom:pi".into(),
+                agent: "pi".into(),
+                state,
+                message: None,
+                seq: Some(seq as u64 + 1),
+                agent_session_id: None,
+                agent_session_path: None,
+            }),
+        );
+    }
+    let notifications = completion_guard_notifications(&mut server, &control_rx);
+    assert!(notifications.contains(&NeedsAttention));
+    assert!(
+        !notifications.contains(&Finished),
+        "startup: {notifications:?}"
+    );
+}
+
+#[test]
+fn completion_guard_api_session_replacement_does_not_notify_finished() {
+    use crate::agent_resume::{AgentSessionRef, PersistedAgentSession};
+    use crate::detect::{Agent, AgentState};
+    use api::schema::{
+        Method, PaneAgentState, PaneReportAgentParams, PaneReportAgentSessionParams,
+    };
+    use protocol::SemanticNotificationKind::Finished;
+
+    for fallback_state in [
+        AgentState::Unknown,
+        AgentState::Working,
+        AgentState::Blocked,
+    ] {
+        for reason in ["new", "resume", "fork"] {
+            let (writer, control_rx, _render_rx) = test_client_writer();
+            let (mut server, pane_id) = completion_guard_server(writer);
+            let terminal_id = server.app.state.workspaces[0].panes[&pane_id]
+                .attached_terminal_id
+                .clone();
+            let terminal = server.app.state.terminals.get_mut(&terminal_id).unwrap();
+            terminal.set_detected_state(Some(Agent::Pi), fallback_state);
+            let cwd = std::env::current_dir().unwrap();
+            let old_session =
+                AgentSessionRef::path(cwd.join("completion-old.jsonl").display().to_string())
+                    .unwrap();
+            let new_session = cwd.join("completion-new.jsonl").display().to_string();
+            terminal.set_persisted_agent_session(PersistedAgentSession {
+                source: "herdr:pi".into(),
+                agent: "pi".into(),
+                session_ref: old_session.clone(),
+            });
+            terminal
+                .set_hook_authority_with_session_ref(
+                    "herdr:pi".into(),
+                    "pi".into(),
+                    AgentState::Idle,
+                    None,
+                    Some(old_session),
+                    Some(10),
+                )
+                .expect("initial authority");
+            let public_pane_id = server.app.public_pane_id(0, pane_id).unwrap();
+            completion_guard_api_report(
+                &mut server,
+                Method::PaneReportAgentSession(PaneReportAgentSessionParams {
+                    pane_id: public_pane_id.clone(),
+                    source: "herdr:pi".into(),
+                    agent: "pi".into(),
+                    seq: Some(11),
+                    agent_session_id: None,
+                    agent_session_path: Some(new_session.clone()),
+                    session_start_source: Some(reason.into()),
+                }),
+            );
+            let mut report = PaneReportAgentParams {
+                pane_id: public_pane_id,
+                source: "herdr:pi".into(),
+                agent: "pi".into(),
+                state: PaneAgentState::Idle,
+                message: None,
+                seq: Some(12),
+                agent_session_id: None,
+                agent_session_path: Some(new_session.clone()),
+            };
+            completion_guard_api_report(&mut server, Method::PaneReportAgent(report.clone()));
+            let terminal = &server.app.state.terminals[&terminal_id];
+            assert_eq!(terminal.state, AgentState::Idle);
+            assert_eq!(
+                terminal
+                    .hook_authority
+                    .as_ref()
+                    .unwrap()
+                    .session_ref
+                    .as_ref()
+                    .unwrap()
+                    .value,
+                new_session
+            );
+            let notifications = completion_guard_notifications(&mut server, &control_rx);
+            assert!(
+                !notifications.contains(&Finished),
+                "{reason}: {notifications:?}"
+            );
+            assert!(server.app.state.workspaces[0].panes[&pane_id].seen);
+            for (seq, state, expected) in [
+                (13, PaneAgentState::Working, AgentState::Working),
+                (14, PaneAgentState::Idle, AgentState::Idle),
+            ] {
+                report.seq = Some(seq);
+                report.state = state;
+                completion_guard_api_report(&mut server, Method::PaneReportAgent(report.clone()));
+                let terminal = &server.app.state.terminals[&terminal_id];
+                assert_eq!(terminal.state, expected);
+                assert_eq!(
+                    terminal.last_agent_completion_seq.is_some(),
+                    expected == AgentState::Idle
+                );
+            }
+            let finished = completion_guard_notifications(&mut server, &control_rx)
+                .into_iter()
+                .filter(|kind| *kind == Finished)
+                .count();
+            assert_eq!(
+                finished, 1,
+                "the first real turn after {reason} must finish"
+            );
+        }
+    }
+}
+
 #[test]
 fn startup_idle_does_not_forward_completion() {
     let mut server = test_headless_server();
diff --git a/src/server/headless/tests/pane_graphics.rs b/src/server/headless/tests/pane_graphics.rs
index 0b7a8095..3e36a8fc 100644
--- a/src/server/headless/tests/pane_graphics.rs
+++ b/src/server/headless/tests/pane_graphics.rs
@@ -56,7 +56,8 @@ async fn client_shell_surface_sends_complete_placements_and_each_live_asset_once
 
 #[tokio::test]
 async fn client_shell_asset_delivery_is_bounded_to_the_current_live_scene() {
-    let (mut server, client_rx, pane_id) = retained_test_server(b"client shell graphics");
+    let (mut server, _control_rx, client_rx, pane_id) =
+        retained_test_server_with_control(b"client shell graphics");
     let client = server.clients.get_mut(&1).unwrap();
     client.mode = ClientConnectionMode::ClientShell;
     client.render_state =
@@ -92,7 +93,8 @@ async fn client_shell_asset_delivery_is_bounded_to_the_current_live_scene() {
 
 #[tokio::test]
 async fn first_kitty_image_updates_retained_surface_without_full_redraw() {
-    let (mut server, client_rx, pane_id) = retained_test_server(b"text before image");
+    let (mut server, _control_rx, client_rx, pane_id) =
+        retained_test_server_with_control(b"text before image");
     let client = server.clients.get_mut(&1).unwrap();
     client.mode = ClientConnectionMode::ClientShell;
     client.render_state =
@@ -183,7 +185,8 @@ async fn first_kitty_image_updates_retained_surface_without_full_redraw() {
 
 #[tokio::test]
 async fn retained_unicode_image_arrives_after_fragmented_upload_without_reupload() {
-    let (mut server, client_rx, pane_id) = retained_test_server(b"\x1b[?1049h");
+    let (mut server, _control_rx, client_rx, pane_id) =
+        retained_test_server_with_control(b"\x1b[?1049h");
     let client = server.clients.get_mut(&1).unwrap();
     client.mode = ClientConnectionMode::ClientShell;
     client.cell_size = crate::kitty_graphics::HostCellSize {
@@ -260,7 +263,8 @@ async fn render_scale_profile_retained_graphics() {
     for retained in [false, true] {
         for with_image in [false, true] {
             for count in [1, 15] {
-                let (mut server, client_rx, root) = retained_test_server(b"populated terminal\r\n");
+                let (mut server, _control_rx, client_rx, root) =
+                    retained_test_server_with_control(b"populated terminal\r\n");
                 let mut pane_ids = vec![root];
                 for index in 1..count {
                     let workspace = &mut server.app.state.workspaces[0];
@@ -336,8 +340,9 @@ async fn render_scale_profile_retained_graphics() {
 
 #[tokio::test]
 async fn client_shell_surface_projects_terminal_kitty_images_from_authoritative_runtime() {
-    let (mut server, client_rx, _pane_id) =
-        retained_test_server(b"\x1b_Ga=T,f=32,t=d,i=7,p=3,s=1,v=1,c=1,r=1,q=2;/wAA/w==\x1b\\");
+    let (mut server, _control_rx, client_rx, _pane_id) = retained_test_server_with_control(
+        b"\x1b_Ga=T,f=32,t=d,i=7,p=3,s=1,v=1,c=1,r=1,q=2;/wAA/w==\x1b\\",
+    );
     let client = server.clients.get_mut(&1).unwrap();
     client.mode = ClientConnectionMode::ClientShell;
     client.render_state =
@@ -366,7 +371,7 @@ async fn client_shell_surface_projects_terminal_kitty_images_from_authoritative_
 
 #[tokio::test]
 async fn client_shell_delivers_equal_pixels_for_distinct_terminal_image_ids() {
-    let (mut server, client_rx, _pane_id) = retained_test_server(
+    let (mut server, _control_rx, client_rx, _pane_id) = retained_test_server_with_control(
         b"\x1b_Ga=T,f=32,t=d,i=7,p=3,s=1,v=1,c=1,r=1,q=2;/wAA/w==\x1b\\\x1b_Ga=T,f=32,t=d,i=8,p=4,s=1,v=1,c=1,r=1,q=2;/wAA/w==\x1b\\",
     );
     let client = server.clients.get_mut(&1).unwrap();
@@ -406,7 +411,8 @@ async fn client_shell_delivers_equal_pixels_for_distinct_terminal_image_ids() {
 
 #[tokio::test]
 async fn full_client_shell_render_lane_does_not_commit_graphics_delivery() {
-    let (mut server, client_rx, pane_id) = retained_test_server(b"client shell graphics");
+    let (mut server, _control_rx, client_rx, pane_id) =
+        retained_test_server_with_control(b"client shell graphics");
     let client = server.clients.get_mut(&1).unwrap();
     client.mode = ClientConnectionMode::ClientShell;
     client.render_state =
@@ -975,7 +981,8 @@ fn resident_direct_stream_survives_non_direct_client_becoming_foreground() {
 #[cfg(unix)]
 #[tokio::test]
 async fn client_shell_direct_graphics_uploads_without_server_authored_coordinates() {
-    let (mut server, client_rx, pane_id) = retained_test_server(b"client shell direct");
+    let (mut server, _control_rx, client_rx, pane_id) =
+        retained_test_server_with_control(b"client shell direct");
     server.app.state.kitty_graphics_enabled = true;
     let client = server.clients.get_mut(&1).unwrap();
     client.mode = ClientConnectionMode::ClientShell;
diff --git a/src/server/headless/tests/pane_move.rs b/src/server/headless/tests/pane_move.rs
index c18d0f1f..789b144f 100644
--- a/src/server/headless/tests/pane_move.rs
+++ b/src/server/headless/tests/pane_move.rs
@@ -55,7 +55,7 @@ async fn public_pane_move_focus_follows_the_moved_pane() {
     let source_id = server.app.public_pane_id(0, source).unwrap();
     let destination_id = server.app.public_workspace_id(1);
     let (control_rx, render_rx) = connect_test_shell(&mut server, 9, 80, 23);
-    let initial = client_shell_snapshot(read_server_message(control_rx.recv().unwrap()));
+    let initial = client_shell_snapshot(&control_rx);
 
     let moved = public_move(
         &mut server,
@@ -81,7 +81,7 @@ async fn public_pane_move_focus_follows_the_moved_pane() {
     assert_eq!(location.focused_tab_id(), Some(moved.pane.tab_id.as_str()));
 
     server.render_and_stream();
-    let snapshot = client_shell_snapshot(read_server_message(control_rx.recv().unwrap()));
+    let snapshot = client_shell_snapshot(&control_rx);
     assert!(snapshot.revision > initial.revision);
     assert_eq!(
         snapshot.focused_workspace_id.as_deref(),
diff --git a/src/server/headless/tests/surface_interest.rs b/src/server/headless/tests/surface_interest.rs
index 7c8b33f8..f3acc27b 100644
--- a/src/server/headless/tests/surface_interest.rs
+++ b/src/server/headless/tests/surface_interest.rs
@@ -79,10 +79,7 @@ async fn metadata_only_shell_is_isolated_until_surface_activation() {
             writer,
         })
     );
-    assert!(matches!(
-        read_server_message(control_rx.recv().expect("metadata snapshot")),
-        ServerMessage::EndpointControl { .. }
-    ));
+    let _ = client_shell_snapshot(&control_rx);
     assert_eq!(server.foreground_client_id, None);
     assert_eq!(server.effective_size, original_size);
 
@@ -410,7 +407,7 @@ async fn presentation_sync_epoch_replays_modes_and_title() {
             writer,
         })
     );
-    let _ = control_rx.recv().expect("initial snapshot");
+    let _ = client_shell_snapshot(&control_rx);
     server.api_window_title = Some("target title".into());
     {
         let client = server.clients.get_mut(&client_id).unwrap();
@@ -526,11 +523,7 @@ async fn two_headless_servers_drive_atomic_endpoint_handoff() {
             writer: source_writer,
         })
     );
-    let source_snapshot = client_shell_snapshot(read_server_message(
-        source_control
-            .recv()
-            .expect("real source metadata snapshot"),
-    ));
+    let source_snapshot = client_shell_snapshot(&source_control);
 
     let mut target_server = test_headless_server();
     let _target_input = install_focused_test_runtime(&mut target_server, b"remote target");
@@ -553,11 +546,7 @@ async fn two_headless_servers_drive_atomic_endpoint_handoff() {
             writer: target_writer,
         })
     );
-    let remote_snapshot = client_shell_snapshot(read_server_message(
-        target_control
-            .recv()
-            .expect("real target metadata snapshot"),
-    ));
+    let remote_snapshot = client_shell_snapshot(&target_control);
 
     let profile = SavedSshEndpoint {
         id: ProfileId::parse("0123456789abcdef0123456789abcdef").unwrap(),
@@ -685,9 +674,7 @@ async fn two_headless_servers_drive_atomic_endpoint_handoff() {
     );
 
     target_server.render_and_stream();
-    let coherent_snapshot = client_shell_snapshot(read_server_message(
-        target_control.recv().expect("target replacement snapshot"),
-    ));
+    let coherent_snapshot = client_shell_snapshot(&target_control);
     let snapshot_progress = activation.receive_snapshot(&target_id, 7, &coherent_snapshot);
     shell.set_endpoint_snapshot(&target_id, coherent_snapshot);
     assert_eq!(