Agent/cyber_agent/trace/goal_tool.py

"""
Goal 工具与目标树持久化入口。

ToolRegistry 向 LLM 暴露 goal()，实际变更由 goal_tool() 执行；
Recursive 还会在此阻止绕过待审核或已批准下一步的 Goal 修改。
"""

import logging
from typing import Optional, List, TYPE_CHECKING

from cyber_agent.tools import tool
from cyber_agent.core.agent_mode import policy_from_context
from cyber_agent.core.task_protocol import ensure_task_protocol

if TYPE_CHECKING:
    from .goal_models import GoalTree, Goal
    from .protocols import TraceStore

logger = logging.getLogger(__name__)


# ===== 知识注入 =====

async def inject_knowledge_for_goal(
    goal: "Goal",
    tree: "GoalTree",
    store: Optional["TraceStore"] = None,
    trace_id: Optional[str] = None,
    knowledge_config: Optional[dict] = None,
    sequence: Optional[int] = None,
) -> Optional[str]:
    """
    为指定 goal 注入相关知识。

    Args:
        goal: 目标对象
        tree: GoalTree
        store: TraceStore（用于持久化）
        trace_id: Trace ID
        knowledge_config: 知识管理配置（KnowledgeConfig 对象）
        sequence: 当前消息序列号（用于记录注入时机）

    Returns:
        注入结果描述（如 "📚 已注入 3 条相关知识"），无结果返回 None
    """
    # 检查是否启用知识注入
    if knowledge_config and not getattr(knowledge_config, 'enable_injection', True):
        logger.debug(f"[Knowledge Inject] 知识注入已禁用，跳过")
        return None

    try:
        import re
        from cyber_agent.tools.builtin.subagent import _run_remote_agent

        logger.info(f"[Knowledge Inject] goal: {goal.id}, query: {goal.description[:80]}")

        # 通过统一 agent 工具调用远端 Librarian，明确走 ask 策略
        result = await _run_remote_agent(
            agent_type="remote_librarian",
            task=goal.description,
            messages=None,
            continue_from=None,  # 知识注入每次独立查询，不续跑
            skills=["ask_strategy"],
        )

        if result.get("status") != "completed":
            logger.warning(f"[Knowledge Inject] 远端 Librarian 调用失败: {result.get('error')}")
            goal.knowledge = []
            return None

        response_text = result.get("summary", "")

        # 从 summary 中解析 source_ids（Librarian prompt 约定以 knowledge-xxx 形式引用）
        source_ids = re.findall(r'\[?(knowledge-[a-zA-Z0-9_-]+)\]?', response_text)
        source_ids = list(dict.fromkeys(source_ids))  # 去重保序

        if source_ids:
            goal.knowledge = [{"id": sid} for sid in source_ids]
            knowledge_count = len(source_ids)
            logger.info(f"[Knowledge Inject] 注入 {knowledge_count} 条知识到 goal {goal.id}")

            if store and trace_id:
                await store.update_goal_tree(trace_id, tree)

                if sequence is not None:
                    await store.append_cognition_event(
                        trace_id=trace_id,
                        event={
                            "type": "query",
                            "sequence": sequence,
                            "goal_id": goal.id,
                            "query": goal.description,
                            "response": response_text[:2000],
                            "source_ids": source_ids,
                            "sources": [],
                        }
                    )
                    logger.info(f"[Knowledge Inject] 已记录 query 事件到 cognition_log")

            return f"📚 已注入 {knowledge_count} 条相关知识"
        else:
            goal.knowledge = []
            logger.info(f"[Knowledge Inject] 未找到相关知识")
            return None

    except Exception as e:
        logger.warning(f"[Knowledge Inject] 知识注入失败: {e}")
        goal.knowledge = []
        return None


# ===== LLM 可调用的 goal 工具 =====

