feat(plugins): add source-aware evidence and recording timelines
This commit is contained in:
chenmingxuan committed
1 parent
72ff342605
commit
ef7af9988e
61 files changed
+3847
-108
No files matched your search
@@ -0,0 +1,391 @@
|
||||
"""Recorder lifecycle and archive limits use fake processes, never a live desktop."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from artemis.config.agent import RecordingConfig
|
||||
from artemis.plugins.models import EvidenceSource
|
||||
from artemis.plugins import recording
|
||||
from artemis.plugins.recording import PluginRecordings, build_capture_command, capture_rectangle
|
||||
|
||||
|
||||
def source(**updates):
|
||||
return EvidenceSource(
|
||||
**{
|
||||
"kind": "computer",
|
||||
"target_id": "123",
|
||||
"label": "Fixture window",
|
||||
"width": 800,
|
||||
"height": 600,
|
||||
"window": {"left": 40, "top": 50, "width": 800, "height": 600},
|
||||
"screen": {"left": -1920, "top": 0, "width": 1920, "height": 1080},
|
||||
**updates,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def context(root, *, enabled=True, **config):
|
||||
return SimpleNamespace(
|
||||
data_engine=SimpleNamespace(base_dir=root, session_start_time=100),
|
||||
agent_config=SimpleNamespace(recording=RecordingConfig(**config)),
|
||||
execution_setup=SimpleNamespace(
|
||||
record_video_enabled=enabled, video_recording_tools_enabled=False
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class FakeRecorderProcess:
|
||||
def __init__(self, *, produces_frames=True):
|
||||
self.pid = 123456
|
||||
self.returncode = None if produces_frames else 1
|
||||
self.stdout = asyncio.StreamReader()
|
||||
if produces_frames:
|
||||
self.stdout.feed_data(b"frame=1\n")
|
||||
self.stdout.feed_eof()
|
||||
self.stdin = SimpleNamespace(write=self._quit, drain=AsyncMock())
|
||||
self.killed = False
|
||||
|
||||
def _quit(self, _data):
|
||||
self.returncode = 0
|
||||
|
||||
async def wait(self):
|
||||
return self.returncode
|
||||
|
||||
def kill(self):
|
||||
self.killed = True
|
||||
self.returncode = -9
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_capture(monkeypatch):
|
||||
state = SimpleNamespace(commands=[], processes=[], fail_hevc=False)
|
||||
|
||||
async def spawn(*command, **_kwargs):
|
||||
state.commands.append(command)
|
||||
encoder = command[command.index("-c:v") + 1]
|
||||
process = FakeRecorderProcess(
|
||||
produces_frames=not (state.fail_hevc and encoder.startswith("hevc"))
|
||||
)
|
||||
state.processes.append(process)
|
||||
Path(command[-1]).write_bytes(b"fake mp4")
|
||||
return process
|
||||
|
||||
local_asyncio = SimpleNamespace(**vars(asyncio))
|
||||
local_asyncio.create_subprocess_exec = spawn
|
||||
monkeypatch.setattr(recording, "asyncio", local_asyncio)
|
||||
monkeypatch.setattr(recording, "os", SimpleNamespace(name="nt", getpid=lambda: 123))
|
||||
monkeypatch.setattr(
|
||||
recording,
|
||||
"CaptureProcessJob",
|
||||
lambda: SimpleNamespace(assign=lambda _pid: None, close=lambda: None),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
recording, "shutil", SimpleNamespace(disk_usage=lambda _path: SimpleNamespace(free=10**10))
|
||||
)
|
||||
monkeypatch.setattr(recording, "time", SimpleNamespace(time=lambda: 123.0))
|
||||
monkeypatch.setattr(recording, "get_ffmpeg_path", lambda: "fake-ffmpeg")
|
||||
monkeypatch.setattr(
|
||||
recording,
|
||||
"probe_video_segment",
|
||||
AsyncMock(return_value={"duration": 2.5, "width": 1280, "height": 720}),
|
||||
)
|
||||
return state
|
||||
|
||||
|
||||
def argument(command, option):
|
||||
return command[command.index(option) + 1]
|
||||
|
||||
|
||||
def test_archive_command_keeps_physical_capture_and_limits_encoded_output(tmp_path):
|
||||
command = build_capture_command(
|
||||
"ffmpeg", source(), RecordingConfig(), tmp_path / "video.mp4", "libx264"
|
||||
)
|
||||
assert argument(command, "-offset_x") == "-1920"
|
||||
assert argument(command, "-video_size") == "1920x1080"
|
||||
assert argument(command, "-vf") == "scale=1280:720"
|
||||
assert argument(command, "-framerate") == "15"
|
||||
assert argument(command, "-b:v") == "1500k"
|
||||
assert argument(command, "-g") == "30"
|
||||
assert "-an" in command
|
||||
|
||||
|
||||
def test_small_window_is_not_upscaled_and_hwnd_capture_is_explicit(tmp_path):
|
||||
command = build_capture_command(
|
||||
"ffmpeg",
|
||||
source(),
|
||||
RecordingConfig(computer_scope="window", codec="h265"),
|
||||
tmp_path / "video.mp4",
|
||||
"hevc_nvenc",
|
||||
)
|
||||
assert argument(command, "-i") == "hwnd=123"
|
||||
assert argument(command, "-vf") == "scale=800:600"
|
||||
assert argument(command, "-b:v") == "1000k"
|
||||
assert argument(command, "-tag:v") == "hvc1"
|
||||
assert "-offset_x" not in command
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"area",
|
||||
[
|
||||
None,
|
||||
{"left": 0},
|
||||
{"left": 0, "top": 0, "width": 0, "height": 600},
|
||||
{"left": 65537, "top": 0, "width": 800, "height": 600},
|
||||
],
|
||||
)
|
||||
def test_invalid_capture_geometry_is_not_sent_to_ffmpeg(area):
|
||||
with pytest.raises(ValueError):
|
||||
capture_rectangle(source(screen=area), "screen")
|
||||
|
||||
|
||||
def test_window_recording_rejects_non_hwnd_targets(tmp_path):
|
||||
with pytest.raises(ValueError, match="HWND"):
|
||||
build_capture_command(
|
||||
"ffmpeg",
|
||||
source(target_id="window-title"),
|
||||
RecordingConfig(computer_scope="window"),
|
||||
tmp_path / "video.mp4",
|
||||
"libx264",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disabled_recording_does_not_spawn_or_create_evidence(tmp_path, fake_capture):
|
||||
recorder = PluginRecordings(context(tmp_path, enabled=False))
|
||||
await recorder.observe("source-1", "visual-test", source())
|
||||
await recorder.stop()
|
||||
assert fake_capture.commands == []
|
||||
assert not (tmp_path / "evidence").exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"codec,hevc_playback,expected_attempts",
|
||||
[("auto", False, 1), ("auto", True, 4), ("h265", False, 4)],
|
||||
)
|
||||
async def test_hardware_hevc_failure_falls_back_to_compatible_archive(
|
||||
tmp_path, fake_capture, codec, hevc_playback, expected_attempts
|
||||
):
|
||||
fake_capture.fail_hevc = True
|
||||
recorder = PluginRecordings(context(tmp_path, codec=codec, hevc_playback=hevc_playback))
|
||||
try:
|
||||
await recorder.observe("source-1", "visual-test", source())
|
||||
assert len(fake_capture.commands) == expected_attempts
|
||||
assert argument(fake_capture.commands[-1], "-c:v") == "libx264"
|
||||
assert len(recorder.active) == 1
|
||||
finally:
|
||||
await recorder.stop()
|
||||
track = next(iter(recorder.sources.values()))
|
||||
assert track["status"] == "ready"
|
||||
assert len(track["segments"]) == 1
|
||||
assert track["segments"][0]["codec"] == "h264"
|
||||
assert track["segments"][0]["duration_ms"] == 2500
|
||||
assert all(process.returncode is not None for process in fake_capture.processes)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_screen_reuses_recording_and_keeps_both_target_ids(tmp_path, fake_capture):
|
||||
recorder = PluginRecordings(context(tmp_path))
|
||||
try:
|
||||
await recorder.observe("window-1", "visual-test", source())
|
||||
await recorder.observe("window-2", "visual-test", source(target_id="456"))
|
||||
assert len(fake_capture.commands) == 1
|
||||
finally:
|
||||
await recorder.stop()
|
||||
assert next(iter(recorder.sources.values()))["target_ids"] == ["window-1", "window-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_low_disk_space_records_failure_without_starting_capture(
|
||||
tmp_path, fake_capture, monkeypatch
|
||||
):
|
||||
monkeypatch.setattr(recording.shutil, "disk_usage", lambda _path: SimpleNamespace(free=1))
|
||||
recorder = PluginRecordings(context(tmp_path))
|
||||
await recorder.observe("source-1", "visual-test", source())
|
||||
await recorder.stop()
|
||||
assert fake_capture.commands == []
|
||||
track = next(iter(recorder.sources.values()))
|
||||
assert track["status"] == "failed"
|
||||
assert "磁盘" in track["message"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_index_write_failure_does_not_leave_capture_processes(
|
||||
tmp_path, fake_capture, monkeypatch
|
||||
):
|
||||
recorder = PluginRecordings(context(tmp_path))
|
||||
|
||||
def unavailable():
|
||||
raise OSError("disk full")
|
||||
|
||||
monkeypatch.setattr(recorder, "_save", unavailable)
|
||||
await recorder.observe("source-1", "visual-test", source())
|
||||
await recorder.stop()
|
||||
assert not recorder.active
|
||||
assert all(process.returncode is not None for process in fake_capture.processes)
|
||||
assert next(iter(recorder.sources.values()))["status"] == "failed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_probe_failure_during_stop_is_recorded_and_capture_is_reaped(
|
||||
tmp_path, fake_capture, monkeypatch
|
||||
):
|
||||
recorder = PluginRecordings(context(tmp_path))
|
||||
await recorder.observe("source-1", "visual-test", source())
|
||||
monkeypatch.setattr(
|
||||
recording, "probe_video_segment", AsyncMock(side_effect=OSError("probe unavailable"))
|
||||
)
|
||||
await recorder.stop()
|
||||
assert not recorder.active
|
||||
assert all(process.returncode is not None for process in fake_capture.processes)
|
||||
track = next(iter(recorder.sources.values()))
|
||||
assert track["status"] == "failed"
|
||||
assert "probe unavailable" in track["message"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_retries_process_cleanup_interrupted_in_watchdog(
|
||||
tmp_path, fake_capture, monkeypatch
|
||||
):
|
||||
recorder = PluginRecordings(context(tmp_path))
|
||||
await recorder.observe("source-1", "visual-test", source())
|
||||
recorder._watchdog.cancel()
|
||||
await asyncio.gather(recorder._watchdog, return_exceptions=True)
|
||||
key = next(iter(recorder.active))
|
||||
entered = asyncio.Event()
|
||||
attempts = 0
|
||||
terminate = recorder._terminate
|
||||
|
||||
async def interrupted_first(process):
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
if attempts == 1:
|
||||
entered.set()
|
||||
await asyncio.Future()
|
||||
await terminate(process)
|
||||
|
||||
monkeypatch.setattr(recorder, "_terminate", interrupted_first)
|
||||
recorder._watchdog = asyncio.create_task(recorder._finish(key))
|
||||
await entered.wait()
|
||||
await recorder.stop()
|
||||
assert attempts == 2
|
||||
assert not recorder.active
|
||||
assert all(process.returncode is not None for process in fake_capture.processes)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_preserves_clip_when_watchdog_is_cancelled_during_probe(
|
||||
tmp_path, fake_capture, monkeypatch
|
||||
):
|
||||
recorder = PluginRecordings(context(tmp_path))
|
||||
await recorder.observe("source-1", "visual-test", source())
|
||||
recorder._watchdog.cancel()
|
||||
await asyncio.gather(recorder._watchdog, return_exceptions=True)
|
||||
key = next(iter(recorder.active))
|
||||
entered = asyncio.Event()
|
||||
attempts = 0
|
||||
|
||||
async def interrupted_probe(_output):
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
if attempts == 1:
|
||||
entered.set()
|
||||
await asyncio.Future()
|
||||
return {"duration": 2.5, "width": 1280, "height": 720}
|
||||
|
||||
monkeypatch.setattr(recording, "probe_video_segment", interrupted_probe)
|
||||
recorder._watchdog = asyncio.create_task(recorder._finish(key))
|
||||
await entered.wait()
|
||||
await recorder.stop()
|
||||
assert attempts == 2
|
||||
assert not recorder.active
|
||||
track = recorder.sources[key]
|
||||
assert track["status"] == "ready"
|
||||
assert track["segments"][0]["duration_ms"] == 2500
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resumed_recorder_appends_and_does_not_overwrite_previous_clips(
|
||||
tmp_path, fake_capture
|
||||
):
|
||||
first = PluginRecordings(context(tmp_path))
|
||||
await first.observe("source-1", "visual-test", source())
|
||||
await first.stop()
|
||||
original = next(iter(first.sources.values()))["segments"][0].copy()
|
||||
original_path = tmp_path / original["file"]
|
||||
original_path.write_bytes(b"immutable original clip")
|
||||
second = PluginRecordings(context(tmp_path))
|
||||
try:
|
||||
await second.observe("source-1", "visual-test", source())
|
||||
finally:
|
||||
await second.stop()
|
||||
segments = next(iter(second.sources.values()))["segments"]
|
||||
assert len(segments) == 2
|
||||
assert segments[0] == original
|
||||
assert segments[1]["file"] != original["file"]
|
||||
assert original_path.read_bytes() == b"immutable original clip"
|
||||
manifest = json.loads(
|
||||
(tmp_path / "evidence/recordings/recordings.json").read_text(encoding="utf-8")
|
||||
)
|
||||
assert len(manifest["sources"][0]["segments"]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resized_window_seals_old_geometry_before_new_capture(tmp_path, fake_capture):
|
||||
recorder = PluginRecordings(context(tmp_path, computer_scope="window"))
|
||||
try:
|
||||
await recorder.observe("source-1", "visual-test", source())
|
||||
await recorder.observe(
|
||||
"source-2",
|
||||
"visual-test",
|
||||
source(window={"left": 40, "top": 50, "width": 1000, "height": 700}),
|
||||
)
|
||||
assert len(fake_capture.commands) == 2
|
||||
assert fake_capture.processes[0].returncode is not None
|
||||
assert len(recorder.active) == 1
|
||||
finally:
|
||||
await recorder.stop()
|
||||
assert len(recorder.sources) == 2
|
||||
assert all(track["status"] == "ready" for track in recorder.sources.values())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dead_worker_recovery_keeps_partial_clip_and_does_not_restart(
|
||||
tmp_path, fake_capture, monkeypatch
|
||||
):
|
||||
recorder = PluginRecordings(context(tmp_path))
|
||||
await recorder.observe("source-1", "visual-test", source())
|
||||
await recorder.stop()
|
||||
track = next(iter(recorder.sources.values()))
|
||||
track.update(status="recording", owner_pid=987654321, owner_created_at=1)
|
||||
track["segments"][0].update(duration=0, duration_ms=0)
|
||||
recorder._save()
|
||||
count = len(fake_capture.commands)
|
||||
assert await recording.recover_interrupted_recordings(tmp_path)
|
||||
recovered = PluginRecordings(context(tmp_path))
|
||||
result = next(iter(recovered.sources.values()))
|
||||
assert result["status"] == "failed"
|
||||
assert result["segments"][0]["duration"] == 2.5
|
||||
assert len(fake_capture.commands) == count
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_worker_recovery_never_finalizes_active_capture(tmp_path, fake_capture):
|
||||
import os
|
||||
import psutil
|
||||
|
||||
recorder = PluginRecordings(context(tmp_path))
|
||||
await recorder.observe("source-1", "visual-test", source())
|
||||
track = next(iter(recorder.sources.values()))
|
||||
track.update(owner_pid=os.getpid(), owner_created_at=psutil.Process().create_time())
|
||||
recorder._save()
|
||||
try:
|
||||
assert not await recording.recover_interrupted_recordings(tmp_path)
|
||||
assert track["status"] == "recording"
|
||||
finally:
|
||||
await recorder.stop()
|
||||
Reference in new issue
Block a user