feat(graphics): add direct pane frame streaming (#2523)

* feat(graphics): add direct pane frame streaming

* feat(graphics): expose pane visibility

* test(graphics): preserve legacy terminal image output

* fix(graphics): harden streaming lifecycle

* fix(input): preserve mouse releases outside panes

1 files changed, 464 insertions, 71 deletions

diff --git a/src/server/headless/pane_graphics.rs b/src/server/headless/pane_graphics.rs
index ed28eab0..d027f79b 100644
--- a/src/server/headless/pane_graphics.rs
+++ b/src/server/headless/pane_graphics.rs
@@ -1,5 +1,3 @@
-use tracing::warn;
-
 use super::{HeadlessServer, RenderImpact};
 use crate::api;
 use crate::protocol::{ServerMessage, MAX_GRAPHICS_FRAME_SIZE};
@@ -12,31 +10,16 @@ pub(super) enum RetainedGraphicsOutcome {
     Fallback,
 }
 
-pub(super) fn frame_pane_graphics_for_client(bytes: Vec<u8>) -> Vec<u8> {
+pub(super) fn frame_pane_graphics(bytes: Vec<u8>) -> Vec<u8> {
     if bytes.is_empty() {
         return bytes;
     }
-    let mut framed = Vec::with_capacity(bytes.len() + 4);
-    framed.extend_from_slice(b"\x1b7");
-    framed.extend_from_slice(&bytes);
-    framed.extend_from_slice(b"\x1b8");
-    framed
+    [b"\x1b7".as_slice(), &bytes, b"\x1b8"].concat()
 }
 
 impl HeadlessServer {
     pub(super) fn pane_graphics_runtime_active(&self) -> bool {
-        !self.app.state.pane_graphics_layers.is_empty()
-            || !self.app.state.pane_graphics_streams.is_empty()
-    }
-
-    pub(super) fn cancel_inactive_pane_graphics_streams(&self) {
-        api::cancel_inactive_pane_graphics_streams(|owner| {
-            self.app
-                .state
-                .pane_graphics_streams
-                .values()
-                .any(|active_owner| active_owner == owner)
-        });
+        !self.app.pane_graphics.slots.is_empty()
     }
 
     pub(super) fn handle_pane_graphics_stream_frame(
@@ -49,12 +32,164 @@ impl HeadlessServer {
         }
 
         let internal_changed = self.drain_all_internal_events_with_forwarding();
+        let direct = match &msg.request.method {
+            api::schema::Method::PaneGraphicsStreamDirect(params) => Some(params.clone()),
+            _ => None,
+        };
+        let direct_key = direct.as_ref().and_then(|params| {
+            self.app.parse_pane_id(&params.pane_id).map(|(_, pane_id)| {
+                (
+                    pane_id,
+                    params
+                        .layer_id
+                        .clone()
+                        .unwrap_or_else(|| api::schema::PANE_GRAPHICS_PRIMARY_LAYER_ID.into()),
+                )
+            })
+        });
+        let direct_client = self
+            .direct_graphics_available()
+            .then_some(self.foreground_client_id)
+            .flatten();
+        let gate_busy = self
+            .app
+            .pane_graphics
+            .slots
+            .values()
+            .any(|slot| slot.direct_gate.is_some());
+        self.app.direct_graphics_available = direct_client.is_some() && !gate_busy;
         let response = self
             .app
             .handle_api_request_after_internal_events_drained(msg.request);
+        self.app.direct_graphics_available = self.direct_graphics_available();
         let succeeded = serde_json::from_str::<api::schema::SuccessResponse>(&response).is_ok();
-        let _ = msg.respond_to.send(response);
 
+        if succeeded {
+            if let (Some(params), Some(client_id)) = (direct, direct_client) {
+                let direct_frame = direct_key.as_ref().and_then(|key| {
+                    self.app.pane_graphics.slots.get(key).and_then(|slot| {
+                        let active_owner = slot.stream_owner.as_deref() == Some(&params.owner)
+                            && slot.stream_is_active();
+                        active_owner.then_some(())?;
+                        slot.layer
+                            .as_ref()
+                            .and_then(crate::app::pane_graphics::Layer::direct_lease)
+                            .map(|lease| {
+                                (
+                                    slot.host_image_id,
+                                    lease.path().to_string_lossy().into_owned(),
+                                    lease.len() as u64,
+                                    lease.fingerprint(),
+                                )
+                            })
+                    })
+                });
+                if let (Some(key), Some((image_id, path, expected_len, transfer_id))) =
+                    (direct_key.clone(), direct_frame)
+                {
+                    let mut next_cache = self
+                        .clients
+                        .get(&client_id)
+                        .map(|client| client.graphics_cache.clone())
+                        .unwrap_or_default();
+                    let command = (!internal_changed)
+                        .then(|| {
+                            crate::kitty_graphics::prepare_direct_file(
+                                &self.app.state,
+                                &self.app.pane_graphics,
+                                self.app.state.view.tab_surface(),
+                                self.clients
+                                    .get(&client_id)
+                                    .map(|client| client.cell_size)
+                                    .unwrap_or_default(),
+                                &mut next_cache,
+                                &key,
+                            )
+                        })
+                        .flatten();
+                    let Some(command) = command else {
+                        if self.install_inline_fallback(&key) {
+                            if msg.respond_to.send(response).is_err() {
+                                self.retire_direct_gate(&key);
+                            }
+                            return if internal_changed {
+                                RenderImpact::Full
+                            } else {
+                                RenderImpact::Graphics
+                            };
+                        }
+                        self.retire_direct_gate(&key);
+                        return if internal_changed {
+                            RenderImpact::Full
+                        } else {
+                            RenderImpact::Graphics
+                        };
+                    };
+                    let message = ServerMessage::GraphicsFile {
+                        path,
+                        expected_len,
+                        image_id,
+                        transfer_id,
+                        leading: command.leading,
+                        control: command.control,
+                    };
+                    let send =
+                        Self::frame_server_message_with_max(&message, MAX_GRAPHICS_FRAME_SIZE)
+                            .map_err(|_| std::sync::mpsc::TrySendError::Full(Vec::new()))
+                            .and_then(|framed| {
+                                let Some(writer) = self
+                                    .clients
+                                    .get(&client_id)
+                                    .and_then(|client| client.writer.as_ref())
+                                else {
+                                    return Err(std::sync::mpsc::TrySendError::Disconnected(
+                                        framed,
+                                    ));
+                                };
+                                writer.render.send_ordered(framed)
+                            });
+                    match send {
+                        Ok(()) => {
+                            if let Some(slot) = self.app.pane_graphics.slots.get_mut(&key) {
+                                slot.direct_gate = Some(crate::app::pane_graphics::DirectGate {
+                                    transfer_id,
+                                    client_id,
+                                    deadline: std::time::Instant::now()
+                                        + crate::app::pane_graphics::DIRECT_DELIVERY_TIMEOUT,
+                                    written: false,
+                                    success_response: response,
+                                    respond_to: msg.respond_to,
+                                });
+                            }
+                            if let Some(client) = self.clients.get_mut(&client_id) {
+                                client.graphics_cache = next_cache;
+                            }
+                            return if internal_changed {
+                                RenderImpact::Full
+                            } else {
+                                RenderImpact::None
+                            };
+                        }
+                        Err(error) => {
+                            self.handle_unwritten_direct_failure(
+                                &key,
+                                response,
+                                msg.respond_to,
+                                error,
+                            );
+                            return RenderImpact::Graphics;
+                        }
+                    }
+                }
+            }
+        }
+
+        let response_failed = msg.respond_to.send(response).is_err();
+        if succeeded && response_failed {
+            if let Some(key) = direct_key {
+                self.retire_direct_gate(&key);
+            }
+        }
         if internal_changed {
             RenderImpact::Full
         } else if succeeded {
@@ -64,6 +199,276 @@ impl HeadlessServer {
         }
     }
 
+    fn install_inline_fallback(&mut self, key: &crate::app::pane_graphics::Key) -> bool {
+        let len = self
+            .app
+            .pane_graphics
+            .slots
+            .get(key)
+            .and_then(|slot| slot.layer.as_ref())
+            .and_then(crate::app::pane_graphics::Layer::direct_lease)
+            .map(crate::pane_graphics_files::Lease::len);
+        let Some(len) = len else {
+            return false;
+        };
+        if !self.app.pane_graphics.can_store_inline(key, len) {
+            return false;
+        }
+        let data = self
+            .app
+            .pane_graphics
+            .slots
+            .get(key)
+            .and_then(|slot| slot.layer.as_ref())
+            .and_then(crate::app::pane_graphics::Layer::direct_lease)
+            .and_then(|lease| lease.copy_rgba().ok());
+        let Some(data) = data else {
+            return false;
+        };
+        let slot = self
+            .app
+            .pane_graphics
+            .slots
+            .get_mut(key)
+            .expect("matched slot");
+        if !slot.stream_is_active() {
+            return false;
+        }
+        let layer = slot.layer.take().expect("direct layer");
+        slot.layer = Some(crate::app::pane_graphics::Layer::inline(
+            layer.format,
+            layer.image_width,
+            layer.image_height,
+            data,
+            layer.render,
+            layer.z_index,
+        ));
+        self.app.pane_graphics.mark_changed();
+        true
+    }
+
+    pub(super) fn handle_unwritten_direct_failure(
+        &mut self,
+        key: &crate::app::pane_graphics::Key,
+        response: String,
+        respond_to: std::sync::mpsc::Sender<String>,
+        error: std::sync::mpsc::TrySendError<Vec<u8>>,
+    ) -> bool {
+        if matches!(error, std::sync::mpsc::TrySendError::Full(_))
+            && self.install_inline_fallback(key)
+            && respond_to.send(response).is_ok()
+        {
+            return true;
+        }
+        self.retire_direct_gate(key);
+        false
+    }
+
+    pub(super) fn retire_direct_gate(&mut self, key: &crate::app::pane_graphics::Key) {
+        if self.app.pane_graphics.slots.remove(key).is_some() {
+            self.app.pane_graphics.mark_changed();
+        }
+    }
+
+    pub(super) fn start_direct_graphics_response(
+        &mut self,
+        client_id: u64,
+        transfer_id: u64,
+        image_id: u32,
+    ) -> bool {
+        if let Some(gate) = self.app.pane_graphics.slots.values_mut().find_map(|slot| {
+            (slot.host_image_id == image_id && slot.stream_is_active())
+                .then_some(slot.direct_gate.as_mut())
+                .flatten()
+                .filter(|gate| gate.client_id == client_id && gate.transfer_id == transfer_id)
+        }) {
+            gate.written = true;
+            gate.deadline =
+                std::time::Instant::now() + crate::app::pane_graphics::DIRECT_RESPONSE_TIMEOUT;
+        }
+        false
+    }
+
+    pub(super) fn complete_direct_graphics(
+        &mut self,
+        client_id: u64,
+        transfer_id: u64,
+        image_id: u32,
+        success: bool,
+    ) -> bool {
+        if self.shutting_down {
+            self.retire_all_direct_graphics();
+            return false;
+        }
+        let key = self.app.pane_graphics.slots.iter().find_map(|(key, slot)| {
+            slot.direct_gate
+                .as_ref()
+                .is_some_and(|gate| {
+                    slot.stream_is_active()
+                        && gate.client_id == client_id
+                        && gate.transfer_id == transfer_id
+                        && slot.host_image_id == image_id
+                        && (!success || gate.written)
+                })
+                .then(|| key.clone())
+        });
+        let Some(key) = key else {
+            return false;
+        };
+        let pane_is_live = self
+            .app
+            .state
+            .workspaces
+            .iter()
+            .any(|workspace| workspace.pane_state(key.0).is_some());
+        if !pane_is_live {
+            self.retire_direct_gate(&key);
+            return false;
+        }
+        if success {
+            let (gate, host_image_id) = {
+                let slot = self
+                    .app
+                    .pane_graphics
+                    .slots
+                    .get_mut(&key)
+                    .expect("matched slot");
+                if !slot.stream_is_active() {
+                    return false;
+                }
+                if let Some(layer) = slot.layer.as_mut() {
+                    layer.mark_resident(client_id);
+                }
+                (
+                    slot.direct_gate.take().expect("matched gate"),
+                    slot.host_image_id,
+                )
+            };
+            if let (Some(client), Some(layer)) = (
+                self.clients.get_mut(&client_id),
+                self.app
+                    .pane_graphics
+                    .slots
+                    .get(&key)
+                    .and_then(|slot| slot.layer.as_ref()),
+            ) {
+                client
+                    .graphics_cache
+                    .trust_pane_layer(&key, host_image_id, layer);
+            }
+            if gate.respond_to.send(gate.success_response).is_err() {
+                self.retire_direct_gate(&key);
+                return true;
+            }
+            return true;
+        }
+
+        if let Some(client) = self.clients.get_mut(&client_id) {
+            client.direct_graphics = false;
+            client.graphics_cache.forget_pane_layer(&key, image_id);
+        }
+        self.app.direct_graphics_available = false;
+        if !self.install_inline_fallback(&key) {
+            self.retire_direct_gate(&key);
+            self.retire_all_direct_graphics();
+            return true;
+        }
+        let gate = self
+            .app
+            .pane_graphics
+            .slots
+            .get_mut(&key)
+            .filter(|slot| slot.stream_is_active())
+            .and_then(|slot| slot.direct_gate.take());
+        let Some(gate) = gate else {
+            self.retire_direct_gate(&key);
+            return true;
+        };
+        if gate.respond_to.send(gate.success_response).is_err() {
+            self.retire_direct_gate(&key);
+            return true;
+        }
+        self.retire_all_direct_graphics();
+        true
+    }
+
+    pub(super) fn expire_direct_graphics(&mut self, now: std::time::Instant) -> bool {
+        let expired = self
+            .app
+            .pane_graphics
+            .slots
+            .iter()
+            .filter_map(|(key, slot)| {
+                slot.direct_gate
+                    .as_ref()
+                    .filter(|gate| gate.deadline <= now)
+                    .map(|gate| (key.clone(), gate.client_id))
+            })
+            .collect::<Vec<_>>();
+        for (key, client_id) in &expired {
+            if let Some(slot) = self.app.pane_graphics.slots.get(key) {
+                if let Some(gate) = &slot.direct_gate {
+                    self.send_to_client(
+                        *client_id,
+                        ServerMessage::GraphicsTransmissionRetired {
+                            transfer_id: gate.transfer_id,
+                            image_id: slot.host_image_id,
+                        },
+                    );
+                }
+            }
+            if let Some(client) = self.clients.get_mut(client_id) {
+                client.direct_graphics = false;
+            }
+            self.retire_direct_gate(key);
+        }
+        if !expired.is_empty() {
+            self.app.direct_graphics_available = false;
+            self.retire_all_direct_graphics();
+        }
+        !expired.is_empty()
+    }
+
+    pub(super) fn retire_all_direct_graphics(&mut self) {
+        let keys = self
+            .app
+            .pane_graphics
+            .slots
+            .iter()
+            .filter(|(_, slot)| {
+                slot.layer
+                    .as_ref()
+                    .is_some_and(crate::app::pane_graphics::Layer::terminal_only)
+            })
+            .map(|(key, _)| key.clone())
+            .collect::<Vec<_>>();
+        for key in keys {
+            self.retire_direct_gate(&key);
+        }
+    }
+
+    pub(super) fn retire_direct_graphics_for_client(&mut self, client_id: u64) {
+        let keys = self
+            .app
+            .pane_graphics
+            .slots
+            .iter()
+            .filter_map(|(key, slot)| {
+                let layer = slot.layer.as_ref()?;
+                let owned = layer.resident_client() == Some(client_id)
+                    || (layer.direct_lease().is_some()
+                        && slot
+                            .direct_gate
+                            .as_ref()
+                            .is_some_and(|gate| gate.client_id == client_id));
+                (layer.terminal_only() && owned).then(|| key.clone())
+            })
+            .collect::<Vec<_>>();
+        for key in keys {
+            self.retire_direct_gate(&key);
+        }
+    }
+
     pub(super) fn render_retained_graphics_update_and_stream(&mut self) -> RetainedGraphicsOutcome {
         crate::render_prof::event("retained_graphics.attempt");
         if self.app.full_redraw_pending {
@@ -117,74 +522,62 @@ impl HeadlessServer {
 
             let mut next_graphics_cache = client.graphics_cache.clone();
             let encode_started = crate::render_prof::timer();
-            let bytes =
-                frame_pane_graphics_for_client(crate::kitty_graphics::encode_local_pane_graphics(
-                    &self.app.state,
-                    &self.app.terminal_runtimes,
-                    self.app.state.view.tab_surface(),
-                    cell_size,
-                    &mut next_graphics_cache,
-                ));
+            let encoded = crate::kitty_graphics::encode_local_pane_graphics(
+                &self.app.state,
+                &self.app.pane_graphics,
+                &self.app.terminal_runtimes,
+                self.app.state.view.tab_surface(),
+                cell_size,
+                Some(crate::kitty_graphics::HEADLESS_GRAPHICS_TRANSACTION_BUDGET),
+                &mut next_graphics_cache,
+            );
             crate::render_prof::duration_since("retained_graphics.graphics_encode", encode_started);
-            if bytes.len() > MAX_GRAPHICS_FRAME_SIZE {
-                warn!(
-                    client_id,
-                    graphics_bytes = bytes.len(),
-                    max = MAX_GRAPHICS_FRAME_SIZE,
-                    "dropping oversized retained graphics payload"
-                );
+            prepared.push((client_id, encoded, next_graphics_cache));
+        }
+
+        let mut broken_clients = Vec::new();
+        for (client_id, encoded, next_graphics_cache) in prepared {
+            let Some(client) = self.clients.get_mut(&client_id) else {
                 continue;
-            }
-            let serialized = if bytes.is_empty() {
+            };
+            let serialized = if encoded.bytes.is_empty() {
                 None
             } else {
-                let serialize_started = crate::render_prof::timer();
-                let framed = match Self::frame_server_message_with_max(
-                    &ServerMessage::Graphics { bytes },
+                match Self::frame_server_message_with_max(
+                    &ServerMessage::Graphics {
+                        bytes: frame_pane_graphics(encoded.bytes),
+                    },
                     MAX_GRAPHICS_FRAME_SIZE,
                 ) {
-                    Ok(framed) => framed,
-                    Err(err) => {
-                        warn!(client_id, err = %err, "failed to serialize retained graphics");
+                    Ok(serialized) => Some(serialized),
+                    Err(_) => {
+                        crate::render_prof::event("retained_graphics_fallback.oversized");
                         return RetainedGraphicsOutcome::Fallback;
                     }
-                };
-                crate::render_prof::duration_since(
-                    "retained_graphics.serialize",
-                    serialize_started,
-                );
-                Some(framed)
-            };
-            prepared.push((client_id, serialized, next_graphics_cache));
-        }
-
-        let mut broken_clients = Vec::new();
-        for (client_id, serialized, next_graphics_cache) in prepared {
-            let Some(client) = self.clients.get_mut(&client_id) else {
-                continue;
+                }
             };
-            let Some(serialized) = serialized else {
-                client.graphics_cache = next_graphics_cache;
-                client.clear_deferred_render();
-                continue;
+            let result = match (serialized, client.writer.as_ref()) {
+                (None, _) => Ok(()),
+                (Some(bytes), Some(writer)) => writer.render.try_send(bytes),
+                (Some(bytes), None) => Err(std::sync::mpsc::TrySendError::Disconnected(bytes)),
             };
-            let Some(writer) = client.writer.as_ref().cloned() else {
-                broken_clients.push(client_id);
-                continue;
-            };
-            match writer.render.try_send(serialized) {
+            match result {
                 Ok(()) => {
                     client.graphics_cache = next_graphics_cache;
-                    client.clear_deferred_render();
+                    if encoded.incomplete {
+                        client.defer_full_render();
+                        deferred = true;
+                    } else {
+                        client.clear_deferred_render();
+                    }
                     crate::render_prof::event("retained_graphics.sent");
                 }
                 Err(std::sync::mpsc::TrySendError::Full(_)) => {
-                    client.defer_pane_graphics_render();
+                    client.defer_full_render();
                     deferred = true;
-                    crate::render_prof::event("retained_graphics.deferred");
                 }
                 Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {
-                    broken_clients.push(client_id);
+                    broken_clients.push(client_id)
                 }
             }
         }