@tool(description="管理执行计划，添加/完成/放弃目标，切换焦点", hidden_params=["context"], groups=["core"])
async def goal(
    add: Optional[str] = None,
    reason: Optional[str] = None,
    after: Optional[str] = None,
    under: Optional[str] = None,
    done: Optional[str] = None,
    abandon: Optional[str] = None,
    focus: Optional[str] = None,
    context: Optional[dict] = None
) -> str:
    """
    LLM 可调用的计划管理工具，由 ToolRegistry 注入当前 GoalTree 和 Trace 上下文。

    它只解析公开参数，再转交 goal_tool() 持久化状态并发送前端事件。

    Args:
        add: 添加目标（逗号分隔多个）
        reason: 创建理由（逗号分隔多个，与 add 一一对应）
        after: 在指定目标后面添加（同层级）
        under: 为指定目标添加子目标
        done: 完成当前目标，值为 summary
        abandon: 放弃当前目标，值为原因
        focus: 切换焦点到指定 ID
        context: 工具执行上下文（包含 store、trace_id、goal_tree）

    Returns:
        str: 更新后的计划状态文本
    """
    # GoalTree 从 context 获取，每个 agent 实例独立，不再依赖全局变量
    tree = context.get("goal_tree") if context else None
    if tree is None:
        return "错误：GoalTree 未初始化"

    # 从 context 获取 store、trace_id 和 knowledge_config
    store = context.get("store") if context else None
    trace_id = context.get("trace_id") if context else None
    knowledge_config = context.get("knowledge_config") if context else None

    return await goal_tool(
        tree=tree,
        store=store,
        trace_id=trace_id,
        add=add,
        reason=reason,
        after=after,
        under=under,
        done=done,
        abandon=abandon,
        focus=focus,
        knowledge_config=knowledge_config,
        context=context
    )


# ===== 核心逻辑函数 =====


