qiuruiandClaude Opus 4.6
5b50a6b1d3
- Add testing methodology and mouse-first architecture docs to CLAUDE.md - Harden safety.py with Unicode NFKC normalization, newline injection detection, and 10K character length limit - Create screen_state.py for detecting BIOS, lock screen, desktop, sleep states via screenshot hash analysis and OCR keyword matching - Integrate ScreenStateDetector into agent.py run_task() loop - Centralize desktop/app keyword constants in screen_state.py - Add 47 new unit tests (test_screen_state, test_mouse_ops, test_safety edge cases) - Migrate test_integration.py cleanup code from combo keys to mouse_ops - Fix pre-existing test_shortcut_action test to match blocked behavior - Add project evaluation table to CLAUDE.md known limitations Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
462 lines
19 KiB
Python
462 lines
19 KiB
Python
"""KVM Agent v2 — Computer Use automation with perception, validation, and memory.
|
|
|
|
Perceive-decide-act loop enhanced with:
|
|
- Local OCR perception for precise click targeting
|
|
- Pre/post screenshot validation with retry
|
|
- Fullscreen window management
|
|
- IME toggle for CJK input
|
|
- Step recording for template creation
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Optional
|
|
|
|
from .actions import Action
|
|
from .kvm_client import KVMClient
|
|
from .llm_planner import LLMPlanner
|
|
from .screen_state import ScreenStateDetector, PCState
|
|
from .validator import ValidationResult
|
|
from . import safety
|
|
from . import mouse_ops
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class StepRecord:
|
|
"""Data captured for each step — used by TemplateRecorder."""
|
|
|
|
step: int
|
|
action: Action
|
|
before_screenshot: bytes
|
|
after_screenshot: bytes
|
|
validation_success: bool
|
|
|
|
|
|
@dataclass
|
|
class TaskResult:
|
|
"""Result of a completed agent task."""
|
|
|
|
success: bool
|
|
steps_taken: int
|
|
final_reason: str
|
|
actions: List[str] = field(default_factory=list)
|
|
step_records: List[StepRecord] = field(default_factory=list)
|
|
|
|
|
|
class KVMAgent:
|
|
"""Main agent orchestrator with perception + validation + retry."""
|
|
|
|
MAX_WAIT_DELAY: float = 30.0
|
|
|
|
def __init__(
|
|
self,
|
|
kvm: KVMClient,
|
|
planner: LLMPlanner,
|
|
max_steps: int = 30,
|
|
step_delay: float = 1.0,
|
|
max_wait_delay: Optional[float] = None,
|
|
click_delay: float = 0.05,
|
|
force_control: bool = True,
|
|
# v2 additions — all optional for backward compat
|
|
perception=None,
|
|
validator=None,
|
|
window_manager=None,
|
|
config=None,
|
|
):
|
|
self.kvm = kvm
|
|
self.planner = planner
|
|
self.max_steps = max_steps
|
|
self.step_delay = step_delay
|
|
self.max_wait_delay = (
|
|
max_wait_delay if max_wait_delay is not None else self.MAX_WAIT_DELAY
|
|
)
|
|
self.click_delay = click_delay
|
|
self.force_control = force_control
|
|
self._running = False
|
|
|
|
# v2 components
|
|
self.perception = perception
|
|
self.validator = validator
|
|
self.window_manager = window_manager
|
|
self._config = config
|
|
self._ime_english = False # Track whether we've toggled IME to English
|
|
self._popup_enabled = True
|
|
self.state_detector = ScreenStateDetector()
|
|
if config:
|
|
self._popup_enabled = config.popup_detection_enabled
|
|
|
|
# ── Main loop ────────────────────────────────────────────────
|
|
|
|
async def run_task(self, task_description: str) -> TaskResult:
|
|
"""Execute a complete task using the perceive-decide-act loop."""
|
|
self._running = True
|
|
self._ime_english = False # Reset IME tracking for new task
|
|
action_history: List[str] = []
|
|
step_records: List[StepRecord] = []
|
|
|
|
try:
|
|
# 1. Enable privacy mode
|
|
logger.info("Enabling privacy mode")
|
|
await self.kvm.set_privacy_mode(True)
|
|
|
|
# 2. Acquire exclusive HID control
|
|
logger.info("Acquiring HID control (force=%s)", self.force_control)
|
|
await self.kvm.acquire_control(force=self.force_control)
|
|
|
|
# 3. Fullscreen the active window
|
|
if self.window_manager and self._should_fullscreen:
|
|
await self.window_manager.ensure_fullscreen()
|
|
|
|
for step in range(self.max_steps):
|
|
if not self._running:
|
|
return TaskResult(
|
|
success=False,
|
|
steps_taken=step,
|
|
final_reason="Stopped by user",
|
|
actions=action_history,
|
|
step_records=step_records,
|
|
)
|
|
|
|
# Safety gate
|
|
if await self.kvm.is_hid_blocked():
|
|
logger.warning("HID blocked by OCR safety — aborting task")
|
|
return TaskResult(
|
|
success=False,
|
|
steps_taken=step,
|
|
final_reason="HID blocked by safety system",
|
|
actions=action_history,
|
|
step_records=step_records,
|
|
)
|
|
|
|
# ── State check ────────────────────────────
|
|
before_screenshot = await self.kvm.screenshot()
|
|
if self.state_detector and self.perception:
|
|
pre_ocr = await self.perception.perceive(before_screenshot)
|
|
detection = self.state_detector.detect(
|
|
before_screenshot,
|
|
pre_ocr.raw_ocr_text,
|
|
bool(pre_ocr.elements),
|
|
)
|
|
if detection.state == PCState.SLEEP:
|
|
logger.info(
|
|
"Target PC sleeping (%s) — attempting wake",
|
|
detection.detail,
|
|
)
|
|
woke = await self.state_detector.wake_from_sleep(self.kvm)
|
|
if woke:
|
|
before_screenshot = await self.kvm.screenshot()
|
|
else:
|
|
logger.warning("Wake failed — aborting")
|
|
return TaskResult(
|
|
success=False,
|
|
steps_taken=step,
|
|
final_reason="Target PC asleep and wake failed",
|
|
actions=action_history,
|
|
step_records=step_records,
|
|
)
|
|
elif detection.state == PCState.LOCK_SCREEN:
|
|
logger.warning(
|
|
"Target PC locked (%s) — cannot proceed without PIN",
|
|
detection.detail,
|
|
)
|
|
|
|
# ── Perceive ─────────────────────────────────
|
|
logger.info(
|
|
"Step %d/%d: capturing screenshot", step + 1, self.max_steps
|
|
)
|
|
|
|
# Auto-dismiss popup if detected (does not consume a step)
|
|
if self.perception and self.validator and self._popup_enabled:
|
|
popup_scene = await self.perception.perceive(before_screenshot)
|
|
if self.validator.detect_popup(popup_scene):
|
|
logger.warning("Popup detected — auto-dismissing")
|
|
if self.window_manager:
|
|
await self.window_manager.dismiss_popup(
|
|
self.perception, before_screenshot)
|
|
await asyncio.sleep(0.5)
|
|
before_screenshot = await self.kvm.screenshot()
|
|
|
|
scene = None
|
|
scene_text = ""
|
|
if self.perception:
|
|
scene = await self.perception.perceive(before_screenshot)
|
|
scene_text = scene.to_text_summary()
|
|
|
|
# ── Decide ───────────────────────────────────
|
|
action = await self.planner.plan_action(
|
|
before_screenshot,
|
|
task_description,
|
|
step,
|
|
action_history,
|
|
max_steps=self.max_steps,
|
|
scene_text=scene_text,
|
|
)
|
|
|
|
action_desc = f"[{step + 1}] {action.type}: {action.reason}"
|
|
action_history.append(action_desc)
|
|
logger.info("Action: %s", action_desc)
|
|
|
|
# ── Done? ────────────────────────────────────
|
|
if action.type == "done":
|
|
return TaskResult(
|
|
success=True,
|
|
steps_taken=step + 1,
|
|
final_reason=action.reason,
|
|
actions=action_history,
|
|
step_records=step_records,
|
|
)
|
|
|
|
# ── Act + Validate ───────────────────────────
|
|
validated = await self._execute_with_validation(
|
|
action, before_screenshot, scene, step,
|
|
)
|
|
after_screenshot = await self.kvm.screenshot()
|
|
|
|
step_records.append(
|
|
StepRecord(
|
|
step=step,
|
|
action=action,
|
|
before_screenshot=before_screenshot,
|
|
after_screenshot=after_screenshot,
|
|
validation_success=validated,
|
|
)
|
|
)
|
|
|
|
await asyncio.sleep(self.step_delay)
|
|
|
|
return TaskResult(
|
|
success=False,
|
|
steps_taken=self.max_steps,
|
|
final_reason="Max steps reached",
|
|
actions=action_history,
|
|
step_records=step_records,
|
|
)
|
|
|
|
finally:
|
|
logger.info("Releasing HID control")
|
|
try:
|
|
await self.kvm.release_control()
|
|
except BaseException as e:
|
|
logger.warning("Failed to release control: %s", e)
|
|
try:
|
|
await self.kvm.set_privacy_mode(False)
|
|
except BaseException as e:
|
|
logger.warning("Failed to restore privacy mode: %s", e)
|
|
|
|
def stop(self):
|
|
"""Signal the agent to stop after the current step."""
|
|
self._running = False
|
|
|
|
# ── Execution with validation ────────────────────────────────
|
|
|
|
async def _execute_with_validation(
|
|
self, action: Action, before_screenshot: bytes, scene, step: int,
|
|
) -> bool:
|
|
"""Execute an action and validate it succeeded. Returns True on success."""
|
|
if action.type == "click":
|
|
return await self._click_with_retry(action, before_screenshot, scene)
|
|
|
|
# IME toggle before typing or IME-sensitive shortcuts (for CJK Windows)
|
|
if self._should_toggle_ime:
|
|
if action.type == "type":
|
|
await self._ensure_english_ime()
|
|
elif (action.type == "shortcut"
|
|
and action.shortcut.lower().strip() in self._IME_SENSITIVE_SHORTCUTS):
|
|
await self._ensure_english_ime()
|
|
|
|
await self._execute_action(action)
|
|
|
|
# Validate non-click actions
|
|
if self.validator and action.type in ("type", "shortcut"):
|
|
after = await self.kvm.screenshot()
|
|
if action.type == "type":
|
|
v = await self.validator.validate_type(
|
|
before_screenshot, after, action.text,
|
|
)
|
|
else:
|
|
v = await self.validator.validate_shortcut(
|
|
before_screenshot, after, action.shortcut,
|
|
)
|
|
return v.result.value == "success"
|
|
|
|
return True # No validator or non-validatable action
|
|
|
|
async def _click_with_retry(
|
|
self, action: Action, before: bytes, scene,
|
|
) -> bool:
|
|
"""Click with OCR-assisted coordinates and fresh-screenshot retry.
|
|
|
|
Strategy:
|
|
Attempt 0: Use OCR from existing scene (zero cost), fall back to LLM coords
|
|
Attempt 1+: Take fresh screenshot + re-OCR (screen may have changed)
|
|
"""
|
|
max_retries = 3
|
|
if self._config:
|
|
max_retries = self._config.click_max_retries
|
|
|
|
for attempt in range(max_retries):
|
|
x, y = action.x, action.y
|
|
|
|
# Always prefer OCR coordinates when available
|
|
if action.reason and self.perception:
|
|
if attempt == 0 and scene:
|
|
# First attempt: use existing scene data (zero cost)
|
|
matches = scene.find_by_text(action.reason)
|
|
if matches and matches[0].bbox != (0.0, 0.0, 0.0, 0.0):
|
|
x, y = matches[0].center_x, matches[0].center_y
|
|
logger.info(
|
|
"OCR-first click (%.3f, %.3f) for '%s'",
|
|
x, y, action.reason,
|
|
)
|
|
else:
|
|
# Retry: fresh screenshot + re-OCR (stale data = stale coords)
|
|
fresh = await self.kvm.screenshot()
|
|
element = await self.perception.find_element(
|
|
fresh, target_text=action.reason,
|
|
)
|
|
if element and element.bbox != (0.0, 0.0, 0.0, 0.0):
|
|
x, y = element.center_x, element.center_y
|
|
logger.info(
|
|
"Retry %d: OCR → (%.3f, %.3f) for '%s'",
|
|
attempt, x, y, action.reason,
|
|
)
|
|
|
|
await self.kvm.mouse_move(x, y)
|
|
await asyncio.sleep(self.click_delay)
|
|
await self.kvm.mouse_click(action.button)
|
|
|
|
if not self.validator:
|
|
return True # No validator — assume success
|
|
|
|
after = await self.kvm.screenshot()
|
|
v = await self.validator.validate_click(
|
|
before, after, target_text=action.reason,
|
|
)
|
|
if v.result.value == "success":
|
|
logger.info("Click validated on attempt %d", attempt + 1)
|
|
return True
|
|
|
|
if v.result == ValidationResult.POPUP_DETECTED:
|
|
logger.warning("Popup after click — dismissing and retrying")
|
|
if self.window_manager:
|
|
await self.window_manager.dismiss_popup(
|
|
self.perception, after)
|
|
before = await self.kvm.screenshot()
|
|
continue # Retry the click
|
|
|
|
logger.info("Click attempt %d: %s", attempt + 1, v.detail)
|
|
before = after # Update baseline for next validation
|
|
await asyncio.sleep(0.3)
|
|
|
|
logger.warning("Click failed after %d retries", max_retries)
|
|
return False
|
|
|
|
# ── IME toggle ───────────────────────────────────────────────
|
|
|
|
# Shortcuts that are intercepted by Chinese IME and need English mode
|
|
_IME_SENSITIVE_SHORTCUTS = {"enter", "return", "space"}
|
|
|
|
async def _ensure_english_ime(self):
|
|
"""Detect IME state via OCR and switch to English if Chinese is active.
|
|
|
|
Delegates to kvm.ensure_english_ime() and updates local tracking.
|
|
"""
|
|
result = await self.kvm.ensure_english_ime()
|
|
self._ime_english = result
|
|
|
|
# ── Legacy execute (safety-checked) ──────────────────────────
|
|
|
|
async def _execute_action(self, action: Action):
|
|
"""Execute a single action via KVM HID API with safety checks."""
|
|
try:
|
|
if action.type == "click":
|
|
await self.kvm.mouse_move(action.x, action.y)
|
|
await asyncio.sleep(self.click_delay)
|
|
await self.kvm.mouse_click(action.button)
|
|
|
|
elif action.type == "type":
|
|
if not safety.validate_type_action(action.text):
|
|
logger.warning("BLOCKED: unsafe text rejected by safety layer")
|
|
return
|
|
await self.kvm.keyboard_type(action.text)
|
|
|
|
elif action.type == "shortcut":
|
|
if not safety.validate_shortcut(action.shortcut):
|
|
# Try translating blocked combo key to mouse operation
|
|
translated = await self._shortcut_to_mouse(action.shortcut)
|
|
if not translated:
|
|
logger.warning(
|
|
"BLOCKED: shortcut '%s' has no mouse equivalent",
|
|
action.shortcut,
|
|
)
|
|
return
|
|
await self.kvm.keyboard_shortcut(action.shortcut)
|
|
|
|
elif action.type == "scroll":
|
|
await self.kvm.mouse_scroll(action.delta)
|
|
|
|
elif action.type == "wait":
|
|
capped_delay = max(0, min(action.delay, self.max_wait_delay))
|
|
await asyncio.sleep(capped_delay)
|
|
|
|
else:
|
|
logger.warning("Unknown action type: %s", action.type)
|
|
except Exception as e:
|
|
logger.warning("Action '%s' failed: %s", action.type, e)
|
|
|
|
# ── Shortcut → mouse translation ─────────────────────────────
|
|
|
|
async def _shortcut_to_mouse(self, shortcut: str) -> bool:
|
|
"""Translate a blocked combo key into an equivalent mouse operation.
|
|
|
|
Returns True if a mouse equivalent was executed, False if no mapping.
|
|
"""
|
|
key = shortcut.lower().strip()
|
|
if key in ("alt+f4",):
|
|
if self.perception:
|
|
return await mouse_ops.close_window(self.kvm, self.perception)
|
|
elif key in ("win+d",):
|
|
await mouse_ops.show_desktop(self.kvm)
|
|
return True
|
|
elif key in ("ctrl+s",):
|
|
if self.perception:
|
|
return await mouse_ops.click_element(
|
|
self.kvm, "保存", self.perception)
|
|
elif key in ("ctrl+w",):
|
|
if self.perception:
|
|
return await mouse_ops.close_window(self.kvm, self.perception)
|
|
elif key in ("win+r", "win+s"):
|
|
# Cannot launch without app name — just log
|
|
logger.info("Shortcut '%s' blocked; use click-based app launch", key)
|
|
return False
|
|
elif key in ("ctrl+c",):
|
|
if self.perception:
|
|
await mouse_ops.right_click(self.kvm, 0.5, 0.5)
|
|
await asyncio.sleep(0.3)
|
|
return await mouse_ops.click_element(
|
|
self.kvm, "复制", self.perception)
|
|
elif key in ("ctrl+v",):
|
|
if self.perception:
|
|
await mouse_ops.right_click(self.kvm, 0.5, 0.5)
|
|
await asyncio.sleep(0.3)
|
|
return await mouse_ops.click_element(
|
|
self.kvm, "粘贴", self.perception)
|
|
return False
|
|
|
|
# ── Property helpers ─────────────────────────────────────────
|
|
|
|
@property
|
|
def _should_fullscreen(self) -> bool:
|
|
if self._config:
|
|
return self._config.fullscreen_enabled
|
|
return True
|
|
|
|
@property
|
|
def _should_toggle_ime(self) -> bool:
|
|
if self._config:
|
|
return self._config.ime_toggle_before_type
|
|
return False
|