1 files changed, 29 insertions, 5 deletions
diff --git a/src/client/shell/graphics.rs b/src/client/shell/graphics.rs
index a7ae81e3..35e61dbb 100644
--- a/src/client/shell/graphics.rs
+++ b/src/client/shell/graphics.rs
@@ -1,12 +1,37 @@
 use super::*;
 
+#[cfg(unix)]
+pub(crate) struct ClientGraphicsCheckpoint(crate::kitty_graphics::surface::ClientState);
+
 impl ClientShellState {
-    #[cfg(unix)]
+    #[cfg(all(test, unix))]
     pub(crate) fn graphics_scope(&self) -> &str {
         self.graphics.scope()
     }
 
     #[cfg(unix)]
+    pub(crate) fn accepts_direct_graphics_asset(
+        &self,
+        key: &crate::protocol::SurfaceGraphicsAssetKey,
+        image_id: u32,
+    ) -> bool {
+        self.graphics.accepts_direct_asset(key, image_id)
+    }
+
+    #[cfg(unix)]
+    pub(crate) fn direct_graphics_checkpoint(&self) -> ClientGraphicsCheckpoint {
+        ClientGraphicsCheckpoint(self.graphics.clone())
+    }
+
+    #[cfg(unix)]
+    pub(crate) fn restore_direct_graphics_checkpoint(
+        &mut self,
+        checkpoint: ClientGraphicsCheckpoint,
+    ) {
+        self.graphics = checkpoint.0;
+    }
+
+    #[cfg(unix)]
     pub(crate) fn trust_direct_graphics_asset(
         &mut self,
         key: &crate::protocol::SurfaceGraphicsAssetKey,
@@ -33,10 +58,9 @@ impl ClientShellState {
 
     pub(super) fn compose_graphics(
         &mut self,
-        frame: &mut FrameData,
         layout: ClientShellLayout,
         occlusion: &crate::kitty_graphics::surface::Occlusion,
-    ) {
+    ) -> crate::kitty_graphics::GraphicsOutput {
         let visibility = if self.endpoint_error.is_some() {
             crate::kitty_graphics::surface::Visibility::Hidden
         } else if self.hits.popup.is_some() {
@@ -49,12 +73,12 @@ impl ClientShellState {
             .popup
             .as_ref()
             .map(|popup| (popup.inner_rect.x, popup.inner_rect.y));
-        frame.graphics = self.graphics.encode(
+        self.graphics.encode_output(
             visibility,
             (layout.pane_surface.x, layout.pane_surface.y),
             popup_origin,
             self.graphics_cell_size,
             occlusion,
-        );
+        )
     }
 }