diff --git a/.agents/skills/ok-script-pr-review/SKILL.md b/.agents/skills/ok-script-pr-review/SKILL.md index 960382f..ad75e7c 100644 --- a/.agents/skills/ok-script-pr-review/SKILL.md +++ b/.agents/skills/ok-script-pr-review/SKILL.md @@ -145,6 +145,8 @@ description: 处理 ok-script-toolkit 主仓及 JetBrains 子仓的 PR 审阅意 `6` 仍需等待（`-Once` 不是超时）、`4` PR 关闭而停止等待、`11` head 变化且本轮数据丢弃、`2` 错误。 等待须在用户能从 Codex 侧边查看输出的终端会话后台运行，复用已有进程；不要默认隐藏 Start-Process，也不要为等待另建定时任务、heartbeat 或 automation，除非用户明确要求安排。 +终端中的等待截止时间和下一次额度查询时间跟随电脑当前时区，显示完整日期、UTC 偏移与剩余时长； +等待已停止时也显示建议的查询时间。它不是审阅预计完成时间。JSON 的 ISO 时间格式保持不变。 本机账本默认在 `%LOCALAPPDATA%\ok-script-pr-review\coderabbit-triggers`，按仓库、PR 与 head 记录 触发，写入先于发送；`-StateDir` 可改位置。修改脚本后运行 `test-coderabbit-helpers.ps1` 与 `test-coderabbit-wait-mock.ps1`、`test-review-threads-mock.ps1`（Windows PowerShell 5.1 与 PowerShell 7 都应通过）。含非 ASCII diff --git a/.agents/skills/ok-script-pr-review/coderabbit-review-helpers.ps1 b/.agents/skills/ok-script-pr-review/coderabbit-review-helpers.ps1 index de67323..a78d5d0 100644 --- a/.agents/skills/ok-script-pr-review/coderabbit-review-helpers.ps1 +++ b/.agents/skills/ok-script-pr-review/coderabbit-review-helpers.ps1 @@ -15,6 +15,12 @@ function Get-CrNow { return [datetimeoffset]::UtcNow } +# Human-readable times follow the computer's current timezone; protocol timestamps stay in UTC. +function Format-CrDisplayTime { + param([datetimeoffset]$Time, [TimeZoneInfo]$TimeZone = [TimeZoneInfo]::Local) + return [TimeZoneInfo]::ConvertTime($Time, $TimeZone).ToString('yyyy-MM-dd HH:mm:ss zzz', [Globalization.CultureInfo]::InvariantCulture) +} + function Wait-CrSeconds { param([double]$Seconds) if ($Seconds -le 0) { return } diff --git a/.agents/skills/ok-script-pr-review/test-coderabbit-helpers.ps1 b/.agents/skills/ok-script-pr-review/test-coderabbit-helpers.ps1 index 72e9cff..ef5792f 100644 --- a/.agents/skills/ok-script-pr-review/test-coderabbit-helpers.ps1 +++ b/.agents/skills/ok-script-pr-review/test-coderabbit-helpers.ps1 @@ -17,6 +17,18 @@ function Assert-Equal { if ($Actual -ne $Expected) { throw "$Message -> got '$Actual', expected '$Expected'" } } +# ---- display timezone: follow the host, including date rollover and fractional offsets ---- +$displayInstant = [datetimeoffset]'2026-10-08T18:04:05Z' +foreach ($case in @( + @{ Minutes = 0; Want = '2026-10-08 18:04:05 +00:00' }, + @{ Minutes = 480; Want = '2026-10-09 02:04:05 +08:00' }, + @{ Minutes = -300; Want = '2026-10-08 13:04:05 -05:00' }, + @{ Minutes = 330; Want = '2026-10-08 23:34:05 +05:30' })) { + $zone = [TimeZoneInfo]::CreateCustomTimeZone("test-$($case.Minutes)", [TimeSpan]::FromMinutes($case.Minutes), 'test', 'test') + Assert-Equal (Format-CrDisplayTime $displayInstant $zone) $case.Want 'display uses the requested timezone' +} +Assert-Equal (Format-CrDisplayTime $displayInstant) ($displayInstant.ToLocalTime().ToString('yyyy-MM-dd HH:mm:ss zzz', [Globalization.CultureInfo]::InvariantCulture)) 'default display follows host timezone' + # ---- quota replies: a countdown is unavailable; only an explicit positive reply is available ---- $plan = '' + "`n" + 'Your [plan](https://docs.coderabbit.ai/management/plans#fair-usage-limits-policy) includes PR reviews subject to [rate limits](https://docs.coderabbit.ai/management/plans#rate-limits).' diff --git a/.agents/skills/ok-script-pr-review/test-coderabbit-wait-mock.ps1 b/.agents/skills/ok-script-pr-review/test-coderabbit-wait-mock.ps1 index c5d2ef5..f72aa1a 100644 --- a/.agents/skills/ok-script-pr-review/test-coderabbit-wait-mock.ps1 +++ b/.agents/skills/ok-script-pr-review/test-coderabbit-wait-mock.ps1 @@ -129,8 +129,8 @@ function Invoke-Wait([hashtable]$Extra = @{}) { function Invoke-Quota([hashtable]$Extra = @{}) { $params = @{ Repo = $repo; PrNumber = $pr; ExpectedHead = $global:W.head; PollSeconds = 10; ReplyTimeoutSeconds = 60 } foreach ($key in $Extra.Keys) { $params[$key] = $Extra[$key] } - $lines = @(& $quotaScript @params 6>$null) - return [pscustomobject]@{ exit = $LASTEXITCODE; result = ([string]$lines[-1] | ConvertFrom-Json) } + $lines = @(& $quotaScript @params 6>&1) + return [pscustomobject]@{ exit = $LASTEXITCODE; result = ([string]$lines[-1] | ConvertFrom-Json); log = ($lines | ForEach-Object { "$_" }) -join "`n" } } function Test-Case([string]$Name, [scriptblock]$Body) { try { & $Body; $script:passed++; Write-Output "ok $Name" } @@ -151,6 +151,26 @@ function Reply-OnProbe([string[]]$Replies) { } } +Test-Case 'quota wait displays local date, offset and time remaining' { + New-World + Reply-OnProbe @('More reviews will be available in 2 minutes.', 'Reviews are available now.') + $r = Invoke-Quota @{ MaxWaitSeconds = 300 } + $expected = $t0.AddSeconds(150).ToLocalTime().ToString('yyyy-MM-dd HH:mm:ss zzz', [Globalization.CultureInfo]::InvariantCulture) + Assert-True ($r.log.Contains("next query at $expected (local time; in 00:02:30)")) 'next query must use the local timezone with a countdown' + Assert-True ($r.log.Contains('review completion time is unknown')) 'do not present quota recheck as review completion' + Assert-True ($r.result.state -eq 'AVAILABLE' -and (Get-Now) -eq $t0.AddSeconds(150)) 'display formatting must not change retry scheduling' +} + +Test-Case 'stopped quota wait still displays suggested local retry time' { + New-World + Reply-OnProbe @('More reviews will be available in 2 minutes.') + $r = Invoke-Quota @{ MaxWaitSeconds = 60 } + $expected = $t0.AddSeconds(150).ToLocalTime().ToString('yyyy-MM-dd HH:mm:ss zzz', [Globalization.CultureInfo]::InvariantCulture) + Assert-True ($r.log.Contains("suggested next query at $expected (local time; in 00:02:30); this wait has stopped")) 'stopped wait must give a readable retry time' + Assert-True ($r.result.nextCheckAt -eq $t0.AddSeconds(150).ToString('o')) 'JSON timestamp must keep its protocol format' + Assert-True ($r.result.state -eq 'UNAVAILABLE' -and (Get-Now) -eq $t0) 'stopped wait must not keep sleeping' +} + Test-Case 'covered head ends REVIEWED without any write' { New-World Add-Status $headA 'success' 'Review completed' diff --git a/.agents/skills/ok-script-pr-review/wait-coderabbit-rate-limit.ps1 b/.agents/skills/ok-script-pr-review/wait-coderabbit-rate-limit.ps1 index 4fc237c..388b56e 100644 --- a/.agents/skills/ok-script-pr-review/wait-coderabbit-rate-limit.ps1 +++ b/.agents/skills/ok-script-pr-review/wait-coderabbit-rate-limit.ps1 @@ -53,6 +53,7 @@ function Wait-WithSessionCheck { try { $deadline = (Get-CrNow).AddSeconds($MaxWaitSeconds) + Write-Host "[rate-limit] waiting deadline: $(Format-CrDisplayTime $deadline) (local time; maximum wait $MaxWaitSeconds seconds; not a review completion estimate)" for ($probeCount = 1; $probeCount -le $MaxProbes; $probeCount++) { $stop = Test-Session if ($stop) { Complete-RateLimit $stop @{ actualHead = $script:ActualHead } } @@ -90,8 +91,13 @@ try { $delay = if ($null -ne $suggested) { $suggested + $RecheckBufferSeconds } else { $DefaultRecheckSeconds } $nextCheck = (Get-CrNow).AddSeconds([Math]::Max(1, $delay)) $detail = @{ probeId = $probe.id; replyId = $reply.comment.id; reply = Get-CrPlainText $reply.comment.body; nextCheckAt = $nextCheck.ToString('o') } - if ($nextCheck -gt $deadline -or $probeCount -eq $MaxProbes) { Complete-RateLimit 'UNAVAILABLE' $detail } - Write-Host "[rate-limit] unavailable; next query at $($nextCheck.ToString('u')) (a countdown is not proof of quota)" + $remaining = [TimeSpan]::FromSeconds([Math]::Max(0, [Math]::Ceiling(($nextCheck - (Get-CrNow)).TotalSeconds))) + $nextQuery = "$(Format-CrDisplayTime $nextCheck) (local time; in $($remaining.ToString('c')))" + if ($nextCheck -gt $deadline -or $probeCount -eq $MaxProbes) { + Write-Host "[rate-limit] unavailable; suggested next query at $nextQuery; this wait has stopped (review completion time is unknown)" + Complete-RateLimit 'UNAVAILABLE' $detail + } + Write-Host "[rate-limit] unavailable; next query at $nextQuery (a countdown is not proof of quota; review completion time is unknown)" $stop = Wait-WithSessionCheck $nextCheck if ($stop) { Complete-RateLimit $stop @{ actualHead = $script:ActualHead } } } diff --git a/.agents/skills/ok-script-pr-review/wait-coderabbit.ps1 b/.agents/skills/ok-script-pr-review/wait-coderabbit.ps1 index 14909d0..d51e0f0 100644 --- a/.agents/skills/ok-script-pr-review/wait-coderabbit.ps1 +++ b/.agents/skills/ok-script-pr-review/wait-coderabbit.ps1 @@ -128,7 +128,7 @@ try { $arrival = $null $unconfirmedSince = $null $restart = $false - Write-Host "[wait] session for $Repo#$PrNumber head=$head" + Write-Host "[wait] session for $Repo#$PrNumber head=$head; waiting deadline: $(Format-CrDisplayTime $globalDeadline) (local time; maximum wait $TimeoutSeconds seconds; not a review completion estimate)" while (-not $restart) { if ((Get-CrNow) -ge $globalDeadline) { Complete-Wait 'TIMEOUT' $head -Extra @{ detail = 'overall timeout' } } diff --git a/.agents/skills/ok-script-pr-review/wait-review-threads.ps1 b/.agents/skills/ok-script-pr-review/wait-review-threads.ps1 index f32c00f..1986572 100644 --- a/.agents/skills/ok-script-pr-review/wait-review-threads.ps1 +++ b/.agents/skills/ok-script-pr-review/wait-review-threads.ps1 @@ -162,6 +162,7 @@ try { } $started = Get-CrNow $deadline = $started.AddSeconds($TimeoutSeconds) + if (-not $Once) { Write-Host "[threads] waiting deadline: $(Format-CrDisplayTime $deadline) (local time; maximum wait $TimeoutSeconds seconds; not a review completion estimate)" } $pr = Get-CrPr $Repo $PrNumber if (-not $ExpectedHead) { $ExpectedHead = $pr.head } $polls = 0 diff --git a/jetbrains b/jetbrains index 0fb95a0..f2f7b97 160000 --- a/jetbrains +++ b/jetbrains @@ -1 +1 @@ -Subproject commit 0fb95a0744e767dc2fb704e79eca7a03622bc375 +Subproject commit f2f7b97bd340e3446bcb0ac18728b05991f3021a diff --git a/package.json b/package.json index 6c52e9b..a329b8d 100644 --- a/package.json +++ b/package.json @@ -435,7 +435,7 @@ "verify:version": "node scripts/release/verify-version.js", "verify:l10n": "node scripts/release/verify-l10n.js", "release": "sh scripts/release.sh", - "test": "npm run verify:version && npm run verify:l10n && npm run test:l10n-literals && npm run compile && npm run test:webview-l10n && npm run test:webview-html && npm run test:design-system && npm run test:thumbnail-actions && npm run test:task-launcher && npm run test:console-accountseg && npm run test:task-launcher-executor && npm run test:task-reconcile && npm run test:temp-screenshots && npm run test:annotation-coords && npm run test:annotation-clipboard && npm run test:annotation-swap && npm run test:annotation-merge && npm run test:annotation-history && npm run test:annotation-mode-cycle && npm run test:position-resource && npm run test:resource-host-contract && npm run test:position-publish-settings && npm run test:point-resource && npm run test:asset-swap-picker && npm run test:thumb-content-hash && npm run test:list-editor && npm run test:project-config && npm run test:project-config-schema && npm run test:convention-sources && npm run test:label-enum-guard && npm run test:save-to-assets-flow && npm run test:save-to-assets && npm run test:asset-pack-worker && npm run test:coco-feature-path && npm run test:box-resource && npm run test:probe-schemas && npm run test:probe-gui-group-names && npm run test:probe-qt-layer && npm run test:probe-multi-account-path && npm run test:probe-window-config && npm run test:run-sandbox && npm run test:startup-patches && npm run test:executor-runtime && npm run test:executor-project-config && npm run test:executor-gconfig && npm run test:project-store && npm run test:project-runtime", + "test": "npm run verify:version && npm run verify:l10n && npm run test:l10n-literals && npm run compile && npm run test:webview-l10n && npm run test:webview-html && npm run test:design-system && npm run test:thumbnail-actions && npm run test:task-launcher && npm run test:console-accountseg && npm run test:task-launcher-executor && npm run test:task-reconcile && npm run test:temp-screenshots && npm run test:annotation-coords && npm run test:annotation-clipboard && npm run test:annotation-swap && npm run test:annotation-merge && npm run test:annotation-history && npm run test:annotation-mode-cycle && npm run test:position-resource && npm run test:resource-host-contract && npm run test:position-publish-settings && npm run test:point-resource && npm run test:asset-swap-picker && npm run test:thumb-content-hash && npm run test:list-editor && npm run test:project-config && npm run test:project-config-schema && npm run test:convention-sources && npm run test:label-enum-guard && npm run test:save-to-assets-flow && npm run test:save-to-assets && npm run test:asset-pack-worker && npm run test:coco-feature-path && npm run test:box-resource && npm run test:probe-schemas && npm run test:probe-gui-group-names && npm run test:probe-qt-layer && npm run test:probe-multi-account-path && npm run test:probe-window-config && npm run test:screenshot-capture && npm run test:screenshot-worker && npm run test:background-screenshot && npm run test:global-hotkey && npm run test:run-sandbox && npm run test:startup-patches && npm run test:executor-runtime && npm run test:executor-project-config && npm run test:executor-gconfig && npm run test:project-store && npm run test:project-runtime", "test:webview-html": "node scripts/test_webview_html.js", "test:design-system": "node scripts/test_design_system.js", "test:thumbnail-actions": "node scripts/test_thumbnail_actions.js", @@ -468,6 +468,10 @@ "test:probe-gui-group-names": "python python/tests/test_probe_gui_group_names.py", "test:probe-qt-layer": "python python/tests/test_probe_qt_layer.py", "test:probe-multi-account-path": "python python/tests/test_probe_multi_account_path.py", + "test:screenshot-worker": "node scripts/test_screenshot_worker.js", + "test:background-screenshot": "node scripts/test_background_screenshot.js", + "test:global-hotkey": "python python/tests/test_global_hotkey.py", + "test:screenshot-capture": "python python/tests/test_capture_game_window.py", "test:probe-window-config": "python python/tests/test_probe_window_config.py", "test:run-sandbox": "python python/tests/test_run_executor_sandbox.py", "test:startup-patches": "python python/tests/test_run_executor_startup_patches.py", diff --git a/python/capture_game_window.py b/python/capture_game_window.py index 2b84728..2e29889 100644 --- a/python/capture_game_window.py +++ b/python/capture_game_window.py @@ -32,6 +32,17 @@ import time import ctypes +_SCRIPT_STARTED = time.perf_counter() +_PROTOCOL_STDOUT = sys.stdout +if '--serve' in sys.argv: + # 项目和框架的普通 print 全部留在日志管道，stdout 只传截图结果。 + sys.stdout = sys.stderr + + +def _log_timing(stage, started): + print(f"capture.timing.{stage}_ms={(time.perf_counter() - started) * 1000:.1f}", + file=sys.stderr, flush=True) + def _set_process_dpi_awareness(): """Declaring DPI awareness must happen before any Win32 window/DC call. @@ -72,6 +83,9 @@ def _set_process_dpi_awareness(): _ctx = BitBltCtxDummy() _desktop_ctx = BitBltCtxDummy() +_REUSE_WGC = False +_wgc_session = None +_log_timing('imports', _SCRIPT_STARTED) def _get_scaling_ratios(hwnd): @@ -175,7 +189,7 @@ def __str__(self): return f"StubHwndWindow(hwnd={self.hwnd}, {self.width}x{self.height})" -def _capture_wgc(hwnd, width, height, attempts=4): +def _capture_wgc(hwnd, width, height, attempts=4, timeout=1.5): """Capture one frame with Windows Graphics Capture. Returns None on failure. WGC reports the window content size through item.Size, which includes the @@ -195,12 +209,30 @@ def _capture_wgc(hwnd, width, height, attempts=4): print(f"WGC unavailable: {e}", file=sys.stderr) return None - stub = _StubHwndWindow(hwnd, width, height) - method = None + global _wgc_session + key = hwnd, width, height + if _wgc_session is not None and _wgc_session[0] != key: + _close_wgc_session() + if _REUSE_WGC and _wgc_session is not None: + _, stub, method = _wgc_session + else: + stub = _StubHwndWindow(hwnd, width, height) + method = None frame = None + # get_frame() 本身会等待首帧；不能把框架的 4 秒等待重复四次。 + # 使用框架已有的退出事件取消请求，不在另一个线程关闭仍在读帧的 D3D 对象。 + deadline = time.monotonic() + timeout + cancel = threading.Timer(timeout, stub.app_exit_event.set) + cancel.daemon = True + cancel.start() try: - method = WindowsGraphicsCaptureMethod(stub) - for _ in range(attempts): + if method is None: + method = WindowsGraphicsCaptureMethod(stub) + if not method.connected(): + return None + for attempt in range(attempts): + if stub.app_exit_event.is_set() or time.monotonic() >= deadline: + break content = getattr(method, 'last_size', None) if content is not None: cw, ch = int(content.Width), int(content.Height) @@ -210,19 +242,39 @@ def _capture_wgc(hwnd, width, height, attempts=4): stub.width, stub.height = cw, ch # content smaller: do not trim frame = method.get_frame() if frame is not None: + if _REUSE_WGC: + _wgc_session = key, stub, method return frame - time.sleep(0.3) + if attempt + 1 < attempts: + stub.app_exit_event.wait(min(0.05, max(0, deadline - time.monotonic()))) except Exception as e: print(f"WGC capture failed: {e}", file=sys.stderr) finally: - if method is not None: - try: - method.close() - except Exception: - pass + cancel.cancel() + if not (_REUSE_WGC and frame is not None and not stub.app_exit_event.is_set()): + stub.app_exit_event.set() + if method is not None: + if _wgc_session is not None and _wgc_session[2] is method: + _wgc_session = None + try: + method.close() + except Exception: + pass return frame if frame is not None else None +def _close_wgc_session(): + global _wgc_session + if _wgc_session is not None: + _, stub, method = _wgc_session + _wgc_session = None + stub.app_exit_event.set() + try: + method.close() + except Exception: + pass + + def _capture_bitblt(hwnd, w, h, ox, oy): """Capture one frame with BitBlt + PrintWindow. Returns None on failure.""" global _ctx @@ -244,8 +296,10 @@ def _capture_bitblt(hwnd, w, h, ox, oy): print(f"bitblt client={w}x{h} offset={ox},{oy} monitor_scaling={m_scaling:.2f} " f"window_scaling={w_scaling:.2f} ratio={ratio:.2f} capture={cap_w}x{cap_h}") - clean_up_bitblt(_ctx) - frame = capture_by_bitblt(_ctx, hwnd, cap_w, cap_h, cap_ox, cap_oy, True) + try: + frame = capture_by_bitblt(_ctx, hwnd, cap_w, cap_h, cap_ox, cap_oy, True) + finally: + clean_up_bitblt(_ctx) if frame is None: return None if ratio != 1.0: @@ -347,7 +401,9 @@ def capture(output_path, exe_names=None, hwnd_class=None, processor=None, method Returns the output path on success, raises on failure. """ + started = time.perf_counter() result = find_hwnd(None, exe_names or [], 0, 0, class_name=hwnd_class) + _log_timing('find_window', started) title, hwnd = result[0], result[1] if not hwnd: raise RuntimeError("No matching window found") @@ -366,17 +422,25 @@ def capture(output_path, exe_names=None, hwnd_class=None, processor=None, method ox, oy = wx - wleft, wy - wtop frame = None + if method not in ('auto', 'wgc'): + _close_wgc_session() if method == 'foreground': + started = time.perf_counter() frame = _capture_foreground(hwnd, w, h) + _log_timing('foreground', started) if frame is not None: print(f"captured via foreground: {frame.shape[1]}x{frame.shape[0]}") else: if method in ('auto', 'wgc'): - frame = _capture_wgc(hwnd, w, h) + started = time.perf_counter() + frame = _capture_wgc(hwnd, w, h, timeout=1.5 if method == 'auto' else 3.0) + _log_timing('wgc', started) if frame is not None: print(f"captured via WGC: {frame.shape[1]}x{frame.shape[0]}") if frame is None and method in ('auto', 'bitblt'): + started = time.perf_counter() frame = _capture_bitblt(hwnd, w, h, ox, oy) + _log_timing('bitblt', started) if frame is not None: print(f"captured via BitBlt: {frame.shape[1]}x{frame.shape[0]}") if frame is None: @@ -389,13 +453,19 @@ def capture(output_path, exe_names=None, hwnd_class=None, processor=None, method # 应用项目配置的 screenshot_processor（如遮挡 UID） if processor: + started = time.perf_counter() try: frame = processor(frame.copy()) except Exception as e: print(f"screenshot_processor error: {e}", file=sys.stderr) + finally: + _log_timing('processor', started) os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) - cv2.imwrite(output_path, frame) + started = time.perf_counter() + if not cv2.imwrite(output_path, frame): + raise RuntimeError("Screenshot file could not be written") + _log_timing('save', started) return output_path @@ -437,7 +507,9 @@ def main(): # 从项目 config 加载 screenshot_processor processor = None if args.project_dir: + started = time.perf_counter() processor = load_screenshot_processor(args.project_dir) + _log_timing('load_project_config', started) print(f"Searching: exe={exe_names}, class={hwnd_class}, method={args.method}") @@ -446,9 +518,118 @@ def main(): except Exception as exc: print(f"Capture failed: {exc}", file=sys.stderr) sys.exit(1) + finally: + _log_timing('total', _SCRIPT_STARTED) print(f"Saved: {args.output_path}") +def _project_sources(project_dir): + """记录本次实际导入的本地源码；变更后由宿主重新启动解释器，避免旧模块缓存。""" + from probe_window_config import _resolve_config_path + root = os.path.normcase(os.path.abspath(project_dir)) + environments = tuple(os.path.join(root, name) + os.sep for name in ('.venv', 'venv')) + paths = {_resolve_config_path(project_dir)} + for loaded in list(sys.modules.values()): + filename = getattr(loaded, '__file__', None) + if not filename: + continue + filename = os.path.normcase(os.path.abspath(filename)) + # 虚拟环境是解释器依赖，不属于项目源码；避免逐张 stat 数千个第三方库文件。 + if filename.startswith(environments): + continue + try: + if os.path.commonpath((root, filename)) == root: + paths.add(filename) + except ValueError: + pass + return {filename: _source_signature(filename) for filename in paths if filename} + + +def _source_signature(filename): + try: + stat = os.stat(filename) + return stat.st_mtime_ns, stat.st_size + except OSError: + return None + + +def warm_capture_session(project_dir, method='auto', load_project=True): + """提前导入项目及初始化后台 WGC，丢弃首帧；不写文件、不激活窗口。""" + if load_project: + load_screenshot_processor(project_dir) + if method not in ('auto', 'wgc'): + return False + # 即使游戏尚未打开，也先完成 WGC 依赖导入。 + try: + from ok.util.window import windows_graphics_available + if not windows_graphics_available(): + return False + importlib.import_module('ok.device.capture_methods.windows_graphics') + except Exception as exc: + print(f"WGC warmup unavailable: {exc}", file=sys.stderr) + return False + from probe_window_config import probe_project_config + config = probe_project_config(project_dir) + # 没有明确目标时，不在后台猜测窗口或弹输入框。 + if not config.get('ok') or not (config.get('exe') or config.get('hwnd_class')): + return False + result = find_hwnd(None, config.get('exe') or [], 0, 0, class_name=config.get('hwnd_class')) + hwnd = result[1] + if not hwnd or win32gui.IsIconic(hwnd): + return False + left, top, right, bottom = win32gui.GetClientRect(hwnd) + width, height = right - left, bottom - top + if width <= 0 or height <= 0: + return False + return _capture_wgc(hwnd, width, height, timeout=1.5) is not None + + +def serve(): + """串行处理截图请求，复用导入和 WGC 会话；每张图读取新帧和当前配置。""" + global _REUSE_WGC + _REUSE_WGC = True + sources = {} + project_warmed = False + for line in sys.stdin: + request = json.loads(line) + response = {'capture_response': request['id'], 'ok': False} + if any(_source_signature(filename) != signature for filename, signature in sources.items()): + _close_wgc_session() + response['restart'] = True + print(json.dumps(response), file=_PROTOCOL_STDOUT, flush=True) + return + started = time.perf_counter() + try: + project_dir = request.get('project_dir') + if request.get('op') == 'probe': + from probe_window_config import probe_project_config + response['config'] = probe_project_config(project_dir) + elif request.get('op') == 'warmup': + response['frame_ready'] = warm_capture_session( + project_dir, request.get('method', 'auto'), load_project=not project_warmed) + project_warmed = True + sources = _project_sources(project_dir) + else: + processor = load_screenshot_processor(project_dir) if project_dir else None + _log_timing('load_project_config', started) + config = request.get('window_config') or {} + capture(request['output_path'], config.get('exe'), config.get('hwnd_class'), + processor=processor, method=request.get('method', 'auto')) + if project_dir: + sources = _project_sources(project_dir) + project_warmed = True + response['ok'] = True + except Exception as exc: + response['error'] = str(exc) + finally: + _log_timing('request', started) + print(json.dumps(response), file=_PROTOCOL_STDOUT, flush=True) + _close_wgc_session() + + if __name__ == "__main__": - main() + if '--serve' in sys.argv: + serve() + else: + main() diff --git a/python/global_hotkey.py b/python/global_hotkey.py new file mode 100644 index 0000000..9baaf31 --- /dev/null +++ b/python/global_hotkey.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import argparse +import ctypes +import sys +import threading +from ctypes import wintypes + +HOTKEY_ID = 0x4F4B +WM_HOTKEY = 0x0312 +WM_QUIT = 0x0012 +MOD_ALT = 0x0001 +MOD_CONTROL = 0x0002 +MOD_NOREPEAT = 0x4000 +VK_S = 0x53 +SYNCHRONIZE = 0x00100000 +INFINITE = 0xFFFFFFFF + + +def start_parent_watch(parent_pid: int, message_thread_id: int, user32: ctypes.WinDLL) -> bool: + if parent_pid <= 0: + return True + + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + open_process = kernel32.OpenProcess + open_process.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + open_process.restype = wintypes.HANDLE + wait_for_single_object = kernel32.WaitForSingleObject + wait_for_single_object.argtypes = [wintypes.HANDLE, wintypes.DWORD] + wait_for_single_object.restype = wintypes.DWORD + close_handle = kernel32.CloseHandle + close_handle.argtypes = [wintypes.HANDLE] + close_handle.restype = wintypes.BOOL + post_thread_message = user32.PostThreadMessageW + post_thread_message.argtypes = [wintypes.DWORD, wintypes.UINT, wintypes.WPARAM, wintypes.LPARAM] + post_thread_message.restype = wintypes.BOOL + + handle = open_process(SYNCHRONIZE, False, parent_pid) + if not handle: + return False + + def wait_for_parent() -> None: + try: + wait_for_single_object(handle, INFINITE) + post_thread_message(message_thread_id, WM_QUIT, 0, 0) + finally: + close_handle(handle) + + threading.Thread(target=wait_for_parent, name="ok-hotkey-parent-watch", daemon=True).start() + return True + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--parent-pid", type=int, default=0) + args = parser.parse_args() + + if sys.platform != "win32": + print("UNSUPPORTED", flush=True) + return 3 + + user32 = ctypes.WinDLL("user32", use_last_error=True) + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + register_hot_key = user32.RegisterHotKey + register_hot_key.argtypes = [wintypes.HWND, ctypes.c_int, wintypes.UINT, wintypes.UINT] + register_hot_key.restype = wintypes.BOOL + unregister_hot_key = user32.UnregisterHotKey + unregister_hot_key.argtypes = [wintypes.HWND, ctypes.c_int] + unregister_hot_key.restype = wintypes.BOOL + get_message = user32.GetMessageW + get_message.argtypes = [ctypes.POINTER(wintypes.MSG), wintypes.HWND, wintypes.UINT, wintypes.UINT] + get_message.restype = ctypes.c_int + peek_message = user32.PeekMessageW + peek_message.argtypes = [ctypes.POINTER(wintypes.MSG), wintypes.HWND, wintypes.UINT, wintypes.UINT, wintypes.UINT] + peek_message.restype = wintypes.BOOL + get_current_thread_id = kernel32.GetCurrentThreadId + get_current_thread_id.restype = wintypes.DWORD + + modifiers = MOD_CONTROL | MOD_ALT | MOD_NOREPEAT + if not register_hot_key(None, HOTKEY_ID, modifiers, VK_S): + print(f"ERROR register_hotkey {ctypes.get_last_error()}", flush=True) + return 2 + + message = wintypes.MSG() + try: + # Create the message queue before the watcher can post WM_QUIT. + peek_message(ctypes.byref(message), None, 0, 0, 0) + if not start_parent_watch(args.parent_pid, get_current_thread_id(), user32): + print(f"ERROR parent_process {ctypes.get_last_error()}", flush=True) + return 5 + print("READY", flush=True) + while True: + result = get_message(ctypes.byref(message), None, 0, 0) + if result == -1: + print(f"ERROR get_message {ctypes.get_last_error()}", flush=True) + return 4 + if result == 0: + return 0 + if message.message == WM_HOTKEY and message.wParam == HOTKEY_ID: + print("TRIGGER", flush=True) + finally: + unregister_hot_key(None, HOTKEY_ID) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/probe_window_config.py b/python/probe_window_config.py index b1d7e2c..5e78dba 100644 --- a/python/probe_window_config.py +++ b/python/probe_window_config.py @@ -205,15 +205,11 @@ def _clean(value): return value -def main(): - if len(sys.argv) < 2: - print(json.dumps({"ok": False, "error": "缺少 project_dir 参数"}, ensure_ascii=False)) - sys.exit(1) - project_dir = sys.argv[1] - +def probe_project_config(project_dir): + """仅解析当前源码，可在截图进程中调用，不执行项目配置。""" config_path = _resolve_config_path(project_dir) if not config_path: - print(json.dumps({ + return { "ok": False, "error": f"找不到 config.py: {project_dir}", "searched": [ @@ -221,25 +217,32 @@ def main(): os.path.join(project_dir, "src", "config.py"), os.path.join(project_dir, "config.py"), ], - }, ensure_ascii=False)) - sys.exit(1) + } - window_config = _extract_config_dict_keys(config_path) - template_matching = _extract_template_matching_keys(config_path) - template_tab = _extract_template_tab_keys(config_path) + with open(config_path, encoding="utf-8") as f: + config_dict = _find_config_dict(ast.parse(f.read(), filename=config_path)) + window_config = _extract_sub_dict(config_dict, "windows", WINDOWS_SUB_KEYS) + template_matching = _extract_sub_dict(config_dict, "template_matching", TEMPLATE_MATCHING_SUB_KEYS) + template_tab = _extract_sub_dict(config_dict, "template_tab", TEMPLATE_TAB_SUB_KEYS) window_config = {k: _clean(v) for k, v in window_config.items()} coco_feature_json = _clean(template_matching.get("coco_feature_json")) label_enum_relative_path = _clean(template_tab.get("label_enum_relative_path")) - print(json.dumps({ + return { "ok": True, "config_path": config_path, **window_config, "coco_feature_json": coco_feature_json, "label_enum_relative_path": label_enum_relative_path, - }, ensure_ascii=False)) - sys.exit(0) + } + + +def main(): + result = probe_project_config(sys.argv[1]) if len(sys.argv) >= 2 else { + "ok": False, "error": "缺少 project_dir 参数"} + print(json.dumps(result, ensure_ascii=False)) + sys.exit(0 if result['ok'] else 1) if __name__ == "__main__": diff --git a/python/tests/test_capture_game_window.py b/python/tests/test_capture_game_window.py new file mode 100644 index 0000000..a890167 --- /dev/null +++ b/python/tests/test_capture_game_window.py @@ -0,0 +1,336 @@ +"""截图等待预算、自动回退与写入失败的回归测试，无需安装 ok 或连接游戏。""" +import importlib.util +import io +import json +import sys +import threading +import time +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +from _test_tmp import make_tmp_dir + + +def module(name, **attributes): + result = types.ModuleType(name) + result.__dict__.update(attributes) + return result + + +window = module('ok.util.window', find_hwnd=Mock(return_value=('game', 123)), + windows_graphics_available=lambda: True) +wgc = module('ok.device.capture_methods.windows_graphics') +bitblt = module('ok.device.capture_methods.bitblt_utils', BitBltCtxDummy=Mock, + capture_by_bitblt=Mock(), capture_desktop_by_bitblt=Mock(), + clean_up_bitblt=Mock(), clean_up_desktop_bitblt=Mock()) +cv2 = module('cv2', imwrite=Mock(return_value=True), resize=Mock(), INTER_LINEAR=1) +gui = module('win32gui', GetClientRect=lambda _: (0, 0, 640, 480), + ClientToScreen=lambda *_: (10, 20), GetWindowRect=lambda _: (0, 0, 650, 500), IsIconic=lambda _: False) +dependencies = {m.__name__: m for m in (window, wgc, bitblt, cv2, gui, + module('win32api'), module('win32con'))} +spec = importlib.util.spec_from_file_location('capture_under_test', + Path(__file__).resolve().parents[1] / 'capture_game_window.py') +capture = importlib.util.module_from_spec(spec) +with patch.dict(sys.modules, dependencies): + spec.loader.exec_module(capture) + + +class FakeWgc: + instances = [] + frame = None + connected_value = True + wait_for_frame = False + raises = False + + def __init__(self, stub): + self.stub = stub + self.closed = False + self.calls = 0 + self.last_size = None + self.instances.append(self) + + def connected(self): + return self.connected_value + + def get_frame(self): + self.calls += 1 + if self.raises: + raise RuntimeError('backend failure') + if self.wait_for_frame: + self.stub.app_exit_event.wait(4) + return self.frame + + def close(self): + self.closed = True + + +class Frame: + shape = (480, 640, 3) + + def copy(self): + return self + + +class CaptureTest(unittest.TestCase): + def setUp(self): + capture._close_wgc_session() + capture._REUSE_WGC = False + self.addCleanup(capture._close_wgc_session) + self.modules = patch.dict(sys.modules, dependencies) + self.modules.start() + self.addCleanup(self.modules.stop) + FakeWgc.instances = [] + FakeWgc.frame = None + FakeWgc.connected_value = True + FakeWgc.wait_for_frame = False + FakeWgc.raises = False + wgc.WindowsGraphicsCaptureMethod = FakeWgc + cv2.imwrite.reset_mock(return_value=True) + cv2.imwrite.return_value = True + bitblt.clean_up_bitblt.reset_mock() + self.output = str(Path(make_tmp_dir('capture')) / 'shot.png') + + def test_frame_returns_immediately_and_closes(self): + FakeWgc.frame = Frame() + started = time.monotonic() + self.assertIs(capture._capture_wgc(123, 640, 480), FakeWgc.frame) + self.assertLess(time.monotonic() - started, 0.5) + self.assertTrue(FakeWgc.instances[0].closed) + self.assertTrue(FakeWgc.instances[0].stub.app_exit_event.is_set()) + + def test_initialization_failure_does_not_retry_or_sleep(self): + FakeWgc.connected_value = False + started = time.monotonic() + self.assertIsNone(capture._capture_wgc(123, 640, 480)) + self.assertLess(time.monotonic() - started, 0.5) + self.assertEqual(FakeWgc.instances[0].calls, 0) + self.assertTrue(FakeWgc.instances[0].closed) + + def test_reused_session_requests_a_new_frame_and_rebuilds_after_resize(self): + capture._REUSE_WGC = True + first, second = Frame(), Frame() + FakeWgc.frame = first + self.assertIs(capture._capture_wgc(123, 640, 480), first) + FakeWgc.frame = second + self.assertIs(capture._capture_wgc(123, 640, 480), second) + self.assertEqual(len(FakeWgc.instances), 1) + self.assertEqual(FakeWgc.instances[0].calls, 2) + self.assertFalse(FakeWgc.instances[0].closed) + capture._capture_wgc(123, 800, 600) + self.assertEqual(len(FakeWgc.instances), 2) + self.assertTrue(FakeWgc.instances[0].closed) + + def test_timeout_discards_reused_session(self): + capture._REUSE_WGC = True + FakeWgc.frame = Frame() + capture._capture_wgc(123, 640, 480) + FakeWgc.frame = None + FakeWgc.wait_for_frame = True + self.assertIsNone(capture._capture_wgc(123, 640, 480, timeout=0.05)) + self.assertIsNone(capture._wgc_session) + self.assertTrue(FakeWgc.instances[0].closed) + + def test_blocked_frame_request_is_cancelled_within_budget(self): + FakeWgc.wait_for_frame = True + started = time.monotonic() + self.assertIsNone(capture._capture_wgc(123, 640, 480, timeout=0.08)) + self.assertLess(time.monotonic() - started, 0.5) + self.assertEqual(FakeWgc.instances[0].calls, 1) + self.assertTrue(FakeWgc.instances[0].closed) + + def test_fast_empty_requests_share_one_deadline(self): + started = time.monotonic() + self.assertIsNone(capture._capture_wgc(123, 640, 480, attempts=100, timeout=0.08)) + self.assertLess(time.monotonic() - started, 0.5) + self.assertLessEqual(FakeWgc.instances[0].calls, 3) + + def test_backend_exception_closes_session(self): + FakeWgc.raises = True + self.assertIsNone(capture._capture_wgc(123, 640, 480)) + self.assertTrue(FakeWgc.instances[0].closed) + + def test_auto_falls_back_and_applies_processor_before_save(self): + frame, processed = Frame(), Frame() + processor = Mock(return_value=processed) + with patch.object(capture, '_capture_wgc', return_value=None) as first, \ + patch.object(capture, '_capture_bitblt', return_value=frame) as fallback, \ + patch.object(capture, '_capture_foreground') as foreground: + capture.capture(self.output, processor=processor) + first.assert_called_once_with(123, 640, 480, timeout=1.5) + fallback.assert_called_once() + foreground.assert_not_called() + processor.assert_called_once_with(frame) + cv2.imwrite.assert_called_once_with(self.output, processed) + + def test_explicit_wgc_keeps_a_larger_budget_and_never_falls_back(self): + with patch.object(capture, '_capture_wgc', return_value=None) as first, \ + patch.object(capture, '_capture_bitblt') as fallback: + with self.assertRaises(RuntimeError): + capture.capture(self.output, method='wgc') + first.assert_called_once_with(123, 640, 480, timeout=3.0) + fallback.assert_not_called() + + def test_failed_write_is_reported_as_failure(self): + cv2.imwrite.return_value = False + with patch.object(capture, '_capture_bitblt', return_value=Frame()): + with self.assertRaisesRegex(RuntimeError, 'could not be written'): + capture.capture(self.output, method='bitblt') + + def test_bitblt_resources_are_released_after_failure(self): + with patch.object(capture, '_get_scaling_ratios', return_value=(1, 1)), \ + patch.object(capture, 'capture_by_bitblt', side_effect=RuntimeError('GDI failure')): + with self.assertRaises(RuntimeError): + capture._capture_bitblt(123, 640, 480, 0, 0) + bitblt.clean_up_bitblt.assert_called_once_with(capture._ctx) + + def test_server_reads_each_request_and_restarts_after_source_change(self): + config = Path(make_tmp_dir('capture-source')) / 'config.py' + config.write_text('version = 1', encoding='utf-8') + protocol = io.StringIO() + def requests(): + yield json.dumps({'id': 1, 'output_path': self.output, 'project_dir': str(config.parent)}) + '\n' + config.write_text('version = 222', encoding='utf-8') + yield json.dumps({'id': 2, 'output_path': self.output, 'project_dir': str(config.parent)}) + '\n' + with patch.object(sys, 'stdin', requests()), \ + patch.object(capture, '_PROTOCOL_STDOUT', protocol), \ + patch.object(capture, 'load_screenshot_processor', return_value=None), \ + patch.object(capture, '_project_sources', side_effect=lambda _: {str(config): capture._source_signature(config)}), \ + patch.object(capture, 'capture') as take: + capture.serve() + replies = [json.loads(line) for line in protocol.getvalue().splitlines()] + self.assertTrue(replies[0]['ok']) + self.assertTrue(replies[1]['restart']) + self.assertEqual(take.call_count, 1) + + def test_server_reloads_processor_for_each_shot(self): + requests = [json.dumps({'id': n, 'output_path': self.output, 'project_dir': 'project'}) + '\n' + for n in (1, 2)] + protocol = io.StringIO() + with patch.object(sys, 'stdin', requests), \ + patch.object(capture, '_PROTOCOL_STDOUT', protocol), \ + patch.object(capture, '_project_sources', return_value={}), \ + patch.object(capture, 'load_screenshot_processor', return_value=None) as load, \ + patch.object(capture, 'capture'): + capture.serve() + self.assertEqual(load.call_count, 2) + self.assertTrue(all(json.loads(line)['ok'] for line in protocol.getvalue().splitlines())) + + def test_probe_reads_current_source_without_importing_config_or_capturing(self): + root = Path(make_tmp_dir('capture-probe')) + config = root / 'config.py' + protocol = io.StringIO() + def requests(): + for n in (1, 2): + config.write_text("raise RuntimeError('must not execute')\n" + + "config = {'windows': {'exe': ['game%d.exe']}}" % n, encoding='utf-8') + yield json.dumps({'id': n, 'op': 'probe', 'project_dir': str(root)}) + '\n' + yield json.dumps({'id': 3, 'op': 'probe', 'project_dir': str(root / 'missing')}) + '\n' + sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + self.addCleanup(lambda: sys.path.pop(0)) + with patch.object(sys, 'stdin', requests()), \ + patch.object(capture, '_PROTOCOL_STDOUT', protocol), \ + patch.object(capture, 'load_screenshot_processor') as load, \ + patch.object(capture, 'capture') as take: + capture.serve() + replies = [json.loads(line) for line in protocol.getvalue().splitlines()] + self.assertEqual(replies[0]['config']['exe'], ['game1.exe']) + self.assertEqual(replies[1]['config']['exe'], ['game2.exe']) + self.assertFalse(replies[2]['config']['ok']) + load.assert_not_called() + take.assert_not_called() + + def test_source_watch_includes_project_modules_and_excludes_virtualenv(self): + root = Path(make_tmp_dir('capture-watch')) + config = root / 'config.py' + local = root / 'src' / 'processor.py' + dependency = root / '.venv' / 'Lib' / 'cv2.py' + for file in (config, local, dependency): + file.parent.mkdir(parents=True, exist_ok=True) + file.write_text('version = 1', encoding='utf-8') + probe = module('probe_window_config', _resolve_config_path=lambda _: str(config)) + with patch.dict(sys.modules, {'probe_window_config': probe, + 'local_processor': module('local_processor', __file__=str(local)), + 'dependency': module('dependency', __file__=str(dependency))}): + sources = capture._project_sources(str(root)) + self.assertIn(str(config), sources) + self.assertIn(capture.os.path.normcase(str(local)), sources) + self.assertNotIn(capture.os.path.normcase(str(dependency)), sources) + + def test_warmup_discards_frame_and_real_shot_requests_a_fresh_frame(self): + capture._REUSE_WGC = True + FakeWgc.frame = Frame() + probe = module('probe_window_config', probe_project_config=lambda _: {'ok': True, 'exe': ['game.exe']}) + processor = Mock(return_value=Frame()) + with patch.dict(sys.modules, {'probe_window_config': probe}), \ + patch.object(capture, 'load_screenshot_processor', return_value=processor), \ + patch.object(capture, '_capture_foreground') as foreground, \ + patch.object(capture, '_capture_bitblt') as bit: + self.assertTrue(capture.warm_capture_session('project')) + cv2.imwrite.assert_not_called() + processor.assert_not_called() + FakeWgc.frame = Frame() + capture.capture(self.output, processor=processor) + self.assertEqual(len(FakeWgc.instances), 1) + self.assertEqual(FakeWgc.instances[0].calls, 2) + cv2.imwrite.assert_called_once() + processor.assert_called_once_with(FakeWgc.frame) + foreground.assert_not_called() + bit.assert_not_called() + + def test_warmup_without_game_or_in_foreground_mode_never_activates_or_saves(self): + probe = module('probe_window_config', probe_project_config=lambda _: {'ok': True, 'exe': ['game.exe']}) + with patch.dict(sys.modules, {'probe_window_config': probe}), \ + patch.object(capture, 'load_screenshot_processor') as load, \ + patch.object(capture, 'find_hwnd', return_value=('', None)) as find, \ + patch.object(capture, '_capture_foreground') as foreground: + self.assertFalse(capture.warm_capture_session('project')) + find.reset_mock() + self.assertFalse(capture.warm_capture_session('project', 'foreground')) + find.assert_not_called() + self.assertEqual(load.call_count, 2) + foreground.assert_not_called() + cv2.imwrite.assert_not_called() + self.assertEqual(FakeWgc.instances, []) + + def test_server_warmup_reuses_project_imports_and_failure_does_not_block_capture(self): + requests = [json.dumps({'id': n, 'op': 'warmup', 'project_dir': 'project'}) + '\n' for n in (1, 2, 3)] + requests.append(json.dumps({'id': 4, 'output_path': self.output, 'project_dir': 'project'}) + '\n') + protocol = io.StringIO() + with patch.object(sys, 'stdin', requests), \ + patch.object(capture, '_PROTOCOL_STDOUT', protocol), \ + patch.object(capture, '_project_sources', return_value={}), \ + patch.object(capture, 'warm_capture_session', side_effect=[False, RuntimeError('warmup failure'), True]) as warm, \ + patch.object(capture, 'load_screenshot_processor'), \ + patch.object(capture, 'capture') as take: + capture.serve() + replies = [json.loads(line) for line in protocol.getvalue().splitlines()] + self.assertTrue(replies[0]['ok']) + self.assertFalse(replies[0]['frame_ready']) + self.assertFalse(replies[1]['ok']) + self.assertTrue(replies[2]['frame_ready']) + self.assertTrue(replies[3]['ok']) + self.assertEqual([call.kwargs['load_project'] for call in warm.call_args_list], [True, False, False]) + take.assert_called_once() + + def test_project_change_after_warmup_requires_restart_before_next_request(self): + config = Path(make_tmp_dir('capture-warm-source')) / 'config.py' + config.write_text('version = 1', encoding='utf-8') + protocol = io.StringIO() + def requests(): + yield json.dumps({'id': 1, 'op': 'warmup', 'project_dir': str(config.parent)}) + '\n' + config.write_text('version = 222', encoding='utf-8') + yield json.dumps({'id': 2, 'op': 'probe', 'project_dir': str(config.parent)}) + '\n' + with patch.object(sys, 'stdin', requests()), \ + patch.object(capture, '_PROTOCOL_STDOUT', protocol), \ + patch.object(capture, 'warm_capture_session', return_value=False), \ + patch.object(capture, '_project_sources', side_effect=lambda _: {str(config): capture._source_signature(config)}): + capture.serve() + replies = [json.loads(line) for line in protocol.getvalue().splitlines()] + self.assertTrue(replies[0]['ok']) + self.assertTrue(replies[1]['restart']) + + +if __name__ == '__main__': + unittest.main() diff --git a/python/tests/test_global_hotkey.py b/python/tests/test_global_hotkey.py new file mode 100644 index 0000000..32616bc --- /dev/null +++ b/python/tests/test_global_hotkey.py @@ -0,0 +1,62 @@ +import contextlib +import importlib.util +import io +from pathlib import Path +import unittest +from unittest.mock import Mock, patch + +spec = importlib.util.spec_from_file_location('global_hotkey', Path(__file__).resolve().parents[1] / 'global_hotkey.py') +hotkey = importlib.util.module_from_spec(spec) +spec.loader.exec_module(hotkey) + + +class ParentLifetimeTests(unittest.TestCase): + def run_helper(self, parent_handle): + events = [] + user = Mock() + kernel = Mock() + user.RegisterHotKey.return_value = True + user.PeekMessageW.side_effect = lambda *args: events.append('queue') or False + user.GetMessageW.return_value = 0 + kernel.GetCurrentThreadId.return_value = 42 + kernel.OpenProcess.side_effect = lambda *args: events.append('parent') or parent_handle + kernel.WaitForSingleObject.return_value = 0 + user.PostThreadMessageW.side_effect = lambda *args: events.append('quit') or True + + class ImmediateThread: + def __init__(self, target, **kwargs): + self.target = target + + def start(self): + self.target() + + output = io.StringIO() + with patch.object(hotkey.sys, 'platform', 'win32'), \ + patch.object(hotkey.sys, 'argv', ['global_hotkey.py', '--parent-pid', '123']), \ + patch.object(hotkey.ctypes, 'WinDLL', side_effect=lambda name, **kwargs: user if name == 'user32' else kernel, create=True), \ + patch.object(hotkey.ctypes, 'get_last_error', return_value=87, create=True), \ + patch.object(hotkey.threading, 'Thread', ImmediateThread), \ + contextlib.redirect_stdout(output): + result = hotkey.main() + return result, output.getvalue(), events, user, kernel + + def test_missing_parent_exits_and_unregisters_hotkey(self): + result, output, events, user, kernel = self.run_helper(0) + self.assertEqual(result, 5) + self.assertNotIn('READY', output) + self.assertIn('ERROR parent_process', output) + user.UnregisterHotKey.assert_called_once() + user.GetMessageW.assert_not_called() + kernel.WaitForSingleObject.assert_not_called() + + def test_already_exited_parent_posts_quit_after_queue_creation(self): + result, output, events, user, kernel = self.run_helper(1234) + self.assertEqual(result, 0) + self.assertEqual(events, ['queue', 'parent', 'quit']) + user.PostThreadMessageW.assert_called_once_with(42, hotkey.WM_QUIT, 0, 0) + kernel.CloseHandle.assert_called_once_with(1234) + user.UnregisterHotKey.assert_called_once() + + +if __name__ == '__main__': + unittest.main() diff --git a/scripts/test_background_screenshot.js b/scripts/test_background_screenshot.js new file mode 100644 index 0000000..362747f --- /dev/null +++ b/scripts/test_background_screenshot.js @@ -0,0 +1,66 @@ +const assert = require('assert'); +const fs = require('fs'); +const path = require('path'); +const vm = require('vm'); +const { makeTmpDir } = require('./test-tmp'); + +const root = makeTmpDir('background-screenshot'); +const outputs = []; +const errors = []; +let behavior = async output => { + fs.writeFileSync(output, 'frame'); + return { ok: true }; +}; +const moduleStub = { exports: {} }; +const vscode = { + workspace: { workspaceFolders: [{ uri: { fsPath: root } }] }, + window: { + showWarningMessage() {}, + showErrorMessage(message) { errors.push(message); }, + showInformationMessage() {}, + }, +}; +const source = fs.readFileSync(path.join(__dirname, '../out/backgroundScreenshot.js'), 'utf8'); +vm.runInNewContext(source, { + exports: moduleStub.exports, + require(name) { + if (name === 'vscode') return vscode; + if (name === './localization') return { tr: (key, args) => `${key} ${args?.error || ''}` }; + if (name === './projectConfig') return { resolveProjectDir: () => root, templatesDirectory: () => 'templates' }; + if (name === './screenshotCapture') return { + captureGameWindow(output) { outputs.push(output); return behavior(output); }, + }; + return require(name); + }, + Date: class extends Date { constructor() { super('2026-10-12T12:00:00Z'); } }, +}, { filename: 'backgroundScreenshot.js' }); +const capture = moduleStub.exports.captureTemplateScreenshotInBackground; + +(async () => { + assert.strictEqual(await capture(), true); + assert.strictEqual(await capture(), true); + assert.notStrictEqual(outputs[0], outputs[1], 'Captures within the same second must keep both images'); + assert.strictEqual(fs.readFileSync(outputs[0], 'utf8'), 'frame'); + assert.strictEqual(fs.readFileSync(outputs[1], 'utf8'), 'frame'); + + let finish; + behavior = () => new Promise(resolve => { finish = resolve; }); + const pending = capture(); + const requestCount = outputs.length; + assert.strictEqual(await capture(), false, 'Repeated hotkeys must not queue overlapping captures'); + assert.strictEqual(outputs.length, requestCount); + finish({ ok: false, reason: 'cancelled' }); + assert.strictEqual(await pending, false); + + behavior = async () => { throw new Error('capture unavailable'); }; + assert.strictEqual(await capture(), false); + assert(errors.some(message => message.includes('capture unavailable'))); + + behavior = async () => ({ ok: true }); + assert.strictEqual(await capture(), true, 'A failed capture must release the busy guard'); + vscode.workspace.workspaceFolders = []; + assert.strictEqual(await capture(), false); + vscode.workspace.workspaceFolders = [{ uri: { fsPath: root } }]; + assert.strictEqual(await capture(), true, 'Missing workspace must not keep the guard locked'); + console.log('test_background_screenshot: OK'); +})().catch(error => { console.error(error); process.exitCode = 1; }); diff --git a/scripts/test_save_to_assets.js b/scripts/test_save_to_assets.js index ea83e47..6ebec4a 100644 --- a/scripts/test_save_to_assets.js +++ b/scripts/test_save_to_assets.js @@ -975,7 +975,33 @@ async function test_deleteImageCleansAllAuthoringSourcesAndRollsBack() { console.log('[PASS] image deletion cleans all three authoring sources and rolls back on failure'); } +function test_tempBmpPromotionPreservesPixelsAndImportSemantics() { + setup(); + try { + const { decodeRgba } = require('../out/pngCrop'); + // 两行不同颜色，包含 BMP 行填充和自底向上的排列。 + const bmp = Buffer.alloc(54 + 24); + bmp.write('BM'); bmp.writeUInt32LE(bmp.length, 2); bmp.writeUInt32LE(54, 10); + bmp.writeUInt32LE(40, 14); bmp.writeInt32LE(3, 18); bmp.writeInt32LE(2, 22); + bmp.writeUInt16LE(1, 26); bmp.writeUInt16LE(24, 28); + for (let y = 0; y < 2; y++) for (let x = 0; x < 3; x++) { + const i = 54 + y * 12 + x * 3; + bmp[i] = 7 + x; bmp[i + 1] = 40 + y; bmp[i + 2] = 120 + x + y; + } + const source = path.join(tmpDir, 'shot.bmp'); fs.writeFileSync(source, bmp); + const data = new TemplateAssetData(tmpDir); + const promoted = data.importImageFile(source, true); + assert(promoted && path.extname(promoted) === '.png', '临时 BMP 入库应转 PNG'); + const before = decodeRgba(bmp), after = decodeRgba(fs.readFileSync(promoted)); + assert(before.width === after.width && before.height === after.height && before.rgba.equals(after.rgba), '像素与尺寸必须完全一致'); + const imported = data.importImageFile(source); + assert(path.extname(imported) === '.bmp' && fs.readFileSync(imported).equals(bmp), '普通导入保留原格式'); + assert(!fs.existsSync(path.join(templateDir, 'coco_annotations.json')), '入库不应提前登记标注'); + } finally { teardown(); } +} + const tests = [ + test_tempBmpPromotionPreservesPixelsAndImportSemantics, test_deleteImageCleansAllAuthoringSourcesAndRollsBack, test_nonOverlappingBboxesPackToSamePage, test_overlappingBboxesSplitToDifferentPages, diff --git a/scripts/test_screenshot_worker.js b/scripts/test_screenshot_worker.js new file mode 100644 index 0000000..5487d62 --- /dev/null +++ b/scripts/test_screenshot_worker.js @@ -0,0 +1,86 @@ +const assert = require('assert'); +const fs = require('fs'); +const path = require('path'); +const { ScreenshotWorker } = require('../out/screenshotWorker'); +const { makeTmpDir } = require('./test-tmp'); + +const tmp = makeTmpDir('capture-worker'); +const script = path.join(tmp, 'worker.js'); +fs.writeFileSync(script, ` +const fs = require('fs'); +const path = require('path'); +const rl = require('readline').createInterface({ input: process.stdin }); +rl.on('line', line => { + const r = JSON.parse(line); + fs.appendFileSync(path.join(r.project_dir, 'requests.jsonl'), JSON.stringify(r) + '\\n'); + if (r.op === 'warmup') { + if (r.method === 'hang') return; + setTimeout(() => console.log(JSON.stringify({ capture_response: r.id, ok: r.method !== 'error', frame_ready: true, error: 'warmup failed' })), 80); + return; + } + if (r.op === 'probe') { + console.log(JSON.stringify({ capture_response: r.id, ok: true, + config: { ok: true, exe: [fs.readFileSync(path.join(r.project_dir, 'exe.txt'), 'utf8')] } })); + return; + } + if (r.output_path === 'hang') return; + if (r.output_path === 'exit') return process.exit(7); + if (r.output_path === 'restart' && !fs.existsSync(path.join(r.project_dir, 'restarted'))) { + fs.writeFileSync(path.join(r.project_dir, 'restarted'), '1'); + const reply = JSON.stringify({ capture_response: r.id, restart: true }) + '\\n'; + const tailCode = 'setTimeout(() => process.stdout.write(' + JSON.stringify(reply) + '), 30)'; + const tail = require('child_process').spawn(process.execPath, ['-e', tailCode], { + stdio: ['ignore', 'inherit', 'ignore'], + }); + tail.unref(); + return process.exit(0); + } + for (let i = 0; i < 1000; i++) process.stderr.write('log\\n'); + console.log('framework output'); + setTimeout(() => console.log(JSON.stringify({ capture_response: r.id, ok: r.output_path !== 'error', error: 'capture failed' })), 10); +}); +`); + +(async () => { + const logs = []; + const worker = new ScreenshotWorker(process.execPath, script, tmp, line => logs.push(line), 2000); + try { + const warming = worker.prewarm('auto'); + assert.strictEqual(worker.prewarm('auto'), warming, 'Concurrent triggers must share one warmup'); + assert.strictEqual(await warming, true); + assert.strictEqual(fs.readFileSync(path.join(tmp, 'requests.jsonl'), 'utf8').trim().split('\n').length, 1); + fs.writeFileSync(path.join(tmp, 'exe.txt'), 'game.exe'); + assert.deepStrictEqual(await worker.probe(), { ok: true, exe: ['game.exe'] }); + await Promise.all([worker.capture('one', {}, 'auto'), worker.capture('two', {}, 'bitblt')]); + fs.writeFileSync(path.join(tmp, 'exe.txt'), 'edited.exe'); + assert.deepStrictEqual(await worker.probe(), { ok: true, exe: ['edited.exe'] }); + assert.strictEqual(logs.filter(line => line.includes('worker_started')).length, 1); + await assert.rejects(worker.prewarm('error'), /warmup failed/); + await worker.prewarm('auto'); + await assert.rejects(worker.capture('error', {}, 'auto'), /capture failed/); + await worker.capture('restart', {}, 'auto'); + assert.strictEqual(logs.filter(line => line.includes('worker_started')).length, 2); + await assert.rejects(worker.capture('exit', {}, 'auto'), /exited/); + await worker.capture('after-exit', {}, 'auto'); + assert.strictEqual(logs.filter(line => line.includes('worker_started')).length, 3); + } finally { worker.dispose(); } + await assert.rejects(worker.capture('disposed', {}, 'auto'), /disposed/); + + const hung = new ScreenshotWorker(process.execPath, script, tmp, () => {}, 200); + try { + const started = Date.now(); + await assert.rejects(hung.prewarm('hang'), /timed out/); + assert(Date.now() - started < 2500); + await hung.capture('recovered', {}, 'auto'); + } finally { hung.dispose(); } + + const idleLogs = []; + const idle = new ScreenshotWorker(process.execPath, script, tmp, line => idleLogs.push(line), 2000, 50); + try { + await idle.prewarm('auto'); + await new Promise(resolve => setTimeout(resolve, 150)); + await idle.capture('second', {}, 'auto'); + assert.strictEqual(idleLogs.filter(line => line.includes('worker_started')).length, 2); + } finally { idle.dispose(); } + console.log('test_screenshot_worker: OK'); +})().catch(error => { console.error(error); process.exitCode = 1; }); diff --git a/src/backgroundScreenshot.ts b/src/backgroundScreenshot.ts new file mode 100644 index 0000000..79e2619 --- /dev/null +++ b/src/backgroundScreenshot.ts @@ -0,0 +1,58 @@ +import * as fs from 'fs'; +import * as path from 'path'; +import { randomUUID } from 'crypto'; +import * as vscode from 'vscode'; +import { tr } from './localization'; +import { resolveProjectDir, templatesDirectory } from './projectConfig'; +import { captureGameWindow } from './screenshotCapture'; + +let captureInProgress = false; + +/** + * Capture directly into the template authoring directory without revealing an IDE panel. + * + * The system-wide hotkey uses this path so pressing Ctrl+Alt+S while the game is focused + * does not bring VS Code to the foreground before capture. The actual window capture still + * goes through the same captureGameWindow implementation used by the normal asset panel. + */ +export async function captureTemplateScreenshotInBackground(): Promise { + if (captureInProgress) return false; + captureInProgress = true; + try { + return await captureScreenshot(); + } catch (error) { + void vscode.window.showErrorMessage(tr('Screenshot failed: {error}', { error: String(error) })); + return false; + } finally { + captureInProgress = false; + } +} + +async function captureScreenshot(): Promise { + const folder = vscode.workspace.workspaceFolders?.[0]; + if (!folder) { + void vscode.window.showWarningMessage(tr('No workspace folder open.')); + return false; + } + + const projectRoot = resolveProjectDir() || folder.uri.fsPath; + const outputDir = path.join(projectRoot, templatesDirectory(projectRoot)); + fs.mkdirSync(outputDir, { recursive: true }); + + const now = new Date(); + const ts = `${now.getFullYear()}${String(now.getMonth() + 1).padStart(2, '0')}${String(now.getDate()).padStart(2, '0')}_${String(now.getHours()).padStart(2, '0')}${String(now.getMinutes()).padStart(2, '0')}${String(now.getSeconds()).padStart(2, '0')}`; + const outputPath = path.join(outputDir, `screenshot_${ts}_${randomUUID()}.png`); + + const outcome = await captureGameWindow(outputPath); + if (!outcome.ok) { + if (outcome.reason === 'noScript') { + void vscode.window.showErrorMessage(tr('Screenshot script not found. capture_game_window.py is missing from the extension.')); + } else if (outcome.reason === 'failed') { + void vscode.window.showErrorMessage(tr('Screenshot failed: {error}', { error: outcome.error || '' })); + } + return false; + } + + void vscode.window.showInformationMessage(tr('Screenshot saved: {name}', { name: path.basename(outputPath) })); + return true; +} diff --git a/src/consolePanel.ts b/src/consolePanel.ts index 71c20f2..dc6c3da 100644 --- a/src/consolePanel.ts +++ b/src/consolePanel.ts @@ -8,6 +8,7 @@ import { loadToolboxState, notifyExecutorRunning, onToolboxStateChange, saveTool import { GameConnectService } from './toolboxConnect'; import { applySharedAssets, errorPage, getNonce } from './webviewHtml'; import { reconcileCachedTaskList, reconcileTaskList } from './taskReconcile'; +import { prewarmScreenshotCapture } from './screenshotCapture'; /** 单个任务的元信息 */ interface TaskInfo { @@ -1562,6 +1563,7 @@ export class ConsoleViewProvider implements vscode.WebviewViewProvider { this.applySnapshot(line.slice(line.indexOf('OK_TOOLKIT_STATE:') + 'OK_TOOLKIT_STATE:'.length).trim()); } else if (line.includes('OK_TOOLKIT_EXECUTOR_READY')) { this.connecting = false; + void prewarmScreenshotCapture(this.executorProjectDir, this.executor?.spawnargs?.[0]); this.appendOutputLine(tr('Executor ready · {count} trigger task(s) enabled', { count: this.enabledTriggers.size, })); diff --git a/src/entrypoint.ts b/src/entrypoint.ts index 606a961..82a1cb4 100644 --- a/src/entrypoint.ts +++ b/src/entrypoint.ts @@ -1,8 +1,12 @@ import * as fs from 'fs'; import * as path from 'path'; import * as vscode from 'vscode'; +import { captureTemplateScreenshotInBackground } from './backgroundScreenshot'; import { activate as activateCore, deactivate as deactivateCore } from './extension'; +import { GlobalScreenshotHotkey } from './globalScreenshotHotkey'; import { resolveProjectDir } from './projectConfig'; +import { getProjectConfig } from './screenshotCapture'; +import { repaintAllAssetGalleries } from './templateAssetPanel'; const WALKTHROUGH_ID = 'okScriptToolkit.getStarted'; const PROJECT_READY_CONTEXT = 'okScriptToolkit.projectReady'; @@ -30,6 +34,16 @@ function isSupportedProject(projectDir: string): boolean { export function activate(context: vscode.ExtensionContext): void { activateCore(context); + const globalScreenshotHotkey = new GlobalScreenshotHotkey( + context.extensionPath, + () => getProjectConfig().pythonPath, + () => { + void captureTemplateScreenshotInBackground().then((captured) => { + if (captured) repaintAllAssetGalleries(); + }); + }, + ); + const updateProjectReady = () => vscode.commands.executeCommand( 'setContext', PROJECT_READY_CONTEXT, @@ -95,6 +109,7 @@ export function activate(context: vscode.ExtensionContext): void { } }), new vscode.Disposable(disposeConfigWatchers), + globalScreenshotHotkey, ); void updateProjectReady(); diff --git a/src/extension.ts b/src/extension.ts index 410593a..4bca420 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -1,6 +1,7 @@ import * as path from 'path'; import * as crypto from 'crypto'; import * as vscode from 'vscode'; +import { setCaptureLogger, disposeCaptureWorker } from './screenshotCapture'; import { LangData, poDirectorySetting } from './langData'; import { tr } from './localization'; import { cocoFeatureRelPaths, refreshCocoFeaturePath } from './cocoFeaturePath'; @@ -64,7 +65,9 @@ export function activate(context: vscode.ExtensionContext): void { // 性能日志输出通道：查看 → 输出 → ok-script-toolkit const cropLog = vscode.window.createOutputChannel('ok-script-toolkit'); setCropLogger((msg) => cropLog.appendLine(msg)); + setCaptureLogger((msg) => cropLog.appendLine(msg)); context.subscriptions.push(cropLog); + context.subscriptions.push({ dispose: disposeCaptureWorker }); // 初始化 worker 线程池（纯 JS 图像处理，主线程零阻塞） initCropWorkerPool(context.extensionPath); diff --git a/src/globalScreenshotHotkey.ts b/src/globalScreenshotHotkey.ts new file mode 100644 index 0000000..8fc742a --- /dev/null +++ b/src/globalScreenshotHotkey.ts @@ -0,0 +1,144 @@ +import * as fs from 'fs'; +import * as path from 'path'; +import { ChildProcess, spawn } from 'child_process'; + +const INITIAL_RETRY_DELAY_MS = 2_000; +const MAX_RETRY_DELAY_MS = 30_000; + +/** + * Windows 系统级 Ctrl+Alt+S 监听器。 + * + * VS Code 的 contributes.keybindings 只能在 VS Code 获得焦点时触发；这里启动一个 + * 仅依赖 Python 标准库的轻量辅助进程，通过 RegisterHotKey 收到系统级快捷键后把 + * TRIGGER 事件转回扩展宿主。截图本身仍由现有 command/面板链路完成。 + */ +export class GlobalScreenshotHotkey { + private child: ChildProcess | undefined; + private generation = 0; + private disposed = false; + private retryTimer: NodeJS.Timeout | undefined; + private retryDelayMs = INITIAL_RETRY_DELAY_MS; + + constructor( + private readonly extensionPath: string, + private readonly resolvePythonPath: () => string, + private readonly onTrigger: () => void, + private readonly log: (message: string) => void = (message) => console.log(`[ok-script Toolkit] ${message}`), + ) { + this.start(); + } + + dispose(): void { + if (this.disposed) return; + this.disposed = true; + this.clearRetry(); + this.stop(); + } + + private start(): void { + if (this.disposed || process.platform !== 'win32') return; + + this.clearRetry(); + const scriptPath = path.join(this.extensionPath, 'python', 'global_hotkey.py'); + if (!fs.existsSync(scriptPath)) { + this.log(`Global screenshot hotkey helper not found: ${scriptPath}`); + return; + } + + const generation = ++this.generation; + let stdoutBuffer = ''; + + try { + const pythonPath = this.resolvePythonPath(); + const child = spawn(pythonPath, [scriptPath, '--parent-pid', String(process.pid)], { + windowsHide: true, + stdio: ['ignore', 'pipe', 'pipe'], + }); + this.child = child; + + child.stdout?.setEncoding('utf8'); + child.stdout?.on('data', (chunk: string) => { + if (generation !== this.generation || this.disposed) return; + stdoutBuffer += chunk; + while (true) { + const newline = stdoutBuffer.indexOf('\n'); + if (newline < 0) break; + const line = stdoutBuffer.slice(0, newline).trim(); + stdoutBuffer = stdoutBuffer.slice(newline + 1); + this.handleLine(line); + } + }); + + child.stderr?.setEncoding('utf8'); + child.stderr?.on('data', (chunk: string) => { + if (generation === this.generation && !this.disposed) { + const text = chunk.trim(); + if (text) this.log(`Global screenshot hotkey helper stderr: ${text}`); + } + }); + + child.on('error', (error) => { + if (generation === this.generation && !this.disposed) { + this.log(`Failed to start global screenshot hotkey helper: ${error.message}`); + this.scheduleRetry(); + } + }); + + child.on('close', (code) => { + if (generation !== this.generation) return; + if (this.child === child) this.child = undefined; + if (!this.disposed && code !== 0) { + this.log(`Global screenshot hotkey helper exited with code ${code ?? 'unknown'}`); + this.scheduleRetry(); + } + }); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + this.log(`Failed to start global screenshot hotkey helper: ${message}`); + this.scheduleRetry(); + } + } + + private stop(): void { + ++this.generation; + const child = this.child; + this.child = undefined; + if (child && !child.killed) child.kill(); + } + + private clearRetry(): void { + if (!this.retryTimer) return; + clearTimeout(this.retryTimer); + this.retryTimer = undefined; + } + + private scheduleRetry(): void { + if (this.disposed || process.platform !== 'win32' || this.retryTimer) return; + + const delayMs = this.retryDelayMs; + this.retryDelayMs = Math.min(this.retryDelayMs * 2, MAX_RETRY_DELAY_MS); + this.retryTimer = setTimeout(() => { + this.retryTimer = undefined; + this.start(); + }, delayMs); + this.retryTimer.unref?.(); + this.log(`Retrying global screenshot hotkey registration after ${delayMs}ms`); + } + + private handleLine(line: string): void { + if (!line) return; + if (line === 'TRIGGER') { + this.onTrigger(); + return; + } + if (line === 'READY') { + this.retryDelayMs = INITIAL_RETRY_DELAY_MS; + this.log('Global screenshot hotkey registered: Ctrl+Alt+S'); + return; + } + if (line === 'UNSUPPORTED') return; + if (line.startsWith('ERROR ')) { + this.log(`Global screenshot hotkey unavailable: ${line.slice('ERROR '.length)}`); + } + } +} diff --git a/src/screenshotCapture.ts b/src/screenshotCapture.ts index f0230af..ed1a6a1 100644 --- a/src/screenshotCapture.ts +++ b/src/screenshotCapture.ts @@ -4,6 +4,21 @@ import * as vscode from 'vscode'; import { execFile } from 'child_process'; import { tr } from './localization'; import { resolveProjectDir } from './projectConfig'; +import { ScreenshotWorker } from './screenshotWorker'; + +let captureWorker: ScreenshotWorker | undefined; +let captureWorkerKey = ''; + +export function disposeCaptureWorker(): void { + captureWorker?.dispose(); + captureWorker = undefined; +} + +let captureLogger: (message: string) => void = () => {}; + +export function setCaptureLogger(logger: (message: string) => void): void { + captureLogger = logger; +} /** 窗口配置（从项目 config.py 的 windows 子字典提取） */ export interface WindowConfig { @@ -35,6 +50,25 @@ interface ProbeWindowConfigResult { label_enum_relative_path?: string | null; } +function windowConfigFromProbe(parsed: ProbeWindowConfigResult | undefined): WindowConfig | undefined { + if (!parsed?.ok) return undefined; + return { + exe: parsed.exe, title: parsed.title, hwnd_class: parsed.hwnd_class, top_hwnd_class: parsed.top_hwnd_class, + cocoFeatureJson: parsed.coco_feature_json ?? undefined, + labelEnumRelativePath: parsed.label_enum_relative_path ?? undefined, + }; +} + +function getCaptureWorker(pythonPath: string, scriptPath: string, projectDir = ''): ScreenshotWorker { + const key = JSON.stringify([pythonPath, scriptPath, projectDir]); + if (!captureWorker || captureWorkerKey !== key) { + disposeCaptureWorker(); + captureWorkerKey = key; + captureWorker = new ScreenshotWorker(pythonPath, scriptPath, projectDir, captureLogger); + } + return captureWorker; +} + /** 读取 okScriptToolkit 扩展配置中的项目路径和 Python 解释器 */ export function getProjectConfig(): { projectDir: string; pythonPath: string } { const cfg = vscode.workspace.getConfiguration('okScriptToolkit'); @@ -71,18 +105,7 @@ export async function probeWindowConfig(projectDir: string, pythonPath: string): for (let i = lines.length - 1; i >= 0; i--) { try { const parsed = JSON.parse(lines[i]) as ProbeWindowConfigResult; - if (parsed.ok) { - return { - exe: parsed.exe, - title: parsed.title, - hwnd_class: parsed.hwnd_class, - top_hwnd_class: parsed.top_hwnd_class, - // 探针解不出来时给的是 `null`（AST 里掺了变量），这里归一成 `undefined` - cocoFeatureJson: parsed.coco_feature_json ?? undefined, - labelEnumRelativePath: parsed.label_enum_relative_path ?? undefined, - }; - } - return undefined; + return windowConfigFromProbe(parsed); } catch { /* 跳过非 JSON 行 */ } } } catch { /* Python 不可用或脚本失败 */ } @@ -93,6 +116,24 @@ export type CaptureOutcome = | { ok: true } | { ok: false; reason: 'noScript' | 'cancelled' | 'failed'; error?: string }; +/** 面板显示或执行器就绪时后台预热；失败只写日志，不影响任务或弹出输入框。 */ +export async function prewarmScreenshotCapture(expectedProjectDir?: string, expectedPythonPath?: string): Promise { + try { + const { projectDir, pythonPath } = getProjectConfig(); + if (!projectDir || !fs.existsSync(projectDir) + || (expectedProjectDir !== undefined && expectedProjectDir !== projectDir) + || (expectedPythonPath !== undefined && expectedPythonPath !== pythonPath)) return; + const script = path.join(path.dirname(__dirname), 'python', 'capture_game_window.py'); + if (!fs.existsSync(script)) return; + const method = vscode.workspace.getConfiguration('okScriptToolkit').get('captureMethod') || 'auto'; + const started = Date.now(); + const frameReady = await getCaptureWorker(pythonPath, script, projectDir).prewarm(method); + captureLogger(`[capture] warmup_ms=${Date.now() - started} frame_ready=${frameReady}`); + } catch (error) { + captureLogger(`[capture] warmup failed: ${String(error)}`); + } +} + /** * 捕获游戏窗口到指定文件。 * @@ -101,10 +142,19 @@ export type CaptureOutcome = * 只负责把图片写到 outputPath，落库/入库由调用方处理。 */ export async function captureGameWindow(outputPath: string, methodOverride?: string): Promise { + const started = Date.now(); const { projectDir, pythonPath } = getProjectConfig(); + const captureScript = path.join(path.dirname(__dirname), 'python', 'capture_game_window.py'); + if (!fs.existsSync(captureScript)) return { ok: false, reason: 'noScript' }; let windowConfig: WindowConfig | undefined; if (projectDir) { - windowConfig = await probeWindowConfig(projectDir, pythonPath); + try { + windowConfig = windowConfigFromProbe(await getCaptureWorker(pythonPath, captureScript, projectDir).probe() as ProbeWindowConfigResult); + } catch (error) { + captureLogger(`[capture] probe failed: ${String(error)}`); + return { ok: false, reason: 'failed', error: String(error) }; + } + captureLogger(`[capture] probe_ms=${Date.now() - started}`); } let titleRegex: string | undefined; @@ -131,16 +181,11 @@ export async function captureGameWindow(outputPath: string, methodOverride?: str fs.mkdirSync(path.dirname(outputPath), { recursive: true }); - const captureScript = path.join(path.dirname(__dirname), 'python', 'capture_game_window.py'); - if (!fs.existsSync(captureScript)) { - return { ok: false, reason: 'noScript' }; - } - const windowConfigJson = (exeNames || hwndClass || titleRegex) ? JSON.stringify({ exe: exeNames, title: titleRegex, hwnd_class: hwndClass }) : undefined; - return runCaptureScript(captureScript, outputPath, titleRegex, windowConfigJson, projectDir || undefined, methodOverride); + return runCaptureScript(captureScript, outputPath, titleRegex, windowConfigJson, pythonPath, projectDir || undefined, methodOverride); } function runCaptureScript( @@ -148,36 +193,24 @@ function runCaptureScript( outputPath: string, titlePattern: string | undefined, configJson?: string, + pythonPath = 'python', projectDir?: string, methodOverride?: string, ): Promise { - return new Promise((resolve) => { - const args: string[] = [scriptPath, outputPath]; - if (configJson) { - args.push('--config-json', configJson); - } else if (titlePattern) { - args.push(titlePattern); - } - if (projectDir) { - args.push('--project-dir', projectDir); + const method = methodOverride + || vscode.workspace.getConfiguration('okScriptToolkit').get('captureMethod') + || 'auto'; + const worker = getCaptureWorker(pythonPath, scriptPath, projectDir); + const started = Date.now(); + const windowConfig = configJson ? JSON.parse(configJson) : { title: titlePattern }; + return worker.capture(outputPath, windowConfig, method).then((): CaptureOutcome => { + captureLogger(`[capture] request_ms=${Date.now() - started}`); + if (!fs.existsSync(outputPath)) { + return { ok: false, reason: 'failed', error: tr('Screenshot failed: file not created') }; } - - const method = methodOverride - || vscode.workspace.getConfiguration('okScriptToolkit').get('captureMethod') - || 'auto'; - args.push('--method', method); - - const { pythonPath } = getProjectConfig(); - execFile(pythonPath, args, { timeout: 30000 }, (error) => { - if (error) { - resolve({ ok: false, reason: 'failed', error: error.message }); - return; - } - if (!fs.existsSync(outputPath)) { - resolve({ ok: false, reason: 'failed', error: tr('Screenshot failed: file not created') }); - return; - } - resolve({ ok: true }); - }); + return { ok: true }; + }, (error: Error): CaptureOutcome => { + captureLogger(`[capture] ${error.message}`); + return { ok: false, reason: 'failed', error: error.message }; }); } diff --git a/src/screenshotWorker.ts b/src/screenshotWorker.ts new file mode 100644 index 0000000..828039c --- /dev/null +++ b/src/screenshotWorker.ts @@ -0,0 +1,131 @@ +import { ChildProcessWithoutNullStreams, spawn } from 'child_process'; +import * as readline from 'readline'; + +interface CaptureReply { + capture_response: number; + ok: boolean; + error?: string; + restart?: boolean; + config?: unknown; + frame_ready?: boolean; +} + +/** 截图专用解释器：串行请求、有限等待、闲置退出；项目源码变更时自动重启。 */ +export class ScreenshotWorker { + private child?: ChildProcessWithoutNullStreams; + private nextId = 0; + private pending?: { id: number; resolve: (reply: CaptureReply) => void; reject: (error: Error) => void; timer: NodeJS.Timeout }; + private idleTimer?: NodeJS.Timeout; + private queue: Promise = Promise.resolve(); + private disposed = false; + private warming?: { method: string; job: Promise }; + + constructor( + private readonly pythonPath: string, + private readonly scriptPath: string, + private readonly projectDir: string, + private readonly log: (line: string) => void, + private readonly timeoutMs = 30000, + private readonly idleMs = 60000, + ) {} + + capture(outputPath: string, windowConfig: unknown, method: string): Promise { + return this.submit({ output_path: outputPath, window_config: windowConfig, method }).then(() => {}); + } + + probe(): Promise { + return this.submit({ op: 'probe' }).then(reply => reply.config); + } + + prewarm(method: string): Promise { + if (this.warming?.method === method) return this.warming.job; + const job = this.submit({ op: 'warmup', method }).then(reply => reply.frame_ready === true); + this.warming = { method, job }; + const clear = () => { if (this.warming?.job === job) this.warming = undefined; }; + void job.then(clear, clear); + return job; + } + + private submit(payload: Record): Promise { + const job = this.queue.then(async () => { + if (this.disposed) throw new Error('Screenshot worker disposed'); + for (let attempt = 0; attempt < 2; attempt++) { + const reply = await this.request(payload); + if (reply.restart) { + this.stop(); + continue; + } + if (!reply.ok) throw new Error(reply.error || 'Screenshot worker failed'); + return reply; + } + throw new Error('Screenshot sources changed repeatedly'); + }); + this.queue = job.catch(() => {}); + return job; + } + + dispose(): void { + this.disposed = true; + this.stop(); + } + + private start(): ChildProcessWithoutNullStreams { + if (this.child) return this.child; + const child = spawn(this.pythonPath, [this.scriptPath, '--serve'], { + cwd: this.projectDir || undefined, + windowsHide: true, + env: { ...process.env, PYTHONIOENCODING: 'utf-8', PYTHONUTF8: '1' }, + stdio: 'pipe', + }); + this.child = child; + this.log(`[capture] worker_started python=${this.pythonPath} project=${this.projectDir}`); + const stdout = readline.createInterface({ input: child.stdout }); + const stderr = readline.createInterface({ input: child.stderr }); + stdout.on('line', line => { + if (this.child !== child) return; + try { + const reply = JSON.parse(line) as CaptureReply; + if (!this.pending || reply.capture_response !== this.pending.id) return; + const pending = this.pending; + this.pending = undefined; + clearTimeout(pending.timer); + this.idleTimer = setTimeout(() => this.stop(), this.idleMs); + this.idleTimer.unref(); + pending.resolve(reply); + } catch { this.log(`[capture] ${line}`); } + }); + stderr.on('line', line => this.log(`[capture] ${line}`)); + child.on('error', error => { if (this.child === child) this.stop(error); }); + child.stdin.on('error', error => { if (this.child === child) this.stop(error); }); + child.on('close', code => { + stdout.close(); + stderr.close(); + if (this.child === child) this.stop(new Error(`Screenshot worker exited (${code})`)); + }); + return child; + } + + private request(payload: Record): Promise { + if (this.idleTimer) clearTimeout(this.idleTimer); + const child = this.start(); + return new Promise((resolve, reject) => { + const id = ++this.nextId; + const timer = setTimeout(() => this.stop(new Error('Screenshot worker timed out')), this.timeoutMs); + this.pending = { id, resolve, reject, timer }; + child.stdin.write(JSON.stringify({ ...payload, id, project_dir: this.projectDir }) + '\n'); + }); + } + + private stop(error = new Error('Screenshot worker stopped')): void { + if (this.idleTimer) clearTimeout(this.idleTimer); + const child = this.child; + this.child = undefined; + child?.kill(); + if (this.pending) { + const pending = this.pending; + this.pending = undefined; + clearTimeout(pending.timer); + pending.reject(error); + } + } +} diff --git a/src/tempScreenshotPanel.ts b/src/tempScreenshotPanel.ts index 18f5d17..806ebfe 100644 --- a/src/tempScreenshotPanel.ts +++ b/src/tempScreenshotPanel.ts @@ -5,7 +5,7 @@ import { execFile } from 'child_process'; import { TempScreenshotStore, MAX_TEMP_SCREENSHOTS } from './tempScreenshotStore'; import { TemplateAssetData } from './templateAssetData'; import { repaintAllAssetGalleries } from './templateAssetPanel'; -import { captureGameWindow } from './screenshotCapture'; +import { captureGameWindow, prewarmScreenshotCapture } from './screenshotCapture'; import { clearPendingDrag, setPendingDrag } from './tempDrag'; import { cropTemplateThumbFileAsync, readImageSize, removeTemplateThumbFile, THUMB_HEIGHT } from './pngCrop'; import { injectWebviewLocalization, tr } from './localization'; @@ -243,7 +243,7 @@ class TempScreenshotController { if (this.busy) return; this.busy = true; try { - const outputPath = this.store.newFilePath(); + const outputPath = this.store.newFilePath('bmp'); const outcome = await captureGameWindow(outputPath, hardForeground ? 'foreground' : undefined); if (outcome.ok) { this.store.register(outputPath); @@ -268,7 +268,7 @@ class TempScreenshotController { async sendToAssets(id: string): Promise { const target = this.store.get(id); if (!target) return false; - const dst = this.assetData.importImageFile(target.filePath); + const dst = this.assetData.importImageFile(target.filePath, true); if (!dst) { this.notice('error', tr('Failed to add to template assets: {error}', { error: target.name })); return false; @@ -326,7 +326,10 @@ export class TempScreenshotViewProvider implements vscode.WebviewViewProvider { () => view.visible, ); controller.attachHtml(); - view.onDidChangeVisibility(() => { if (view.visible) void controller.refresh(); }); + if (view.visible) void prewarmScreenshotCapture(); + view.onDidChangeVisibility(() => { + if (view.visible) { void prewarmScreenshotCapture(); void controller.refresh(); } + }); view.onDidDispose(() => controller.dispose()); } } diff --git a/src/tempScreenshotStore.ts b/src/tempScreenshotStore.ts index 35dd151..7fdda93 100644 --- a/src/tempScreenshotStore.ts +++ b/src/tempScreenshotStore.ts @@ -17,13 +17,13 @@ const IMAGE_EXT_RE = /\.(png|jpe?g|bmp)$/i; /** 生成单调递增的文件名（时间戳 + 自增序号，保证按名字排序即为时间顺序）。 */ let seq = 0; -function nextName(): string { +function nextName(extension = 'png'): string { seq = (seq + 1) % 1000; - return `shot_${Date.now()}_${String(seq).padStart(3, '0')}.png`; + return `shot_${Date.now()}_${String(seq).padStart(3, '0')}.${extension}`; } /** - * 临时截图存储：把截图以 PNG 落盘到扩展 globalStorage 的工作区子目录， + * 临时截图存储：截图以无损 BMP 落盘，粘贴图片保持 PNG， * 列表直接从磁盘派生（无需额外索引文件），最多保留 {@link MAX_TEMP_SCREENSHOTS} 张。 * * 该存储被侧边栏视图与标注管理面板共享：侧边栏负责增删，标注管理面板负责 @@ -64,9 +64,9 @@ export class TempScreenshotStore { } /** 为外部写入者（如截图脚本）预留一个落盘路径，随后调用 register() 生效。 */ - newFilePath(): string { + newFilePath(extension: 'png' | 'bmp' = 'png'): string { this.ensure(); - return path.join(this.dir, nextName()); + return path.join(this.dir, nextName(extension)); } /** 写入 base64 PNG（可带 data URI 前缀）。 */ diff --git a/src/templateAssetData.ts b/src/templateAssetData.ts index 2890ff0..0e1d732 100644 --- a/src/templateAssetData.ts +++ b/src/templateAssetData.ts @@ -484,13 +484,19 @@ export class TemplateAssetData extends CocoAnnotationData { * 否则"导进来但一张框都没标"的图会立刻在标注文件里占一条空记录。 * 模板序号占位同时看磁盘与 COCO（见 `nextImageName`），所以不登记也不会撞名。 */ - importImageFile(srcPath: string): string | undefined { + importImageFile(srcPath: string, convertBmp = false): string | undefined { try { this.ensureTemplateFolder(); - const ext = path.extname(srcPath); + const convert = convertBmp && path.extname(srcPath).toLowerCase() === '.bmp'; + const ext = convert ? '.png' : path.extname(srcPath); const name = this.nextImageName() + ext; const dst = path.join(this.templateFolder, name); - fs.copyFileSync(srcPath, dst); + if (convert) { + const image = decodeRgba(fs.readFileSync(srcPath)); + fs.writeFileSync(dst, encodePngRgb(image.width, image.height, image.rgba)); + } else { + fs.copyFileSync(srcPath, dst); + } return dst; } catch { return undefined; diff --git a/src/templateAssetPanel.ts b/src/templateAssetPanel.ts index 95de5ea..04eb77d 100644 --- a/src/templateAssetPanel.ts +++ b/src/templateAssetPanel.ts @@ -6,7 +6,7 @@ import { AnnotationPanel } from './annotationPanel'; import { cropTemplateThumbFileAsync, THUMB_HEIGHT } from './pngCrop'; import { injectWebviewLocalization, tr } from './localization'; import { TempScreenshotStore } from './tempScreenshotStore'; -import { captureGameWindow, getProjectConfig, probeWindowConfig } from './screenshotCapture'; +import { captureGameWindow, getProjectConfig, probeWindowConfig, prewarmScreenshotCapture } from './screenshotCapture'; import { takePendingDrag } from './tempDrag'; import { ideSetting, labelEnumClassName, labelEnumPathInputError, labelEnumPathSetting, normalizeLabelEnumPathInput, setIdeSetting, templatesDirectory } from './projectConfig'; import { labelEnumRenameImpact, labelEnumRenameMessage, referencingFiles, writableClassName } from './labelEnumGuard'; @@ -809,7 +809,7 @@ class AssetGalleryController { void vscode.window.showWarningMessage(tr('The dragged screenshot is no longer available.')); return; } - const dst = this.data.importImageFile(target.filePath); + const dst = this.data.importImageFile(target.filePath, true); if (!dst) { void vscode.window.showErrorMessage(tr('Failed to add to template assets: {error}', { error: target.name })); return; @@ -853,7 +853,10 @@ export class TemplateAssetViewProvider implements vscode.WebviewViewProvider { this.tempStore, ); controller.attachHtml(); - view.onDidChangeVisibility(() => { if (view.visible) void controller.update(); }); + if (view.visible) void prewarmScreenshotCapture(); + view.onDidChangeVisibility(() => { + if (view.visible) { void prewarmScreenshotCapture(); void controller.update(); } + }); view.onDidDispose(() => controller.dispose()); } } @@ -866,6 +869,7 @@ export class TemplateAssetPanel { static show(data: TemplateAssetData, thumbDir: string, extensionUri: vscode.Uri, tempStore?: TempScreenshotStore): void { if (TemplateAssetPanel.current) { TemplateAssetPanel.current.panel.reveal(); + void prewarmScreenshotCapture(); void TemplateAssetPanel.current.controller.update(); return; } @@ -905,8 +909,9 @@ export class TemplateAssetPanel { readonly controller: AssetGalleryController, ) { controller.attachHtml(); + if (panel.visible) void prewarmScreenshotCapture(); panel.onDidChangeViewState((e) => { - if (e.webviewPanel.visible) void this.controller.update(); + if (e.webviewPanel.visible) { void prewarmScreenshotCapture(); void this.controller.update(); } }); panel.onDidDispose(() => { this.controller.dispose();