# core/control_plane.py
from dataclasses import dataclass, field
from typing import Any, Dict, Optional, List
from enum import Enum
from datetime import datetime
import asyncio
import json
import uuid
class AgentState(Enum):
"""Agent执行状态"""
RUNNING = "running"
PAUSED = "paused"
RESUMING = "resuming"
COMPLETED = "completed"
FAILED = "failed"
TERMINATED = "terminated"
@dataclass
class Checkpoint:
"""检查点:包含Agent的完整可恢复状态"""
checkpoint_id: str
agent_state: Dict[str, Any] # Agent的内部状态快照
message_history: List[Dict[str, Any]] # 完整消息历史
tool_context: Dict[str, Any] # 当前工具调用上下文
execution_context: Dict[str, Any] # 任务执行上下文
timestamp: datetime
iteration: int
safe_point: str # 检查点所在的安全位置标识
metadata: Dict[str, Any] = field(default_factory=dict)
def serialize(self) -> bytes:
"""序列化检查点为 JSON 字节,避免反序列化执行任意代码"""
payload = {
"checkpoint_id": self.checkpoint_id,
"agent_state": self.agent_state,
"message_history": self.message_history,
"tool_context": self.tool_context,
"execution_context": self.execution_context,
"timestamp": self.timestamp.isoformat(),
"iteration": self.iteration,
"safe_point": self.safe_point,
"metadata": self.metadata,
}
return json.dumps(payload, ensure_ascii=False).encode("utf-8")
@classmethod
def deserialize(cls, data: bytes) -> 'Checkpoint':
"""从 JSON 字节恢复检查点;只接受显式字段"""
payload = json.loads(data.decode("utf-8"))
return cls(
checkpoint_id=payload["checkpoint_id"],
agent_state=payload["agent_state"],
message_history=payload["message_history"],
tool_context=payload["tool_context"],
execution_context=payload["execution_context"],
timestamp=datetime.fromisoformat(payload["timestamp"]),
iteration=payload["iteration"],
safe_point=payload["safe_point"],
metadata=payload.get("metadata", {}),
)
class CheckpointManager:
"""检查点管理器"""
def __init__(self, storage_backend=None):
self.checkpoints: Dict[str, Checkpoint] = {}
self.storage_backend = storage_backend # 可选的持久化存储
self.latest_checkpoint_id: Optional[str] = None
async def create_checkpoint(
self,
agent: 'ControlledAgent',
safe_point: str
) -> Checkpoint:
"""在指定的安全点创建检查点"""
checkpoint = Checkpoint(
checkpoint_id=str(uuid.uuid4()),
agent_state=agent._serialize_state(),
message_history=agent.message_history.copy(),
tool_context=agent.tool_context.copy(),
execution_context=agent.execution_context.copy(),
timestamp=datetime.now(),
iteration=agent.iteration,
safe_point=safe_point,
metadata={
"task_id": agent.task_id,
"user_id": agent.user_id,
}
)
# 如果有持久化存储,异步保存
if self.storage_backend:
await self.storage_backend.save(checkpoint)
self.checkpoints[checkpoint.checkpoint_id] = checkpoint
self.latest_checkpoint_id = checkpoint.checkpoint_id
return checkpoint
async def restore_from_checkpoint(
self,
agent: 'ControlledAgent',
checkpoint_id: str
) -> bool:
"""从检查点恢复Agent状态"""
checkpoint = self.checkpoints.get(checkpoint_id)
if not checkpoint:
if self.storage_backend:
checkpoint = await self.storage_backend.load(checkpoint_id)
else:
return False
# 恢复Agent状态
agent._restore_state(checkpoint.agent_state)
agent.message_history = checkpoint.message_history.copy()
agent.tool_context = checkpoint.tool_context.copy()
agent.execution_context = checkpoint.execution_context.copy()
agent.iteration = checkpoint.iteration
return True
async def list_checkpoints(self) -> List[Dict[str, Any]]:
"""列出所有检查点摘要"""
return [
{
"id": cp.checkpoint_id,
"timestamp": cp.timestamp.isoformat(),
"iteration": cp.iteration,
"safe_point": cp.safe_point,
"task_id": cp.metadata.get("task_id")
}
for cp in self.checkpoints.values()
]
async def delete_checkpoint(self, checkpoint_id: str) -> bool:
"""删除检查点"""
if checkpoint_id in self.checkpoints:
del self.checkpoints[checkpoint_id]
if self.storage_backend:
await self.storage_backend.delete(checkpoint_id)
return True
return False