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, {