从工具使用到系统设计，
继续学习 Agent 与软硬件协作。

"""
Minimal Autonomous Agent Runtime Skeleton (可终止进程执行与契约校验参考实现)
Author: Yvain Zhang
Architecture: ReAct FSM + Context Budget + Schema Validation + Process Watchdog + External Verifier
Note: 用于阅读控制流程；权限隔离、完整预算与超时处理需另行实现。
"""
import os
import sys
import json
import re
import time
import signal
import multiprocessing as mp
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, Dict, Any, List, Optional

class ToolStatus(Enum):
    SUCCESS = "SUCCESS"
    FAILURE = "FAILURE"
    TIMEOUT = "TIMEOUT"
    VALIDATION_ERROR = "VALIDATION_ERROR"
    SECURITY_BLOCK = "SECURITY_BLOCK"

@dataclass
class ToolResult:
    status: ToolStatus
    output: str
    exit_code: int = 0
    error: Optional[str] = None

@dataclass
class ToolSpec:
    name: str
    description: str
    schema: Dict[str, Any]  # JSON Schema 参数契约
    handler: Callable[..., Any]
    is_destructive: bool = False

def sanitize_timeout(timeout_val: Any, default_sec: float = 30.0, min_sec: float = 0.001, max_sec: float = 300.0) -> float:
    """校验并归一化超时参数，防止 null、非法字符串或负值导致看门狗崩溃"""
    if timeout_val is None:
        return default_sec
    try:
        t = float(timeout_val)
        if t <= 0:
            return default_sec
        return min(max(t, min_sec), max_sec)
    except (ValueError, TypeError):
        return default_sec

def validate_schema(schema: Dict[str, Any], args: Any) -> Optional[str]:
    """轻量级参数契约校验器：先验证 args 为字典对象，再检查必填项与字段类型，拦截非法入参"""
    if args is None:
        return "参数对象为空 (null)，预期为包含工具入参的 JSON 字典对象"
    if not isinstance(args, dict):
        return f"参数格式非法: 预期为字典 (JSON Object)，实际传入 {type(args).__name__}"

    for field in schema.get("required", []):
        if field not in args:
            return f"缺少必填参数 '{field}'"
    properties = schema.get("properties", {})
    type_map = {"string": str, "integer": int, "number": (int, float), "boolean": bool, "array": list, "object": dict}
    for key, val in args.items():
        if key in properties:
            expected_type_name = properties[key].get("type")
            expected_type = type_map.get(expected_type_name)
            if expected_type_name == "integer" and isinstance(val, bool):
                return f"参数 '{key}' 应为 integer，不能传入 bool"
            if expected_type and not isinstance(val, expected_type):
                return f"参数 '{key}' 类型错误: 预期 {expected_type_name}，实际为 {type(val).__name__}"
    return None

