516 lines
21 KiB
Python
516 lines
21 KiB
Python
"""Bounded Windows plugin recording, with one task-local manifest and no scheduler."""
|
||
|
||
import asyncio
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import os
|
||
from pathlib import Path
|
||
import shutil
|
||
import sys
|
||
import time
|
||
import uuid
|
||
|
||
import psutil
|
||
|
||
from artemis.config.agent import RecordingConfig
|
||
from artemis.plugins.models import EvidenceSource
|
||
from artemis.plugins.recording_process import CaptureProcessJob
|
||
from artemis.utils.video import get_ffmpeg_path, probe_video_segment
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
async def recover_interrupted_recordings(root: Path) -> bool:
|
||
"""Publish usable fragments after their owning worker died; never restart capture."""
|
||
manifest = root / "evidence" / "recordings" / "recordings.json"
|
||
if not manifest.is_file() or manifest.resolve() != manifest.absolute():
|
||
return False
|
||
try:
|
||
if manifest.stat().st_size > 4 * 1024 * 1024:
|
||
return False
|
||
data = json.loads(manifest.read_text(encoding="utf-8"))
|
||
if data.get("version") != 1 or not isinstance(data.get("sources"), list):
|
||
return False
|
||
changed = False
|
||
async with asyncio.timeout(30):
|
||
for track in data["sources"]:
|
||
if not isinstance(track, dict) or track.get("status") != "recording":
|
||
continue
|
||
try:
|
||
owner = psutil.Process(int(track["owner_pid"]))
|
||
if abs(owner.create_time() - float(track["owner_created_at"])) < 0.01:
|
||
continue
|
||
except psutil.AccessDenied:
|
||
continue # An inaccessible worker is not proof it has exited.
|
||
except (psutil.NoSuchProcess, KeyError, ValueError, TypeError):
|
||
pass
|
||
for segment in track.get("segments", []):
|
||
video = root / str(segment.get("file") or "")
|
||
if (
|
||
not video.is_file()
|
||
or video.resolve() != video.absolute()
|
||
or not video.resolve().is_relative_to(root.resolve())
|
||
):
|
||
continue
|
||
try:
|
||
metadata = await asyncio.wait_for(probe_video_segment(video), 10)
|
||
duration = max(0.0, float(metadata.get("duration") or 0))
|
||
segment.update(
|
||
duration=duration,
|
||
duration_ms=round(duration * 1000),
|
||
width=int(metadata.get("width") or 0),
|
||
height=int(metadata.get("height") or 0),
|
||
bytes=video.stat().st_size,
|
||
)
|
||
except (OSError, ValueError, RuntimeError, TimeoutError):
|
||
segment.update(duration=0, duration_ms=0)
|
||
track.update(
|
||
status="failed",
|
||
message="录制工作进程已退出；已保留可读取的片段，末尾可能缺失。",
|
||
)
|
||
changed = True
|
||
if changed:
|
||
temporary = manifest.with_suffix(".recovery.tmp")
|
||
temporary.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
|
||
temporary.replace(manifest)
|
||
return changed
|
||
except (OSError, ValueError, TypeError, KeyError, TimeoutError):
|
||
logger.warning("无法恢复电脑录像索引，原始片段已保留")
|
||
return False
|
||
|
||
|
||
def recording_config(ctx) -> RecordingConfig:
|
||
value = getattr(getattr(ctx, "agent_config", None), "recording", None)
|
||
return value if isinstance(value, RecordingConfig) else RecordingConfig()
|
||
|
||
|
||
def capture_rectangle(source: EvidenceSource, scope: str) -> dict[str, int]:
|
||
value = source.window if scope == "window" else source.screen
|
||
if not value or set(value) != {"left", "top", "width", "height"}:
|
||
raise ValueError("插件未提供完整的物理屏幕/窗口范围")
|
||
if not 1 <= value["width"] <= 32768 or not 1 <= value["height"] <= 32768:
|
||
raise ValueError("插件返回了无效的录制范围")
|
||
if any(abs(value[key]) > 65536 for key in ("left", "top")):
|
||
raise ValueError("插件返回了无效的屏幕原点")
|
||
return value
|
||
|
||
|
||
def build_capture_command(
|
||
ffmpeg: str, source: EvidenceSource, config: RecordingConfig, output: Path, encoder: str
|
||
) -> list[str]:
|
||
area = capture_rectangle(source, config.computer_scope)
|
||
if config.computer_scope == "window":
|
||
# GDI/BitBlt can return black for GPU-rendered Chromium windows.
|
||
# Frames come from an isolated Windows Graphics Capture worker instead.
|
||
if not source.target_id.isdecimal():
|
||
raise ValueError("窗口录制需要真实 HWND")
|
||
capture = []
|
||
else:
|
||
capture = [
|
||
"-offset_x",
|
||
str(area["left"]),
|
||
"-offset_y",
|
||
str(area["top"]),
|
||
"-video_size",
|
||
f"{area['width']}x{area['height']}",
|
||
"-i",
|
||
"desktop",
|
||
]
|
||
hevc = encoder.startswith("hevc")
|
||
bitrate = config.h265_kbps if hevc else config.h264_kbps
|
||
# Never upscale. Even dimensions are required by yuv420p encoders.
|
||
factor = min(1.0, config.max_long_edge / max(area["width"], area["height"]))
|
||
width = max(2, int(area["width"] * factor) // 2 * 2)
|
||
height = max(2, int(area["height"] * factor) // 2 * 2)
|
||
input_args = ["-f", "gdigrab", "-framerate", str(config.fps), "-draw_mouse", "1", *capture]
|
||
if config.computer_scope == "window":
|
||
input_args = [
|
||
"-f",
|
||
"rawvideo",
|
||
"-pixel_format",
|
||
"bgra",
|
||
"-video_size",
|
||
f"{width}x{height}",
|
||
"-framerate",
|
||
str(config.fps),
|
||
"-use_wallclock_as_timestamps",
|
||
"1",
|
||
"-i",
|
||
"pipe:0",
|
||
]
|
||
command = [
|
||
ffmpeg,
|
||
"-hide_banner",
|
||
"-loglevel",
|
||
"error",
|
||
"-y",
|
||
*input_args,
|
||
"-an",
|
||
"-vf",
|
||
f"scale={width}:{height}",
|
||
"-c:v",
|
||
encoder,
|
||
"-pix_fmt",
|
||
"yuv420p",
|
||
"-b:v",
|
||
f"{bitrate}k",
|
||
"-maxrate",
|
||
f"{bitrate * 2}k",
|
||
"-bufsize",
|
||
f"{bitrate * 2}k",
|
||
"-g",
|
||
str(config.fps * 2),
|
||
]
|
||
if encoder == "libx264":
|
||
# Live capture must produce a first encoded frame before readiness.
|
||
# Lookahead/thread buffering can exceed that deadline on slow desktops.
|
||
command += ["-preset", "veryfast", "-tune", "zerolatency"]
|
||
if hevc:
|
||
command += ["-tag:v", "hvc1"]
|
||
command += [
|
||
"-movflags",
|
||
"+frag_keyframe+empty_moov+default_base_moof",
|
||
"-progress",
|
||
"pipe:1",
|
||
"-stats_period",
|
||
"0.5",
|
||
str(output),
|
||
]
|
||
if config.computer_scope == "window":
|
||
return [
|
||
sys.executable,
|
||
"-m",
|
||
"artemis.plugins.window_capture",
|
||
"--hwnd",
|
||
source.target_id,
|
||
"--width",
|
||
str(width),
|
||
"--height",
|
||
str(height),
|
||
"--fps",
|
||
str(config.fps),
|
||
"--",
|
||
*command,
|
||
]
|
||
return command
|
||
|
||
|
||
class PluginRecordings:
|
||
def __init__(self, ctx, *, folder: Path | None = None, force: bool = False):
|
||
self.ctx = ctx
|
||
self.config = recording_config(ctx)
|
||
self.root = Path(ctx.data_engine.base_dir)
|
||
self.folder = folder or self.root / "evidence" / "recordings"
|
||
if not self.folder.resolve().is_relative_to(self.root.resolve()):
|
||
raise ValueError("录像目录必须属于当前任务")
|
||
self.force = force
|
||
self.sources: dict[str, dict] = {}
|
||
self.active: dict[str, tuple] = {}
|
||
self.failed: set[str] = set()
|
||
self._lock = asyncio.Lock()
|
||
self._watchdog = None
|
||
self._job = None
|
||
manifest = self.folder / "recordings.json"
|
||
if manifest.is_file():
|
||
# Resumed workers append; prior finalized segments are immutable.
|
||
data = json.loads(manifest.read_text(encoding="utf-8"))
|
||
if data.get("version") != 1 or not isinstance(data.get("sources"), list):
|
||
raise ValueError("已有电脑录像索引无效，已保留原文件")
|
||
if any(
|
||
not isinstance(item, dict)
|
||
or not isinstance(item.get("id"), str)
|
||
or not isinstance(item.get("segments"), list)
|
||
for item in data["sources"]
|
||
):
|
||
raise ValueError("已有电脑录像来源无效，已保留原文件")
|
||
self.sources = {item["id"]: item for item in data["sources"]}
|
||
|
||
@property
|
||
def enabled(self):
|
||
setup = getattr(self.ctx, "execution_setup", None)
|
||
return self.force or bool(setup and setup.record_video_enabled)
|
||
|
||
def _save(self):
|
||
self.folder.mkdir(parents=True, exist_ok=True)
|
||
target = self.folder / "recordings.json"
|
||
temporary = target.with_suffix(".tmp")
|
||
temporary.write_text(
|
||
json.dumps({"version": 1, "sources": list(self.sources.values())}, ensure_ascii=False),
|
||
encoding="utf-8",
|
||
)
|
||
temporary.replace(target)
|
||
|
||
async def observe(
|
||
self, source_id: str, plugin_id: str, source: EvidenceSource, *, exclusive: bool = False
|
||
):
|
||
if not self.enabled:
|
||
return
|
||
async with self._lock:
|
||
try:
|
||
area = capture_rectangle(source, self.config.computer_scope)
|
||
key = (
|
||
source.target_id + json.dumps(area, sort_keys=True)
|
||
if self.config.computer_scope == "window"
|
||
else json.dumps(area, sort_keys=True)
|
||
)
|
||
track_id = (
|
||
"computer-" + hashlib.sha256(f"{plugin_id}:{key}".encode()).hexdigest()[:16]
|
||
)
|
||
track = self.sources.setdefault(
|
||
track_id,
|
||
{
|
||
"id": track_id,
|
||
"kind": "computer",
|
||
"plugin_id": plugin_id,
|
||
"label": f"电脑 · {source.label or source.target_id}"
|
||
if self.config.computer_scope == "window"
|
||
else f"电脑屏幕 · {area['left']},{area['top']}",
|
||
"scope": self.config.computer_scope,
|
||
"area": area,
|
||
"target_ids": [],
|
||
"status": "unavailable",
|
||
"message": "",
|
||
"segments": [],
|
||
},
|
||
)
|
||
if source_id not in track["target_ids"]:
|
||
track["target_ids"].append(source_id)
|
||
if self.config.computer_scope == "window":
|
||
for other in list(self.active):
|
||
previous = self.sources[other]
|
||
if other != track_id and (
|
||
exclusive or previous.get("target_id") == source.target_id
|
||
):
|
||
await self._finish(other)
|
||
track["target_id"] = source.target_id
|
||
if track_id in self.active or track_id in self.failed:
|
||
self._save()
|
||
return
|
||
if os.name != "nt":
|
||
raise OSError("当前电脑录制仅支持 Windows")
|
||
self.folder.mkdir(parents=True, exist_ok=True)
|
||
if shutil.disk_usage(self.folder).free < self.config.minimum_disk_mb * 1024 * 1024:
|
||
raise OSError("磁盘可用空间不足，已停止录像")
|
||
await self._start(track_id, source)
|
||
if self._watchdog is None or self._watchdog.done():
|
||
self._watchdog = asyncio.create_task(self._watch())
|
||
except (OSError, ValueError, RuntimeError, TimeoutError) as exc:
|
||
# A recording failure never changes the result of a completed tool action.
|
||
failed_id = locals().get("track_id", source_id)
|
||
if failed_id not in self.failed:
|
||
logger.warning("电脑录像不可用：%s", exc)
|
||
self.failed.add(failed_id)
|
||
track = self.sources.setdefault(
|
||
failed_id,
|
||
{
|
||
"id": failed_id,
|
||
"kind": "computer",
|
||
"plugin_id": plugin_id,
|
||
"label": "电脑录像",
|
||
"target_ids": [source_id],
|
||
"segments": [],
|
||
},
|
||
)
|
||
track.update(status="failed", message=str(exc))
|
||
try:
|
||
self._save()
|
||
except OSError:
|
||
logger.warning("电脑录像状态无法写入任务目录")
|
||
|
||
async def _start(self, key, source):
|
||
ffmpeg = get_ffmpeg_path()
|
||
if not ffmpeg:
|
||
raise OSError("未找到 FFmpeg，无法录制电脑画面")
|
||
encoders = []
|
||
if self.config.codec == "h265" or (
|
||
self.config.codec == "auto" and self.config.hevc_playback
|
||
):
|
||
encoders += ["hevc_nvenc", "hevc_qsv", "hevc_amf"]
|
||
encoders += ["libx264"]
|
||
track = self.sources[key]
|
||
errors = []
|
||
if self._job is None:
|
||
self._job = CaptureProcessJob()
|
||
for encoder in encoders:
|
||
output = self.folder / f"{key}-{uuid.uuid4().hex[:12]}.mp4"
|
||
log_path = output.with_suffix(".log")
|
||
spawned = time.time()
|
||
with log_path.open("wb") as error_log:
|
||
process = await asyncio.create_subprocess_exec(
|
||
*build_capture_command(ffmpeg, source, self.config, output, encoder),
|
||
stdin=asyncio.subprocess.PIPE,
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=error_log,
|
||
creationflags=0x08000000,
|
||
)
|
||
try:
|
||
self._job.assign(process.pid)
|
||
if self.config.computer_scope == "window":
|
||
process.stdin.write(b"start\n")
|
||
await process.stdin.drain()
|
||
first_capture = spawned
|
||
async with asyncio.timeout(12):
|
||
while line := await process.stdout.readline():
|
||
if line.startswith(b"capture_started_at="):
|
||
first_capture = float(line.partition(b"=")[2])
|
||
if line.startswith(b"frame=") and int(line.partition(b"=")[2]) > 0:
|
||
break
|
||
else:
|
||
raise OSError(
|
||
log_path.read_text(errors="replace")[-1000:] or "录制进程未产生画面"
|
||
)
|
||
segment = {
|
||
"file": output.relative_to(self.root).as_posix(),
|
||
"offset_ms": max(
|
||
0, round((first_capture - self.ctx.data_engine.session_start_time) * 1000)
|
||
),
|
||
"duration": 0,
|
||
"duration_ms": 0,
|
||
"start": 0,
|
||
"codec": "h265" if encoder.startswith("hevc") else "h264",
|
||
"encoder": encoder,
|
||
"capture_backend": "windows_graphics_capture"
|
||
if self.config.computer_scope == "window"
|
||
else "gdigrab",
|
||
}
|
||
track["segments"].append(segment)
|
||
track.update(
|
||
status="recording",
|
||
codec=segment["codec"],
|
||
message="；".join(errors),
|
||
owner_pid=os.getpid(),
|
||
owner_created_at=psutil.Process().create_time(),
|
||
)
|
||
self.active[key] = (
|
||
process,
|
||
output,
|
||
segment,
|
||
asyncio.create_task(self._drain(process.stdout, segment)),
|
||
)
|
||
try:
|
||
self._save()
|
||
except OSError:
|
||
await self._finish(key, "录像索引写入失败")
|
||
# Retain the completed partial clip rather than deleting evidence.
|
||
raise
|
||
return
|
||
except (OSError, ValueError, TimeoutError) as exc:
|
||
await self._terminate(process)
|
||
if not any(
|
||
item["file"] == output.relative_to(self.root).as_posix()
|
||
for item in track["segments"]
|
||
):
|
||
output.unlink(missing_ok=True)
|
||
else:
|
||
raise
|
||
detail = "等待首帧超时" if isinstance(exc, TimeoutError) else str(exc)[:160]
|
||
errors.append(f"{encoder} 不可用，已尝试兼容编码：{detail}")
|
||
except asyncio.CancelledError:
|
||
await self._terminate(process)
|
||
raise
|
||
raise OSError("；".join(errors))
|
||
|
||
@staticmethod
|
||
async def _drain(stream, segment):
|
||
while line := await stream.readline():
|
||
if line.startswith(b"out_time_us="):
|
||
try:
|
||
duration = max(0, int(line.partition(b"=")[2])) / 1_000_000
|
||
except ValueError:
|
||
continue
|
||
segment.update(duration=duration, duration_ms=round(duration * 1000))
|
||
|
||
@staticmethod
|
||
async def _terminate(process):
|
||
if process.returncode is None:
|
||
try:
|
||
process.stdin.write(b"q\n")
|
||
await process.stdin.drain()
|
||
await asyncio.wait_for(process.wait(), 5)
|
||
except (OSError, ConnectionError, TimeoutError):
|
||
if process.returncode is None:
|
||
try:
|
||
process.kill()
|
||
except ProcessLookupError:
|
||
pass
|
||
await asyncio.wait_for(process.wait(), 5)
|
||
|
||
async def _finish(self, key, error=""):
|
||
process, output, segment, drain = self.active[key]
|
||
track = self.sources[key]
|
||
interrupted = False
|
||
try:
|
||
await self._terminate(process)
|
||
if process.returncode:
|
||
detail = output.with_suffix(".log").read_text(errors="replace")[-1000:]
|
||
error = error or f"录像进程异常退出 ({process.returncode})：{detail}"
|
||
await asyncio.wait_for(drain, 2)
|
||
metadata = await asyncio.wait_for(probe_video_segment(output), 10)
|
||
duration = float(metadata.get("duration") or 0)
|
||
segment.update(
|
||
duration=duration,
|
||
duration_ms=round(duration * 1000),
|
||
width=int(metadata.get("width") or 0),
|
||
height=int(metadata.get("height") or 0),
|
||
)
|
||
if duration > 0:
|
||
segment["bytes"] = output.stat().st_size
|
||
else:
|
||
error = error or "未取得有效的录像片段"
|
||
except asyncio.CancelledError:
|
||
interrupted = True
|
||
raise
|
||
except (OSError, ValueError, RuntimeError, TimeoutError) as exc:
|
||
error = error or f"录像收尾失败：{exc}"
|
||
finally:
|
||
# Cancellation while finalizing keeps ownership for stop() to retry.
|
||
if process.returncode is not None and not interrupted:
|
||
self.active.pop(key, None)
|
||
drain.cancel()
|
||
await asyncio.gather(drain, return_exceptions=True)
|
||
track.update(status="failed" if error else "ready", message=error or track["message"])
|
||
if error:
|
||
self.failed.add(key)
|
||
logger.warning("电脑录像中断：%s", error)
|
||
|
||
async def _watch(self):
|
||
while self.active:
|
||
await asyncio.sleep(5)
|
||
async with self._lock:
|
||
for key, (process, _, _, _) in list(self.active.items()):
|
||
reason = "录制进程意外退出" if process.returncode is not None else ""
|
||
try:
|
||
if (
|
||
shutil.disk_usage(self.folder).free
|
||
< self.config.minimum_disk_mb * 1024 * 1024
|
||
):
|
||
reason = "磁盘可用空间不足"
|
||
except OSError:
|
||
reason = "录像目录不可用"
|
||
if reason:
|
||
await self._finish(key, reason)
|
||
self._save_safely()
|
||
|
||
def _save_safely(self):
|
||
try:
|
||
self._save()
|
||
except OSError:
|
||
logger.warning("电脑录像状态无法写入任务目录")
|
||
|
||
async def stop(self):
|
||
if self._watchdog is not None:
|
||
self._watchdog.cancel()
|
||
await asyncio.gather(self._watchdog, return_exceptions=True)
|
||
self._watchdog = None
|
||
async with self._lock:
|
||
try:
|
||
for key in list(self.active):
|
||
await self._finish(key)
|
||
if self.sources:
|
||
self._save_safely()
|
||
finally:
|
||
if self._job is not None:
|
||
self._job.close()
|
||
self._job = None
|