From c11beb172877ba12ab76b429778a4be865b07aa4 Mon Sep 17 00:00:00 2001
From: PJ <ipriyanshujain@gmail.com>
Date: Sun, 31 May 2026 15:25:58 +0530
Subject: [PATCH] feat(driver): expose Snapshot on DeviceDriver and sidecar
 client

Snapshot wraps the new atomic-snapshot gRPC: the runner gets hierarchy
and screenshot from one round-trip whose two reads are serialized on
the sidecar side.
---
 internal/driver/driver.go         |  5 +++++
 internal/driver/sidecar/client.go | 17 +++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/internal/driver/driver.go b/internal/driver/driver.go
index 6bde5e0..6a0b24a 100644
--- a/internal/driver/driver.go
+++ b/internal/driver/driver.go
@@ -22,6 +22,11 @@ type DeviceDriver interface {
 
 	Hierarchy(ctx context.Context) (string, error)
 	Screenshot(ctx context.Context) (Image, error)
+	// Snapshot returns the hierarchy and screenshot captured back-to-back
+	// under a backend-side mutex, so the pair describes the same on-device
+	// frame. Prefer this over calling Hierarchy and Screenshot separately:
+	// independent reads can land on different frames during transitions.
+	Snapshot(ctx context.Context) (string, Image, error)
 	// RecentLogs returns log entries at or after `since`, filtered to
 	// `minLevel` or above. An empty minLevel defaults to "E".
 	RecentLogs(ctx context.Context, since time.Time, minLevel string) ([]LogEntry, error)
diff --git a/internal/driver/sidecar/client.go b/internal/driver/sidecar/client.go
index f356c9a..272d365 100644
--- a/internal/driver/sidecar/client.go
+++ b/internal/driver/sidecar/client.go
@@ -162,6 +162,23 @@ func (c *Client) Screenshot(ctx context.Context) (driver.Image, error) {
 	}, nil
 }
 
+// Snapshot fetches hierarchy and screenshot in a single sidecar round-trip.
+// The sidecar serializes the two reads behind a mutex so the returned pair
+// describes the same on-device frame, removing the cross-fade race the
+// runner used to see when fetching them as independent goroutines.
+func (c *Client) Snapshot(ctx context.Context) (string, driver.Image, error) {
+	response, err := c.stub.Snapshot(ctx, &driverpb.Empty{})
+	if err != nil {
+		return "", driver.Image{}, err
+	}
+	image := response.GetScreenshot()
+	return response.GetHierarchy().GetJson(), driver.Image{
+		PNG:    image.GetPng(),
+		Width:  int(image.GetWidth()),
+		Height: int(image.GetHeight()),
+	}, nil
+}
+
 func (c *Client) WaitForIdle(ctx context.Context, duration time.Duration) error {
 	_, err := c.stub.WaitForIdle(ctx, &driverpb.Duration{Millis: duration.Milliseconds()})
 	return err