Agent Harness：把多步智能体变成可暂停、可恢复、可审计的状态机

import uuid
from typing import Dict, Any, List, Optional
from pydantic import BaseModel, Field

class ToolCall(BaseModel):
    tool_name: str
    arguments: Dict[str, Any]
    output: Optional[str] = None
    status: str = "pending"  # pending, success, failed

class AgentStep(BaseModel):
    step_number: int
    thought: str
    tool_calls: List[ToolCall] = []
    timestamp: float

class AgentState(BaseModel):
    session_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    status: str = "PENDING"  # PENDING, RUNNING, PAUSED, COMPLETED, FAILED
    current_step: int = 0
    history: List[AgentStep] = []
    context: Dict[str, Any] = {}
    error_message: Optional[str] = None
    
    # 运行预算控制
    max_steps: int = 10
    accumulated_cost: float = 0.0
    max_cost_usd: float = 0.50

class AgentStateMachineHarness:
    def __init__(self, state: AgentState):
        self.state = state

    def load_state(self, serialized_state: dict):
        self.state = AgentState(**serialized_state)

    def dump_state(self) -> dict:
        return self.state.model_dump()

    def transition_to(self, next_status: str):
        valid_transitions = {
            "PENDING": ["RUNNING"],
            "RUNNING": ["PAUSED", "COMPLETED", "FAILED"],
            "PAUSED": ["RUNNING", "FAILED"],
            "FAILED": ["PENDING", "RUNNING"],  # 允许从失败状态重启或恢复
            "COMPLETED": []
        }
        if next_status not in valid_transitions[self.state.status]:
            raise ValueError(f"Invalid transition from {self.state.status} to {next_status}")
        self.state.status = next_status

class BudgetExceededException(Exception):
    pass

class AgentStateMachineHarness(AgentStateMachineHarness):
    def enforce_budget(self, estimated_step_cost: float):
        # 检查步骤数超限
        if self.state.current_step >= self.state.max_steps:
            self.transition_to("FAILED")
            self.state.error_message = f"Max steps ({self.state.max_steps}) exceeded."
            self.send_to_dead_letter_queue("MAX_STEPS_EXCEEDED")
            raise BudgetExceededException(self.state.error_message)
        
        # 检查资金预算超限
        new_cost = self.state.accumulated_cost + estimated_step_cost
        if new_cost > self.state.max_cost_usd:
            self.transition_to("FAILED")
            self.state.error_message = f"Max cost budget ({self.state.max_cost_usd} USD) exceeded."
            self.send_to_dead_letter_queue("BUDGET_EXCEEDED")
            raise BudgetExceededException(self.state.error_message)
            
        self.state.accumulated_cost = new_cost

    def send_to_dead_letter_queue(self, reason: str):
        # 将当前状态快照序列化后发送至死信队列（DLQ），供开发人员离线分析
        dead_letter_payload = {
            "reason": reason,
            "snapshot": self.dump_state()
        }
        # 实际生产中此处对接消息队列，如 RabbitMQ, AWS SQS 或数据库归档表
        print(f"[DLQ ALERT] Session {self.state.session_id} sent to DLQ. Reason: {reason}")

{
  "trace_id": "trace_882a8f3b_901c",
  "session_id": "session_abc123",
  "step_number": 3,
  "inputs": {
    "query": "查询用户 2026 年 Q1 的消费总额，并计算折合美元"
  },
  "thought": "我已经获取了人民币消费总额为 7000 元。现在我需要调用货币转换工具将其转换为美元。当前汇率需要通过实时 API 获取。",
  "tool_calls": [
    {
      "tool_name": "get_exchange_rate",
      "arguments": {
        "from_currency": "CNY",
        "to_currency": "USD"
      },
      "output": "{\"rate\": 0.14}",
      "status": "success"
    }
  ],
  "outputs": {
    "intermediate_result": "980 USD"
  },
  "metrics": {
    "step_latency_seconds": 1.45,
    "prompt_tokens": 850,
    "completion_tokens": 120,
    "step_cost_usd": 0.015
  }
}

