|
|
@@ -13,6 +13,13 @@ from datetime import datetime
|
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
|
|
from cyber_agent.tools import tool
|
|
|
+from cyber_agent.core.agent_mode import (
|
|
|
+ AgentMode,
|
|
|
+ AgentPolicy,
|
|
|
+ apply_policy_to_context,
|
|
|
+ assert_removed_config_absent,
|
|
|
+ policy_from_context,
|
|
|
+)
|
|
|
from cyber_agent.trace.models import Trace, Messages
|
|
|
from cyber_agent.trace.trace_id import generate_sub_trace_id
|
|
|
from cyber_agent.trace.goal_models import GoalTree
|
|
|
@@ -22,9 +29,6 @@ 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()
|
|
|
@@ -39,13 +43,6 @@ 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()
|
|
|
@@ -158,27 +155,32 @@ async def _update_collaborator(
|
|
|
async def _update_goal_start(
|
|
|
store, trace_id: str, goal_id: str, mode: str,
|
|
|
sub_trace_ids: List[Dict[str, str]],
|
|
|
+ *,
|
|
|
+ accumulate_sub_trace_ids: bool,
|
|
|
) -> 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
|
|
|
+ next_entries = sub_trace_ids
|
|
|
+ if accumulate_sub_trace_ids:
|
|
|
+ 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
|
|
|
+ next_entries = list(merged_entries.values())
|
|
|
|
|
|
await store.update_goal(
|
|
|
trace_id, goal_id,
|
|
|
type="agent_call",
|
|
|
agent_call_mode=mode,
|
|
|
status="in_progress",
|
|
|
- sub_trace_ids=list(merged_entries.values()),
|
|
|
+ sub_trace_ids=next_entries,
|
|
|
)
|
|
|
|
|
|
|
|
|
@@ -216,12 +218,16 @@ def _aggregate_stats(results: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
|
}
|
|
|
|
|
|
|
|
|
-def _get_allowed_tools(context: dict, agent_depth: int) -> Optional[List[str]]:
|
|
|
+def _get_allowed_tools(
|
|
|
+ context: dict,
|
|
|
+ agent_depth: int,
|
|
|
+ policy: AgentPolicy,
|
|
|
+) -> Optional[List[str]]:
|
|
|
"""按子 Agent 深度生成唯一的工具权限列表。"""
|
|
|
runner = context.get("runner")
|
|
|
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:
|
|
|
+ if policy.mode is AgentMode.LEGACY or agent_depth >= policy.max_depth:
|
|
|
blocked_tools.add("agent")
|
|
|
return [
|
|
|
name for name in runner.tools.get_tool_names()
|
|
|
@@ -481,13 +487,17 @@ async def _run_agents(
|
|
|
single (len(tasks)==1): delegate 模式,全量工具
|
|
|
multi (len(tasks)>1): explore 模式,并行执行
|
|
|
|
|
|
- AGENT_RECURSION_ENABLED=true 时,depth < MAX_AGENT_DEPTH 的本地
|
|
|
- Sub-Agent 可继续使用 agent 工具。
|
|
|
+ 本地 Sub-Agent 的深度、直接孩子配额和工具权限由父 Trace
|
|
|
+ 已持久化的 AgentPolicy 决定,不重新读取环境变量。
|
|
|
"""
|
|
|
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}"}
|
|
|
+ try:
|
|
|
+ policy = policy_from_context(parent_trace.context)
|
|
|
+ except ValueError as exc:
|
|
|
+ return {"status": "failed", "error": str(exc)}
|
|
|
|
|
|
# continue_from: 复用已有 trace(仅 single)
|
|
|
sub_trace_id = None
|
|
|
@@ -510,68 +520,82 @@ async def _run_agents(
|
|
|
"status": "failed",
|
|
|
"error": "continue_from trace owner does not match the current trace owner",
|
|
|
}
|
|
|
+ if existing.context.get("created_by_tool") != "agent":
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "continue_from must reference a child created by the agent tool",
|
|
|
+ }
|
|
|
+ try:
|
|
|
+ child_policy = policy_from_context(existing.context)
|
|
|
+ except ValueError as exc:
|
|
|
+ return {"status": "failed", "error": str(exc)}
|
|
|
+ if (
|
|
|
+ child_policy.mode is not policy.mode
|
|
|
+ or child_policy.revision != policy.revision
|
|
|
+ ):
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "continue_from trace Agent mode does not match the current trace",
|
|
|
+ }
|
|
|
|
|
|
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]
|
|
|
+ parent_depth, expected_root_trace_id = await _resolve_trace_lineage(
|
|
|
+ store, parent_trace
|
|
|
+ )
|
|
|
child_depth, root_trace_id = await _resolve_trace_lineage(store, existing)
|
|
|
+ if (
|
|
|
+ child_depth != parent_depth + 1
|
|
|
+ or root_trace_id != expected_root_trace_id
|
|
|
+ ):
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "continue_from trace lineage does not match the current trace",
|
|
|
+ }
|
|
|
+ if child_depth > policy.max_depth:
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": (
|
|
|
+ "continue_from trace exceeds the persisted Agent mode depth: "
|
|
|
+ f"depth={child_depth}, max={policy.max_depth}"
|
|
|
+ ),
|
|
|
+ }
|
|
|
all_sub_trace_ids = [{"trace_id": sub_trace_id, "mission": mission}]
|
|
|
child_records = [{
|
|
|
"trace_id": sub_trace_id,
|
|
|
"depth": child_depth,
|
|
|
- "context": {
|
|
|
+ "context": apply_policy_to_context({
|
|
|
**existing.context,
|
|
|
"agent_depth": child_depth,
|
|
|
"root_trace_id": root_trace_id,
|
|
|
- },
|
|
|
+ }, policy),
|
|
|
"is_new": False,
|
|
|
}]
|
|
|
else:
|
|
|
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:
|
|
|
+ if parent_depth >= policy.max_depth:
|
|
|
return {
|
|
|
"status": "failed",
|
|
|
"error": (
|
|
|
f"Local Sub-Agent depth limit reached: "
|
|
|
- f"depth={parent_depth}, max={MAX_AGENT_DEPTH}"
|
|
|
+ f"depth={parent_depth}, max={policy.max_depth}, "
|
|
|
+ f"mode={policy.mode.value}"
|
|
|
),
|
|
|
}
|
|
|
|
|
|
- # 将配额检查和 Trace 创建放在同一进程内临界区,
|
|
|
- # 保证两个并发 agent() 调用不会共同突破 6 个直接孩子。
|
|
|
- async with _LOCAL_AGENT_SPAWN_LOCK:
|
|
|
- existing_children = await store.list_traces(
|
|
|
- parent_trace_id=trace_id,
|
|
|
- created_by_tool="agent",
|
|
|
- limit=MAX_LOCAL_CHILDREN_PER_TRACE + 1,
|
|
|
- )
|
|
|
- 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}"
|
|
|
- ),
|
|
|
- }
|
|
|
-
|
|
|
+ async def create_child_traces() -> None:
|
|
|
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 = {
|
|
|
+ child_context = apply_policy_to_context({
|
|
|
"created_by_tool": "agent",
|
|
|
"agent_depth": child_depth,
|
|
|
"root_trace_id": root_trace_id,
|
|
|
- }
|
|
|
+ }, policy)
|
|
|
sub_trace = Trace(
|
|
|
trace_id=stid,
|
|
|
mode="agent",
|
|
|
@@ -597,14 +621,40 @@ async def _run_agents(
|
|
|
})
|
|
|
|
|
|
if single:
|
|
|
+ nonlocal sub_trace_id
|
|
|
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
|
|
|
+ child_limit = policy.max_children_per_parent
|
|
|
+ if child_limit is None:
|
|
|
+ await create_child_traces()
|
|
|
+ else:
|
|
|
+ # Recursive 模式在单进程临界区内完成“计数 -> 创建”,
|
|
|
+ # 保证并发批次不会共同突破六个直接孩子。
|
|
|
+ async with _LOCAL_AGENT_SPAWN_LOCK:
|
|
|
+ existing_children = await store.list_traces(
|
|
|
+ parent_trace_id=trace_id,
|
|
|
+ created_by_tool="agent",
|
|
|
+ limit=child_limit + 1,
|
|
|
+ )
|
|
|
+ requested_children = len(tasks)
|
|
|
+ if len(existing_children) + requested_children > child_limit:
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": (
|
|
|
+ "Local child Agent limit exceeded: "
|
|
|
+ f"existing={len(existing_children)}, "
|
|
|
+ f"requested={requested_children}, max={child_limit}"
|
|
|
+ ),
|
|
|
+ }
|
|
|
+ await create_child_traces()
|
|
|
+
|
|
|
+ await _update_goal_start(
|
|
|
+ store, trace_id, goal_id,
|
|
|
+ "delegate" if single else "explore",
|
|
|
+ all_sub_trace_ids,
|
|
|
+ accumulate_sub_trace_ids=policy.accumulate_sub_trace_ids,
|
|
|
+ )
|
|
|
+ goal_started = True
|
|
|
|
|
|
# 创建执行协程。Trace 在上面已经按批次预创建,
|
|
|
# 因此孩子配额不会被并发调用绕过。
|
|
|
@@ -632,7 +682,7 @@ async def _run_agents(
|
|
|
|
|
|
# 构建消息
|
|
|
agent_msgs = list(msgs) + [{"role": "user", "content": task_item}]
|
|
|
- allowed_tools = _get_allowed_tools(context, child_depth)
|
|
|
+ allowed_tools = _get_allowed_tools(context, child_depth, policy)
|
|
|
|
|
|
debug = getattr(runner, 'debug', False)
|
|
|
agent_label = (agent_type or ("delegate" if single else f"explore-{i+1}"))
|
|
|
@@ -667,6 +717,7 @@ async def _run_agents(
|
|
|
store, trace_id, goal_id,
|
|
|
"delegate" if single else "explore",
|
|
|
all_sub_trace_ids,
|
|
|
+ accumulate_sub_trace_ids=policy.accumulate_sub_trace_ids,
|
|
|
)
|
|
|
|
|
|
# 执行
|
|
|
@@ -868,6 +919,11 @@ async def agent(
|
|
|
- 远端:由服务器按 agent_type 白名单过滤(如 remote_librarian 允许 ask_strategy / upload_strategy)
|
|
|
context: 框架自动注入的上下文
|
|
|
"""
|
|
|
+ try:
|
|
|
+ assert_removed_config_absent()
|
|
|
+ except ValueError as exc:
|
|
|
+ return {"status": "failed", "error": str(exc)}
|
|
|
+
|
|
|
# 远端路由:agent_type 以 remote_ 开头
|
|
|
if agent_type and agent_type.startswith(REMOTE_PREFIX):
|
|
|
if not isinstance(task, str):
|
|
|
@@ -997,18 +1053,71 @@ async def evaluate(
|
|
|
|
|
|
# 获取父 Trace 信息
|
|
|
parent_trace = await store.get_trace(trace_id)
|
|
|
+ if not parent_trace:
|
|
|
+ return {"status": "failed", "error": f"Parent trace not found: {trace_id}"}
|
|
|
+ try:
|
|
|
+ policy = policy_from_context(parent_trace.context)
|
|
|
+ except ValueError as exc:
|
|
|
+ return {"status": "failed", "error": str(exc)}
|
|
|
|
|
|
# 处理 continue_from 或创建新 Sub-Trace
|
|
|
if continue_from:
|
|
|
existing_trace = await store.get_trace(continue_from)
|
|
|
if not existing_trace:
|
|
|
return {"status": "failed", "error": f"Continue-from trace not found: {continue_from}"}
|
|
|
+ if existing_trace.parent_trace_id != trace_id:
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "continue_from must reference a direct child of the current trace",
|
|
|
+ }
|
|
|
+ if existing_trace.uid != parent_trace.uid:
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "continue_from trace owner does not match the current trace owner",
|
|
|
+ }
|
|
|
+ if existing_trace.context.get("created_by_tool") != "evaluate":
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "continue_from must reference a child created by evaluate",
|
|
|
+ }
|
|
|
+ try:
|
|
|
+ existing_policy = policy_from_context(existing_trace.context)
|
|
|
+ except ValueError as exc:
|
|
|
+ return {"status": "failed", "error": str(exc)}
|
|
|
+ if (
|
|
|
+ existing_policy.mode is not policy.mode
|
|
|
+ or existing_policy.revision != policy.revision
|
|
|
+ ):
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "continue_from trace Agent mode does not match the current trace",
|
|
|
+ }
|
|
|
+ parent_depth, expected_root_trace_id = await _resolve_trace_lineage(
|
|
|
+ store, parent_trace
|
|
|
+ )
|
|
|
+ child_depth, root_trace_id = await _resolve_trace_lineage(
|
|
|
+ store, existing_trace
|
|
|
+ )
|
|
|
+ if (
|
|
|
+ child_depth != parent_depth + 1
|
|
|
+ or root_trace_id != expected_root_trace_id
|
|
|
+ ):
|
|
|
+ return {
|
|
|
+ "status": "failed",
|
|
|
+ "error": "continue_from trace lineage does not match the current trace",
|
|
|
+ }
|
|
|
sub_trace_id = continue_from
|
|
|
goal_tree = await store.get_goal_tree(continue_from)
|
|
|
mission = goal_tree.mission if goal_tree else eval_prompt
|
|
|
sub_trace_ids = [{"trace_id": sub_trace_id, "mission": mission}]
|
|
|
else:
|
|
|
sub_trace_id = generate_sub_trace_id(trace_id, "evaluate")
|
|
|
+ parent_depth, root_trace_id = await _resolve_trace_lineage(store, parent_trace)
|
|
|
+ evaluator_context = apply_policy_to_context({
|
|
|
+ "created_by_tool": "evaluate",
|
|
|
+ "agent_depth": parent_depth + 1,
|
|
|
+ "root_trace_id": root_trace_id,
|
|
|
+ }, policy)
|
|
|
sub_trace = Trace(
|
|
|
trace_id=sub_trace_id,
|
|
|
mode="agent",
|
|
|
@@ -1019,7 +1128,7 @@ async def evaluate(
|
|
|
uid=parent_trace.uid if parent_trace else None,
|
|
|
model=parent_trace.model if parent_trace else None,
|
|
|
status="running",
|
|
|
- context={"created_by_tool": "evaluate"},
|
|
|
+ context=evaluator_context,
|
|
|
created_at=datetime.now(),
|
|
|
)
|
|
|
await store.create_trace(sub_trace)
|
|
|
@@ -1033,7 +1142,14 @@ async def evaluate(
|
|
|
)
|
|
|
|
|
|
# 更新主 Goal 为 in_progress
|
|
|
- await _update_goal_start(store, trace_id, current_goal_id, "evaluate", sub_trace_ids)
|
|
|
+ await _update_goal_start(
|
|
|
+ store,
|
|
|
+ trace_id,
|
|
|
+ current_goal_id,
|
|
|
+ "evaluate",
|
|
|
+ sub_trace_ids,
|
|
|
+ accumulate_sub_trace_ids=policy.accumulate_sub_trace_ids,
|
|
|
+ )
|
|
|
|
|
|
# 注册为活跃协作者
|
|
|
eval_name = f"评估: {(goal_id or 'unknown')[:20]}"
|