coding-agent/s11_error_recovery/code.py

#!/usr/bin/env python3
"""
s11: 错误恢复 — 三条恢复路径 + 指数退避。

运行:  python s11_error_recovery/code.py
需要: pip install anthropic python-dotenv + .env 中配置 ANTHROPIC_API_KEY

相对 s10 的变化：
  - LLM 调用包在 try/except 中，并提供三条恢复路径
  - 路径 1：max_tokens -> 从 8K 升级到 64K（第一次升级不追加截断输出），
             然后发送续写提示（最多 3 次）
  - 路径 2：prompt_too_long -> reactive compact -> 重试（一次）
  - 路径 3：429/529 -> 带抖动的指数退避（最多 10 次），
             连续 529 时切换备用模型
  - with_retry 包装器处理瞬时错误
  - RecoveryState 跟踪升级 / 压缩 / 529 / 模型

ASCII 流程：
  messages -> 提示词组装 -> 压缩+加载 -> [try] LLM [except] -> tools -> 循环
                                                   |          |
                                             stop_reason   错误类型
                                             max_tokens?   prompt_too_long? -> 压缩
                                             升级 /        429/529? -> 退避
                                             继续          其他? -> 记录 + 退出
"""

import os, subprocess, time, random, json
from pathlib import Path

try:
    import readline
    readline.parse_and_bind('set bind-tty-special-chars off')
except ImportError:
    pass

from anthropic import Anthropic
from dotenv import load_dotenv

load_dotenv(override=True)
if os.getenv("ANTHROPIC_BASE_URL"):
    os.environ.pop("ANTHROPIC_AUTH_TOKEN", None)

WORKDIR = Path.cwd()
MEMORY_DIR = WORKDIR / ".memory"
MEMORY_INDEX = MEMORY_DIR / "MEMORY.md"
client = Anthropic(base_url=os.getenv("ANTHROPIC_BASE_URL"))
PRIMARY_MODEL = os.environ["MODEL_ID"]
FALLBACK_MODEL = os.getenv("FALLBACK_MODEL_ID")

# ── 常量 ──

ESCALATED_MAX_TOKENS = 64000
DEFAULT_MAX_TOKENS = 8000
MAX_RECOVERY_RETRIES = 3
MAX_RETRIES = 10
BASE_DELAY_MS = 500
MAX_CONSECUTIVE_529 = 3
CONTINUATION_PROMPT = (
    "已触发输出 token 上限。请直接继续 — "
    "不要道歉，不要回顾，从中断处接着写。"
)

# ── 提示词组装 （来自 s10，已同步） ──

PROMPT_SECTIONS = {
    "identity": "你是一个编码 Agent。直接行动，不要只解释。",
    "tools": "可用工具：bash, read_file, write_file.",
    "workspace": f"工作目录：{WORKDIR}",
    "memory": "有可用的相关记忆时，会在下方注入。",
}


def assemble_system_prompt(context: dict) -> str:
    sections = [PROMPT_SECTIONS["identity"],
                PROMPT_SECTIONS["tools"],
                PROMPT_SECTIONS["workspace"]]
    memories = context.get("memories", "")
    if memories:
        sections.append(f"相关记忆：\n{memories}")
    return "\n\n".join(sections)


_last_context_key, _last_prompt = None, None


def get_system_prompt(context: dict) -> str:
    global _last_context_key, _last_prompt
    key = json.dumps(context, sort_keys=True, ensure_ascii=False, default=str)
    if key == _last_context_key and _last_prompt:
        print("  \033[90m[缓存命中] 系统提示词未变化\033[0m")
        return _last_prompt
    _last_context_key = key
    _last_prompt = assemble_system_prompt(context)

    loaded = ["identity", "tools", "workspace"]
    if context.get("memories"):
        loaded.append("memory")
    print(f"  \033[32m[已组装] 片段：{', '.join(loaded)}\033[0m")
    return _last_prompt


# ── 工具 （未改动） ──

def safe_path(p: str) -> Path:
    path = (WORKDIR / p).resolve()
    if not path.is_relative_to(WORKDIR):
        raise ValueError(f"路径逃逸出工作区：{p}")
    return path


def run_bash(command: str) -> str:
    try:
        r = subprocess.run(command, shell=True, cwd=WORKDIR,
                           capture_output=True, text=True, timeout=120)
        out = (r.stdout + r.stderr).strip()
        return out[:50000] if out else "（无输出）"
    except subprocess.TimeoutExpired:
        return "错误：执行超时（120 秒）"


