From 16cd67e68040d52a647108c5d11e14502be471de Mon Sep 17 00:00:00 2001
From: Romil Bijarnia <romilbij@gmail.com>
Date: Wed, 11 Nov 2026 19:41:18 +1100
Subject: [PATCH] fix(server): capture device screenshots through their
 selected host

---
 apps/server/src/device/DeviceHubProxy.test.ts | 70 +++++++++++++++++--
 apps/server/src/device/DeviceHubProxy.ts      |  5 +-
 2 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/apps/server/src/device/DeviceHubProxy.test.ts b/apps/server/src/device/DeviceHubProxy.test.ts
index 72c8f9e0e5d56..3daa305bdfb79 100644
--- a/apps/server/src/device/DeviceHubProxy.test.ts
+++ b/apps/server/src/device/DeviceHubProxy.test.ts
@@ -27,7 +27,10 @@ const fixture = (
 ) => {
   let finalized = 0;
   const requests: string[] = [];
+  const screenshots: Array<Parameters<DeviceService.DeviceService["Service"]["screenshot"]>[0]> =
+    [];
   const headers: Array<Readonly<Record<string, string>>> = [];
+  const unexpectedOperation = Effect.die("Unexpected device operation");
   const client = HttpClient.make((request, _url, signal) =>
     Effect.gen(function* () {
       requests.push(request.url);
@@ -45,7 +48,7 @@ const fixture = (
   const { handler, dispose } = HttpRouter.toWebHandler(
     DeviceHubProxy.layer.pipe(
       Layer.provideMerge(
-        Layer.succeed(EnvironmentAuth.EnvironmentAuth, {
+        Layer.mock(EnvironmentAuth.EnvironmentAuth)({
           authenticateWebSocketUpgrade: () =>
             authError
               ? Effect.fail(authError)
@@ -55,24 +58,50 @@ const fixture = (
                   method: "bearer-access-token",
                   scopes,
                 }),
-        } as unknown as EnvironmentAuth.EnvironmentAuth["Service"]),
+        }),
       ),
       Layer.provideMerge(
-        Layer.succeed(DeviceService.DeviceService, {
-          currentReadiness: () =>
+        Layer.mock(DeviceService.DeviceService)({
+          currentReadiness: (hostId) =>
             Effect.succeed({
-              hostId: LOCAL_DEVICE_HOST_ID,
+              hostId: hostId ?? LOCAL_DEVICE_HOST_ID,
+              nodePath: "node",
               hub: { origin: "http://hub.test" },
+              devices: {
+                list: unexpectedOperation,
+                boot: () => unexpectedOperation,
+                attach: () => unexpectedOperation,
+                shutdown: () => unexpectedOperation,
+                screenshot: () => unexpectedOperation,
+              },
+              run: () => unexpectedOperation,
+              helpers: { serveSimAxSettings: null, serveSimCli: null },
               ...(upstream ? { upstream } : {}),
             }),
-        } as DeviceService.DeviceService["Service"]),
+          screenshot: (input) =>
+            Effect.sync(() => {
+              screenshots.push(input);
+              return {
+                device: {
+                  hostId: input.hostId ?? LOCAL_DEVICE_HOST_ID,
+                  id: input.deviceId,
+                  platform: "ios",
+                  name: "Phone",
+                  version: "27.0",
+                  booted: true,
+                  physical: false,
+                },
+                png: new Uint8Array([137, 80, 78, 71]),
+              };
+            }),
+        }),
       ),
       Layer.provideMerge(Layer.succeed(HttpClient.HttpClient, client)),
     ),
     { disableLogger: true },
   );
   disposers.push(dispose);
-  return { handler, requests, headers, finalized: () => finalized };
+  return { handler, requests, screenshots, headers, finalized: () => finalized };
 };
 
 describe("device hub proxy", () => {
@@ -239,6 +268,33 @@ describe("device hub proxy", () => {
     ).toBe(404);
     expect(requests).toEqual([]);
   });
+
+  it.each(["serve-sim", "serve-emu"])(
+    "captures %s screenshots through hosts with their own stream routing",
+    async (vendor) => {
+      const { handler, requests, screenshots } = fixture(
+        [AuthOrchestrationReadScope],
+        false,
+        undefined,
+        () => Effect.succeed(null),
+      );
+      const response = await handler(
+        new Request(
+          `http://t3.test/api/device-hub/vendor/${vendor}/api/screenshot?device=udid-1&hostId=pool`,
+          {
+            method: "POST",
+          },
+        ),
+      );
+      expect(response.status).toBe(200);
+      expect(response.headers.get("content-type")).toBe("image/png");
+      expect(new Uint8Array(await response.arrayBuffer())).toEqual(
+        new Uint8Array([137, 80, 78, 71]),
+      );
+      expect(requests).toEqual([]);
+      expect(screenshots).toEqual([{ hostId: "pool", deviceId: "udid-1" }]);
+    },
+  );
 });
 
 it.each([
diff --git a/apps/server/src/device/DeviceHubProxy.ts b/apps/server/src/device/DeviceHubProxy.ts
index ff7d292564fba..56132433c8869 100644
--- a/apps/server/src/device/DeviceHubProxy.ts
+++ b/apps/server/src/device/DeviceHubProxy.ts
@@ -214,7 +214,10 @@ const handler = Effect.gen(function* () {
   upstreamSearch.delete("hostId");
   const search = upstreamSearch.size > 0 ? `?${upstreamSearch.toString()}` : "";
   const screenshotDevice = upstreamSearch.get("device");
-  if (!ready.upstream && hubPath === "/vendor/serve-sim/api/screenshot" && screenshotDevice) {
+  if (
+    screenshotDevice &&
+    ["/vendor/serve-sim/api/screenshot", "/vendor/serve-emu/api/screenshot"].includes(hubPath)
+  ) {
     const { png } = yield* devices.screenshot({ hostId: ready.hostId, deviceId: screenshotDevice });
     return HttpServerResponse.uint8Array(png, {
       contentType: "image/png",