DevFlow/src-tauri/src/commands/ai/tools/spawn_agent.rs

//! spawn_agent AI 工具声明式注册(派发子代理,隔离上下文执行专项任务,报告经返回值交付)。
//!
//! 架构:子代理 = 纯内存子会话 + 独立 agentic 循环(agentic/sub_agent.rs),本工具是唯一入口:
//! 1. 解析参数(task/tools/max_rounds/model),白名单钳制写死在入口(tools ∩ Low 风险 ∖ spawn_agent);
//! 2. 组装任务书(父上下文头 + 白名单说明 + 被裁剪尾注),注册级联取消条目(parent→children);
//! 3. 同步等待子循环跑完,报告作为 tool_result 返回值交付(Claude Code Task 语义同构)。
//!
//! ## 自引用绕法(两条进程级 OnceLock,先例 GLOBAL_TOOL_META)
//! - 工具注册表:白名单钳制需要 AiToolRegistry 自身,而注册期无法捕获自身 Arc,
//!   state.rs init 在 build_ai_tool_registry 返回后调 init_global_tool_registry 注入;
//! - AppHandle:子循环 emit/guard 需要,lib.rs setup 阶段调 init_app_handle 注入。
//!
//! ## 父会话定位
//! 工具 handler 签名只收 args,父 loop 的 conv_id 经 task-local 注入(复用 command_stream
//! 的 task-local 范式):process_tool_calls / ai_approve 执行工具前 scope_caller_conv 注入,
//! handler 内 caller_conv() 读取;缺省兜底 active_conversation_id(直调防御)。

use std::sync::Arc;

use tokio::sync::{Mutex, RwLock};

use df_ai::ai_tools::{
    AiToolRegistry, GoalMeta, GoalStyle, RiskLevel, SigMeta, SigValKind, ToolMeta,
};
use df_ai::declare_tool;
use df_ai::provider::ToolDefinition;

use crate::commands::ai::agentic::conv_state::ConvStateStore;
use crate::commands::ai::agentic::sub_agent::{
    clamp_tool_whitelist, next_child_id, run_sub_agent_loop, sub_agent_system_prompt,
    sub_agents_register, sub_agents_remove, SubAgentRegistry,
    SUB_AGENT_DEFAULT_MAX_ROUNDS, SUB_AGENT_MAX_ROUNDS_CAP,
};
use crate::commands::ai::AiSession;
use crate::commands::ai::conversation::SaveMetaOverride;
use crate::state::LlmConcurrency;

use df_storage::db::Database;

/// spawn_agent 工具捕获上下文(Arc 全字段,Clone 廉价;先例 GetAppConfigCtx)。
///
/// 各 Arc 由 state.rs init 在 build_ai_tool_registry **之前**提前构建并 clone 注入,
/// AppState struct 字段复用同一 Arc(注册先于 ai_session 构建的自引用绕法)。
/// derive Clone:declare_tool! 把捕获变量 clone 进工具闭包。
#[derive(Clone)]
pub struct SpawnAgentCtx {
    /// 父/子会话状态(读父 model_override / active_conversation_id,子会话建条)
    pub session: Arc<Mutex<AiSession>>,
    /// 会话生成态存储(子循环 guard 经 AppHandle 侧使用,此处预留一致生命周期;
    /// 与 AppState.conv_states 共享同一 Arc,随子代理消费场景接入启用)
    #[allow(dead_code)]
    pub conv_states: Arc<ConvStateStore>,
    /// LLM 并发控制(子循环 global + per_conv(child) 双份额)
    pub llm_concurrency: LlmConcurrency,
    /// 文件访问授权白名单(子代理投影段定位工作区:首个持久授权目录;
    /// 与 AppState.allowed_dirs 共享同一 Arc)
    pub allowed_dirs: Arc<RwLock<crate::state::AllowedDirs>>,
    /// 数据库句柄(子循环 provider 解析)
    pub db: Arc<Database>,
    /// 级联取消注册表(parent→children,与 AppState.sub_agent_registry 共享同一 Arc)
    pub sub_agents: SubAgentRegistry,
}

// ── 进程级 OnceLock:工具注册表 + AppHandle(注册期无法捕获,运行期注入,见模块头) ──

static GLOBAL_SPAWN_TOOL_REGISTRY: std::sync::OnceLock<Arc<AiToolRegistry>> =
    std::sync::OnceLock::new();
static GLOBAL_APP: std::sync::OnceLock<tauri::AppHandle> = std::sync::OnceLock::new();

/// state.rs init 注入工具注册表(白名单钳制数据源;重复注入忽略,首次生效)。
pub(crate) fn init_global_tool_registry(registry: Arc<AiToolRegistry>) {
    let _ = GLOBAL_SPAWN_TOOL_REGISTRY.set(registry);
}

