iAs/src/worker.rs

#![deprecated = "Worker 进程模式已废弃，请使用 daemon 模式（统一 tokio task 架构）"]

//! Worker — 无状态执行进程（已废弃）
//!
//! 此模块保留供编译参考，不再被 daemon 使用。
//! 新架构使用 daemon 内的 LLM/Tool/Send 三个消费者 tokio task。
//!   1. 连接 daemon 的 Unix Domain Socket
//!   2. 读取 TaskFrame (用户消息 + 上下文 + 环境变量)
//!   3. 注入环境变量，构建 Conversation，执行 LLM 对话 + 工具调用
//!   4. 逐帧输出结果 (usage, reply, bye)
//!   5. 高风险工具遇审批 → 发送 RequestApproval → 退出，由 daemon 重新 spawn

use crate::ipc::{OutputFrame, TaskFrame, recv_task, send_frame};
use crate::llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT};
use crate::tools::types::SkillResult;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::net::UnixStream;
use tracing::info;

/// Worker 与工具执行器之间的共享状态
struct WorkerShared {
    /// 待发送的审批请求 (tool, reason)
    pending_approval: tokio::sync::Mutex<Option<(String, String)>>,
    /// 本次对话中已由用户批准的工具名（避免重复审批）
    approved_tool: tokio::sync::Mutex<Option<String>>,
    /// 本地记忆缓存
    memories: tokio::sync::Mutex<HashMap<String, Vec<String>>>,
    /// 本次写入的新记忆（退出前发送 DbWrite 帧同步到 daemon）
    new_memories: tokio::sync::Mutex<Vec<String>>,
    /// 预载的摘要
    summaries: Vec<String>,
    /// 用户 ID
    user_id: String,
}

/// Worker 入口: 连接 daemon，处理一条消息后退出
pub async fn run(sock_path: &str) -> Result<(), String> {
    let mut stream = UnixStream::connect(sock_path)
        .await
        .map_err(|e| format!("连接 daemon 失败 ({}): {}", sock_path, e))?;

    info!("Worker 已连接 daemon");

    // 1. 读 task
    let task: TaskFrame = recv_task(&mut stream).await?;
    info!(
        "Worker 收到任务: user={} msg={:.60}",
        task.user_id, task.msg.text
    );

    // 2. 注入环境变量
    for (k, v) in &task.env {
        unsafe {
            std::env::set_var(k, v);
        }
    }

    // 3. 构建 ConversationConfig
    let system_prompt = if task.system_prompt.is_empty() {
        DEFAULT_SYSTEM_PROMPT.to_string()
    } else {
        task.system_prompt.clone()
    };
    let model = if task.model.is_empty() {
        std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string())
    } else {
        task.model.clone()
    };

    let cfg = ConversationConfig {
        system_prompt,
        model,
        tools: Some(task.tools.clone()),
        ..Default::default()
    };

    // 4. 构建 Conversation 并加载上下文
    let mut conv =
        Conversation::new(cfg).map_err(|e| format!("初始化 Conversation 失败: {}", e))?;
    info!("Worker LLM: {}/{}", conv.provider_name(), conv.model());

    {
        let session = conv.session();
        let mut s = session.lock().await;
        s.load_from_history(&task.history);
        s.current_user_id = task.user_id.clone();
    }

    // 5. 共享状态
    let approved_tool = task.approved_tool.clone().filter(|t| !t.is_empty());
    let shared = Arc::new(WorkerShared {
        pending_approval: tokio::sync::Mutex::new(None),
        approved_tool: tokio::sync::Mutex::new(approved_tool),
        memories: tokio::sync::Mutex::new(HashMap::from([(
            task.user_id.clone(),
            task.memories.clone(),
        )])),
        new_memories: tokio::sync::Mutex::new(Vec::new()),
        summaries: task.summaries.clone(),
        user_id: task.user_id.clone(),
    });

    // 6. 构建工具执行器
    let executor = build_worker_executor(conv.session(), shared.clone());
    conv.set_tool_executor(executor);

    // 7. 执行 LLM 对话（带审批检测）
    use crate::llm::ChatResult;
    let ChatResult { reply, used_tools: _used_tools, usage, .. } = conv
        .chat_with_tools(task.msg.text.clone())
        .await
        .map_err(|e| format!("LLM 对话失败: {}", e))?;

    // 8. 检查是否有审批请求（由工具执行器设置）
    let approval_req = shared.pending_approval.lock().await.take();
    if let Some((tool, reason)) = approval_req {
        info!("Worker 发送审批请求: tool={}", tool);
        let _ = send_frame(&mut stream, &OutputFrame::RequestApproval { tool, reason }).await;
        let _ = send_frame(&mut stream, &OutputFrame::Bye).await;
        return Ok(());
    }

    // 9. 发送结果帧
    // 先同步新写入的记忆到 daemon
    let new_mems = shared.new_memories.lock().await;
    for content in new_mems.iter() {
        let _ = send_frame(
            &mut stream,
            &OutputFrame::DbWrite {
                table: "user_memories".into(),
                row: serde_json::json!({"content": content}),
            },
        )
        .await;
    }
    drop(new_mems);

    if let Some(u) = &usage {
        let _ = send_frame(
            &mut stream,
            &OutputFrame::UsageReport {
                model: conv.model().to_string(),
                provider: conv.provider_name().to_string(),
                prompt_tokens: u.prompt_tokens,
                completion_tokens: u.completion_tokens,
                total_tokens: u.total_tokens,
                cache_hit_tokens: u.prompt_cache_hit_tokens,
                cache_miss_tokens: u.prompt_cache_miss_tokens,
                user_id: task.user_id.clone(),
            },
        )
        .await;
    }

    let _ = send_frame(
        &mut stream,
        &OutputFrame::Reply {
            text: reply.clone(),
        },
    )
    .await;

    let _ = send_frame(&mut stream, &OutputFrame::Bye).await;

    info!(
        "Worker 完成: user={} reply_len={}",
        task.user_id,
        reply.len()
    );
    Ok(())
}

