"""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): 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 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