class ContextBudgetManager:
    """按近期轮数裁剪历史；未实现 Token 计数或历史内容摘要"""
    def __init__(self, max_tokens: int = 128000, keep_recent_turns: int = 4):
        self.max_tokens = max_tokens
        self.keep_recent_turns = keep_recent_turns

    def prune(self, messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
        if len(messages) <= (self.keep_recent_turns * 2 + 2):
            return messages
        prefix = messages[:2]  # 保持 System Prompt + 根目标不可变，最大化 Prefix Cache 命中
        recent = messages[-(self.keep_recent_turns * 2):]
        compact_summary = {
            "role": "system",
            "content": "[SYSTEM COMPACT: 前序探索步骤已折叠。请专注于当前激活的任务分支与测试断言。]"
        }
        return prefix + [compact_summary] + recent

def _kill_process_tree(proc: mp.Process):
    """尝试终止 worker 及同组进程；不覆盖脱离该进程组的后代"""
    pid = proc.pid
    if not pid:
        return

    parent_pgid = os.getpgrp()

    # 1. 尝试向 worker 建立的独立进程组发送 SIGTERM（优雅终止）
    # 若 worker 的 os.setsid() 成功，其 PGID 等于 pid
    if pid != parent_pgid:
        try:
            os.killpg(pid, signal.SIGTERM)
        except (ProcessLookupError, PermissionError):
            pass

    # 兜底直接向 worker 进程发送 SIGTERM
    try:
        proc.terminate()
    except Exception:
        pass

    # 给予短暂优雅退出窗口
    proc.join(timeout=0.05)

    # 2. 再向该进程组发送 SIGKILL
    # 同组内忽略 SIGTERM 的后代进程也会收到该信号
    if pid != parent_pgid:
        try:
            os.killpg(pid, signal.SIGKILL)
        except (ProcessLookupError, PermissionError):
            pass

    # 兜底对 worker 发送 SIGKILL
    try:
        if proc.is_alive():
            proc.kill()
    except Exception:
        pass

    try:
        proc.join(timeout=0.05)
    except Exception:
        pass

class ToolDispatcher:
    """部分参数类型校验、进程组执行与超时检查示例；未实现权限沙箱"""
    def __init__(self):
        self.registry: Dict[str, ToolSpec] = {}
        self.dangerous_patterns = [
            re.compile(r"\brm\s+-rf\s+/", re.IGNORECASE),
            re.compile(r"\bmkfs(\.\w+)?\b", re.IGNORECASE),
            re.compile(r"\bdd\s+if=", re.IGNORECASE),
            re.compile(r":\(\)\{.*\}")
        ]

    def register(self, spec: ToolSpec):
        self.registry[spec.name] = spec

    def dispatch(self, name: str, args: Any, timeout_sec: Any = 30.0, allow_destructive: bool = False) -> ToolResult:
        if name not in self.registry:
            return ToolResult(status=ToolStatus.FAILURE, output="", exit_code=1, error=f"未注册的工具: {name}")

        spec = self.registry[name]

        # 1. 契约校验：在调用前拦截 null、非对象或缺失必填项的参数
        schema_err = validate_schema(spec.schema, args)
        if schema_err:
            return ToolResult(status=ToolStatus.VALIDATION_ERROR, output="", exit_code=2, error=f"参数契约校验失败: {schema_err}")

        # 2. 超时参数归一化与破坏性操作安全检查
        clean_timeout = sanitize_timeout(timeout_sec)

        if spec.is_destructive and not allow_destructive:
            return ToolResult(status=ToolStatus.SECURITY_BLOCK, output="", exit_code=126, error=f"安全阻断: 工具 '{name}' 涉及破坏性操作，未获显式授权。")

        for val in args.values():
            if isinstance(val, str):
                for pat in self.dangerous_patterns:
                    if pat.search(val):
                        return ToolResult(status=ToolStatus.SECURITY_BLOCK, output="", exit_code=126, error=f"安全阻断: 检测到高危指令模式 '{pat.pattern}'")

        # 3. 在 worker 进程中执行工具，并尝试建立独立进程组
        ctx = mp.get_context("fork" if hasattr(os, "fork") else None)
        parent_conn, child_conn = ctx.Pipe()

        def _worker(conn, fn, kwargs):
            # 脱离父会话成为独立 Process Group Leader，使后续派生的 shell/子命令归入同组
            try:
                os.setsid()
            except Exception:
                pass
            try:
                r = fn(**kwargs)
                conn.send((True, r))
            except Exception as e:
                conn.send((False, str(e)))
            finally:
                try:
                    conn.close()
                except Exception:
                    pass

        proc = ctx.Process(target=_worker, args=(child_conn, spec.handler, args))
        proc.start()
        child_conn.close()

        # 等待管道可读；后续 recv 的完整读取还需单独处理超时与输出上限
        if not parent_conn.poll(clean_timeout):
            # 等待超时：尝试终止 worker 及其同组进程
            _kill_process_tree(proc)
            parent_conn.close()
            return ToolResult(status=ToolStatus.TIMEOUT, output="", exit_code=124, error=f"工具执行超时 (超过 {clean_timeout} 秒)，已尝试终止工具进程及其同进程组内的子进程。")

        # 管道就绪后立即读取排空数据，解除 worker 端 send 阻塞，随后回收进程
        try:
            ok, val = parent_conn.recv()
            parent_conn.close()
            proc.join(timeout=0.5)
            if proc.is_alive():
                _kill_process_tree(proc)

            if ok:
                return val if isinstance(val, ToolResult) else ToolResult(status=ToolStatus.SUCCESS, output=str(val))
            return ToolResult(status=ToolStatus.FAILURE, output="", exit_code=1, error=str(val))
        except EOFError:
            parent_conn.close()
            _kill_process_tree(proc)
            return ToolResult(status=ToolStatus.FAILURE, output="", exit_code=proc.exitcode or 1, error=f"工具进程异常退出 (exit code: {proc.exitcode})，未接收到返回数据。")
        except Exception as e:
            parent_conn.close()
            _kill_process_tree(proc)
            return ToolResult(status=ToolStatus.FAILURE, output="", exit_code=1, error=str(e))

class AutonomousRunner:
    """Agent 主控制循环：感知 -> 推理 -> 执行 -> 观察 -> 外部独立验收"""
    def __init__(self, llm_client, dispatcher: ToolDispatcher, verifier: Optional[Callable[[], ToolResult]] = None, max_steps: int = 25):
        self.llm = llm_client
        self.dispatcher = dispatcher
        self.verifier = verifier  # 外部验收器，检查结果是否符合任务要求
        self.budget = ContextBudgetManager()
        self.max_steps = max_steps
        self.history: List[Dict[str, str]] = []

    def execute_goal(self, goal: str, system_prompt: str) -> Dict[str, Any]:
        self.history = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"目标任务: {goal}\n规则: 请逐步执行。只有当客观验收断言真正通过后，任务方可结束。"}
        ]

        for step in range(1, self.max_steps + 1):
            # 1. 缓存对齐与上下文裁剪装配
            active_context = self.budget.prune(self.history)

            # 2. 大模型决策推理
            decision = self.llm.generate_decision(messages=active_context)

            # 3. 模型声明完成后，调用外部验收器；未配置时返回未验证状态
            if decision.get("is_complete"):
                if self.verifier is None:
                    return {"status": "SUCCESS_UNVERIFIED", "steps": step, "output": decision.get("final_summary"), "warning": "未配置独立验证器，结果未经验收断言"}

                verify_res = self.verifier()
                if verify_res.status == ToolStatus.SUCCESS:
                    return {"status": "SUCCESS", "steps": step, "output": decision.get("final_summary")}
                else:
                    obs_entry = f"[第 {step} 步独立验收失败] 退出码: {verify_res.exit_code}\n断言报错: {verify_res.error or verify_res.output}\n系统判定: 任务尚未完成，模型不可提前声明成功。请继续修复。"
                    self.history.append({"role": "assistant", "content": json.dumps(decision)})
                    self.history.append({"role": "user", "content": obs_entry})
                    continue

            # 4. 工具分发执行
            tool_name = decision.get("tool_name", "")
            tool_args = decision.get("tool_args")  # 保留原始值，交由 dispatcher 严格校验
            timeout_sec = decision.get("timeout_sec", 30.0)
            tool_result = self.dispatcher.dispatch(tool_name, tool_args, timeout_sec=timeout_sec)

            # 5. 观察清洗与结构化记忆追加
            obs_preview = tool_result.output[:800] if tool_result.output else "[无标准输出]"
            obs_entry = f"[第 {step} 步观察] 状态: {tool_result.status.value}, 返回码: {tool_result.exit_code}\n输出截取: {obs_preview}"
            if tool_result.error:
                obs_entry += f"\n错误跟踪: {tool_result.error}"

            self.history.append({"role": "assistant", "content": json.dumps(decision)})
            self.history.append({"role": "user", "content": obs_entry})

        return {"status": "BUDGET_EXHAUSTED", "steps": self.max_steps, "error": "超出最大步数预算，系统未能在限制步数内收敛。"}