/// handler 读取工具注册表(未初始化返 None,handler 报明确错误)。
fn tool_registry() -> Option<Arc<AiToolRegistry>> {
    GLOBAL_SPAWN_TOOL_REGISTRY.get().cloned()
}

/// lib.rs setup 注入 AppHandle(子循环 emit/guard 所需;重复注入忽略,首次生效)。
pub(crate) fn init_app_handle(app: tauri::AppHandle) {
    let _ = GLOBAL_APP.set(app);
}

/// handler 读取 AppHandle(未初始化返 None,handler 报明确错误)。
/// pub(crate):编排会话启动通道(chat.rs orchestration_session_spawner)同款复用。
pub(crate) fn app_handle() -> Option<tauri::AppHandle> {
    GLOBAL_APP.get().cloned()
}

// ── 父会话 task-local(复用 command_stream 的 task-local 范式) ──

tokio::task_local! {
    static CALLER_CONV: String;
}

/// 在注入父会话 id 的作用域内执行 future(process_tool_calls / ai_approve 调用工具前包裹)。
pub(crate) async fn scope_caller_conv<F, R>(conv_id: String, fut: F) -> R
where
    F: std::future::Future<Output = R>,
{
    CALLER_CONV.scope(conv_id, fut).await
}

/// handler 读取发起调用的父会话 id(未经 scope 直调时返 None,调用方走兜底)。
pub(crate) fn caller_conv() -> Option<String> {
    CALLER_CONV.try_with(|c| c.clone()).ok()
}

// ── 参数解析与任务书组装(纯函数,单测锁定) ──

/// 解析 max_rounds:缺省默认 12,合法域钳制 [1, 20](防参数放大预算失控)。
fn parse_max_rounds(raw: Option<u64>) -> u32 {
    raw.map(|v| v as u32)
        .unwrap_or(SUB_AGENT_DEFAULT_MAX_ROUNDS)
        .clamp(1, SUB_AGENT_MAX_ROUNDS_CAP)
}

/// 组装子代理任务书:父上下文头 + 任务正文 + 工具白名单说明 + 被裁剪项尾注。
fn build_task_prompt(task: &str, allowed: &[String], trimmed: &[String]) -> String {
    let mut p = String::new();
    p.push_str("[父任务上下文] 你是主会话派生的子代理,以下是父任务交办的任务书。\n\n<task>\n");
    p.push_str(task);
    p.push_str("\n</task>\n\n");
    if allowed.is_empty() {
        p.push_str("可用工具:无(白名单为空),请直接基于推理输出结论报告。\n");
    } else {
        p.push_str(&format!(
            "可用工具白名单(仅限以下工具,且需授权/高风险操作会被自动拒绝):{}\n",
            allowed.join(", ")
        ));
    }
    if !trimmed.is_empty() {
        p.push_str(&format!(
            "(以下请求的工具因超出子代理权限或未注册被剔除,如确需请写入报告交由主会话完成:{})\n",
            trimmed.join(", ")
        ));
    }
    p
}

/// 注册 spawn_agent 工具(RiskLevel::Medium,派发需用户审批,与 create_task 同档)。
pub fn register(registry: &mut AiToolRegistry, ctx: SpawnAgentCtx) {
    // schema 手写 JSON(task 必填 + 三个可选参数,描述面向 LLM)
    let schema = serde_json::json!({
        "type": "object",
        "properties": {
            "task": {
                "type": "string",
                "description": "自包含任务书:背景/目标/完成标准。子代理看不到父对话,任务书必须自带全部所需上下文"
            },
            "tools": {
                "type": "array",
                "items": { "type": "string" },
                "description": "可选,工具名白名单。仅低风险只读工具可保留,其余被剔除并在任务书尾注说明;缺省=全部低风险工具"
            },
            "max_rounds": {
                "type": "integer",
                "description": "可选,工具循环轮数上限,默认 12,最大 20"
            },
            "model": {
                "type": "string",
                "description": "可选,子代理用模型 id;缺省继承当前会话模型"
            },
            "temperature": {
                "type": "number",
                "description": "可选,子代理采样温度(0~2);缺省沿用 provider 配置/默认 0.7,不影响主会话"
            }
        },
        "required": ["task"]
    });
    declare_tool!(
        registry,
        ctx: SpawnAgentCtx,
        "spawn_agent",
        "派发一个子代理在隔离上下文中独立完成专项任务(不与用户交互),最终返回结论报告。适用:探索检索/多文件阅读汇总/资料核对等可独立收敛的子任务。参数 task 为自包含任务书(子代理看不到父对话);tools 为可选工具白名单;max_rounds 可选轮数上限(默认 12);model 可选模型(缺省继承当前会话);temperature 可选采样温度(缺省沿用 provider 配置)。注意:子代理无审批通道,写文件/命令执行等高风险工具会被自动拒绝,结论以 report 字段文本返回",
        RiskLevel::Medium,
        meta: ToolMeta {
            // 审批卡文案模板(模板 {} 按位置填充,键取 task)
            approval_hint: Some(("派发子代理: {}", &["task"])),
            display_verb: Some("派发"),
            // G1 目标推理:任务书截断展示(Truncate 25 字符,对齐 run_command 风格)
            goal: Some(GoalMeta {
                verb: "派发子代理",
                keys: &["task"],
                style: GoalStyle::Truncate,
                fallback: Some("派发子代理执行专项任务"),
            }),
            // G2 探索熔断签名:task 进签名,同任务书重复派发会被判重复打转熔断
            sig: Some(SigMeta {
                groups: &[(&["task"], SigValKind::Str)],
            }),
            ..Default::default()
        },
        schema: schema,
        args => {
            spawn_agent_handler(&ctx, &args).await
        }
    );
}

