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

This commit is contained in:
KingKDot
2026-07-19 02:33:30 -05:00
parent 1a68b5ddf1
commit 698bcf681e
14 changed files with 360 additions and 64 deletions
+48 -4
View File
@@ -68,6 +68,7 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
const streamProfileSelect = document.getElementById("streamProfileSelect");
const streamProfileDetail = document.getElementById("streamProfileDetail");
const bitrateSelect = document.getElementById("bitrateSelect");
const requestKeyframeBtn = document.getElementById("requestKeyframeBtn");
const smoothingSlider = document.getElementById("smoothingSlider");
const smoothingValue = document.getElementById("smoothingValue");
const qualitySlider = document.getElementById("qualitySlider");
@@ -143,6 +144,7 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
let h264KeyframeErrorStreak = 0;
let h264RecoveryAttempts = 0;
let h264LastDecodeWarnAt = 0;
let lastAutomaticKeyframeAt = 0;
const H264_LOW_FPS_THRESHOLD = 6;
const H264_FALLBACK_WARMUP_MS = 10000;
const H264_MIN_FRAMES_BEFORE_FALLBACK = 120;
@@ -150,6 +152,8 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
const H264_KEYFRAME_ERROR_RESTART_THRESHOLD = 24;
const H264_MAX_RECOVERY_ATTEMPTS = 1;
const H264_DECODE_WARN_THROTTLE_MS = 2000;
const VIDEO_FRAME_GAP_KEYFRAME_MS = 500;
const AUTOMATIC_KEYFRAME_COOLDOWN_MS = 2000;
const inputBackpressureBytes = 256 * 1024;
const diagnostics = {
agent: null,
@@ -621,6 +625,9 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
if (stopBtn) {
stopBtn.disabled = !wsOpen || (!isStarting && !isStreaming && !isStopping && !isStalled);
}
if (requestKeyframeBtn) {
requestKeyframeBtn.disabled = !wsOpen || (!isStreaming && !isStalled);
}
if (recordBtn) {
recordBtn.disabled =
!recordingActive &&
@@ -1733,6 +1740,13 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
sharedSettingsSaver.scheduleSave();
});
}
if (requestKeyframeBtn) {
requestKeyframeBtn.addEventListener("click", function () {
if (streamState !== "streaming" && streamState !== "stalled") return;
sendCmd("desktop_request_keyframe", { reason: "manual_viewer" });
});
}
function needsWebrtcFirewallWarning(mode) {
if (firewallWarningAcked) return false;
@@ -2012,8 +2026,32 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
return buf.length >= 8 && buf[0] === 0x46 && buf[1] === 0x52 && buf[2] === 0x4d;
}
function requestAutomaticKeyframe(reason, details = {}) {
if (!desiredStreaming || getWebrtcMode() !== "off") return false;
if (streamState !== "streaming" && streamState !== "stalled" && streamState !== "starting") {
return false;
}
const codec = String(diagnostics.codec || negotiatedCodec || "").toLowerCase();
if (codec !== "h264" && codec !== "hevc") return false;
if (!ws || ws.readyState !== WebSocket.OPEN) return false;
const now = performance.now();
if (lastAutomaticKeyframeAt && now - lastAutomaticKeyframeAt < AUTOMATIC_KEYFRAME_COOLDOWN_MS) {
return false;
}
lastAutomaticKeyframeAt = now;
rdDebug("automatic keyframe request", { reason, codec, ...details });
sendCmd("desktop_request_keyframe", { reason, codec });
return true;
}
function markFrameReceived() {
lastFrameAt = performance.now();
const now = performance.now();
const frameGapMs = lastFrameAt ? now - lastFrameAt : 0;
lastFrameAt = now;
if (frameGapMs >= VIDEO_FRAME_GAP_KEYFRAME_MS) {
requestAutomaticKeyframe("viewer_frame_gap", { frameGapMs: Math.round(frameGapMs) });
}
clearOfflineTimer();
if (!firstFrameLogged) {
firstFrameLogged = true;
@@ -2512,17 +2550,23 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
h264KeyframeErrorStreak = 0;
} catch (err) {
h264PendingTimings.pop();
if (codecName === "h264" && isKeyframeRequiredError(err)) {
if (isKeyframeRequiredError(err)) {
h264KeyframeErrorStreak += 1;
requestAutomaticKeyframe(`${codecName}_decoder_keyframe_required`, {
streak: h264KeyframeErrorStreak,
});
const now = Date.now();
if (now - h264LastDecodeWarnAt >= H264_DECODE_WARN_THROTTLE_MS) {
h264LastDecodeWarnAt = now;
console.warn("rd: h264 decode waiting for keyframe", {
console.warn(`rd: ${codecName} decode waiting for keyframe`, {
streak: h264KeyframeErrorStreak,
recoveries: h264RecoveryAttempts,
});
}
if (h264KeyframeErrorStreak >= H264_KEYFRAME_ERROR_RESTART_THRESHOLD) {
if (
codecName === "h264" &&
h264KeyframeErrorStreak >= H264_KEYFRAME_ERROR_RESTART_THRESHOLD
) {
const restarted = tryRecoverH264Stream("h264_keyframe_required");
if (!restarted) {
fallbackToJpegCodec("h264_keyframe_required_loop");
@@ -219,6 +219,15 @@
<label class="mt-2 flex items-center gap-2 text-sm text-slate-200" title="Use GPU desktop duplication (Windows only)">
<input type="checkbox" id="duplicationCtrl" /> GPU Capture
</label>
<button
id="requestKeyframeBtn"
type="button"
class="button ghost w-full justify-center mt-3"
title="Request a fresh video keyframe from the remote client"
disabled
>
<i class="fa-solid fa-film"></i><span>Request Keyframe</span>
</button>
</div>
<div class="px-4 py-3 border-b border-slate-800">
@@ -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);
@@ -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);
@@ -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, {
+39 -1
View File
@@ -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;
}
});
});
+36 -17
View File
@@ -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();
@@ -64,6 +64,7 @@ test("WebRTC input stays inactive until start and then forwards mouse and keyboa
});
await expect(page.locator("#streamProfileSelect option[value='1440:60']")).toHaveCount(1);
await expect(page.locator("#streamProfileSelect option[value='2160:60']")).toHaveCount(1);
await expect(page.locator("#requestKeyframeBtn")).toBeDisabled();
await page.evaluate(() => {
const sent = (window as typeof window & { __rdSent: ArrayBuffer[] }).__rdSent;
@@ -198,4 +199,37 @@ test("WebRTC input stays inactive until start and then forwards mouse and keyboa
expect.objectContaining({ type: "key_down", key: "Enter", code: "Enter" }),
expect.objectContaining({ type: "key_up", key: "Enter", code: "Enter" }),
]));
await page.evaluate(() => {
const socket = (window as typeof window & { __rdSocket: EventTarget }).__rdSocket;
const framePacket = Uint8Array.from([0x46, 0x52, 0x4d, 0, 0, 0, 0, 0]).buffer;
socket.dispatchEvent(new MessageEvent("message", { data: framePacket }));
});
await expect(page.locator("#requestKeyframeBtn")).toBeEnabled();
await page.waitForTimeout(550);
await page.evaluate(() => {
const sent = (window as typeof window & { __rdSent: ArrayBuffer[] }).__rdSent;
sent.length = 0;
const socket = (window as typeof window & { __rdSocket: EventTarget }).__rdSocket;
const framePacket = Uint8Array.from([0x46, 0x52, 0x4d, 0, 0, 0, 0, 0]).buffer;
socket.dispatchEvent(new MessageEvent("message", { data: framePacket }));
});
await expect.poll(async () => page.evaluate(async () => {
const sent = (window as typeof window & { __rdSent: ArrayBuffer[] }).__rdSent;
const { decodeMsgpack } = await import("/assets/msgpack-helpers.js");
return sent.map((message) => decodeMsgpack(message))
.filter((message) => message.type === "desktop_request_keyframe");
})).toEqual([
expect.objectContaining({ type: "desktop_request_keyframe", reason: "viewer_frame_gap" }),
]);
await page.locator("#rdSettingsBtn").click();
await page.locator("#requestKeyframeBtn").click();
await expect.poll(async () => page.evaluate(async () => {
const sent = (window as typeof window & { __rdSent: ArrayBuffer[] }).__rdSent;
const { decodeMsgpack } = await import("/assets/msgpack-helpers.js");
return sent.map((message) => decodeMsgpack(message))
.filter((message) => message.type === "desktop_request_keyframe" && message.reason === "manual_viewer");
})).toEqual([
expect.objectContaining({ type: "desktop_request_keyframe", reason: "manual_viewer" }),
]);
});