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
@@ -645,6 +645,7 @@ func ResetPrevbackstage() {
func RequestDesktopFullFrame() { func RequestDesktopFullFrame() {
requestFullFrames(2) requestFullFrames(2)
webrtcpub.RequestKeyframe()
RequestDesktopH264Keyframe() RequestDesktopH264Keyframe()
} }
@@ -11,6 +11,7 @@ import (
"overlord-client/cmd/agent/config" "overlord-client/cmd/agent/config"
rt "overlord-client/cmd/agent/runtime" rt "overlord-client/cmd/agent/runtime"
"overlord-client/cmd/agent/webrtcpub"
"overlord-client/cmd/agent/wire" "overlord-client/cmd/agent/wire"
"github.com/vmihailenco/msgpack/v5" "github.com/vmihailenco/msgpack/v5"
@@ -476,3 +477,13 @@ func TestSetFrameFlowTargetFPSEnvOverride(t *testing.T) {
t.Fatalf("expected env slot limit 12, got %d", got) t.Fatalf("expected env slot limit 12, got %d", got)
} }
} }
func TestRequestDesktopFullFrameRequestsWebRTCKeyframe(t *testing.T) {
_ = webrtcpub.ConsumeKeyframeRequest()
RequestDesktopFullFrame()
if !webrtcpub.ConsumeKeyframeRequest() {
t.Fatal("expected desktop full-frame request to force the next WebRTC video frame to be a keyframe")
}
}
@@ -9,12 +9,24 @@ import (
"overlord-client/cmd/agent/wire" "overlord-client/cmd/agent/wire"
) )
func directVideoKeyframeDue(requested bool, nowNs, lastKeyframeNs int64) bool {
if requested || lastKeyframeNs == 0 {
return true
}
return time.Duration(nowNs-lastKeyframeNs) >= keyframeEvery
}
func tryBuildDirectH264Frame(display int) (wire.Frame, time.Duration, time.Duration, bool, error) { func tryBuildDirectH264Frame(display int) (wire.Frame, time.Duration, time.Duration, bool, error) {
codec := blockCodec() codec := blockCodec()
if (codec != "h264" && codec != "hevc") || (codec == "h264" && useDesktopSoftwareH264()) || !useDesktopDuplication() { if (codec != "h264" && codec != "hevc") || (codec == "h264" && useDesktopSoftwareH264()) || !useDesktopDuplication() {
return wire.Frame{}, 0, 0, false, nil return wire.Frame{}, 0, 0, false, nil
} }
forceKeyframe := webrtcpub.ConsumeKeyframeRequest() now := time.Now()
forceKeyframe := directVideoKeyframeDue(
webrtcpub.ConsumeKeyframeRequest(),
now.UnixNano(),
lastKeyframe.Load(),
)
var data []byte var data []byte
var width, height int var width, height int
var captureDur, encodeDur time.Duration var captureDur, encodeDur time.Duration
@@ -31,8 +43,9 @@ func tryBuildDirectH264Frame(display int) (wire.Frame, time.Duration, time.Durat
if len(data) == 0 { if len(data) == 0 {
return wire.Frame{}, captureDur, encodeDur, true, nil return wire.Frame{}, captureDur, encodeDur, true, nil
} }
now := time.Now() if forceKeyframe {
lastKeyframe.Store(now.UnixNano()) lastKeyframe.Store(now.UnixNano())
}
statFullFrames.Add(1) statFullFrames.Add(1)
return wire.Frame{Type: "frame", Header: wire.FrameHeader{Monitor: display, FPS: 0, Format: codec, Width: width, Height: height}, Data: data}, captureDur, encodeDur, true, nil return wire.Frame{Type: "frame", Header: wire.FrameHeader{Monitor: display, FPS: 0, Format: codec, Width: width, Height: height}, Data: data}, captureDur, encodeDur, true, nil
} }
@@ -90,3 +90,28 @@ func TestCaptureDisplayDXGI(t *testing.T) {
t.Fatalf("invalid capture bounds: %dx%d", bounds.Dx(), bounds.Dy()) t.Fatalf("invalid capture bounds: %dx%d", bounds.Dx(), bounds.Dy())
} }
} }
func TestDirectVideoKeyframeDue(t *testing.T) {
now := (10 * time.Second).Nanoseconds()
last := now - keyframeEvery.Nanoseconds()
tests := []struct {
name string
requested bool
last int64
want bool
}{
{name: "explicit request", requested: true, last: now, want: true},
{name: "first frame", last: 0, want: true},
{name: "before periodic interval", last: last + 1, want: false},
{name: "at periodic interval", last: last, want: true},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if got := directVideoKeyframeDue(test.requested, now, test.last); got != test.want {
t.Fatalf("directVideoKeyframeDue() = %v, want %v", got, test.want)
}
})
}
}
@@ -20,13 +20,11 @@ import (
var ( var (
monitorCountFn = capture.MonitorCount monitorCountFn = capture.MonitorCount
displayBoundsFn = capture.DisplayBounds displayBoundsFn = capture.DisplayBounds
captureDisplayRGBABitBltFn = capture.CaptureDisplayRGBABitBlt captureRectFn = screenshot.CaptureRect
) )
func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplays bool) error { func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplays bool) error {
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2 //garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
restore := capture.BypassResolutionCap()
defer restore()
if allDisplays { if allDisplays {
log.Printf("screenshot: capturing all displays") log.Printf("screenshot: capturing all displays")
@@ -86,8 +84,6 @@ func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplay
return err return err
} }
capture.RequestDesktopH264Keyframe()
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{ return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{
Type: "command_result", Type: "command_result",
CommandID: cmdID, CommandID: cmdID,
@@ -157,7 +153,8 @@ func captureScreenshotImageWindows(allDisplays bool) (*image.RGBA, int, image.Re
} }
if !allDisplays { if !allDisplays {
img, err := captureDisplayRGBABitBltFn(0) bounds := displayBoundsFn(0)
img, err := captureRectFn(bounds)
if err != nil || img == nil { if err != nil || img == nil {
return nil, 0, image.Rectangle{}, err return nil, 0, image.Rectangle{}, err
} }
@@ -174,7 +171,7 @@ func captureScreenshotImageWindows(allDisplays bool) (*image.RGBA, int, image.Re
for i := 0; i < monitorCount; i++ { for i := 0; i < monitorCount; i++ {
bounds := displayBoundsFn(i) bounds := displayBoundsFn(i)
img, err := captureDisplayRGBABitBltFn(i) img, err := captureRectFn(bounds)
if err != nil || img == nil { if err != nil || img == nil {
return nil, 0, image.Rectangle{}, err return nil, 0, image.Rectangle{}, err
} }
@@ -16,14 +16,17 @@ import (
func TestCaptureScreenshotImageWindows_PrimaryOnly(t *testing.T) { func TestCaptureScreenshotImageWindows_PrimaryOnly(t *testing.T) {
origMonitorCountFn := monitorCountFn origMonitorCountFn := monitorCountFn
origCaptureFn := captureDisplayRGBABitBltFn origBoundsFn := displayBoundsFn
origCaptureFn := captureRectFn
t.Cleanup(func() { t.Cleanup(func() {
monitorCountFn = origMonitorCountFn monitorCountFn = origMonitorCountFn
captureDisplayRGBABitBltFn = origCaptureFn displayBoundsFn = origBoundsFn
captureRectFn = origCaptureFn
}) })
monitorCountFn = func() int { return 2 } monitorCountFn = func() int { return 2 }
captureDisplayRGBABitBltFn = func(display int) (*image.RGBA, error) { displayBoundsFn = func(int) image.Rectangle { return image.Rect(0, 0, 3, 2) }
captureRectFn = func(image.Rectangle) (*image.RGBA, error) {
img := image.NewRGBA(image.Rect(0, 0, 3, 2)) img := image.NewRGBA(image.Rect(0, 0, 3, 2))
img.SetRGBA(0, 0, color.RGBA{R: 255, A: 255}) img.SetRGBA(0, 0, color.RGBA{R: 255, A: 255})
return img, nil return img, nil
@@ -47,11 +50,11 @@ func TestCaptureScreenshotImageWindows_PrimaryOnly(t *testing.T) {
func TestCaptureScreenshotImageWindows_AllDisplaysStitched(t *testing.T) { func TestCaptureScreenshotImageWindows_AllDisplaysStitched(t *testing.T) {
origMonitorCountFn := monitorCountFn origMonitorCountFn := monitorCountFn
origBoundsFn := displayBoundsFn origBoundsFn := displayBoundsFn
origCaptureFn := captureDisplayRGBABitBltFn origCaptureFn := captureRectFn
t.Cleanup(func() { t.Cleanup(func() {
monitorCountFn = origMonitorCountFn monitorCountFn = origMonitorCountFn
displayBoundsFn = origBoundsFn displayBoundsFn = origBoundsFn
captureDisplayRGBABitBltFn = origCaptureFn captureRectFn = origCaptureFn
}) })
monitorCountFn = func() int { return 2 } monitorCountFn = func() int { return 2 }
@@ -61,9 +64,9 @@ func TestCaptureScreenshotImageWindows_AllDisplaysStitched(t *testing.T) {
} }
return image.Rect(2, 0, 4, 2) return image.Rect(2, 0, 4, 2)
} }
captureDisplayRGBABitBltFn = func(display int) (*image.RGBA, error) { captureRectFn = func(bounds image.Rectangle) (*image.RGBA, error) {
img := image.NewRGBA(image.Rect(0, 0, 2, 2)) img := image.NewRGBA(image.Rect(0, 0, 2, 2))
if display == 0 { if bounds.Min.X == 0 {
img.SetRGBA(0, 0, color.RGBA{R: 255, A: 255}) img.SetRGBA(0, 0, color.RGBA{R: 255, A: 255})
} else { } else {
img.SetRGBA(0, 0, color.RGBA{G: 255, A: 255}) img.SetRGBA(0, 0, color.RGBA{G: 255, A: 255})
@@ -112,16 +115,19 @@ func TestCaptureScreenshotImageWindows_NoDisplays(t *testing.T) {
} }
} }
func TestHandleScreenshot_DoesNotEmitLiveFrame(t *testing.T) { func TestHandleScreenshot_RepeatedRequestsDoNotEmitLiveFrames(t *testing.T) {
origMonitorCountFn := monitorCountFn origMonitorCountFn := monitorCountFn
origCaptureFn := captureDisplayRGBABitBltFn origBoundsFn := displayBoundsFn
origCaptureFn := captureRectFn
t.Cleanup(func() { t.Cleanup(func() {
monitorCountFn = origMonitorCountFn monitorCountFn = origMonitorCountFn
captureDisplayRGBABitBltFn = origCaptureFn displayBoundsFn = origBoundsFn
captureRectFn = origCaptureFn
}) })
monitorCountFn = func() int { return 1 } monitorCountFn = func() int { return 1 }
captureDisplayRGBABitBltFn = func(display int) (*image.RGBA, error) { displayBoundsFn = func(int) image.Rectangle { return image.Rect(0, 0, 2, 2) }
captureRectFn = func(image.Rectangle) (*image.RGBA, error) {
img := image.NewRGBA(image.Rect(0, 0, 2, 2)) img := image.NewRGBA(image.Rect(0, 0, 2, 2))
img.SetRGBA(0, 0, color.RGBA{B: 255, A: 255}) img.SetRGBA(0, 0, color.RGBA{B: 255, A: 255})
return img, nil return img, nil
@@ -129,26 +135,30 @@ func TestHandleScreenshot_DoesNotEmitLiveFrame(t *testing.T) {
ctx := context.Background() ctx := context.Background()
writer := &testWriter{} writer := &testWriter{}
if err := HandleScreenshot(ctx, &rt.Env{Conn: writer}, "cmd-1", false); err != nil { for i, commandID := range []string{"cmd-1", "cmd-2"} {
t.Fatalf("HandleScreenshot failed: %v", err) if err := HandleScreenshot(ctx, &rt.Env{Conn: writer}, commandID, false); err != nil {
t.Fatalf("HandleScreenshot request %d failed: %v", i+1, err)
} }
if len(writer.msgs) != 2 { }
t.Fatalf("expected 2 messages without live frame emission, got %d", len(writer.msgs)) if len(writer.msgs) != 4 {
t.Fatalf("expected 4 result messages without live frame emission, got %d", len(writer.msgs))
} }
var first wire.ScreenshotResult for i := range 2 {
if err := msgpack.Unmarshal(writer.msgs[0], &first); err != nil { var screenshotResult wire.ScreenshotResult
t.Fatalf("unmarshal screenshot_result: %v", err) if err := msgpack.Unmarshal(writer.msgs[i*2], &screenshotResult); err != nil {
t.Fatalf("unmarshal screenshot_result %d: %v", i+1, err)
} }
if first.Type != "screenshot_result" { if screenshotResult.Type != "screenshot_result" {
t.Fatalf("expected screenshot_result, got %q", first.Type) t.Fatalf("expected screenshot_result for request %d, got %q", i+1, screenshotResult.Type)
} }
var second wire.CommandResult var commandResult wire.CommandResult
if err := msgpack.Unmarshal(writer.msgs[1], &second); err != nil { if err := msgpack.Unmarshal(writer.msgs[i*2+1], &commandResult); err != nil {
t.Fatalf("unmarshal command_result: %v", err) t.Fatalf("unmarshal command_result %d: %v", i+1, err)
}
if commandResult.Type != "command_result" || !commandResult.OK {
t.Fatalf("expected successful command_result for request %d, got %+v", i+1, commandResult)
} }
if second.Type != "command_result" || !second.OK {
t.Fatalf("expected successful command_result, got %+v", second)
} }
} }
+48 -4
View File
@@ -68,6 +68,7 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
const streamProfileSelect = document.getElementById("streamProfileSelect"); const streamProfileSelect = document.getElementById("streamProfileSelect");
const streamProfileDetail = document.getElementById("streamProfileDetail"); const streamProfileDetail = document.getElementById("streamProfileDetail");
const bitrateSelect = document.getElementById("bitrateSelect"); const bitrateSelect = document.getElementById("bitrateSelect");
const requestKeyframeBtn = document.getElementById("requestKeyframeBtn");
const smoothingSlider = document.getElementById("smoothingSlider"); const smoothingSlider = document.getElementById("smoothingSlider");
const smoothingValue = document.getElementById("smoothingValue"); const smoothingValue = document.getElementById("smoothingValue");
const qualitySlider = document.getElementById("qualitySlider"); const qualitySlider = document.getElementById("qualitySlider");
@@ -143,6 +144,7 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
let h264KeyframeErrorStreak = 0; let h264KeyframeErrorStreak = 0;
let h264RecoveryAttempts = 0; let h264RecoveryAttempts = 0;
let h264LastDecodeWarnAt = 0; let h264LastDecodeWarnAt = 0;
let lastAutomaticKeyframeAt = 0;
const H264_LOW_FPS_THRESHOLD = 6; const H264_LOW_FPS_THRESHOLD = 6;
const H264_FALLBACK_WARMUP_MS = 10000; const H264_FALLBACK_WARMUP_MS = 10000;
const H264_MIN_FRAMES_BEFORE_FALLBACK = 120; 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_KEYFRAME_ERROR_RESTART_THRESHOLD = 24;
const H264_MAX_RECOVERY_ATTEMPTS = 1; const H264_MAX_RECOVERY_ATTEMPTS = 1;
const H264_DECODE_WARN_THROTTLE_MS = 2000; 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 inputBackpressureBytes = 256 * 1024;
const diagnostics = { const diagnostics = {
agent: null, agent: null,
@@ -621,6 +625,9 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
if (stopBtn) { if (stopBtn) {
stopBtn.disabled = !wsOpen || (!isStarting && !isStreaming && !isStopping && !isStalled); stopBtn.disabled = !wsOpen || (!isStarting && !isStreaming && !isStopping && !isStalled);
} }
if (requestKeyframeBtn) {
requestKeyframeBtn.disabled = !wsOpen || (!isStreaming && !isStalled);
}
if (recordBtn) { if (recordBtn) {
recordBtn.disabled = recordBtn.disabled =
!recordingActive && !recordingActive &&
@@ -1733,6 +1740,13 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
sharedSettingsSaver.scheduleSave(); sharedSettingsSaver.scheduleSave();
}); });
} }
if (requestKeyframeBtn) {
requestKeyframeBtn.addEventListener("click", function () {
if (streamState !== "streaming" && streamState !== "stalled") return;
sendCmd("desktop_request_keyframe", { reason: "manual_viewer" });
});
}
function needsWebrtcFirewallWarning(mode) { function needsWebrtcFirewallWarning(mode) {
if (firewallWarningAcked) return false; 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; 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() { 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(); clearOfflineTimer();
if (!firstFrameLogged) { if (!firstFrameLogged) {
firstFrameLogged = true; firstFrameLogged = true;
@@ -2512,17 +2550,23 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s
h264KeyframeErrorStreak = 0; h264KeyframeErrorStreak = 0;
} catch (err) { } catch (err) {
h264PendingTimings.pop(); h264PendingTimings.pop();
if (codecName === "h264" && isKeyframeRequiredError(err)) { if (isKeyframeRequiredError(err)) {
h264KeyframeErrorStreak += 1; h264KeyframeErrorStreak += 1;
requestAutomaticKeyframe(`${codecName}_decoder_keyframe_required`, {
streak: h264KeyframeErrorStreak,
});
const now = Date.now(); const now = Date.now();
if (now - h264LastDecodeWarnAt >= H264_DECODE_WARN_THROTTLE_MS) { if (now - h264LastDecodeWarnAt >= H264_DECODE_WARN_THROTTLE_MS) {
h264LastDecodeWarnAt = now; h264LastDecodeWarnAt = now;
console.warn("rd: h264 decode waiting for keyframe", { console.warn(`rd: ${codecName} decode waiting for keyframe`, {
streak: h264KeyframeErrorStreak, streak: h264KeyframeErrorStreak,
recoveries: h264RecoveryAttempts, recoveries: h264RecoveryAttempts,
}); });
} }
if (h264KeyframeErrorStreak >= H264_KEYFRAME_ERROR_RESTART_THRESHOLD) { if (
codecName === "h264" &&
h264KeyframeErrorStreak >= H264_KEYFRAME_ERROR_RESTART_THRESHOLD
) {
const restarted = tryRecoverH264Stream("h264_keyframe_required"); const restarted = tryRecoverH264Stream("h264_keyframe_required");
if (!restarted) { if (!restarted) {
fallbackToJpegCodec("h264_keyframe_required_loop"); 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)"> <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 <input type="checkbox" id="duplicationCtrl" /> GPU Capture
</label> </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>
<div class="px-4 py-3 border-b border-slate-800"> <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 * as sessionManager from "../../sessions/sessionManager";
import type { SocketData } from "../../sessions/types"; import type { SocketData } from "../../sessions/types";
import type { ClientInfo } from "../../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 { queueClientDbUpdate, scheduleQueuedClientDbFlush } from "../../client-db-sync";
import { getMaxPayloadLimit, getMessageByteLength, isAllowedClientMessageType } from "../../wsValidation"; import { getMaxPayloadLimit, getMessageByteLength, isAllowedClientMessageType } from "../../wsValidation";
import { stopAllProxiesForClient } from "../socks5-proxy-manager"; import { stopAllProxiesForClient } from "../socks5-proxy-manager";
import { verifyBuildToken, isBuildBanned } from "../build-signing"; import { verifyBuildToken, isBuildBanned } from "../build-signing";
import { clearThumbnail } from "../../thumbnails"; import { clearThumbnail } from "../../thumbnails";
import { stopRemoteDesktopRecording } from "../rd-recording"; import { stopRemoteDesktopRecording } from "../rd-recording";
import { requestRemoteDesktopKeyframeAfterScreenshot } from "../ws-console-rd-backstage";
import { import {
isAuthenticatedViewerRole, isAuthenticatedViewerRole,
registerViewerSocket, 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 {} try { ws.send(encodeMessage({ type: "frame_ack" })); } catch {}
} }
break; break;
case "screenshot_result": case "screenshot_result":
handleScreenshotThumbnailResult(client, payload); handleScreenshotThumbnailResult(client, payload);
deps.handleNotificationScreenshotResult(client.id, payload); deps.handleNotificationScreenshotResult(client.id, payload);
requestRemoteDesktopKeyframeAfterScreenshot(client.id);
break; break;
case "console_output": case "console_output":
deps.handleConsoleOutput(client.id, payload); deps.handleConsoleOutput(client.id, payload);
@@ -14,6 +14,7 @@ import {
handleDesktopEncoderCapabilities, handleDesktopEncoderCapabilities,
backstageStreamingState, backstageStreamingState,
rdStreamingState, rdStreamingState,
requestRemoteDesktopKeyframeAfterScreenshot,
} from "./ws-console-rd-backstage"; } from "./ws-console-rd-backstage";
type MockWs = { type MockWs = {
@@ -118,6 +119,72 @@ describe("remote desktop viewer control", () => {
expect(rdStreamingState.get(clientId)?.isStreaming).toBe(false); 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", () => { test("does not forward desktop_start when a macOS client is missing required permissions", () => {
const clientId = `rd-mac-perms-${Date.now().toString(36)}`; const clientId = `rd-mac-perms-${Date.now().toString(36)}`;
const { info, agentWs } = createClient(clientId); const { info, agentWs } = createClient(clientId);
@@ -178,6 +178,11 @@ function broadcastFrameToViewers(
const rdSendStats = { lastLog: 0, frames: 0, sendMs: 0, bytes: 0 }; const rdSendStats = { lastLog: 0, frames: 0, sendMs: 0, bytes: 0 };
const rdDebugFrameLogAt = new Map<string, number>(); const rdDebugFrameLogAt = new Map<string, number>();
const rdCanvasFrameAckPending = new Map<string, string>(); 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, { export const rdStreamingState = new Map<string, {
isStreaming: boolean; isStreaming: boolean;
display: number; 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) { export function handleRemoteDesktopViewerMessage(ws: ServerWebSocket<SocketData>, raw: string | ArrayBuffer | Uint8Array) {
const payload = decodeViewerPayload(raw); const payload = decodeViewerPayload(raw);
if (!payload) return; if (!payload) return;
@@ -489,6 +506,15 @@ export function handleRemoteDesktopViewerMessage(ws: ServerWebSocket<SocketData>
safeSendViewer(ws, { type: "status", status: "stopped" }); safeSendViewer(ws, { type: "status", status: "stopped" });
break; 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": { case "desktop_record_start": {
if (!state.isStreaming) { if (!state.isStreaming) {
safeSendViewer(ws, { safeSendViewer(ws, {
+39 -1
View File
@@ -1,7 +1,8 @@
import { describe, expect, test } from "bun:test"; import { describe, expect, test } from "bun:test";
import type { ClientInfo } from "./types";
import { decodeMessage } from "./protocol"; import { decodeMessage } from "./protocol";
import { metrics } from "./metrics"; 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 }; type MockWs = { sent: Uint8Array[]; send: (msg: Uint8Array) => void };
@@ -178,3 +179,40 @@ describe("wsHandlers ping/pong", () => {
expect(snapshot.ping.count).toBe(0); 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;
}
});
});
+34 -15
View File
@@ -273,30 +273,49 @@ export function handleScreenshotThumbnailResult(info: ClientInfo, payload: any):
return true; return true;
} }
export function handleFrame(info: ClientInfo, payload: any): boolean { export function shouldRelayFrameToViewers(payload: unknown): boolean {
const bytes = payload.data as unknown as Uint8Array; if (!payload || typeof payload !== "object" || !("header" in payload)) return false;
const header = (payload as any).header; 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 allowedFormats = ["jpeg", "jpg", "webp"];
const fmt = String(header?.format || "").toLowerCase(); const fmt = String(header.format || "").toLowerCase();
const safeFormat = allowedFormats.includes(fmt) ? fmt : ""; const safeFormat = allowedFormats.includes(fmt) ? fmt : "";
metrics.recordBytesReceived(bytes.length); metrics.recordBytesReceived(bytes.length);
let handledByViewerRelay = false; let handledByViewerRelay = false;
if (relayToViewers) {
try { try {
const globalAny: any = globalThis as any; const runtime = globalThis as unknown as {
if (header?.webcam) { __webcamBroadcast?: (clientId: string, data: Uint8Array, frameHeader: Record<string, unknown>) => boolean;
if (globalAny.__webcamBroadcast) { __backstageBroadcast?: (clientId: string, data: Uint8Array, frameHeader: Record<string, unknown>) => boolean;
handledByViewerRelay = globalAny.__webcamBroadcast(info.id, bytes, header); __rdBroadcast?: (clientId: string, data: Uint8Array, frameHeader: Record<string, unknown>) => boolean;
} };
} else if (header?.backstage) { if (header.webcam) {
if (globalAny.__backstageBroadcast) { handledByViewerRelay = runtime.__webcamBroadcast?.(info.id, bytes, header) ?? false;
handledByViewerRelay = globalAny.__backstageBroadcast(info.id, bytes, header); } else if (header.backstage) {
} handledByViewerRelay = runtime.__backstageBroadcast?.(info.id, bytes, header) ?? false;
} else if (globalAny.__rdBroadcast) { } else {
handledByViewerRelay = globalAny.__rdBroadcast(info.id, bytes, header); handledByViewerRelay = runtime.__rdBroadcast?.(info.id, bytes, header) ?? false;
} }
} catch {} } catch {}
}
if (safeFormat) { if (safeFormat) {
const now = Date.now(); 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='1440:60']")).toHaveCount(1);
await expect(page.locator("#streamProfileSelect option[value='2160:60']")).toHaveCount(1); await expect(page.locator("#streamProfileSelect option[value='2160:60']")).toHaveCount(1);
await expect(page.locator("#requestKeyframeBtn")).toBeDisabled();
await page.evaluate(() => { await page.evaluate(() => {
const sent = (window as typeof window & { __rdSent: ArrayBuffer[] }).__rdSent; 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_down", key: "Enter", code: "Enter" }),
expect.objectContaining({ type: "key_up", 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" }),
]);
}); });