/// spawn_agent handler:钳制白名单 → 注册级联条目 → 同步跑子循环 → 报告作返回值。
async fn spawn_agent_handler(
    ctx: &SpawnAgentCtx,
    args: &serde_json::Value,
) -> anyhow::Result<serde_json::Value> {
    // 参数解析
    let task = args
        .get("task")
        .and_then(|v| v.as_str())
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
        .ok_or_else(|| anyhow::anyhow!("参数 task 不能为空:请传自包含任务书(背景/目标/完成标准)"))?;
    let max_rounds = parse_max_rounds(args.get("max_rounds").and_then(|v| v.as_u64()));
    let model = args
        .get("model")
        .and_then(|v| v.as_str())
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty());
    // 采样温度独立覆盖(缺省 None=沿用 provider 配置/默认 0.7,不影响主会话)
    let temperature = args.get("temperature").and_then(|v| v.as_f64()).map(|t| t as f32);
    let requested: Vec<String> = args
        .get("tools")
        .and_then(|v| v.as_array())
        .map(|a| a.iter().filter_map(|x| x.as_str().map(|s| s.to_string())).collect())
        .unwrap_or_default();

    // 白名单钳制(写死在入口,不依赖 LLM 自觉)+ 工具定义投影
    let registry = tool_registry()
        .ok_or_else(|| anyhow::anyhow!("子代理工具注册表未就绪(应用初始化未完成)"))?;
    let (allowed, trimmed) = clamp_tool_whitelist(&registry, &requested);
    let tool_defs: Vec<ToolDefinition> = registry
        .tool_definitions()
        .into_iter()
        .filter(|d| allowed.contains(&d.function.name))
        .collect();

    // 父会话定位:task-local 注入优先,兜底 active(直调防御)
    let parent_conv_id = match caller_conv() {
        Some(p) if !p.is_empty() => p,
        _ => ctx
            .session
            .lock()
            .await
            .active_conversation_id
            .clone()
            .unwrap_or_default(),
    };
    if parent_conv_id.is_empty() {
        return Err(anyhow::anyhow!("spawn_agent 无法定位父会话:只能在会话内调用"));
    }

    // model 不传 → 继承父会话模型 override(与父对话同模型语义);
    // lang 恒沿用父会话语言(agent_language,缺省中文),子代理主体与工具纪律段随之同语,
    // 否则英文会话派生的子代理仍用中文提示词(报告语言/术语断层)。一次锁内取两者。
    let (child_model, parent_lang) = {
        let session = ctx.session.lock().await;
        let parent = session.conv_read(&parent_conv_id);
        let model = match model {
            Some(m) => Some(m),
            None => parent.and_then(|c| c.model_override.clone()),
        };
        let lang = parent
            .and_then(|c| c.agent_language.clone())
            .filter(|s| !s.is_empty())
            .unwrap_or_else(|| "zh-CN".to_string());
        (model, lang)
    };

    let child_id = next_child_id(&parent_conv_id);
    let task_prompt = build_task_prompt(&task, &allowed, &trimmed);
    // P2-3:系统提示词注入实际生效的预算值(父可传 1..20,子代理按真实预算规划)
    let system_prompt = sub_agent_system_prompt(max_rounds, &parent_lang);
    // 项目规则文件投影与主会话同源:父会话工作区的 devflow.md/CLAUDE.md 对子代理同样生效
    // (以父会话为缓存键,复用父会话已读内容零重复读盘;无授权目录/无文件时原样返回)。
    let work_dir = {
        let snap = ctx.allowed_dirs.read().await;
        crate::commands::ai::devflow_md::work_dir_of(&snap)
    };
    let system_prompt = crate::commands::ai::devflow_md::append_projection(
        system_prompt, work_dir.as_deref(), &parent_conv_id,
    ).await;

    // P3-5:spawn 时刻快照父会话 save 元数据(active provider/created_at)。子代理后台
    // 长跑收尾时用户可能已切走会话,出口落库若读 active_* 会继承别的会话元数据;
    // 快照随子循环传入,出口经 save_conversation_meta 显式落库。
    let parent_save_meta = {
        let session = ctx.session.lock().await;
        SaveMetaOverride {
            provider_id: session.active_provider_id.clone(),
            created_at: session.active_conv_created_at.clone(),
        }
    };

    // 级联注册(父停止/强制发送/删除会遍历 children 级联取消)
    sub_agents_register(&ctx.sub_agents, &parent_conv_id, &child_id);

    let app = match app_handle() {
        Some(h) => h,
        None => {
            sub_agents_remove(&ctx.sub_agents, &parent_conv_id, &child_id);
            return Err(anyhow::anyhow!("应用句柄未就绪(启动初始化未完成)"));
        }
    };

    tracing::info!(
        parent = %parent_conv_id,
        agent = %child_id,
        task_chars = task.chars().count(),
        tools = allowed.len(),
        max_rounds,
        model = ?child_model,
        lang = %parent_lang,
        temperature = ?temperature,
        "[sub-agent] 派发子代理"
    );

    // 同步等待子循环(报告经返回值交付;外层超时对 spawn_agent 取子代理挂钟预算同款上限)
    let outcome = run_sub_agent_loop(
        ctx.session.clone(),
        registry.clone(),
        ctx.db.clone(),
        app,
        parent_conv_id.clone(),
        child_id.clone(),
        system_prompt,
        task_prompt,
        tool_defs,
        allowed,
        max_rounds,
        child_model,
        temperature,
        ctx.llm_concurrency.clone(),
        parent_save_meta,
    )
    .await;

    // 级联注销(空父条目整条清除)
    sub_agents_remove(&ctx.sub_agents, &parent_conv_id, &child_id);

    Ok(serde_json::json!({
        "status": outcome.status.as_str(),
        "agent_id": child_id,
        "rounds": outcome.rounds,
        "tokens": outcome.tokens,
        // 实际生效模型(回退/候选切换后的最终值):父会话据此知晓子代理真用了哪个模型
        "model": outcome.model,
        "report": outcome.summary,
    }))
}

