|
|
@@ -22,6 +22,12 @@ from cyber_agent.trace.websocket import broadcast_sub_trace_started, broadcast_s
|
|
|
# ===== 远端路由常量 =====
|
|
|
|
|
|
REMOTE_PREFIX = "remote_"
|
|
|
+MAX_AGENT_DEPTH = 5
|
|
|
+MAX_LOCAL_CHILDREN_PER_TRACE = 6
|
|
|
+
|
|
|
+# POC 阶段使用进程内锁保护“检查配额 -> 创建 Trace”。
|
|
|
+# 多 worker 部署时需要换成跨进程的原子配额。
|
|
|
+_LOCAL_AGENT_SPAWN_LOCK = asyncio.Lock()
|
|
|
|
|
|
|
|
|
def _knowhub_api() -> str:
|
|
|
@@ -33,6 +39,13 @@ def _remote_agent_timeout() -> float:
|
|
|
return float(os.getenv("REMOTE_AGENT_TIMEOUT", "600"))
|
|
|
|
|
|
|
|
|
+def _recursion_enabled() -> bool:
|
|
|
+ """是否允许本地 Sub-Agent 继续创建 Sub-Agent。"""
|
|
|
+ return os.getenv("AGENT_RECURSION_ENABLED", "false").strip().lower() in {
|
|
|
+ "1", "true", "yes", "on",
|
|
|
+ }
|
|
|
+
|
|
|
+
|
|
|
# 兼容旧代码对 module-level 常量的引用(运行时值 = 首次 import 时的快照)
|
|
|
KNOWHUB_API = _knowhub_api()
|
|
|
REMOTE_AGENT_TIMEOUT = _remote_agent_timeout()
|
|
|
@@ -143,23 +156,35 @@ async def _update_collaborator(
|
|
|
|
|
|
|
|
|
async def _update_goal_start(
|
|
|
- store, trace_id: str, goal_id: str, mode: str, sub_trace_ids: List[str]
|
|
|
+ store, trace_id: str, goal_id: str, mode: str,
|
|
|
+ sub_trace_ids: List[Dict[str, str]],
|
|
|
) -> None:
|
|
|
"""标记 Goal 开始执行"""
|
|
|
if not goal_id:
|
|
|
return
|
|
|
+
|
|
|
+ tree = await store.get_goal_tree(trace_id)
|
|
|
+ goal = tree.find(goal_id) if tree else None
|
|
|
+ existing_entries = goal.sub_trace_ids if goal and goal.sub_trace_ids else []
|
|
|
+ merged_entries: Dict[str, Dict[str, str]] = {}
|
|
|
+ for entry in [*existing_entries, *sub_trace_ids]:
|
|
|
+ if isinstance(entry, str):
|
|
|
+ merged_entries[entry] = {"trace_id": entry, "mission": entry}
|
|
|
+ elif isinstance(entry, dict) and entry.get("trace_id"):
|
|
|
+ merged_entries[entry["trace_id"]] = entry
|
|
|
+
|
|
|
await store.update_goal(
|
|
|
trace_id, goal_id,
|
|
|
type="agent_call",
|
|
|
agent_call_mode=mode,
|
|
|
status="in_progress",
|
|
|
- sub_trace_ids=sub_trace_ids
|
|
|
+ sub_trace_ids=list(merged_entries.values()),
|
|
|
)
|
|
|
|
|
|
|
|
|
async def _update_goal_complete(
|
|
|
store, trace_id: str, goal_id: str,
|
|
|
- status: str, summary: str, sub_trace_ids: List[str]
|
|
|
+ status: str, summary: str,
|
|
|
) -> None:
|
|
|
"""标记 Goal 完成"""
|
|
|
if not goal_id:
|
|
|
@@ -168,7 +193,6 @@ async def _update_goal_complete(
|
|
|
trace_id, goal_id,
|
|
|
status=status,
|
|
|
summary=summary,
|
|
|
- sub_trace_ids=sub_trace_ids
|
|
|
)
|
|
|
|
|
|
|
|
|
@@ -192,15 +216,50 @@ def _aggregate_stats(results: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
|
}
|
|
|
|
|
|
|
|
|
-def _get_allowed_tools(single: bool, context: dict) -> Optional[List[str]]:
|
|
|
- """获取允许工具列表。获取所有工具,排除 agent 和 evaluate"""
|
|
|
+def _get_allowed_tools(context: dict, agent_depth: int) -> Optional[List[str]]:
|
|
|
+ """按子 Agent 深度生成唯一的工具权限列表。"""
|
|
|
runner = context.get("runner")
|
|
|
- if runner and hasattr(runner, "tools") and hasattr(runner.tools, "_tools"):
|
|
|
- all_tools = list(runner.tools._tools.keys())
|
|
|
- return [t for t in all_tools if t not in ("agent", "evaluate")]
|
|
|
+ if runner and hasattr(runner, "tools") and hasattr(runner.tools, "get_tool_names"):
|
|
|
+ blocked_tools = {"evaluate", "bash_command"}
|
|
|
+ if not _recursion_enabled() or agent_depth >= MAX_AGENT_DEPTH:
|
|
|
+ blocked_tools.add("agent")
|
|
|
+ return [
|
|
|
+ name for name in runner.tools.get_tool_names()
|
|
|
+ if name not in blocked_tools
|
|
|
+ ]
|
|
|
return None
|
|
|
|
|
|
|
|
|
+async def _resolve_trace_lineage(store, trace: Trace) -> tuple[int, str]:
|
|
|
+ """返回 Trace 的 (agent_depth, root_trace_id)。
|
|
|
+
|
|
|
+ 新 Trace 直接读 context;旧 Trace 缺少字段时沿 parent_trace_id
|
|
|
+ 回溯,不依赖 Trace ID 字符串格式。
|
|
|
+ """
|
|
|
+ depth = trace.context.get("agent_depth")
|
|
|
+ root_trace_id = trace.context.get("root_trace_id")
|
|
|
+ if isinstance(depth, int) and depth >= 0 and isinstance(root_trace_id, str):
|
|
|
+ return depth, root_trace_id
|
|
|
+
|
|
|
+ depth = 0
|
|
|
+ current = trace
|
|
|
+ root_trace_id = trace.trace_id
|
|
|
+ visited = {trace.trace_id}
|
|
|
+ while current.parent_trace_id:
|
|
|
+ parent_id = current.parent_trace_id
|
|
|
+ if parent_id in visited:
|
|
|
+ break
|
|
|
+ visited.add(parent_id)
|
|
|
+ parent = await store.get_trace(parent_id)
|
|
|
+ if not parent:
|
|
|
+ break
|
|
|
+ depth += 1
|
|
|
+ root_trace_id = parent.trace_id
|
|
|
+ current = parent
|
|
|
+
|
|
|
+ return depth, root_trace_id
|
|
|
+
|
|
|
+
|
|
|
def _format_single_result(result: Dict[str, Any], sub_trace_id: str, continued: bool) -> Dict[str, Any]:
|
|
|
"""格式化单任务(delegate)结果"""
|
|
|
lines = [DELEGATE_RESULT_HEADER]
|
|
|
@@ -419,69 +478,149 @@ async def _run_agents(
|
|
|
"""
|
|
|
统一 agent 执行逻辑。
|
|
|
|
|
|
- single (len(tasks)==1): delegate 模式,全量工具(排除 agent/evaluate)
|
|
|
- multi (len(tasks)>1): explore 模式,只读工具,并行执行
|
|
|
+ single (len(tasks)==1): delegate 模式,全量工具
|
|
|
+ multi (len(tasks)>1): explore 模式,并行执行
|
|
|
+
|
|
|
+ AGENT_RECURSION_ENABLED=true 时,depth < MAX_AGENT_DEPTH 的本地
|
|
|
+ Sub-Agent 可继续使用 agent 工具。
|
|
|
"""
|
|
|
single = len(tasks) == 1
|
|
|
parent_trace = await store.get_trace(trace_id)
|
|
|
+ if not parent_trace:
|
|
|
+ return {"status": "failed", "error": f"Parent trace not found: {trace_id}"}
|
|
|
|
|
|
# continue_from: 复用已有 trace(仅 single)
|
|
|
sub_trace_id = None
|
|
|
continued = False
|
|
|
+ goal_started = False
|
|
|
+ child_records: List[Dict[str, Any]] = []
|
|
|
+ all_sub_trace_ids: List[Dict[str, str]] = []
|
|
|
+
|
|
|
if single and continue_from:
|
|
|
existing = await store.get_trace(continue_from)
|
|
|
if not existing:
|
|
|
return {"status": "failed", "error": f"Continue-from trace not found: {continue_from}"}
|
|
|
+ if existing.parent_trace_id != trace_id:
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "continue_from must reference a direct child of the current trace",
|
|
|
+ }
|
|
|
+ if existing.uid != parent_trace.uid:
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "continue_from trace owner does not match the current trace owner",
|
|
|
+ }
|
|
|
+
|
|
|
sub_trace_id = continue_from
|
|
|
continued = True
|
|
|
goal_tree = await store.get_goal_tree(continue_from)
|
|
|
mission = goal_tree.mission if goal_tree else tasks[0]
|
|
|
- sub_trace_ids = [{"trace_id": sub_trace_id, "mission": mission}]
|
|
|
+ child_depth, root_trace_id = await _resolve_trace_lineage(store, existing)
|
|
|
+ all_sub_trace_ids = [{"trace_id": sub_trace_id, "mission": mission}]
|
|
|
+ child_records = [{
|
|
|
+ "trace_id": sub_trace_id,
|
|
|
+ "depth": child_depth,
|
|
|
+ "context": {
|
|
|
+ **existing.context,
|
|
|
+ "agent_depth": child_depth,
|
|
|
+ "root_trace_id": root_trace_id,
|
|
|
+ },
|
|
|
+ "is_new": False,
|
|
|
+ }]
|
|
|
else:
|
|
|
- sub_trace_ids = []
|
|
|
-
|
|
|
- # 创建 sub-traces 和执行协程
|
|
|
- coros = []
|
|
|
- all_sub_trace_ids = list(sub_trace_ids) # copy for continue_from case
|
|
|
+ parent_depth, root_trace_id = await _resolve_trace_lineage(store, parent_trace)
|
|
|
+ if parent_depth > 0 and not _recursion_enabled():
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "Local Sub-Agent recursion is disabled",
|
|
|
+ }
|
|
|
+ if parent_depth >= MAX_AGENT_DEPTH:
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": (
|
|
|
+ f"Local Sub-Agent depth limit reached: "
|
|
|
+ f"depth={parent_depth}, max={MAX_AGENT_DEPTH}"
|
|
|
+ ),
|
|
|
+ }
|
|
|
|
|
|
- for i, (task_item, msgs) in enumerate(zip(tasks, per_agent_msgs)):
|
|
|
- if single and continued:
|
|
|
- # continue_from 已经设置了 sub_trace_id
|
|
|
- pass
|
|
|
- else:
|
|
|
- resolved_agent_type = agent_type or ("delegate" if single else "explore")
|
|
|
- suffix = "delegate" if single else f"explore-{i+1:03d}"
|
|
|
- stid = generate_sub_trace_id(trace_id, suffix)
|
|
|
-
|
|
|
- sub_trace = Trace(
|
|
|
- trace_id=stid,
|
|
|
- mode="agent",
|
|
|
- task=task_item,
|
|
|
+ # 将配额检查和 Trace 创建放在同一进程内临界区,
|
|
|
+ # 保证两个并发 agent() 调用不会共同突破 6 个直接孩子。
|
|
|
+ async with _LOCAL_AGENT_SPAWN_LOCK:
|
|
|
+ existing_children = await store.list_traces(
|
|
|
parent_trace_id=trace_id,
|
|
|
- parent_goal_id=goal_id,
|
|
|
- agent_type=resolved_agent_type,
|
|
|
- uid=parent_trace.uid if parent_trace else None,
|
|
|
- model=parent_trace.model if parent_trace else None,
|
|
|
- status="running",
|
|
|
- context={"created_by_tool": "agent"},
|
|
|
- created_at=datetime.now(),
|
|
|
+ created_by_tool="agent",
|
|
|
+ limit=MAX_LOCAL_CHILDREN_PER_TRACE + 1,
|
|
|
)
|
|
|
- await store.create_trace(sub_trace)
|
|
|
- await store.update_goal_tree(stid, GoalTree(mission=task_item))
|
|
|
+ requested_children = len(tasks)
|
|
|
+ if len(existing_children) + requested_children > MAX_LOCAL_CHILDREN_PER_TRACE:
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": (
|
|
|
+ f"Local child Agent limit exceeded: existing={len(existing_children)}, "
|
|
|
+ f"requested={requested_children}, "
|
|
|
+ f"max={MAX_LOCAL_CHILDREN_PER_TRACE}"
|
|
|
+ ),
|
|
|
+ }
|
|
|
|
|
|
- all_sub_trace_ids.append({"trace_id": stid, "mission": task_item})
|
|
|
+ child_depth = parent_depth + 1
|
|
|
+ for i, task_item in enumerate(tasks):
|
|
|
+ resolved_agent_type = agent_type or ("delegate" if single else "explore")
|
|
|
+ suffix = "delegate" if single else f"explore-{i+1:03d}"
|
|
|
+ stid = generate_sub_trace_id(trace_id, suffix)
|
|
|
+ child_context = {
|
|
|
+ "created_by_tool": "agent",
|
|
|
+ "agent_depth": child_depth,
|
|
|
+ "root_trace_id": root_trace_id,
|
|
|
+ }
|
|
|
+ sub_trace = Trace(
|
|
|
+ trace_id=stid,
|
|
|
+ mode="agent",
|
|
|
+ task=task_item,
|
|
|
+ parent_trace_id=trace_id,
|
|
|
+ parent_goal_id=goal_id,
|
|
|
+ agent_type=resolved_agent_type,
|
|
|
+ uid=parent_trace.uid,
|
|
|
+ model=parent_trace.model,
|
|
|
+ status="running",
|
|
|
+ context=child_context,
|
|
|
+ created_at=datetime.now(),
|
|
|
+ )
|
|
|
+ await store.create_trace(sub_trace)
|
|
|
+ await store.update_goal_tree(stid, GoalTree(mission=task_item))
|
|
|
+ all_sub_trace_ids.append({"trace_id": stid, "mission": task_item})
|
|
|
+ child_records.append({
|
|
|
+ "trace_id": stid,
|
|
|
+ "depth": child_depth,
|
|
|
+ "context": child_context,
|
|
|
+ "is_new": True,
|
|
|
+ "agent_type": resolved_agent_type,
|
|
|
+ })
|
|
|
|
|
|
- # 广播 sub_trace_started
|
|
|
- await broadcast_sub_trace_started(
|
|
|
- trace_id, stid, goal_id or "",
|
|
|
- resolved_agent_type, task_item,
|
|
|
+ if single:
|
|
|
+ sub_trace_id = child_records[0]["trace_id"]
|
|
|
+
|
|
|
+ await _update_goal_start(
|
|
|
+ store, trace_id, goal_id,
|
|
|
+ "delegate" if single else "explore",
|
|
|
+ all_sub_trace_ids,
|
|
|
)
|
|
|
+ goal_started = True
|
|
|
|
|
|
- if single:
|
|
|
- sub_trace_id = stid
|
|
|
+ # 创建执行协程。Trace 在上面已经按批次预创建,
|
|
|
+ # 因此孩子配额不会被并发调用绕过。
|
|
|
+ coros = []
|
|
|
+ for i, (task_item, msgs, child_record) in enumerate(
|
|
|
+ zip(tasks, per_agent_msgs, child_records)
|
|
|
+ ):
|
|
|
+ cur_stid = child_record["trace_id"]
|
|
|
+ child_depth = child_record["depth"]
|
|
|
+ if child_record["is_new"]:
|
|
|
+ await broadcast_sub_trace_started(
|
|
|
+ trace_id, cur_stid, goal_id or "",
|
|
|
+ child_record["agent_type"], task_item,
|
|
|
+ )
|
|
|
|
|
|
# 注册为活跃协作者
|
|
|
- cur_stid = sub_trace_id if single else all_sub_trace_ids[-1]["trace_id"]
|
|
|
collab_name = task_item[:30] if single and not continued else (
|
|
|
f"delegate-{cur_stid[:8]}" if single else f"explore-{i+1}"
|
|
|
)
|
|
|
@@ -493,12 +632,12 @@ async def _run_agents(
|
|
|
|
|
|
# 构建消息
|
|
|
agent_msgs = list(msgs) + [{"role": "user", "content": task_item}]
|
|
|
- allowed_tools = _get_allowed_tools(single, context)
|
|
|
+ allowed_tools = _get_allowed_tools(context, child_depth)
|
|
|
|
|
|
debug = getattr(runner, 'debug', False)
|
|
|
agent_label = (agent_type or ("delegate" if single else f"explore-{i+1}"))
|
|
|
debug_printer = _make_event_printer(agent_label) if debug else None
|
|
|
- # 为了彻底封死逃逸风险,禁用由于缺省 tool_groups 注入进来的 agent 和 bash_command
|
|
|
+ # allowed_tools 已是当前深度的精确白名单。
|
|
|
on_event = _make_interactive_handler(
|
|
|
runner, cur_stid, trace_id, debug_printer=debug_printer
|
|
|
)
|
|
|
@@ -512,22 +651,23 @@ async def _run_agents(
|
|
|
model=parent_trace.model if parent_trace else "gpt-4o",
|
|
|
uid=parent_trace.uid if parent_trace else None,
|
|
|
tools=allowed_tools,
|
|
|
- tool_groups=[], # 清空默认的 "core" 组,防止把 agent 和 bash_command 偷渡进来
|
|
|
- exclude_tools=["agent", "evaluate", "bash_command"], # 严格锁死危险元工具
|
|
|
+ tool_groups=[], # tools 是精确白名单,不再合并默认 core 组
|
|
|
name=task_item[:50],
|
|
|
skills=skills,
|
|
|
knowledge=context.get("knowledge_config"),
|
|
|
+ context=child_record["context"],
|
|
|
),
|
|
|
on_event=on_event,
|
|
|
)
|
|
|
coros.append((i, cur_stid, collab_name, coro))
|
|
|
|
|
|
- # 更新主 Goal 为 in_progress
|
|
|
- await _update_goal_start(
|
|
|
- store, trace_id, goal_id,
|
|
|
- "delegate" if single else "explore",
|
|
|
- all_sub_trace_ids,
|
|
|
- )
|
|
|
+ # continue_from 不进入新建临界区,但仍需要恢复 Goal 的运行状态。
|
|
|
+ if not goal_started:
|
|
|
+ await _update_goal_start(
|
|
|
+ store, trace_id, goal_id,
|
|
|
+ "delegate" if single else "explore",
|
|
|
+ all_sub_trace_ids,
|
|
|
+ )
|
|
|
|
|
|
# 执行
|
|
|
if single:
|
|
|
@@ -555,7 +695,6 @@ async def _run_agents(
|
|
|
store, trace_id, goal_id,
|
|
|
result.get("status", "completed"),
|
|
|
formatted["summary"],
|
|
|
- all_sub_trace_ids,
|
|
|
)
|
|
|
return formatted
|
|
|
|
|
|
@@ -572,7 +711,6 @@ async def _run_agents(
|
|
|
await _update_goal_complete(
|
|
|
store, trace_id, goal_id,
|
|
|
"failed", f"委托任务失败: {error_msg}",
|
|
|
- all_sub_trace_ids,
|
|
|
)
|
|
|
return {
|
|
|
"mode": "delegate",
|
|
|
@@ -639,7 +777,6 @@ async def _run_agents(
|
|
|
store, trace_id, goal_id,
|
|
|
formatted["status"],
|
|
|
formatted["summary"],
|
|
|
- all_sub_trace_ids,
|
|
|
)
|
|
|
return formatted
|
|
|
|
|
|
@@ -790,6 +927,14 @@ async def agent(
|
|
|
if messages is None:
|
|
|
per_agent_msgs: List[Messages] = [[] for _ in tasks]
|
|
|
elif messages and isinstance(messages[0], list):
|
|
|
+ if len(messages) != len(tasks):
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": (
|
|
|
+ "2D messages must contain exactly one message list per task: "
|
|
|
+ f"tasks={len(tasks)}, messages={len(messages)}"
|
|
|
+ ),
|
|
|
+ }
|
|
|
per_agent_msgs = messages # 2D: per-agent
|
|
|
else:
|
|
|
per_agent_msgs = [messages] * len(tasks) # 1D: 共享
|
|
|
@@ -910,6 +1055,8 @@ async def evaluate(
|
|
|
model=parent_trace.model if parent_trace else "gpt-4o",
|
|
|
uid=parent_trace.uid if parent_trace else None,
|
|
|
tools=allowed_tools,
|
|
|
+ tool_groups=[],
|
|
|
+ exclude_tools=["agent", "evaluate", "bash_command"],
|
|
|
name=f"评估: {goal_id}",
|
|
|
),
|
|
|
on_event=_make_interactive_handler(
|
|
|
@@ -937,7 +1084,6 @@ async def evaluate(
|
|
|
store, trace_id, current_goal_id,
|
|
|
result.get("status", "completed"),
|
|
|
formatted_summary,
|
|
|
- sub_trace_ids,
|
|
|
)
|
|
|
|
|
|
return {
|
|
|
@@ -961,7 +1107,6 @@ async def evaluate(
|
|
|
await _update_goal_complete(
|
|
|
store, trace_id, current_goal_id,
|
|
|
"failed", f"评估任务失败: {error_msg}",
|
|
|
- sub_trace_ids,
|
|
|
)
|
|
|
return {
|
|
|
"mode": "evaluate",
|