fix: keep direct graphics bound to its client (#3549)

6 files changed, 250 insertions, 27 deletions

diff --git a/src/app/pane_graphics.rs b/src/app/pane_graphics.rs
index 7b516116..e6c43fa1 100644
--- a/src/app/pane_graphics.rs
+++ b/src/app/pane_graphics.rs
@@ -65,6 +65,13 @@ impl Slot {
             .as_ref()
             .is_some_and(|active| active.load(std::sync::atomic::Ordering::Acquire))
     }
+
+    pub(crate) fn direct_client(&self) -> Option<u64> {
+        self.direct_gate
+            .as_ref()
+            .map(|gate| gate.client_id)
+            .or_else(|| self.layer.as_ref().and_then(Layer::resident_client))
+    }
 }
 
 impl Drop for Slot {
diff --git a/src/kitty_graphics.rs b/src/kitty_graphics.rs
index 37bae34e..2eca21f5 100644
--- a/src/kitty_graphics.rs
+++ b/src/kitty_graphics.rs
@@ -502,6 +502,7 @@ fn collect_visible_placements(
     cell_size: HostCellSize,
     uploaded_images: &HashMap<u32, ImageSignature>,
     oversized_images: &HashMap<HostSourceKey, ImageSignature>,
+    client_id: u64,
 ) -> Vec<HostPlacement> {
     let Some(target) = surface.target else {
         tracing::debug!("collect_visible_placements: no tab surface target");
@@ -536,9 +537,10 @@ fn collect_visible_placements(
             .filter_map(|((pane_id, layer_id), slot)| {
                 (*pane_id == info.id)
                     .then(|| {
-                        slot.layer
-                            .as_ref()
-                            .map(|layer| (layer_id, slot.host_image_id, layer))
+                        slot.layer.as_ref().and_then(|layer| {
+                            (!layer.terminal_only() || slot.direct_client() == Some(client_id))
+                                .then_some((layer_id, slot.host_image_id, layer))
+                        })
                     })
                     .flatten()
             })
diff --git a/src/kitty_graphics/surface.rs b/src/kitty_graphics/surface.rs
index 0e69a145..bf2482b6 100644
--- a/src/kitty_graphics/surface.rs
+++ b/src/kitty_graphics/surface.rs
@@ -339,6 +339,7 @@ pub(crate) fn collect_scene(
         cell_size,
         &uploaded_images,
         &delivered_terminal_images,
+        client_id,
     );
 
     if let (Some(popup), Some((width, height)), Some((_, target))) = (
@@ -476,12 +477,8 @@ pub(crate) fn collect_scene(
         .iter()
         .filter_map(|(key, slot)| {
             let layer = slot.layer.as_ref()?;
-            let retained_for_client = layer.resident_client() == Some(client_id)
-                || slot
-                    .direct_gate
-                    .as_ref()
-                    .is_some_and(|gate| gate.client_id == client_id);
-            retained_for_client.then(|| pane_layer_asset_key(app, key, layer))?
+            (slot.direct_client() == Some(client_id))
+                .then(|| pane_layer_asset_key(app, key, layer))?
         })
         .collect::<Vec<_>>();
     retained_assets.sort_by_key(|key| format!("{:?}", key.source));
diff --git a/src/server/headless.rs b/src/server/headless.rs
index 1efa2d1a..04a1b17b 100644
--- a/src/server/headless.rs
+++ b/src/server/headless.rs
@@ -933,16 +933,54 @@ impl HeadlessServer {
             .count()
     }
 
-    fn direct_graphics_available(&self) -> bool {
-        self.app_client_count() == 1
-            && self.foreground_client_id.is_some_and(|id| {
-                self.clients.get(&id).is_some_and(|client| {
-                    client.is_shell_client()
-                        && client.writer.is_some()
-                        && client.direct_graphics
-                        && client.pixel_mouse
-                })
+    fn client_supports_direct_graphics(&self, client_id: u64) -> bool {
+        self.clients.get(&client_id).is_some_and(|client| {
+            client.is_shell_client()
+                && client.writer.is_some()
+                && client.direct_graphics
+                && client.pixel_mouse
+        })
+    }
+
+    fn existing_direct_graphics_client(&self) -> Option<u64> {
+        self.app
+            .pane_graphics
+            .slots
+            .values()
+            .filter_map(crate::app::pane_graphics::Slot::direct_client)
+            .filter_map(|client_id| {
+                let client = self.clients.get(&client_id)?;
+                self.client_supports_direct_graphics(client_id)
+                    .then_some((client.last_activity, client_id))
             })
+            .max()
+            .map(|(_, client_id)| client_id)
+    }
+
+    fn direct_graphics_client(&self) -> Option<u64> {
+        self.existing_direct_graphics_client().or_else(|| {
+            self.foreground_client_id
+                .filter(|client_id| self.client_supports_direct_graphics(*client_id))
+        })
+    }
+
+    fn direct_graphics_client_for_key(&self, key: &crate::app::pane_graphics::Key) -> Option<u64> {
+        match self
+            .app
+            .pane_graphics
+            .slots
+            .get(key)
+            .and_then(crate::app::pane_graphics::Slot::direct_client)
+        {
+            Some(client_id) => self
+                .client_supports_direct_graphics(client_id)
+                .then_some(client_id),
+            None => self.direct_graphics_client(),
+        }
+    }
+
+    fn direct_graphics_available(&self) -> bool {
+        self.direct_graphics_client().is_some()
     }
 
     fn has_app_client(&self) -> bool {
@@ -987,6 +1025,7 @@ impl HeadlessServer {
                 self.send_shell_focus_target(target, crate::ghostty::FocusEvent::Lost);
             }
         }
+        self.app.direct_graphics_available = self.direct_graphics_available();
         if was_foreground {
             self.promote_latest_remaining_client()
         } else {
diff --git a/src/server/headless/pane_graphics.rs b/src/server/headless/pane_graphics.rs
index bce53b72..e1410be1 100644
--- a/src/server/headless/pane_graphics.rs
+++ b/src/server/headless/pane_graphics.rs
@@ -33,10 +33,9 @@ impl HeadlessServer {
                 )
             })
         });
-        let direct_client = self
-            .direct_graphics_available()
-            .then_some(self.foreground_client_id)
-            .flatten();
+        let direct_client = direct_key
+            .as_ref()
+            .and_then(|key| self.direct_graphics_client_for_key(key));
         let gate_busy = self
             .app
             .pane_graphics
@@ -259,6 +258,7 @@ impl HeadlessServer {
         if self.app.pane_graphics.slots.remove(key).is_some() {
             self.app.pane_graphics.mark_changed();
         }
+        self.app.direct_graphics_available = self.direct_graphics_available();
     }
 
     pub(super) fn start_direct_graphics_response(
@@ -340,6 +340,7 @@ impl HeadlessServer {
                 self.retire_direct_gate(&key);
                 return true;
             }
+            self.app.direct_graphics_available = self.direct_graphics_available();
             return true;
         }
 
diff --git a/src/server/headless/tests/pane_graphics.rs b/src/server/headless/tests/pane_graphics.rs
index b5938d56..6c0cd8b0 100644
--- a/src/server/headless/tests/pane_graphics.rs
+++ b/src/server/headless/tests/pane_graphics.rs
@@ -304,6 +304,7 @@ fn direct_stream_message(
     path: String,
     image_width: u32,
     image_height: u32,
+    sequence: u64,
 ) -> (api::ApiRequestMessage, std::sync::mpsc::Receiver<String>) {
     let (respond_to, response_rx) = std::sync::mpsc::channel();
     (
@@ -320,7 +321,7 @@ fn direct_stream_message(
                         image_height,
                         format: api::schema::PaneGraphicsFormat::Rgba,
                         path,
-                        sequence: 1,
+                        sequence,
                         revision: 1,
                         placement: Default::default(),
                     },
@@ -612,6 +613,114 @@ fn rejected_or_stale_requests_do_not_schedule_rendering() {
 
 #[cfg(unix)]
 #[tokio::test]
+async fn direct_graphics_availability_follows_foreground_client_with_background_clients() {
+    let (mut server, _client_rx, _pane_id) = retained_test_server(b"direct eligibility");
+    let foreground = server.clients.get_mut(&1).unwrap();
+    foreground.direct_graphics = true;
+    foreground.pixel_mouse = true;
+
+    let (background_writer, _background_control_rx, _background_render_rx) = test_client_writer();
+    server.clients.insert(
+        2,
+        ClientConnection::new(
+            (80, 24),
+            crate::kitty_graphics::HostCellSize::default(),
+            2,
+            RenderEncoding::SemanticFrame,
+            Some(background_writer),
+        ),
+    );
+
+    server.sync_foreground_client_state();
+    assert!(server.direct_graphics_available());
+    assert!(server.app.direct_graphics_available);
+
+    server.foreground_client_id = Some(2);
+    server.sync_foreground_client_state();
+    assert!(!server.direct_graphics_available());
+    assert!(!server.app.direct_graphics_available);
+}
+
+#[cfg(unix)]
+#[test]
+fn direct_graphics_routing_prefers_the_target_stream_owner() {
+    let (mut server, first_key, response_rx) = direct_gate_server(&[1, 2, 3, 4]);
+    add_direct_client(&mut server, 7);
+    add_direct_client(&mut server, 8);
+    server.clients.get_mut(&8).unwrap().last_activity = 2;
+    let (transfer_id, image_id) = direct_ids(&server, &first_key);
+    assert!(server.complete_direct_graphics(7, transfer_id, image_id, true));
+    assert_eq!(response_rx.recv().unwrap(), "ack");
+
+    let second_key = (first_key.0, "second".into());
+    let second_layer = crate::app::pane_graphics::Layer {
+        format: crate::api::schema::PaneGraphicsFormat::Rgba,
+        image_width: 1,
+        image_height: 1,
+        backing: crate::app::pane_graphics::Backing::Resident {
+            len: 4,
+            client_id: 8,
+        },
+        data_fingerprint: 2,
+        render: Default::default(),
+        z_index: 0,
+    };
+    server.app.pane_graphics.slots.insert(
+        second_key.clone(),
+        crate::app::pane_graphics::Slot::test((1 << 31) | 901, Some(second_layer)),
+    );
+
+    assert_eq!(server.direct_graphics_client_for_key(&first_key), Some(7));
+    assert_eq!(server.direct_graphics_client_for_key(&second_key), Some(8));
+    assert_eq!(
+        server.direct_graphics_client_for_key(&(first_key.0, "new".into())),
+        Some(8)
+    );
+}
+
+#[cfg(unix)]
+#[test]
+fn resident_direct_stream_survives_non_direct_client_becoming_foreground() {
+    let (mut server, key, response_rx) = direct_gate_server(&[1, 2, 3, 4]);
+    add_direct_client(&mut server, 7);
+    server.foreground_client_id = Some(7);
+    server.sync_foreground_client_state();
+    let (transfer_id, image_id) = direct_ids(&server, &key);
+    assert!(server.complete_direct_graphics(7, transfer_id, image_id, true));
+    assert_eq!(response_rx.recv().unwrap(), "ack");
+
+    let (writer, _control_rx, _render_rx) = test_client_writer();
+    server.clients.insert(
+        8,
+        ClientConnection::new(
+            (80, 24),
+            crate::kitty_graphics::HostCellSize::default(),
+            2,
+            RenderEncoding::SemanticFrame,
+            Some(writer),
+        ),
+    );
+    server.foreground_client_id = Some(8);
+    server.sync_foreground_client_state();
+
+    assert!(server.direct_graphics_available());
+    assert!(server.app.direct_graphics_available);
+    assert_eq!(
+        server.app.pane_graphics.slots[&key]
+            .layer
+            .as_ref()
+            .and_then(crate::app::pane_graphics::Layer::resident_client),
+        Some(7)
+    );
+
+    server.remove_client_and_resize_if_needed(7);
+    assert!(!server.app.pane_graphics.slots.contains_key(&key));
+    assert!(!server.direct_graphics_available());
+    assert!(!server.app.direct_graphics_available);
+}
+
+#[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");
     server.app.state.kitty_graphics_enabled = true;
@@ -626,11 +735,25 @@ async fn client_shell_direct_graphics_uploads_without_server_authored_coordinate
     client.direct_graphics = true;
     client.pixel_mouse = true;
     server.app.direct_graphics_available = true;
+    let (background_writer, _background_control_rx, background_render_rx) = test_client_writer();
+    server.clients.insert(
+        2,
+        ClientConnection::new(
+            (80, 24),
+            crate::kitty_graphics::HostCellSize {
+                width_px: 10,
+                height_px: 20,
+            },
+            2,
+            RenderEncoding::SemanticFrame,
+            Some(background_writer),
+        ),
+    );
     set_stream_owner(&mut server, pane_id, "browser");
     let public_pane_id = server.app.public_pane_id(0, pane_id).unwrap();
     let path = sparse_direct_frame(&server, "client-shell-direct.rgba", 1, 1);
     let (message, response_rx) =
-        direct_stream_message("shell-direct", &public_pane_id, "browser", path, 1, 1);
+        direct_stream_message("shell-direct", &public_pane_id, "browser", path, 1, 1, 1);
 
     assert_eq!(
         server.handle_pane_graphics_stream_frame(message),
@@ -656,6 +779,7 @@ async fn client_shell_direct_graphics_uploads_without_server_authored_coordinate
         }
         other => panic!("expected client shell graphics file, got {other:?}"),
     };
+    assert!(background_render_rx.try_recv().is_err());
     assert_eq!(
         image_id,
         crate::kitty_graphics::surface::host_image_id(&server.client_shell_boot_id, &asset)
@@ -684,6 +808,47 @@ async fn client_shell_direct_graphics_uploads_without_server_authored_coordinate
         &response_rx.recv_timeout(Duration::from_secs(1)).unwrap()
     )
     .is_ok());
+
+    server.foreground_client_id = Some(2);
+    server.sync_foreground_client_state();
+    let next_path = sparse_direct_frame(&server, "client-shell-direct-next.rgba", 1, 1);
+    let (next_message, next_response_rx) = direct_stream_message(
+        "shell-direct-next",
+        &public_pane_id,
+        "browser",
+        next_path,
+        1,
+        1,
+        2,
+    );
+    assert_eq!(
+        server.handle_pane_graphics_stream_frame(next_message),
+        RenderImpact::None
+    );
+    let (next_transfer_id, next_image_id, next_asset) = match read_server_message(
+        client_rx
+            .recv_timeout(Duration::from_secs(1))
+            .expect("sticky owner direct upload"),
+    ) {
+        ServerMessage::GraphicsFile {
+            transfer_id,
+            image_id,
+            surface_asset: Some(asset),
+            ..
+        } => (transfer_id, image_id, asset),
+        other => panic!("expected sticky owner graphics file, got {other:?}"),
+    };
+    assert!(background_render_rx.try_recv().is_err());
+    server.start_direct_graphics_response(1, next_transfer_id, next_image_id);
+    assert!(server.complete_direct_graphics(1, next_transfer_id, next_image_id, true));
+    assert!(serde_json::from_str::<api::schema::SuccessResponse>(
+        &next_response_rx
+            .recv_timeout(Duration::from_secs(1))
+            .unwrap()
+    )
+    .is_ok());
+    assert!(server.app.direct_graphics_available);
+
     server.clients.get_mut(&1).unwrap().request_repaint();
     server.render_and_stream();
     let surface = read_server_message(
@@ -696,8 +861,20 @@ async fn client_shell_direct_graphics_uploads_without_server_authored_coordinate
     };
     assert_eq!(surface.graphics.placements.len(), 1);
     assert!(surface.graphics.assets.is_empty());
-    assert_eq!(surface.graphics.placements[0].asset, asset);
-    assert_eq!(surface.graphics.retained_assets, vec![asset.clone()]);
+    assert_eq!(surface.graphics.placements[0].asset, next_asset);
+    assert_eq!(surface.graphics.retained_assets, vec![next_asset.clone()]);
+
+    let background_surface = read_server_message(
+        background_render_rx
+            .recv_timeout(Duration::from_secs(1))
+            .expect("non-owning client shell scene"),
+    );
+    let ServerMessage::PaneSurface(background_surface) = background_surface else {
+        panic!("expected non-owning client shell pane surface");
+    };
+    assert!(background_surface.graphics.assets.is_empty());
+    assert!(background_surface.graphics.placements.is_empty());
+    assert!(background_surface.graphics.retained_assets.is_empty());
 
     let (hidden, _) = crate::server::client_shell_graphics::collect(
         &server.app,
@@ -716,7 +893,7 @@ async fn client_shell_direct_graphics_uploads_without_server_authored_coordinate
         1,
     );
     assert!(hidden.placements.is_empty());
-    assert_eq!(hidden.retained_assets, vec![asset]);
+    assert_eq!(hidden.retained_assets, vec![next_asset]);
 }
 
 #[cfg(unix)]