/// 构建 Worker 侧的工具执行器（无 DB 连接）
fn build_worker_executor(
    _session: Arc<tokio::sync::Mutex<crate::context::types::ChatSession>>,
    shared: Arc<WorkerShared>,
) -> crate::llm::conversation::ToolExecutor {
    Arc::new(move |name: &str, args_json: &str| {
        let shared = shared.clone();
        let n = name.to_string();
        let args = args_json.to_string();

        Box::pin(async move {
            match n.as_str() {
                "read_memories" => {
                    let cache = shared.memories.lock().await;
                    let mems = cache.get(&shared.user_id);
                    match mems {
                        Some(m) if !m.is_empty() => {
                            let lines: Vec<String> = m
                                .iter()
                                .enumerate()
                                .map(|(i, v)| format!("{}. {}", i + 1, v))
                                .collect();
                            Ok(lines.join("\n"))
                        }
                        _ => Ok("暂无长期记忆".to_string()),
                    }
                }
                "write_memory" => {
                    let params: serde_json::Value = serde_json::from_str(&args).unwrap_or_default();
                    let content = params.get("content").and_then(|v| v.as_str()).unwrap_or("");
                    if content.is_empty() {
                        return Ok("请提供 content 参数".to_string());
                    }
                    shared
                        .memories
                        .lock()
                        .await
                        .entry(shared.user_id.clone())
                        .or_default()
                        .push(content.to_string());
                    shared.new_memories.lock().await.push(content.to_string());
                    Ok(format!("已添加长期记忆: {}", content))
                }
                "read_summaries" => {
                    if shared.summaries.is_empty() {
                        Ok("暂无历史摘要".to_string())
                    } else {
                        Ok(shared
                            .summaries
                            .iter()
                            .enumerate()
                            .map(|(i, s)| format!("{}. {}", i + 1, s))
                            .collect::<Vec<_>>()
                            .join("\n---\n"))
                    }
                }
                "query_capabilities" => {
                    let specs = crate::tools::builtin::BuiltinRegistry::specs();
                    Ok(crate::tools::build_capability_guide(&specs))
                }
                // call_capability 或直接调用
                _ => {
                    let (target_name, target_args) = if n == "call_capability" {
                        let cp: serde_json::Value = serde_json::from_str(&args).unwrap_or_default();
                        let name = cp
                            .get("name")
                            .and_then(|v| v.as_str())
                            .unwrap_or("")
                            .to_string();
                        let unpacked = crate::tools::unpack_call_params(&cp);
                        (name, serde_json::to_string(&unpacked).unwrap_or_default())
                    } else {
                        (n.clone(), args.clone())
                    };

                    if target_name.is_empty() {
                        return Ok(format!("未知工具: {}", n));
                    }

                    // 高风险工具 → 请求审批（除非已被用户批准）
                    if crate::tools::builtin::BuiltinRegistry::is_high_risk(&target_name) {
                        let is_approved = shared.approved_tool.lock().await.as_deref() == Some(&target_name);
                        if !is_approved {
                            let reason = format!("LLM 请求执行高风险工具: {}", target_name);
                            *shared.pending_approval.lock().await = Some((target_name.clone(), reason));
                            return Ok(SkillResult::rejected(&target_name).output);
                        }
                        // 批准过 → 执行一次后清除（同一对话内同一工具不会再被拦截）
                        *shared.approved_tool.lock().await = None;
                    }

                    // 执行内置工具
                    if let Some(result) =
                        crate::tools::builtin::BuiltinRegistry::execute(&target_name, &target_args)
                            .await
                    {
                        if target_name == "manage_memos" {
                            let cp: serde_json::Value =
                                serde_json::from_str(&target_args).unwrap_or_default();
                            if cp.get("action").and_then(|v| v.as_str()) == Some("add") {
                                if let Some(content) = cp.get("content").and_then(|v| v.as_str()) {
                                    shared
                                        .memories
                                        .lock()
                                        .await
                                        .entry(shared.user_id.clone())
                                        .or_default()
                                        .push(content.to_string());
                                    shared.new_memories.lock().await.push(content.to_string());
                                }
                            }
                        }
                        return Ok(result.output);
                    }

                    Ok(format!("未知工具: {}", target_name))
                }
            }
        })
    })
}