def run_read(path: str, limit: int | None = None) -> str:
    try:
        lines = safe_path(path).read_text().splitlines()
        if limit and limit < len(lines):
            lines = lines[:limit] + [f"... ({len(lines) - limit} 行更多内容）"]
        return "\n".join(lines)
    except Exception as e:
        return f"错误：{e}"


def run_write(path: str, content: str) -> str:
    try:
        file_path = safe_path(path)
        file_path.parent.mkdir(parents=True, exist_ok=True)
        file_path.write_text(content)
        return f"已写入 {len(content)} 字节到 {path}"
    except Exception as e:
        return f"错误：{e}"


TOOLS = [
    {"name": "bash", "description": "运行一条 shell 命令。",
     "input_schema": {"type": "object",
                      "properties": {"command": {"type": "string"}},
                      "required": ["command"]}},
    {"name": "read_file", "description": "读取文件内容。",
     "input_schema": {"type": "object",
                      "properties": {"path": {"type": "string"},
                                     "limit": {"type": "integer"}},
                      "required": ["path"]}},
    {"name": "write_file", "description": "向文件写入内容。",
     "input_schema": {"type": "object",
                      "properties": {"path": {"type": "string"},
                                     "content": {"type": "string"}},
                      "required": ["path", "content"]}},
]

TOOL_HANDLERS = {"bash": run_bash, "read_file": run_read, "write_file": run_write}


# ── 错误恢复 (s11 new) ──

class RecoveryState:
    """在循环中跟踪恢复尝试。"""
    def __init__(self):
        self.has_escalated = False
        self.recovery_count = 0
        self.consecutive_529 = 0
        self.has_attempted_reactive_compact = False
        self.current_model = PRIMARY_MODEL


def retry_delay(attempt, retry_after=None):
    """带抖动的指数退避。Retry-After 优先。"""
    if retry_after:
        return retry_after
    base = min(BASE_DELAY_MS * (2 ** attempt), 32000) / 1000
    jitter = random.uniform(0, base * 0.25)
    return base + jitter


def with_retry(fn, state: RecoveryState):
    """针对瞬时错误（429/529）的指数退避。
    非瞬时错误会重新抛给外层处理器。"""
    for attempt in range(MAX_RETRIES):
        try:
            result = fn()
            state.consecutive_529 = 0
            return result
        except Exception as e:
            name = type(e).__name__
            msg = str(e).lower()

            # 429 速率限制 -> 指数退避
            if "ratelimit" in name.lower() or "429" in msg:
                delay = retry_delay(attempt)
                print(f"  \033[33m[429 速率限制] 重试 {attempt+1}/{MAX_RETRIES},"
                      f" 等待 {delay:.1f}s\033[0m")
                time.sleep(delay)
                continue

            # 529 过载 -> 指数退避 + 备用模型
            if "过载" in name.lower() or "529" in msg or "过载" in msg:
                state.consecutive_529 += 1
                if state.consecutive_529 >= MAX_CONSECUTIVE_529:
                    if FALLBACK_MODEL:
                        state.current_model = FALLBACK_MODEL
                        state.consecutive_529 = 0
                        print(f"  \033[31m[529 x{MAX_CONSECUTIVE_529}]"
                              f" 切换到 {FALLBACK_MODEL}\033[0m")
                    else:
                        state.consecutive_529 = 0
                        print(f"  \033[31m[529 x{MAX_CONSECUTIVE_529}]"
                              f" 未配置 FALLBACK_MODEL_ID，继续重试\033[0m")
                delay = retry_delay(attempt)
                print(f"  \033[33m[529 过载] 重试 {attempt+1}/{MAX_RETRIES},"
                      f" 等待 {delay:.1f}s\033[0m")
                time.sleep(delay)
                continue

            # 非瞬时错误 -> 重新抛给外层 try/except
            raise
    raise RuntimeError(f"最大重试次数（{MAX_RETRIES}）已超出")


def is_prompt_too_long_error(e: Exception) -> bool:
    """检查 API 错误是否表示提示词/上下文过长。"""
    msg = str(e).lower()
    return (("prompt" in msg and "long" in msg)
            or "提示词过长" in msg
            or "context_length_exceeded" in msg
            or "max_context_window" in msg)


def reactive_compact(messages: list) -> list:
    """应急压缩 — 教学版本保留最后 N 条消息。
    真实 CC 会通过 LLM 生成压缩摘要，然后用压缩后的消息列表重试。
    因为 s08/s09 已覆盖基于 LLM 的压缩，教学版本简化为保留尾部消息。"""
    print("  \033[31m[响应式压缩] 裁剪到最后 5 条消息\033[0m")
    tail = messages[-5:]
    return [{"role": "user",
             "content": "[响应式压缩] 早前对话已被裁剪。请从中断处继续。"}, *tail]