#[cfg(test)]
mod tests {
    use super::*;

    /// parse_max_rounds:缺省 6,合法域钳制 [1,20]。
    #[test]
    fn parse_max_rounds_clamps() {
        assert_eq!(parse_max_rounds(None), SUB_AGENT_DEFAULT_MAX_ROUNDS);
        assert_eq!(parse_max_rounds(Some(1)), 1);
        assert_eq!(parse_max_rounds(Some(6)), 6);
        assert_eq!(parse_max_rounds(Some(0)), 1, "0 钳到下界 1");
        assert_eq!(parse_max_rounds(Some(999)), SUB_AGENT_MAX_ROUNDS_CAP, "超上限钳到 20");
    }

    /// 任务书:含父上下文头、任务正文、白名单说明与被裁剪尾注。
    #[test]
    fn task_prompt_carries_context_and_notes() {
        let p = build_task_prompt(
            "统计 crates 下 TODO 数量",
            &["read_file".to_string(), "grep".to_string()],
            &["run_command".to_string()],
        );
        assert!(p.contains("[父任务上下文]"), "应有父上下文头");
        assert!(p.contains("<task>\n统计 crates 下 TODO 数量\n</task>"), "任务正文应入 task 标签");
        assert!(p.contains("read_file, grep"), "应列白名单");
        assert!(p.contains("run_command"), "被裁剪项应有尾注说明");
        // 空白名单走"无工具"说明
        let empty = build_task_prompt("t", &[], &[]);
        assert!(empty.contains("白名单为空"));
    }

    /// task-local 父会话注入/读取:scope 内可读,scope 外读不到。
    #[tokio::test]
    async fn caller_conv_scoped() {
        assert!(caller_conv().is_none(), "scope 外应读不到");
        let seen = scope_caller_conv("conv-p::agent-9".to_string(), async {
            caller_conv()
        }).await;
        assert_eq!(seen.as_deref(), Some("conv-p::agent-9"));
    }
}