365 lines
16 KiB
Python
365 lines
16 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
orchestrator.py — Pi Agent 兜底能力 P0 PoC：Python 编排器
|
||
=========================================================
|
||
|
||
职责（对应 `docs/architecture/fallback.md` 的 FallbackLane 编排与熔断设计）：
|
||
|
||
1. 以子进程拉起 Pi headless（`node .../cli.js -p --mode json`），逐行解析 JSONL 事件流；
|
||
2. 三重熔断：超时 / 步数上限 / 输出体量上限，触发即杀进程树（Windows taskkill /T /F）
|
||
并把运行显式标记为 failed（绝不把失败包装成成功 —— GOAL.md 价值观）；
|
||
3. 环境变量白名单清洗后传给子进程（剥离 CONDA_*/PYTHONHOME/PYTHONPATH 等），
|
||
并设置 PI_CODING_AGENT_DIR 指向 poc 内配置目录（NOTES-pi-runtime.md §3）；
|
||
4. 成败判定基于事件流里的 stopReason（"stop"=正常 / "error"=失败），
|
||
**绝不相信进程退出码**（pi headless 失败也退 0 —— Agent-A 实测坑 #1）；
|
||
5. 全程写运行日志到 runs/<runId>/：events.jsonl（原始事件流）、orchestrator.log、result.json。
|
||
|
||
mock 模式说明：AgentRunner 是协议抽象，mock executor（demo_task.py）产出与真实
|
||
pi 相同结构的事件序列，因此熔断/日志/stopReason 判定走**同一条代码路径**。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import queue
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Callable, Iterable, Iterator, Optional
|
||
|
||
from sandbox import POC_ROOT, RUNS_ROOT, PI_HOME, PI_CLI, create_run_dirs
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 集中配置（改行为只改这里）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@dataclass
|
||
class BreakerConfig:
|
||
"""三重熔断阈值（对应方案 §4.6「熔断三闸」）。"""
|
||
timeout_sec: float = 600.0 # 闸 1：单次运行超时（默认 10min，方案 §S10）
|
||
max_steps: int = 50 # 闸 2：工具调用步数上限（方案 §4.6 步数 50）
|
||
max_output_bytes: int = 10 * 1024 * 1024 # 闸 3：assistant 输出累计体量上限（10MiB）
|
||
poll_interval_sec: float = 1.0 # 读事件流的轮询间隔（也用于检测进程挂起）
|
||
|
||
|
||
# 子进程环境变量白名单（L4 进程层围墙，对应方案 §4.2 buildChildEnv 语义）
|
||
ENV_WHITELIST = (
|
||
"PATH", "PATHEXT", "SYSTEMROOT", "SYSTEMDRIVE", "WINDIR", "COMSPEC",
|
||
"TEMP", "TMP", "USERPROFILE", "APPDATA", "LOCALAPPDATA", "HOME",
|
||
"NODE_OPTIONS", # node 运行时可能需要；其余一律不传
|
||
)
|
||
|
||
# 显式剥离名单（即使同名出现在白名单也再剥一层，双保险）
|
||
ENV_STRIP_PREFIXES = ("CONDA_", "PYTHON", "PIP_", "VIRTUAL_ENV")
|
||
|
||
|
||
def build_child_env(extra: Optional[dict] = None) -> dict:
|
||
"""
|
||
构造传给 Pi 子进程的清洗后环境（白名单制）。
|
||
|
||
- 只放行 ENV_WHITELIST 里的变量；
|
||
- 剥离 CONDA_*/PYTHONHOME/PYTHONPATH/PIP_*/VIRTUAL_ENV*（防止 Anaconda 污染 node 子进程）；
|
||
- 强制设置 PI_CODING_AGENT_DIR 指向 poc 内 runtime/pi-home（配置圈禁，Agent-A 实测机制）。
|
||
"""
|
||
env: dict = {}
|
||
for key in ENV_WHITELIST:
|
||
if key in os.environ and not key.upper().startswith(ENV_STRIP_PREFIXES):
|
||
env[key] = os.environ[key]
|
||
env["PI_CODING_AGENT_DIR"] = str(PI_HOME)
|
||
if extra:
|
||
for k, v in extra.items():
|
||
if not k.upper().startswith(ENV_STRIP_PREFIXES):
|
||
env[k] = v
|
||
return env
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 运行结果
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@dataclass
|
||
class RunResult:
|
||
"""一次兜底运行的最终判定。success 只由 stopReason 决定。"""
|
||
run_id: str
|
||
success: bool
|
||
stop_reason: str = "" # "stop" / "error" / "breaker:<哪种熔断>" / "harness_error"
|
||
error_message: str = ""
|
||
steps: int = 0 # 工具调用步数（tool_execution_start 计数）
|
||
output_bytes: int = 0 # assistant 输出累计字节
|
||
elapsed_sec: float = 0.0
|
||
final_text: str = "" # 最后一条 assistant 消息文本（即 Pi 的报告）
|
||
run_dir: str = ""
|
||
extra: dict = field(default_factory=dict)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 事件流来源抽象
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# AgentRunner：给定 (task, work_dir, event_sink) 产出 JSONL 事件 dict 迭代器。
|
||
# 真实实现 = pi headless 子进程；mock 实现 = demo_task.py 的纯 Python executor。
|
||
AgentRunner = Callable[[str, Path], Iterator[dict]]
|
||
|
||
|
||
class Orchestrator:
|
||
"""编排器本体：消费事件流，执行熔断与判定，落日志。"""
|
||
|
||
def __init__(self, config: Optional[BreakerConfig] = None):
|
||
self.cfg = config or BreakerConfig()
|
||
|
||
# -- 主入口 -----------------------------------------------------------
|
||
|
||
def run(self, task: str, runner: AgentRunner, run_id: Optional[str] = None,
|
||
on_tool_event: Optional[Callable[[dict], None]] = None) -> RunResult:
|
||
"""
|
||
执行一次兜底运行。
|
||
|
||
参数:
|
||
task: 任务简报（自然语言）。
|
||
runner: AgentRunner，产出事件流（真实或 mock 共用本条熔断路径）。
|
||
run_id: 缺省自动生成；目录 runs/<run_id>/ 由 sandbox.create_run_dirs 建好。
|
||
on_tool_event: 每个 tool_execution_start/end 事件的回调（tool_bridge 借此签发 callId）。
|
||
|
||
返回 RunResult；任何熔断/错误都会显式写进 result.json 与 orchestrator.log。
|
||
"""
|
||
run_id = run_id or time.strftime("run-%Y%m%d-%H%M%S-") + uuid.uuid4().hex[:6]
|
||
dirs = create_run_dirs(run_id)
|
||
run_dir = dirs["root"]
|
||
log_path = run_dir / "orchestrator.log"
|
||
events_path = run_dir / "events.jsonl"
|
||
|
||
def log(msg: str) -> None:
|
||
line = f"[{time.strftime('%H:%M:%S')}] {msg}\n"
|
||
with open(log_path, "a", encoding="utf-8") as f:
|
||
f.write(line)
|
||
|
||
res = RunResult(run_id=run_id, success=False, run_dir=str(run_dir))
|
||
t0 = time.monotonic()
|
||
log(f"run_id={run_id} task={task[:120]!r}")
|
||
log(f"breaker config: {self.cfg}")
|
||
|
||
breaker_tripped: Optional[str] = None
|
||
last_assistant_text = ""
|
||
stop_reason = ""
|
||
error_message = ""
|
||
|
||
try:
|
||
with open(events_path, "w", encoding="utf-8") as evf:
|
||
for event in runner(task, dirs["work"]):
|
||
elapsed = time.monotonic() - t0
|
||
|
||
# —— 闸 1：超时 ——
|
||
if elapsed > self.cfg.timeout_sec:
|
||
breaker_tripped = f"breaker:timeout({elapsed:.1f}s>{self.cfg.timeout_sec}s)"
|
||
self._safe_write_event(evf, {"type": "breaker", "reason": breaker_tripped})
|
||
break
|
||
|
||
self._safe_write_event(evf, event)
|
||
etype = event.get("type", "")
|
||
|
||
# —— 步数统计 + 闸 2 ——
|
||
if etype == "tool_execution_start":
|
||
res.steps += 1
|
||
if on_tool_event:
|
||
on_tool_event(event)
|
||
if res.steps > self.cfg.max_steps:
|
||
breaker_tripped = f"breaker:max_steps({res.steps}>{self.cfg.max_steps})"
|
||
self._safe_write_event(evf, {"type": "breaker", "reason": breaker_tripped})
|
||
break
|
||
elif etype == "tool_execution_end" and on_tool_event:
|
||
on_tool_event(event)
|
||
|
||
# —— 输出体量统计 + 闸 3 ——
|
||
text_delta = self._extract_text_delta(event)
|
||
if text_delta:
|
||
res.output_bytes += len(text_delta.encode("utf-8"))
|
||
if res.output_bytes > self.cfg.max_output_bytes:
|
||
breaker_tripped = (
|
||
f"breaker:max_output({res.output_bytes}>{self.cfg.max_output_bytes})"
|
||
)
|
||
self._safe_write_event(evf, {"type": "breaker", "reason": breaker_tripped})
|
||
break
|
||
|
||
# —— stopReason 判定（成败唯一权威，Agent-A 坑 #1）——
|
||
sr, em, txt = self._extract_stop(event)
|
||
if sr:
|
||
stop_reason, error_message = sr, em
|
||
if txt:
|
||
last_assistant_text = txt
|
||
|
||
# 自动重试事件：只记录（Agent-A 坑 #4，重试 ~14s 已计入超时预算）
|
||
if etype == "auto_retry_start":
|
||
log(f"auto_retry_start attempt={event.get('attempt')}")
|
||
|
||
if etype == "agent_end":
|
||
break
|
||
|
||
except Exception as exc: # runner 自身抛错（如子进程 spawn 失败）
|
||
stop_reason = "harness_error"
|
||
error_message = f"{type(exc).__name__}: {exc}"
|
||
log(f"HARNESS ERROR: {error_message}")
|
||
|
||
res.elapsed_sec = time.monotonic() - t0
|
||
res.final_text = last_assistant_text
|
||
|
||
if breaker_tripped:
|
||
res.stop_reason = breaker_tripped
|
||
res.error_message = "熔断触发，运行显式标记失败"
|
||
log(f"BREAKER TRIPPED: {breaker_tripped} -> failed")
|
||
elif stop_reason == "stop":
|
||
res.success = True
|
||
res.stop_reason = "stop"
|
||
log(f"OK stopReason=stop steps={res.steps} out={res.output_bytes}B "
|
||
f"elapsed={res.elapsed_sec:.1f}s")
|
||
else:
|
||
res.stop_reason = stop_reason or "error:no_stop_reason"
|
||
res.error_message = error_message or "事件流未给出 stopReason=stop，按失败处理"
|
||
log(f"FAILED stopReason={res.stop_reason} err={res.error_message}")
|
||
|
||
with open(run_dir / "result.json", "w", encoding="utf-8") as f:
|
||
json.dump(res.__dict__, f, ensure_ascii=False, indent=2)
|
||
return res
|
||
|
||
# -- 事件解析辅助 ------------------------------------------------------
|
||
|
||
@staticmethod
|
||
def _safe_write_event(evf, event: dict) -> None:
|
||
try:
|
||
evf.write(json.dumps(event, ensure_ascii=False) + "\n")
|
||
except Exception:
|
||
evf.write(json.dumps({"type": "unserializable_event"}) + "\n")
|
||
|
||
@staticmethod
|
||
def _extract_text_delta(event: dict) -> str:
|
||
"""从 message_update/message_end 事件取 assistant 文本增量（用于输出体量闸）。"""
|
||
if event.get("type") == "message_update":
|
||
delta = event.get("delta") or {}
|
||
if isinstance(delta, dict):
|
||
return str(delta.get("text") or "")
|
||
if event.get("type") == "message_end":
|
||
msg = event.get("message") or {}
|
||
for part in (msg.get("content") or []):
|
||
if isinstance(part, dict) and part.get("type") == "text":
|
||
return str(part.get("text") or "")
|
||
return ""
|
||
|
||
@staticmethod
|
||
def _extract_stop(event: dict):
|
||
"""
|
||
从事件里提取 (stopReason, errorMessage, 文本)。
|
||
兼容两个位置：message_end.message.stopReason 与 agent_end.messages 最后一条。
|
||
"""
|
||
etype = event.get("type")
|
||
msg = None
|
||
if etype == "message_end":
|
||
msg = event.get("message") or {}
|
||
elif etype == "agent_end":
|
||
msgs = event.get("messages") or []
|
||
assistants = [m for m in msgs if m.get("role") == "assistant"]
|
||
msg = assistants[-1] if assistants else None
|
||
if not msg:
|
||
return None, None, None
|
||
txt = ""
|
||
for part in (msg.get("content") or []):
|
||
if isinstance(part, dict) and part.get("type") == "text":
|
||
txt += str(part.get("text") or "")
|
||
return msg.get("stopReason"), msg.get("errorMessage"), txt
|
||
|
||
# -- 真实 pi headless runner ------------------------------------------
|
||
|
||
def pi_headless_runner(self, model: str = "aps-local/qwen3-32b-local",
|
||
tools: str = "read,grep,find,ls",
|
||
guard_ext: Optional[Path] = None,
|
||
extra_args: Optional[list] = None) -> AgentRunner:
|
||
"""
|
||
构造真实 pi headless 的 AgentRunner（供 --real 模式，Agent-C 使用）。
|
||
|
||
命令形态（Agent-A 实测）：
|
||
node <PI_CLI> -p --mode json --model <m> --tools <白名单> [-e guard.ts] "task"
|
||
- cwd 圈禁在 runs/<runId>/work（L2 文件层，pi 自身无沙箱）；
|
||
- 环境经 build_child_env 清洗（L4）；
|
||
- 子进程 stdout 由读线程转 queue，主循环带 poll 超时 —— 进程挂起也能被闸 1 抓到；
|
||
- 熔断/结束时 taskkill /T /F 杀整棵进程树。
|
||
"""
|
||
cmd_tail = extra_args or []
|
||
|
||
def runner(task: str, work_dir: Path) -> Iterator[dict]:
|
||
cmd = [
|
||
"node", str(PI_CLI),
|
||
"-p", "--mode", "json",
|
||
"--model", model,
|
||
"--tools", tools,
|
||
]
|
||
if guard_ext:
|
||
cmd += ["-e", str(guard_ext)]
|
||
cmd += cmd_tail + [task]
|
||
|
||
env = build_child_env()
|
||
proc = subprocess.Popen(
|
||
cmd, cwd=str(work_dir), env=env,
|
||
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||
creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
|
||
)
|
||
q: "queue.Queue[Optional[str]]" = queue.Queue()
|
||
|
||
def reader():
|
||
try:
|
||
for raw in proc.stdout:
|
||
q.put(raw.decode("utf-8", errors="replace"))
|
||
finally:
|
||
q.put(None) # EOF 哨兵
|
||
|
||
t = threading.Thread(target=reader, daemon=True)
|
||
t.start()
|
||
|
||
killed = False
|
||
try:
|
||
while True:
|
||
try:
|
||
line = q.get(timeout=self.cfg.poll_interval_sec)
|
||
except queue.Empty:
|
||
if proc.poll() is not None and t.is_alive() is False:
|
||
break
|
||
# 让上层熔断循环有机会检查超时：产出一个心跳事件
|
||
yield {"type": "harness_heartbeat"}
|
||
continue
|
||
if line is None:
|
||
break
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
yield json.loads(line)
|
||
except json.JSONDecodeError:
|
||
yield {"type": "harness_nonjson", "raw": line[:500]}
|
||
finally:
|
||
# 进程回收（L4）：编排器熔断 break 时会 close 本生成器，
|
||
# 这里只杀进程树、不许再 yield（closing 中的生成器 yield 会 RuntimeError）。
|
||
if proc.poll() is None:
|
||
killed = True
|
||
kill_process_tree(proc.pid)
|
||
try:
|
||
proc.wait(timeout=15)
|
||
except Exception:
|
||
pass
|
||
# 以下仅正常 EOF 退出时可达（被 close 时 GeneratorExit 会直接传播出去）
|
||
if killed:
|
||
yield {"type": "harness_killed", "pid": proc.pid}
|
||
err = proc.stderr.read().decode("utf-8", errors="replace") if proc.stderr else ""
|
||
if err.strip():
|
||
yield {"type": "harness_stderr", "text": err[-2000:]}
|
||
|
||
return runner
|
||
|
||
|
||
def kill_process_tree(pid: int) -> None:
|
||
"""Windows 下杀整棵进程树（node 起的孙进程一并回收，对应方案 §4.2 L4 进程回收）。"""
|
||
subprocess.run(
|
||
["taskkill", "/PID", str(pid), "/T", "/F"],
|
||
capture_output=True, timeout=30,
|
||
)
|