async def goal_tool(
    tree: "GoalTree",
    store: Optional["TraceStore"] = None,
    trace_id: Optional[str] = None,
    add: Optional[str] = None,
    reason: Optional[str] = None,
    after: Optional[str] = None,
    under: Optional[str] = None,
    done: Optional[str] = None,
    abandon: Optional[str] = None,
    focus: Optional[str] = None,
    knowledge_config: Optional[object] = None,
    context: Optional[dict] = None,
) -> str:
    """
    Goal 工具的共用实现，更新 GoalTree、TraceStore 和实时事件。

    Recursive 会先读持久化协议状态，关闭自动父级完成与兄弟切换，把后续决策留给 TaskReview。

    Args:
        tree: GoalTree 实例
        store: TraceStore 实例（用于推送事件）
        trace_id: 当前 Trace ID
        add: 添加目标（逗号分隔多个）
        reason: 创建理由（逗号分隔多个，与 add 一一对应）
        after: 在指定目标后面添加（同层级）
        under: 为指定目标添加子目标
        done: 完成当前目标，值为 summary
        abandon: 放弃当前目标，值为原因
        focus: 切换焦点到指定 ID
        knowledge_config: 知识管理配置（KnowledgeConfig 对象）

    Returns:
        更新后的计划状态文本
    """
    changes = []
    recursive_protocol = False
    if store and trace_id:
        trace = await store.get_trace(trace_id)
        if trace:
            recursive_protocol = policy_from_context(
                trace.context
            ).requires_task_protocol
            if recursive_protocol:
                state = ensure_task_protocol(trace.context)
                has_goal_mutation = any(
                    value is not None
                    for value in (add, done, abandon, focus)
                )
                if state["pending_reviews"] and has_goal_mutation:
                    return (
                        "错误：当前有子任务报告待审核，只允许调用 "
                        "review_task_result。"
                    )
                if state["next_actions"] and has_goal_mutation:
                    return (
                        "错误：当前有已批准的下一步待执行，只允许调用 "
                        "agent 执行该 action。"
                    )

    # 1. 处理 done（完成当前目标）
    if done is not None:
        if not tree.current_id:
            return f"错误：没有当前目标可以完成。当前焦点为空，请先使用 focus 参数切换到要完成的目标。\n\n当前计划：\n{tree.to_prompt()}"

        # 完成当前目标
        # 如果同时指定了 focus，则不清空焦点（后面会切换到新目标）
        # 如果只有 done，则清空焦点
        clear_focus = (focus is None)
        goal = tree.complete(
            tree.current_id,
            done,
            clear_focus=clear_focus,
            auto_advance=not recursive_protocol,
            cascade_parent=not recursive_protocol,
        )
        display_id = tree._generate_display_id(goal)
        changes.append(f"已完成: {display_id}. {goal.description}")

        # 推送事件
        if store and trace_id:
            await store.update_goal(
                trace_id,
                goal.id,
                cascade_completion=not recursive_protocol,
                status="completed",
                summary=done,
            )

        # 检查是否有级联完成的父目标（complete方法已经处理，这里只需要记录）
        if goal.parent_id and not recursive_protocol:
            parent = tree.find(goal.parent_id)
            if parent and parent.status == "completed":
                parent_display_id = tree._generate_display_id(parent)
                changes.append(f"自动完成: {parent_display_id}. {parent.description}（所有子目标已完成）")

    # 2. 处理 focus（切换焦点到新目标）
    if focus is not None:
        goal = tree.find_by_display_id(focus)

        if not goal:
            return f"错误：找不到目标 {focus}\n\n当前计划：\n{tree.to_prompt()}"

        tree.focus(goal.id)
        display_id = tree._generate_display_id(goal)
        changes.append(f"切换焦点: {display_id}. {goal.description}")

        # 自动注入知识
        inject_msg = await inject_knowledge_for_goal(
            goal, tree, store, trace_id, knowledge_config, sequence=context.get("sequence")
        )
        if inject_msg:
            changes.append(inject_msg)

    # 3. 处理 abandon（放弃当前目标）
    if abandon is not None:
        if not tree.current_id:
            return f"错误：没有当前目标可以放弃。当前焦点为空。\n\n当前计划：\n{tree.to_prompt()}"
        goal = tree.abandon(
            tree.current_id,
            abandon,
            auto_advance=not recursive_protocol,
        )
        display_id = tree._generate_display_id(goal)
        changes.append(f"已放弃: {display_id}. {goal.description}")

        # 推送事件
        if store and trace_id:
            await store.update_goal(trace_id, goal.id, status="abandoned", summary=abandon)

    # 4. 处理 add
    if add is not None:
        # 检查 after 和 under 互斥
        if after is not None and under is not None:
            return "错误：after 和 under 参数不能同时指定"

        descriptions = [d.strip() for d in add.split(",") if d.strip()]
        if descriptions:
            # 解析 reasons（与 descriptions 一一对应）
            reasons = None
            if reason:
                reasons = [r.strip() for r in reason.split(",")]
                # 如果 reasons 数量少于 descriptions，补空字符串
                while len(reasons) < len(descriptions):
                    reasons.append("")

            # 确定添加位置
            if after is not None:
                # 在指定 goal 后面添加（同层级）
                target_goal = tree.find_by_display_id(after)

                if not target_goal:
                    return f"错误：找不到目标 {after}\n\n当前计划：\n{tree.to_prompt()}"

                new_goals = tree.add_goals_after(target_goal.id, descriptions, reasons=reasons)
                changes.append(f"在 {tree._generate_display_id(target_goal)} 后面添加 {len(new_goals)} 个同级目标")

            elif under is not None:
                # 为指定 goal 添加子目标
                parent_goal = tree.find_by_display_id(under)

                if not parent_goal:
                    return f"错误：找不到目标 {under}\n\n当前计划：\n{tree.to_prompt()}"

                new_goals = tree.add_goals(descriptions, reasons=reasons, parent_id=parent_goal.id)
                changes.append(f"在 {tree._generate_display_id(parent_goal)} 下添加 {len(new_goals)} 个子目标")

            else:
                # 默认行为：添加到当前焦点下（如果有焦点），否则添加到顶层
                parent_id = tree.current_id
                new_goals = tree.add_goals(descriptions, reasons=reasons, parent_id=parent_id)

                if parent_id:
                    parent_display_id = tree._generate_display_id(tree.find(parent_id))
                    changes.append(f"在 {parent_display_id} 下添加 {len(new_goals)} 个子目标")
                else:
                    changes.append(f"添加 {len(new_goals)} 个顶层目标")

            # 推送事件
            if store and trace_id:
                for goal in new_goals:
                    await store.add_goal(trace_id, goal)

    # 将完整内存树状态（含 current_id）同步到存储，
    # 因为 store.add_goal / update_goal 各自从磁盘加载，不包含 focus 等内存变更
    if store and trace_id and changes:
        await store.update_goal_tree(trace_id, tree)

    # 返回当前状态
    result = []
    if changes:
        result.append("## 更新")
        result.extend(f"- {c}" for c in changes)
        result.append("")

    result.append("## Current Plan")
    result.append(tree.to_prompt())

    return "\n".join(result)


