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() {
requestFullFrames(2)
webrtcpub.RequestKeyframe()
RequestDesktopH264Keyframe()
}
@@ -11,6 +11,7 @@ import (
"overlord-client/cmd/agent/config"
rt "overlord-client/cmd/agent/runtime"
"overlord-client/cmd/agent/webrtcpub"
"overlord-client/cmd/agent/wire"
"github.com/vmihailenco/msgpack/v5"
@@ -476,3 +477,13 @@ func TestSetFrameFlowTargetFPSEnvOverride(t *testing.T) {
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"
)
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) {
codec := blockCodec()
if (codec != "h264" && codec != "hevc") || (codec == "h264" && useDesktopSoftwareH264()) || !useDesktopDuplication() {
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 width, height int
var captureDur, encodeDur time.Duration
@@ -31,8 +43,9 @@ func tryBuildDirectH264Frame(display int) (wire.Frame, time.Duration, time.Durat
if len(data) == 0 {
return wire.Frame{}, captureDur, encodeDur, true, nil
}
now := time.Now()
lastKeyframe.Store(now.UnixNano())
if forceKeyframe {
lastKeyframe.Store(now.UnixNano())
}
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
}
@@ -90,3 +90,28 @@ func TestCaptureDisplayDXGI(t *testing.T) {
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)
}
})
}
}
@@ -18,15 +18,13 @@ import (
)
var (
monitorCountFn = capture.MonitorCount
displayBoundsFn = capture.DisplayBounds
captureDisplayRGBABitBltFn = capture.CaptureDisplayRGBABitBlt
monitorCountFn = capture.MonitorCount
displayBoundsFn = capture.DisplayBounds
captureRectFn = screenshot.CaptureRect
)
func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplays bool) error {
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
restore := capture.BypassResolutionCap()
defer restore()
if allDisplays {
log.Printf("screenshot: capturing all displays")
@@ -86,8 +84,6 @@ func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplay
return err
}
capture.RequestDesktopH264Keyframe()
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{
Type: "command_result",
CommandID: cmdID,
@@ -157,7 +153,8 @@ func captureScreenshotImageWindows(allDisplays bool) (*image.RGBA, int, image.Re
}
if !allDisplays {
img, err := captureDisplayRGBABitBltFn(0)
bounds := displayBoundsFn(0)
img, err := captureRectFn(bounds)
if err != nil || img == nil {
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++ {
bounds := displayBoundsFn(i)
img, err := captureDisplayRGBABitBltFn(i)
img, err := captureRectFn(bounds)
if err != nil || img == nil {
return nil, 0, image.Rectangle{}, err
}
@@ -16,14 +16,17 @@ import (
func TestCaptureScreenshotImageWindows_PrimaryOnly(t *testing.T) {
origMonitorCountFn := monitorCountFn
origCaptureFn := captureDisplayRGBABitBltFn
origBoundsFn := displayBoundsFn
origCaptureFn := captureRectFn
t.Cleanup(func() {
monitorCountFn = origMonitorCountFn
captureDisplayRGBABitBltFn = origCaptureFn
displayBoundsFn = origBoundsFn
captureRectFn = origCaptureFn
})
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.SetRGBA(0, 0, color.RGBA{R: 255, A: 255})
return img, nil
@@ -47,11 +50,11 @@ func TestCaptureScreenshotImageWindows_PrimaryOnly(t *testing.T) {
func TestCaptureScreenshotImageWindows_AllDisplaysStitched(t *testing.T) {
origMonitorCountFn := monitorCountFn
origBoundsFn := displayBoundsFn
origCaptureFn := captureDisplayRGBABitBltFn
origCaptureFn := captureRectFn
t.Cleanup(func() {
monitorCountFn = origMonitorCountFn
displayBoundsFn = origBoundsFn
captureDisplayRGBABitBltFn = origCaptureFn
captureRectFn = origCaptureFn
})
monitorCountFn = func() int { return 2 }
@@ -61,9 +64,9 @@ func TestCaptureScreenshotImageWindows_AllDisplaysStitched(t *testing.T) {
}
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))
if display == 0 {
if bounds.Min.X == 0 {
img.SetRGBA(0, 0, color.RGBA{R: 255, A: 255})
} else {
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
origCaptureFn := captureDisplayRGBABitBltFn
origBoundsFn := displayBoundsFn
origCaptureFn := captureRectFn
t.Cleanup(func() {
monitorCountFn = origMonitorCountFn
captureDisplayRGBABitBltFn = origCaptureFn
displayBoundsFn = origBoundsFn
captureRectFn = origCaptureFn
})
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.SetRGBA(0, 0, color.RGBA{B: 255, A: 255})
return img, nil
@@ -129,26 +135,30 @@ func TestHandleScreenshot_DoesNotEmitLiveFrame(t *testing.T) {
ctx := context.Background()
writer := &testWriter{}
if err := HandleScreenshot(ctx, &rt.Env{Conn: writer}, "cmd-1", false); err != nil {
t.Fatalf("HandleScreenshot failed: %v", err)
for i, commandID := range []string{"cmd-1", "cmd-2"} {
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
if err := msgpack.Unmarshal(writer.msgs[0], &first); err != nil {
t.Fatalf("unmarshal screenshot_result: %v", err)
}
if first.Type != "screenshot_result" {
t.Fatalf("expected screenshot_result, got %q", first.Type)
}
for i := range 2 {
var screenshotResult wire.ScreenshotResult
if err := msgpack.Unmarshal(writer.msgs[i*2], &screenshotResult); err != nil {
t.Fatalf("unmarshal screenshot_result %d: %v", i+1, err)
}
if screenshotResult.Type != "screenshot_result" {
t.Fatalf("expected screenshot_result for request %d, got %q", i+1, screenshotResult.Type)
}
var second wire.CommandResult
if err := msgpack.Unmarshal(writer.msgs[1], &second); err != nil {
t.Fatalf("unmarshal command_result: %v", err)
}
if second.Type != "command_result" || !second.OK {
t.Fatalf("expected successful command_result, got %+v", second)
var commandResult wire.CommandResult
if err := msgpack.Unmarshal(writer.msgs[i*2+1], &commandResult); err != nil {
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)
}
}
}