Agent 状态管理：从无状态到有状态的架构演进

┌─────────────────────────────────────────────────────┐
│                 Agent 状态全景图                     │
├──────────────┬──────────────────┬───────────────────┤
│   状态类型    │    生命周期       │    存储介质       │
├──────────────┼──────────────────┼───────────────────┤
│ 会话状态     │ 单次会话          │ 内存 / Redis     │
│ (消息历史)   │ 30分钟-24小时     │                   │
├──────────────┼──────────────────┼───────────────────┤
│ 工作流状态   │ 任务执行期间       │ Redis / 数据库   │
│ (执行进度)   │ 分钟-天           │                   │
├──────────────┼──────────────────┼───────────────────┤
│ 用户状态     │ 用户生命周期       │ 数据库           │
│ (偏好/画像)  │ 永久              │                   │
├──────────────┼──────────────────┼───────────────────┤
│ 检查点状态   │ 可恢复期间        │ 对象存储/数据库  │
│ (快照)       │ 可配置            │                   │
├──────────────┼──────────────────┼───────────────────┤
│ 共享状态     │ 多Agent协作期间   │ Redis/共享存储   │
│ (黑板/消息)  │ 会话级           │                   │
└──────────────┴──────────────────┴───────────────────┘

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum

class MessageRole(Enum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"

@dataclass
class Message:
    role: MessageRole
    content: str
    tool_calls: list | None = None
    tool_call_id: str | None = None
    timestamp: datetime = field(default_factory=datetime.now)
    metadata: dict = field(default_factory=dict)

@dataclass
class SessionState:
    """完整的会话状态"""
    session_id: str
    user_id: str
    agent_id: str
    messages: list[Message] = field(default_factory=list)
    context: dict = field(default_factory=dict)  # 上下文变量
    active_tools: list[str] = field(default_factory=list)
    pending_tool_calls: list[dict] = field(default_factory=list)
    created_at: datetime = field(default_factory=datetime.now)
    updated_at: datetime = field(default_factory=datetime.now)
    expires_at: datetime | None = None
    status: str = "active"  # active / paused / completed / error

class SessionStore:
    """会话状态存储——分层缓存策略"""
    
    def __init__(self, redis, postgres):
        self.redis = redis       # 热数据
        self.postgres = postgres # 冷数据/持久化
        self.ttl = 86400 * 7     # 7天过期
    
    async def get(self, session_id: str) -> SessionState | None:
        # L1: Redis 缓存
        cached = await self.redis.get(f"session:{session_id}")
        if cached:
            return SessionState.from_json(cached)
        
        # L2: PostgreSQL
        row = await self.postgres.fetchrow(
            "SELECT * FROM agent_sessions WHERE session_id = $1",
            session_id
        )
        if not row:
            return None
        
        state = SessionState.from_db(row)
        
        # 回填缓存
        await self.redis.setex(
            f"session:{session_id}",
            3600,  # 缓存1小时
            state.to_json()
        )
        
        return state
    
    async def save(self, state: SessionState):
        state.updated_at = datetime.now()
        
        # 写入 Redis（热路径）
        await self.redis.setex(
            f"session:{state.session_id}",
            3600,
            state.to_json()
        )
        
        # 异步写入 PostgreSQL（冷路径）
        asyncio.create_task(self._persist_to_db(state))
    
    async def _persist_to_db(self, state: SessionState):
        await self.postgres.execute("""
            INSERT INTO agent_sessions (session_id, user_id, agent_id, state, updated_at)
            VALUES ($1, $2, $3, $4, $5)
            ON CONFLICT (session_id) 
            DO UPDATE SET state = $4, updated_at = $5
        """, state.session_id, state.user_id, state.agent_id,
            state.to_json(), state.updated_at)

class ContextWindowManager:
    """管理 Agent 的上下文窗口"""
    
    def __init__(self, max_tokens: int = 128000):
        self.max_tokens = max_tokens
        self.reserved_for_output = 4096
        self.available = max_tokens - self.reserved_for_output
    
    def prepare_context(
        self,
        system_prompt: str,
        messages: list[Message],
        tools_schema: list[dict]
    ) -> list[dict]:
        """在 Context Window 内准备消息"""
        
        # 计算各部分 Token
        system_tokens = self._count_tokens(system_prompt)
        tools_tokens = self._count_tokens(json.dumps(tools_schema))
        remaining = self.available - system_tokens - tools_tokens
        
        # 从最新消息向前保留
        prepared = []
        total = 0
        
        for msg in reversed(messages):
            msg_tokens = self._count_tokens(msg.content)
            if total + msg_tokens > remaining:
                break
            prepared.insert(0, msg)
            total += msg_tokens
        
        # 如果截断了，添加摘要提示
        if len(prepared) < len(messages):
            summary = self._generate_summary(messages[:len(messages) - len(prepared)])
            prepared.insert(0, Message(
                role=MessageRole.SYSTEM,
                content=f"[Earlier conversation summary: {summary}]"
            ))
        
        return prepared
    
    def _count_tokens(self, text: str) -> int:
        # 使用 tiktoken 精确计算
        import tiktoken
        enc = tiktoken.encoding_for_model("gpt-5")
        return len(enc.encode(text))

class CheckpointManager:
    """Agent 执行检查点管理"""
    
    def __init__(self, storage):
        self.storage = storage
    
    async def save_checkpoint(
        self,
        execution_id: str,
        state: WorkflowState,
        step_index: int,
        step_name: str
    ):
        """保存执行检查点"""
        checkpoint = Checkpoint(
            execution_id=execution_id,
            step_index=step_index,
            step_name=step_name,
            state=state,
            timestamp=datetime.now()
        )
        
        await self.storage.save(checkpoint)
        
        # 保留最近 N 个检查点
        await self._prune_old_checkpoints(execution_id, keep=10)
    
    async def restore(self, execution_id: str) -> tuple[WorkflowState, int]:
        """从最新检查点恢复"""
        latest = await self.storage.get_latest(execution_id)
        if not latest:
            raise NoCheckpointError(execution_id)
        
        logger.info(
            f"Restoring from checkpoint: step={latest.step_index}, "
            f"name={latest.step_name}"
        )
        
        return latest.state, latest.step_index
    
    async def list_checkpoints(self, execution_id: str) -> list[Checkpoint]:
        """列出所有检查点（用于调试）"""
        return await self.storage.list(execution_id)


@dataclass
class WorkflowState:
    """工作流状态——可序列化"""
    execution_id: str
    current_step: str
    step_index: int
    results: dict         # {step_name: result}
    variables: dict       # 工作流变量
    pending_actions: list # 待执行操作
    error: str | None     # 错误信息（如果有的话）
    iteration: int        # 循环计数
    
    def serialize(self) -> bytes:
        return pickle.dumps(self)  # 或使用 JSON
    
    @classmethod
    def deserialize(cls, data: bytes) -> "WorkflowState":
        return pickle.loads(data)

class ResumableAgentExecutor:
    """支持断点续传的 Agent 执行器"""
    
    def __init__(self, checkpoint_mgr: CheckpointManager):
        self.checkpoints = checkpoint_mgr
    
    async def execute(
        self,
        workflow: Workflow,
        initial_state: WorkflowState,
        execution_id: str | None = None
    ) -> WorkflowState:
        
        execution_id = execution_id or str(uuid.uuid4())
        
        # 尝试从检查点恢复
        try:
            state, start_step = await self.checkpoints.restore(execution_id)
            logger.info(f"Resuming from step {start_step}")
        except NoCheckpointError:
            state = initial_state
            start_step = 0
        
        # 获取工作流步骤
        steps = workflow.get_steps()
        
        for i, step in enumerate(steps[start_step:], start=start_step):
            try:
                # 执行前保存检查点
                state.current_step = step.name
                state.step_index = i
                await self.checkpoints.save_checkpoint(
                    execution_id, state, i, step.name
                )
                
                # 执行步骤
                result = await step.execute(state)
                
                # 更新状态
                state.results[step.name] = result
                state.variables.update(result.get("variables", {}))
                
                # 条件分支
                if step.condition:
                    next_step = step.condition(result)
                    if next_step:
                        state.pending_actions = [next_step]
                
            except Exception as e:
                state.error = str(e)
                await self.checkpoints.save_checkpoint(
                    execution_id, state, i, step.name
                )
                
                # 重试逻辑
                if step.retry_policy:
                    retry_count = state.variables.get(f"retry_{step.name}", 0)
                    if retry_count < step.retry_policy.max_attempts:
                        state.variables[f"retry_{step.name}"] = retry_count + 1
                        await asyncio.sleep(step.retry_policy.backoff(retry_count))
                        # 重新执行当前步骤
                        continue
                
                raise
        
        return state

class SharedBlackboard:
    """多 Agent 共享黑板"""
    
    def __init__(self, redis_client):
        self.redis = redis_client
        self.namespace = "blackboard"
    
    async def write(
        self,
        key: str,
        value: any,
        agent_id: str,
        ttl: int = 3600
    ):
        """写入共享状态"""
        entry = {
            "value": value,
            "writer": agent_id,
            "timestamp": time.time(),
        }
        await self.redis.hset(
            f"{self.namespace}:{key}",
            mapping={k: json.dumps(v) for k, v in entry.items()}
        )
        await self.redis.expire(f"{self.namespace}:{key}", ttl)
        
        # 通知订阅者
        await self.redis.publish(
            f"{self.namespace}:updates",
            json.dumps({"key": key, "writer": agent_id})
        )
    
    async def read(self, key: str) -> any:
        """读取共享状态"""
        data = await self.redis.hgetall(f"{self.namespace}:{key}")
        if not data:
            return None
        return json.loads(data.get("value", "null"))
    
    async def subscribe(
        self,
        key_pattern: str,
        callback: callable
    ):
        """订阅状态变更"""
        pubsub = self.redis.pubsub()
        await pubsub.subscribe(f"{self.namespace}:updates")
        
        async for message in pubsub.listen():
            if message["type"] == "message":
                data = json.loads(message["data"])
                if fnmatch.fnmatch(data["key"], key_pattern):
                    await callback(data)

class AgentMessageBus:
    """Agent 间异步消息总线"""
    
    def __init__(self, redis_client):
        self.redis = redis_client
        self.queues = {}  # {agent_id: Queue}
    
    async def send(
        self,
        from_agent: str,
        to_agent: str,
        message_type: str,
        payload: dict,
        reply_to: str | None = None
    ):
        """发送消息给另一个 Agent"""
        msg = AgentMessage(
            id=str(uuid.uuid4()),
            from_agent=from_agent,
            to_agent=to_agent,
            type=message_type,
            payload=payload,
            reply_to=reply_to,
            timestamp=datetime.now()
        )
        
        # 推入接收者的队列
        await self.redis.lpush(
            f"agent:inbox:{to_agent}",
            msg.to_json()
        )
    
    async def receive(
        self,
        agent_id: str,
        timeout: int = 30
    ) -> AgentMessage | None:
        """接收消息"""
        result = await self.redis.brpop(
            f"agent:inbox:{agent_id}",
            timeout=timeout
        )
        if result:
            return AgentMessage.from_json(result[1])
        return None
    
    async def request_reply(
        self,
        from_agent: str,
        to_agent: str,
        message_type: str,
        payload: dict,
        timeout: int = 60
    ) -> dict | None:
        """请求-回复模式"""
        reply_channel = f"reply:{uuid.uuid4()}"
        
        await self.send(
            from_agent, to_agent, message_type,
            payload, reply_to=reply_channel
        )
        
        # 等待回复
        result = await self.redis.brpop(reply_channel, timeout=timeout)
        if result:
            return json.loads(result[1])
        return None

class StateSerializer:
    """状态序列化器"""
    
    SCHEMA_VERSION = "2.0"
    
    def serialize(self, state: any) -> str:
        """序列化状态为 JSON"""
        data = {
            "schema_version": self.SCHEMA_VERSION,
            "type": type(state).__name__,
            "data": self._to_dict(state),
            "timestamp": datetime.now().isoformat()
        }
        return json.dumps(data, ensure_ascii=False, default=str)
    
    def deserialize(self, raw: str) -> any:
        """反序列化"""
        data = json.loads(raw)
        
        # 版本迁移
        if data["schema_version"] != self.SCHEMA_VERSION:
            data = self._migrate(data)
        
        return self._from_dict(data["type"], data["data"])
    
    def _migrate(self, data: dict) -> dict:
        """状态版本迁移"""
        migrations = [
            ("1.0", "1.1", self._migrate_1_0_to_1_1),
            ("1.1", "2.0", self._migrate_1_1_to_2_0),
        ]
        
        current = data["schema_version"]
        for from_v, to_v, migrator in migrations:
            if current == from_v:
                data = migrator(data)
                current = to_v
        
        return data

□ 会话状态分层存储（内存 → Redis → 数据库）
□ 上下文窗口管理策略（滑动窗口 + 摘要）
□ 工作流检查点定期保存
□ 检查点支持断点续传
□ 多 Agent 共享状态通过消息总线
□ 状态序列化支持版本迁移
□ 过期状态自动清理
□ 状态加密敏感字段
□ 状态变更审计日志
□ 状态一致性测试（并发读写）