def create_goal_tool_schema() -> dict:
    """创建 goal 工具的 JSON Schema"""
    return {
        "name": "goal",
        "description": """管理执行计划。目标工具是灵活的支持系统，帮助你组织和追踪工作进度。

使用策略（按需选择）：
- 全局规划：先规划所有目标，再逐个执行
- 渐进规划：走一步看一步，每次只创建下一个目标
- 动态调整：行动中随时 abandon 不可行的目标，创建新目标

参数：
- add: 添加目标（逗号分隔多个）
- reason: 创建理由（逗号分隔，与 add 一一对应）
- after: 在指定目标后面添加同级目标。使用目标 ID。
- under: 为指定目标添加子目标。使用目标 ID。如已有子目标，追加到最后。
- done: 完成当前目标，值为 summary（记录关键结论）
- abandon: 放弃当前目标，值为原因
- focus: 切换焦点到指定目标。使用目标 ID。

位置控制（优先使用 after）：
- 不指定 after/under: 添加到当前 focus 下作为子目标（无 focus 时添加到顶层）
- after="X": 在目标 X 后面添加兄弟节点（同层级）
- under="X": 为目标 X 添加子目标
- after 和 under 不能同时指定

执行顺序：
- done → focus → abandon → add
- 如果同时指定 done 和 focus，会先完成当前目标，再切换焦点到新目标

示例：
- goal(add="分析代码, 实现功能, 测试") - 添加顶层目标
- goal(add="设计接口, 实现代码", under="2") - 为目标2添加子目标
- goal(add="编写文档", after="3") - 在目标3后面添加同级任务
- goal(add="集成测试", after="2.2") - 在目标2.2后面添加同级任务
- goal(done="发现用户模型在 models/user.py") - 完成当前目标
- goal(done="已完成调研", focus="2") - 完成当前目标，切换到目标2
- goal(abandon="方案A需要Redis，环境没有") - 放弃当前目标

注意：
- 目标 ID 的格式为 "1", "2", "2.1", "2.2" 等，在计划视图中可以看到
- reason 应该与 add 的目标数量一致，如果数量不一致，缺少的 reason 将为空
""",
        "parameters": {
            "type": "object",
            "properties": {
                "add": {
                    "type": "string",
                    "description": "添加目标（逗号分隔多个）"
                },
                "reason": {
                    "type": "string",
                    "description": "创建理由（逗号分隔多个，与 add 一一对应）。说明为什么要做这些目标。"
                },
                "after": {
                    "type": "string",
                    "description": "在指定目标后面添加（同层级）。使用目标的 ID，如 \"2\" 或 \"2.1\"。"
                },
                "under": {
                    "type": "string",
                    "description": "为指定目标添加子目标。使用目标的 ID，如 \"2\" 或 \"2.1\"。"
                },
                "done": {
                    "type": "string",
                    "description": "完成当前目标，值为 summary"
                },
                "abandon": {
                    "type": "string",
                    "description": "放弃当前目标，值为原因"
                },
                "focus": {
                    "type": "string",
                    "description": "切换焦点到指定目标。使用目标的 ID，如 \"2\" 或 \"2.1\"。"
                }
            },
            "required": []
        }
    }