# ── 上下文 ──

def update_context(context: dict, messages: list) -> dict:
    """从真实状态推导上下文：有哪些工具、是否存在记忆文件。"""
    memories = ""
    if MEMORY_INDEX.exists():
        content = MEMORY_INDEX.read_text().strip()
        if content:
            memories = content
    return {
        "enabled_tools": list(TOOL_HANDLERS.keys()),
        "workspace": str(WORKDIR),
        "memories": memories,
    }


# ── Agent 循环 ──

def agent_loop(messages: list, context: dict):
    """主循环：为 LLM 调用包上错误恢复逻辑。"""
    system = get_system_prompt(context)
    state = RecoveryState()
    max_tokens = DEFAULT_MAX_TOKENS

    while True:
        # ── LLM 调用：with_retry 处理 429/529，外层处理其余错误 ──
        try:
            response = with_retry(
                lambda: client.messages.create(
                    model=state.current_model, system=system,
                    messages=messages, tools=TOOLS,
                    max_tokens=max_tokens),
                state)
        except Exception as e:
            # 路径 2：prompt_too_long -> 响应式压缩（一次）
            if is_prompt_too_long_error(e):
                if not state.has_attempted_reactive_compact:
                    messages[:] = reactive_compact(messages)
                    state.has_attempted_reactive_compact = True
                    continue
                print("  \033[31m[不可恢复] 压缩后仍然过长\033[0m")
                messages.append({"role": "assistant", "content": [
                    {"type": "text",
                     "text": "[错误] 上下文过大，无法继续。"}]})
                return

            # 不可恢复
            name = type(e).__name__
            print(f"  \033[31m[不可恢复] {name}: {str(e)[:100]}\033[0m")
            messages.append({"role": "assistant", "content": [
                {"type": "text", "text": f"[错误] {name}: {str(e)[:200]}"}]})
            return

        # ── 路径 1：max_tokens -> 升级或继续 ──
        if response.stop_reason == "max_tokens":
            # 第一次升级：不追加截断输出，重试同一个请求
            if not state.has_escalated:
                max_tokens = ESCALATED_MAX_TOKENS
                state.has_escalated = True
                print(f"  \033[33m[max_tokens] 升级"
                      f" {DEFAULT_MAX_TOKENS} -> {ESCALATED_MAX_TOKENS}\033[0m")
                continue
            # 64K 仍被截断：保存截断输出 + 续写提示
            messages.append({"role": "assistant", "content": response.content})
            if state.recovery_count < MAX_RECOVERY_RETRIES:
                messages.append({"role": "user", "content": CONTINUATION_PROMPT})
                state.recovery_count += 1
                print(f"  \033[33m[max_tokens] 续写"
                      f" {state.recovery_count}/{MAX_RECOVERY_RETRIES}\033[0m")
                continue
            print("  \033[31m[max_tokens] 恢复次数已达上限\033[0m")
            return

        # 正常完成：追加 assistant 响应
        messages.append({"role": "assistant", "content": response.content})

        if response.stop_reason != "tool_use":
            return

        # ── 工具执行 ──
        results = []
        for block in response.content:
            if block.type != "tool_use":
                continue
            print(f"\033[36m> {block.name}\033[0m")
            handler = TOOL_HANDLERS.get(block.name)
            output = handler(**block.input) if handler else f"未知工具：{block.name}"
            print(str(output)[:200])
            results.append({"type": "tool_result",
                            "tool_use_id": block.id, "content": output})
        messages.append({"role": "user", "content": results})

        context = update_context(context, messages)
        system = get_system_prompt(context)


if __name__ == "__main__":
    print("s11: 错误恢复")
    print("输入问题后按回车发送。输入 q 退出。\n")
    history = []
    context = update_context({}, [])
    while True:
        try:
            query = input("\033[36ms11 >> \033[0m")
        except (EOFError, KeyboardInterrupt):
            break
        if query.strip().lower() in ("q", "exit", ""):
            break
        turn_start = len(history)
        history.append({"role": "user", "content": query})
        agent_loop(history, context)
        context = update_context(context, history)
        for msg in history[turn_start:]:
            if msg.get("role") != "assistant":
                continue
            for block in msg["content"]:
                if getattr(block, "type", None) == "text":
                    print(block.text)
        print()