class AgentStateMachineHarness(AgentStateMachineHarness):
    def request_human_approval(self, action_details: dict) -> str:
        # 将状态挂起
        self.transition_to("PAUSED")
        # 保存需要审核的动作详情到上下文
        self.state.context["pending_approval_action"] = action_details
        # 实际生产中，这里会向 Web 前端或 Slack 发送一条通知
        return f"Session {self.state.session_id} paused. Awaiting approval for: {action_details['tool_name']}"

    def resume_with_approval(self, approved: bool, feedback: Optional[str] = None):
        if self.state.status != "PAUSED":
            raise ValueError("Can only resume from PAUSED status")
        
        if approved:
            # 清理审核挂载点，恢复运行状态
            self.state.context.pop("pending_approval_action", None)
            self.transition_to("RUNNING")
            print("Approval granted. Resuming Agent execution...")
        else:
            self.transition_to("FAILED")
            self.state.error_message = f"Rejected by human auditor. Feedback: {feedback}"
            print(f"Execution rejected. Feedback: {feedback}")

def replay_agent_at_step(failed_session_state: dict, test_prompt_patch: str) -> dict:
    """
    加载失败的 Agent 状态快照，替换引发错误的 Prompt，并从出错的步骤开始单步重试。
    """
    harness = AgentStateMachineHarness(AgentState(**failed_session_state))
    
    # 1. 验证状态是否可恢复（应该处于 FAILED 或 PAUSED 状态）
    if harness.state.status not in ["FAILED", "PAUSED"]:
        print("Warning: Replaying a running or completed session.")
    
    # 2. 注入修改后的 Prompt/系统指令（模拟 Patch 操作）
    harness.state.context["system_instruction_override"] = test_prompt_patch
    
    # 3. 将状态重置为 RUNNING，准备单步执行
    harness.state.status = "RUNNING"
    harness.state.error_message = None
    
    print(f"Replaying session {harness.state.session_id} from step {harness.state.current_step}...")
    return harness.dump_state()

{
  "session_id": "session_tx_99281",
  "status": "FAILED",
  "current_step": 3,
  "history": [
    {
      "step_number": 1,
      "thought": "先获取客户 ID",
      "tool_calls": [{"tool_name": "get_user_id", "arguments": {"name": "Alice"}, "output": "usr_001", "status": "success"}],
      "timestamp": 1716900000.0
    },
    {
      "step_number": 2,
      "thought": "锁定该客户的对账账户",
      "tool_calls": [{"tool_name": "lock_account", "arguments": {"account_id": "acc_99"}, "output": "locked_success", "status": "success"}],
      "timestamp": 1716900005.0
    },
    {
      "step_number": 3,
      "thought": "获取当前余额",
      "tool_calls": [{"tool_name": "get_bank_balance", "arguments": {"account_id": "acc_99"}, "output": null, "status": "failed"}],
      "timestamp": 1716900010.0
    }
  ],
  "context": {
    "user_id": "usr_001",
    "account_id": "acc_99"
  },
  "error_message": "HTTP 503 Service Unavailable on get_bank_balance",
  "accumulated_cost": 0.045
}

# 模拟从 JSON 文件读取故障快照
fault_snapshot = { ... } # 填入上述第一步的 JSON 数据

harness = AgentStateMachineHarness(AgentState(**fault_snapshot))

# 修复第 3 步的输出（人工注入备用渠道获取到的真实结果）
harness.state.history[-1].tool_calls[0].output = "{\"balance\": 50000, \"currency\": \"CNY\"}"
harness.state.history[-1].tool_calls[0].status = "success"

# 清除错误信息，将状态设回等待运行
harness.state.status = "PENDING" # 允许转移到 RUNNING
harness.state.error_message = None

# 将 current_step 指向第 4 步，表明下一步直接基于已修复的第 3 步结果继续执行
harness.state.current_step = 3  # 注：如果 current_step 表示已完成的步骤索引，更新它

# 导出就绪状态
restored_state = harness.dump_state()
print("Restored State successfully. Ready to resume:")
print(restored_state)