From 61c3d3e05924a8539ea71d4195f553569fdb0afd Mon Sep 17 00:00:00 2001 From: fusheng Date: Mon, 12 Oct 2026 22:50:09 +0800 Subject: [PATCH 1/3] =?UTF-8?q?feat(agent):=20=E5=AE=9E=E7=8E=B0=20Agent?= =?UTF-8?q?=20=E7=B3=BB=E7=BB=9F=E5=B9=B6=E6=8E=A5=E5=85=A5=20TUI=20?= =?UTF-8?q?=E8=81=8A=E5=A4=A9=E9=A1=B5=E9=9D=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 Agent 核心模块，支持流式对话与基于提示词的 ReAct 工具调用循环 - 实现 SessionStore 会话持久化，将对话消息写入 SQLite 数据库 - 新增 AppConfig 模块，从环境变量读取供应商、模型、工作目录等配置 - 新增 Runtime 模块，提供工作目录上下文与路径越界校验 - 新增 ToolRegistry 工具系统，内置文件读写与目录列表工具 - 为 HttpTransport 添加 execute_sse 方法，支持原始 SSE 事件流 - 为 AIClient 添加 post_sse 便捷方法，供供应商自定义解析 - 重写聊天页面，将 Agent 接入 TUI 实现实时流式渲染与交互 --- src/core/agent/mod.rs | 538 ++++++++++++++++++ src/core/client/ai_client.rs | 25 + src/core/config/mod.rs | 139 +++++ src/core/mod.rs | 6 +- src/core/providers/config.rs | 13 + src/core/providers/mod.rs | 7 +- .../providers/openai_compatible/provider.rs | 71 ++- .../providers/openai_compatible/response.rs | 20 + src/core/providers/provider.rs | 21 +- src/core/providers/types.rs | 17 + src/core/runtime/mod.rs | 121 ++++ src/core/tools/mod.rs | 310 ++++++++++ src/core/transport/http_transoprt.rs | 19 + src/pages/chat.rs | 296 +++++++++- 14 files changed, 1582 insertions(+), 21 deletions(-) diff --git a/src/core/agent/mod.rs b/src/core/agent/mod.rs index e69de29..22a13f3 100644 --- a/src/core/agent/mod.rs +++ b/src/core/agent/mod.rs @@ -0,0 +1,538 @@ +//! Agent 系统 +//! +//! 顶层 Agent 负责把底层能力编排成一次完整的对话轮次： +//! - **流式对话**：通过 [`Provider::complete_stream`] 增量获取模型输出，并以 +//! [`AgentEvent::Token`] 事件实时转发给上层（如 TUI）。 +//! - **工具调用循环**：采用基于提示词的 ReAct 方式，解析模型输出中的工具调用意图， +//! 执行工具（见 [`crate::core::tools`]），把结果回灌后继续对话，直到模型给出最终答复。 +//! - **会话持久化**：可选地把用户/助手/工具消息写入 `chat_message` 表，并支持加载历史。 +//! +//! 使用方式： +//! ```ignore +//! let mut agent = Agent::from_config(&config)?; +//! agent.enable_persistence(1, Some("New Chat"))?; +//! let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); +//! tokio::spawn(async move { agent.run_turn("hello", tx).await }); +//! while let Some(event) = rx.recv().await { /* ... */ } +//! ``` + +use std::sync::{Arc, Mutex}; + +use anyhow::{Context, Result}; +use futures_util::{Stream, StreamExt}; +use log::warn; +use rusqlite::{params, Connection, OptionalExtension}; +use serde_json::{json, Value}; +use tokio::sync::mpsc; + +use crate::core::config::AppConfig; +use crate::core::providers::{ + create_provider, CompletionRequest, Message, Provider, Role, +}; +use crate::core::runtime::Runtime; +use crate::core::tools::ToolRegistry; +use crate::platform::get_database_path; + +/// 默认最大工具调用迭代次数，防止无限循环。 +pub const DEFAULT_MAX_ITERATIONS: usize = 8; + +/// Agent 运行过程中向上层发送的事件。 +#[derive(Debug, Clone)] +pub enum AgentEvent { + /// 模型输出的一个增量文本片段。 + Token(String), + /// Agent 决定调用某个工具。 + ToolCall { name: String, arguments: Value }, + /// 工具执行完成并返回结果。 + ToolResult { name: String, output: String }, + /// 本轮对话结束。 + Done, + /// 发生错误（发送后本轮通常随之结束）。 + Error(String), +} + +/// 解析出的一次工具调用。 +#[derive(Debug, Clone)] +pub struct ToolCall { + pub name: String, + pub arguments: Value, +} + +/// 从模型输出中解析 `tool` 代码块形式的工具调用。 +/// +/// 约定格式： +/// ````text +/// ```tool +/// {"name": "read_file", "arguments": {"path": "src/main.rs"}} +/// ``` +/// ```` +pub fn parse_tool_call(content: &str) -> Option { + const MARKER: &str = "```tool"; + + let start = content.find(MARKER)?; + let after_marker = &content[start + MARKER.len()..]; + let end = after_marker.find("```")?; + let json_str = after_marker[..end].trim(); + + let value: Value = serde_json::from_str(json_str).ok()?; + let name = value.get("name")?.as_str()?.to_string(); + let arguments = value.get("arguments").cloned().unwrap_or_else(|| json!({})); + + Some(ToolCall { name, arguments }) +} + +// ============================================================ +// 会话持久化 +// ============================================================ + +/// 定位/创建 provider→model 记录所需的信息。 +#[derive(Debug, Clone)] +pub struct ModelRef { + /// 供应商名称（在 `ai_provider` 中按名匹配），如 "compatible"。 + pub provider_name: String, + /// 供应商 API Base URL。 + pub provider_api: String, + /// 供应商 API 密钥（可选）。 + pub provider_key: Option, + /// 模型代码（`llm_model.model_code`），如 "gpt-4o-mini"。 + pub model_code: String, + /// 模型展示名（可选，默认取 `model_code`）。 + pub model_name: Option, +} + +/// 会话存储：把对话消息持久化到 SQLite。 +/// +/// 内部使用 [`Mutex`] 保证 `Send + Sync`，可安全地被 Agent 在异步任务中共享。 +pub struct SessionStore { + conn: Mutex, +} + +impl SessionStore { + /// 打开（或创建）应用数据库连接。 + pub fn new() -> Result { + let path = get_database_path()?; + let conn = Connection::open(&path) + .with_context(|| format!("Failed to open database: {}", path.display()))?; + Ok(Self { + conn: Mutex::new(conn), + }) + } + + fn lock(&self) -> Result> { + self.conn + .lock() + .map_err(|_| anyhow::anyhow!("SessionStore mutex poisoned")) + } + + /// 确保 `ai_provider`→`llm_model` 记录存在，返回 `llm_model` 的 id。 + /// + /// 供应商按 `name` 匹配、模型按 `(provider_id, model_code)` 匹配；缺失时以真实配置创建， + /// 从而建立正确的 provider→model 关联（而非占位的 `provider_id = 0`）。 + pub fn ensure_model(&self, m: &ModelRef) -> Result { + let conn = self.lock()?; + + // 1) 查/建供应商。 + let provider_id: i32 = match conn + .query_row( + "SELECT id FROM ai_provider WHERE name = ?1 ORDER BY id ASC LIMIT 1", + params![m.provider_name], + |row| row.get(0), + ) + .optional()? + { + Some(id) => id, + None => { + conn.execute( + "INSERT INTO ai_provider (name, avatar, api, key, timeout, status) \ + VALUES (?1, '', ?2, ?3, 5000, 1)", + params![m.provider_name, m.provider_api, m.provider_key], + )?; + conn.last_insert_rowid() as i32 + } + }; + + // 2) 查/建模型。 + let model_name = m.model_name.clone().unwrap_or_else(|| m.model_code.clone()); + let model_id: i32 = match conn + .query_row( + "SELECT id FROM llm_model WHERE provider_id = ?1 AND model_code = ?2 \ + ORDER BY id ASC LIMIT 1", + params![provider_id, m.model_code], + |row| row.get(0), + ) + .optional()? + { + Some(id) => id, + None => { + conn.execute( + "INSERT INTO llm_model (provider_id, name, avatar, model_code, max_tokens, max_response_tokens, weight, call_count, status) \ + VALUES (?1, ?2, '', ?3, 0, 0, 0, 0, 1)", + params![provider_id, model_name, m.model_code], + )?; + conn.last_insert_rowid() as i32 + } + }; + + Ok(model_id) + } + + /// 创建一条会话记录。 + pub fn create_session(&self, uuid: &str, title: Option<&str>, llm_id: i32) -> Result<()> { + let conn = self.lock()?; + conn.execute( + "INSERT INTO chat_session (chat_uuid, title, llm_id, status) VALUES (?1, ?2, ?3, 1)", + params![uuid, title, llm_id], + )?; + Ok(()) + } + + /// 保存一条消息。 + pub fn save_message( + &self, + session_uuid: &str, + role: &str, + content: &str, + model: Option<&str>, + ) -> Result<()> { + let conn = self.lock()?; + conn.execute( + "INSERT INTO chat_message (session_id, role, content, model, status) VALUES (?1, ?2, ?3, ?4, 1)", + params![session_uuid, role, content, model], + )?; + Ok(()) + } + + /// 加载某个会话的历史消息（按时间升序）。 + pub fn load_history(&self, session_uuid: &str) -> Result> { + let conn = self.lock()?; + let mut stmt = conn.prepare( + "SELECT role, content FROM chat_message WHERE session_id = ?1 ORDER BY id ASC", + )?; + + let rows = stmt.query_map(params![session_uuid], |row| { + let role: String = row.get(0)?; + let content: String = row.get(1)?; + Ok((role, content)) + })?; + + let mut history = Vec::new(); + for row in rows { + let (role, content) = row?; + let role = match role.as_str() { + "system" => Role::System, + "assistant" => Role::Assistant, + _ => Role::User, + }; + history.push(Message { role, content }); + } + + Ok(history) + } +} + +// ============================================================ +// Agent +// ============================================================ + +/// 顶层对话 Agent。 +pub struct Agent { + provider: Box, + model: String, + system_prompt: String, + tools: ToolRegistry, + history: Vec, + max_iterations: usize, + store: Option>, + session_uuid: String, + /// 供应商元数据（用于持久化时建立 provider→model 关联）。 + provider_name: String, + provider_api: String, + provider_key: Option, +} + +impl Agent { + /// 使用给定组件构造 Agent。 + pub fn new( + provider: Box, + model: String, + system_prompt: String, + tools: ToolRegistry, + ) -> Self { + Self { + provider, + model, + system_prompt, + tools, + history: Vec::new(), + max_iterations: DEFAULT_MAX_ITERATIONS, + store: None, + session_uuid: generate_uuid(), + provider_name: String::new(), + provider_api: String::new(), + provider_key: None, + } + } + + /// 基于 [`AppConfig`] 构造 Agent，自动创建供应商并注册默认内置工具。 + pub fn from_config(config: &AppConfig) -> Result { + let provider = create_provider(config.provider_config())?; + let runtime = Arc::new(Runtime::new(config.working_dir.clone())); + let tools = ToolRegistry::with_defaults(runtime); + + let key = if config.api_key.is_empty() { + None + } else { + Some(config.api_key.clone()) + }; + + Ok(Self::new( + provider, + config.model.clone(), + config.system_prompt.clone(), + tools, + ) + .with_provider_meta(config.provider_kind.as_str(), config.base_url.clone(), key)) + } + + /// 设置供应商元数据（用于持久化时建立 provider→model 关联）。 + pub fn with_provider_meta( + mut self, + name: impl Into, + api: impl Into, + key: Option, + ) -> Self { + self.provider_name = name.into(); + self.provider_api = api.into(); + self.provider_key = key; + self + } + + /// 设置最大工具调用迭代次数。 + pub fn with_max_iterations(mut self, max_iterations: usize) -> Self { + self.max_iterations = max_iterations.max(1); + self + } + + /// 设置自定义工具注册表。 + pub fn with_tools(mut self, tools: ToolRegistry) -> Self { + self.tools = tools; + self + } + + /// 当前会话 UUID。 + pub fn session_uuid(&self) -> &str { + &self.session_uuid + } + + /// 当前对话历史。 + pub fn history(&self) -> &[Message] { + &self.history + } + + /// 开启持久化：创建新会话并绑定存储。 + pub fn enable_persistence(&mut self, llm_id: i32, title: Option<&str>) -> Result<()> { + let store = Arc::new(SessionStore::new()?); + let uuid = generate_uuid(); + store.create_session(&uuid, title, llm_id)?; + self.store = Some(store); + self.session_uuid = uuid; + Ok(()) + } + + /// 基于当前模型配置开启持久化：自动确保 provider→model 记录存在并创建新会话。 + /// + /// 适用于配置来自环境变量、没有显式 `llm_id` 的场景（如 TUI）。 + pub fn enable_persistence_auto(&mut self, title: Option<&str>) -> Result<()> { + let store = Arc::new(SessionStore::new()?); + let model_ref = ModelRef { + provider_name: if self.provider_name.is_empty() { + "custom".to_string() + } else { + self.provider_name.clone() + }, + provider_api: self.provider_api.clone(), + provider_key: self.provider_key.clone(), + model_code: self.model.clone(), + model_name: None, + }; + let llm_id = store.ensure_model(&model_ref)?; + let uuid = generate_uuid(); + store.create_session(&uuid, title, llm_id)?; + self.store = Some(store); + self.session_uuid = uuid; + Ok(()) + } + + /// 加载已有会话：读取历史消息并绑定存储。 + pub fn load_session(&mut self, session_uuid: String) -> Result<()> { + let store = Arc::new(SessionStore::new()?); + self.history = store.load_history(&session_uuid)?; + self.store = Some(store); + self.session_uuid = session_uuid; + Ok(()) + } + + /// 执行一轮完整对话（含流式输出与工具调用循环），过程事件通过 `tx` 发送。 + pub async fn run_turn( + &mut self, + user_input: &str, + tx: mpsc::UnboundedSender, + ) -> Result<()> { + self.history.push(Message { + role: Role::User, + content: user_input.to_string(), + }); + self.persist("user", user_input, None); + + let mut iteration = 0; + loop { + iteration += 1; + + let mut messages = Vec::with_capacity(self.history.len() + 1); + messages.push(Message { + role: Role::System, + content: self.build_system_prompt(), + }); + messages.extend(self.history.iter().cloned()); + + let request = CompletionRequest { + model: self.model.clone(), + messages, + temperature: None, + max_tokens: None, + stream: true, + }; + + let mut stream = match self.provider.complete_stream(request).await { + Ok(stream) => stream, + Err(e) => { + let _ = tx.send(AgentEvent::Error(e.to_string())); + return Err(e); + } + }; + + let mut accumulated = String::new(); + while let Some(item) = stream.next().await { + match item { + Ok(delta) => { + if !delta.content.is_empty() { + accumulated.push_str(&delta.content); + let _ = tx.send(AgentEvent::Token(delta.content)); + } + } + Err(e) => { + let _ = tx.send(AgentEvent::Error(e.to_string())); + return Err(e); + } + } + } + + self.history.push(Message { + role: Role::Assistant, + content: accumulated.clone(), + }); + self.persist("assistant", &accumulated, Some(self.model.as_str())); + + // 若模型请求调用工具且尚未超出迭代上限，则执行并把结果回灌。 + if iteration < self.max_iterations { + if let Some(call) = parse_tool_call(&accumulated) { + let _ = tx.send(AgentEvent::ToolCall { + name: call.name.clone(), + arguments: call.arguments.clone(), + }); + + let output = match self.tools.execute(&call.name, call.arguments.clone()).await + { + Ok(output) => output, + Err(e) => format!("Tool '{}' failed: {}", call.name, e), + }; + + let _ = tx.send(AgentEvent::ToolResult { + name: call.name.clone(), + output: output.clone(), + }); + + let tool_message = format!("[Tool result for '{}']\n{}", call.name, output); + self.history.push(Message { + role: Role::User, + content: tool_message.clone(), + }); + self.persist("tool", &tool_message, None); + continue; + } + } + + let _ = tx.send(AgentEvent::Done); + break; + } + + Ok(()) + } + + /// 便捷方法：消费自身执行单轮对话，返回一个事件流。 + /// + /// 适用于「一次性」调用；多轮对话请保留 Agent 并重复调用 [`Agent::run_turn`]。 + pub fn into_stream(mut self, user_input: String) -> impl Stream { + let (tx, rx) = mpsc::unbounded_channel(); + tokio::spawn(async move { + if let Err(e) = self.run_turn(&user_input, tx.clone()).await { + let _ = tx.send(AgentEvent::Error(e.to_string())); + } + }); + tokio_stream::wrappers::UnboundedReceiverStream::new(rx) + } + + /// 组合系统提示词与工具说明。 + fn build_system_prompt(&self) -> String { + let tools = self.tools.definitions_prompt(); + if tools.is_empty() { + self.system_prompt.clone() + } else { + format!("{}\n\n{}", self.system_prompt, tools) + } + } + + /// 持久化一条消息（若已开启持久化）。 + fn persist(&self, role: &str, content: &str, model: Option<&str>) { + if let Some(store) = &self.store { + if let Err(e) = store.save_message(&self.session_uuid, role, content, model) { + warn!("Failed to persist '{}' message: {}", role, e); + } + } + } +} + +/// 生成一个会话唯一标识（无外部 uuid 依赖）。 +fn generate_uuid() -> String { + use rand::RngExt; + + let ts = chrono::Utc::now() + .timestamp_nanos_opt() + .unwrap_or_default(); + let rand_part: u32 = rand::rng().random_range(0..u32::MAX); + + format!("{:x}-{:08x}", ts, rand_part) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_tool_call_valid() { + let content = "Let me read the file.\n```tool\n{\"name\":\"read_file\",\"arguments\":{\"path\":\"a.rs\"}}\n```"; + let call = parse_tool_call(content).unwrap(); + assert_eq!(call.name, "read_file"); + assert_eq!(call.arguments["path"], "a.rs"); + } + + #[test] + fn parse_tool_call_none() { + assert!(parse_tool_call("just a normal answer").is_none()); + } + + #[test] + fn uuid_is_unique_enough() { + assert_ne!(generate_uuid(), generate_uuid()); + } +} diff --git a/src/core/client/ai_client.rs b/src/core/client/ai_client.rs index daf2e6c..ec221e2 100644 --- a/src/core/client/ai_client.rs +++ b/src/core/client/ai_client.rs @@ -10,6 +10,7 @@ use crate::core::client::base_client::BaseClient; use crate::core::client::config::ClientConfig; use crate::core::transport::http_transoprt::HttpTransport; use crate::core::transport::retry::RetryPolicy; +use crate::core::transport::sse::SseStream; use crate::core::transport::TransportRequest; #[derive(Clone)] @@ -106,4 +107,28 @@ impl AiClient { pub fn transport(&self) -> &HttpTransport { &self.transport } + + /// 发送 POST 请求并返回原始 SSE 事件流，供需要自定义解析 + /// （如处理 `data: [DONE]` 结束标记）的供应商使用。 + pub async fn post_sse( + &self, + url: &str, + headers: HashMap, + body: &impl Serialize, + ) -> Result { + let body_bytes = serde_json::to_vec(body) + .map_err(|e| anyhow!("Failed to serialize request body: {}", e))?; + + let request = TransportRequest { + method: Method::POST, + url: url.to_string(), + headers, + body: Some(Bytes::from(body_bytes)), + }; + + self.transport + .execute_sse(request) + .await + .map_err(|e| anyhow!(e)) + } } \ No newline at end of file diff --git a/src/core/config/mod.rs b/src/core/config/mod.rs index e69de29..25fa858 100644 --- a/src/core/config/mod.rs +++ b/src/core/config/mod.rs @@ -0,0 +1,139 @@ +//! 应用配置 +//! +//! 负责从环境变量（`.env` 文件由 `dotenv` 加载）中读取运行 Agent 所需的配置， +//! 并转换为 [`ProviderConfig`] 交给供应商工厂使用。 +//! +//! 支持的环境变量： +//! | 变量 | 说明 | 默认值 | +//! |------|------|--------| +//! | `OMEGA_PROVIDER` | 供应商类型：compatible/claude/gemini/deepseek/ollama | compatible | +//! | `OMEGA_API_KEY` (回退 `API_KEY`) | API 密钥 | 空 | +//! | `OMEGA_BASE_URL` | API Base URL | https://api.openai.com/v1 | +//! | `OMEGA_MODEL` | 模型名称 | gpt-4o-mini | +//! | `OMEGA_TEMPERATURE` | 采样温度 | 0.7 | +//! | `OMEGA_MAX_TOKENS` | 单次最大输出 token | 不限制 | +//! | `OMEGA_SYSTEM_PROMPT` | 系统提示词 | 内置默认 | +//! | `OMEGA_WORKING_DIR` | 工作目录 | 当前目录 | + +use std::path::PathBuf; + +use anyhow::{Context, Result}; + +use crate::core::providers::{ProviderConfig, ProviderKind}; + +/// 默认系统提示词。 +pub const DEFAULT_SYSTEM_PROMPT: &str = "You are Omega Code, an AI coding assistant running in \ +the user's terminal. Be concise, accurate and helpful. When you need to inspect or modify the \ +project, use the available tools instead of guessing."; + +/// 应用配置。 +#[derive(Debug, Clone)] +pub struct AppConfig { + pub provider_kind: ProviderKind, + pub api_key: String, + pub base_url: String, + pub model: String, + pub temperature: Option, + pub max_tokens: Option, + pub system_prompt: String, + pub working_dir: PathBuf, +} + +impl Default for AppConfig { + fn default() -> Self { + Self { + provider_kind: ProviderKind::Compatible, + api_key: String::new(), + base_url: "https://api.openai.com/v1".to_string(), + model: "gpt-4o-mini".to_string(), + temperature: Some(0.7), + max_tokens: None, + system_prompt: DEFAULT_SYSTEM_PROMPT.to_string(), + working_dir: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")), + } + } +} + +impl AppConfig { + /// 从环境变量加载配置。会先尝试加载 `.env` 文件（若存在）。 + pub fn from_env() -> Result { + // 忽略 .env 不存在的错误。 + let _ = dotenv::dotenv(); + + let mut config = AppConfig::default(); + + if let Ok(kind) = std::env::var("OMEGA_PROVIDER") { + config.provider_kind = parse_provider_kind(&kind) + .with_context(|| format!("Invalid OMEGA_PROVIDER: {}", kind))?; + } + + if let Ok(key) = std::env::var("OMEGA_API_KEY").or_else(|_| std::env::var("API_KEY")) { + config.api_key = key; + } + + if let Ok(url) = std::env::var("OMEGA_BASE_URL") { + config.base_url = url; + } + + if let Ok(model) = std::env::var("OMEGA_MODEL") { + config.model = model; + } + + if let Ok(temp) = std::env::var("OMEGA_TEMPERATURE") { + config.temperature = temp.parse::().ok(); + } + + if let Ok(max) = std::env::var("OMEGA_MAX_TOKENS") { + config.max_tokens = max.parse::().ok(); + } + + if let Ok(prompt) = std::env::var("OMEGA_SYSTEM_PROMPT") { + if !prompt.trim().is_empty() { + config.system_prompt = prompt; + } + } + + if let Ok(dir) = std::env::var("OMEGA_WORKING_DIR") { + config.working_dir = PathBuf::from(dir); + } + + Ok(config) + } + + /// 转换为供应商配置。 + pub fn provider_config(&self) -> ProviderConfig { + ProviderConfig { + kind: self.provider_kind.clone(), + api_key: self.api_key.clone(), + base_url: self.base_url.clone(), + model: self.model.clone(), + } + } +} + +/// 将字符串解析为 [`ProviderKind`]。 +pub fn parse_provider_kind(value: &str) -> Result { + match value.trim().to_lowercase().as_str() { + "compatible" | "openai" => Ok(ProviderKind::Compatible), + "claude" | "anthropic" => Ok(ProviderKind::Claude), + "gemini" | "google" => Ok(ProviderKind::Gemini), + "deepseek" => Ok(ProviderKind::Deepseek), + "ollama" => Ok(ProviderKind::Ollama), + other => anyhow::bail!( + "unknown provider kind '{}', expected one of: compatible, claude, gemini, deepseek, ollama", + other + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_kinds() { + assert!(matches!(parse_provider_kind("Claude").unwrap(), ProviderKind::Claude)); + assert!(matches!(parse_provider_kind("deepseek").unwrap(), ProviderKind::Deepseek)); + assert!(parse_provider_kind("unknown").is_err()); + } +} diff --git a/src/core/mod.rs b/src/core/mod.rs index 2f611c3..cec949c 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -4,6 +4,6 @@ pub mod config; pub mod runtime; pub mod db; pub mod agent; -mod providers; -mod errors; -mod client; \ No newline at end of file +pub mod providers; +pub mod client; +mod errors; \ No newline at end of file diff --git a/src/core/providers/config.rs b/src/core/providers/config.rs index 8f47572..cad677b 100644 --- a/src/core/providers/config.rs +++ b/src/core/providers/config.rs @@ -9,6 +9,19 @@ pub enum ProviderKind { Ollama, } +impl ProviderKind { + /// 稳定的小写标识，用作供应商名称（与 `parse_provider_kind` 互逆）。 + pub fn as_str(&self) -> &'static str { + match self { + ProviderKind::Compatible => "compatible", + ProviderKind::Claude => "claude", + ProviderKind::Gemini => "gemini", + ProviderKind::Deepseek => "deepseek", + ProviderKind::Ollama => "ollama", + } + } +} + #[derive(Debug, Clone)] pub struct ProviderConfig { pub kind: ProviderKind, diff --git a/src/core/providers/mod.rs b/src/core/providers/mod.rs index 4b01473..7210394 100644 --- a/src/core/providers/mod.rs +++ b/src/core/providers/mod.rs @@ -1,6 +1,7 @@ -use crate::core::providers::config::{ProviderConfig, ProviderKind}; -use crate::core::providers::provider::Provider; -pub use types::{CompletionRequest, CompletionResponse, Message, Role, FinishReason, Usage}; +pub use types::{CompletionRequest, CompletionResponse, CompletionStream, Message, Role, FinishReason, StreamDelta, Usage}; + +pub use config::{ProviderConfig, ProviderKind}; +pub use provider::Provider; mod types; mod provider; diff --git a/src/core/providers/openai_compatible/provider.rs b/src/core/providers/openai_compatible/provider.rs index 2a04174..e90b25a 100644 --- a/src/core/providers/openai_compatible/provider.rs +++ b/src/core/providers/openai_compatible/provider.rs @@ -1,12 +1,15 @@ use anyhow::Result; use async_trait::async_trait; +use futures_util::StreamExt; use std::collections::HashMap; use crate::core::client::{AiClient, ClientConfig}; use crate::core::transport::retry::RetryPolicy; use crate::core::providers::openai_compatible::request::OpenAiRequest; -use crate::core::providers::openai_compatible::response::OpenAiResponse; +use crate::core::providers::openai_compatible::response::{OpenAiResponse, OpenAiStreamChunk}; use crate::core::providers::provider::Provider; -use crate::core::providers::types::{CompletionRequest, CompletionResponse}; +use crate::core::providers::types::{ + CompletionRequest, CompletionResponse, CompletionStream, FinishReason, StreamDelta, +}; pub struct OpenAiCompatibleProvider { client: AiClient, @@ -50,4 +53,68 @@ impl Provider for OpenAiCompatibleProvider { Ok(CompletionResponse::from(response)) } + + async fn complete_stream( + &self, + mut request: CompletionRequest, + ) -> Result { + request.stream = true; + let openai_request = OpenAiRequest::from_completion_request(request); + + let mut headers = HashMap::new(); + headers.insert( + "Authorization".to_string(), + format!("Bearer {}", self.api_key), + ); + headers.insert("Content-Type".to_string(), "application/json".to_string()); + + let url = format!("{}/chat/completions", self.base_url); + + let sse = self + .client + .post_sse(&url, headers, &openai_request) + .await + .map_err(|e| anyhow::anyhow!("OpenAI stream error: {}", e))?; + + let stream = sse.filter_map(|event_result| async move { + match event_result { + Ok(event) => { + let data = event.data.trim(); + if data.is_empty() || data == "[DONE]" { + return None; + } + match serde_json::from_str::(data) { + Ok(chunk) => { + let choice = chunk.choices.first(); + let content = choice + .and_then(|c| c.delta.content.clone()) + .unwrap_or_default(); + let finish_reason = choice + .and_then(|c| c.finish_reason.as_ref()) + .map(|r| match r.as_str() { + "stop" => FinishReason::Stop, + "length" => FinishReason::Length, + _ => FinishReason::Error, + }); + if content.is_empty() && finish_reason.is_none() { + None + } else { + Some(Ok(StreamDelta { + content, + finish_reason, + })) + } + } + Err(e) => Some(Err(anyhow::anyhow!( + "Failed to parse stream chunk: {}", + e + ))), + } + } + Err(e) => Some(Err(anyhow::anyhow!("SSE error: {}", e))), + } + }); + + Ok(Box::pin(stream)) + } } \ No newline at end of file diff --git a/src/core/providers/openai_compatible/response.rs b/src/core/providers/openai_compatible/response.rs index a56aa84..5814d02 100644 --- a/src/core/providers/openai_compatible/response.rs +++ b/src/core/providers/openai_compatible/response.rs @@ -7,6 +7,26 @@ pub struct OpenAiResponse { pub usage: Option, } +/// OpenAI 兼容流式响应块（SSE `data` 负载）。 +#[derive(Deserialize)] +pub struct OpenAiStreamChunk { + #[serde(default)] + pub choices: Vec, +} + +#[derive(Deserialize)] +pub struct OpenAiStreamChoice { + #[serde(default)] + pub delta: OpenAiStreamDelta, + pub finish_reason: Option, +} + +#[derive(Default, Deserialize)] +pub struct OpenAiStreamDelta { + #[serde(default)] + pub content: Option, +} + #[derive(Deserialize)] pub struct OpenAiChoice { pub message: OpenAiAssistantMessage, diff --git a/src/core/providers/provider.rs b/src/core/providers/provider.rs index 263a179..169b300 100644 --- a/src/core/providers/provider.rs +++ b/src/core/providers/provider.rs @@ -1,11 +1,28 @@ use anyhow::Result; use async_trait::async_trait; +use futures_util::stream; use crate::core::client::AiClient; -use crate::core::providers::types::{CompletionRequest, CompletionResponse}; +use crate::core::providers::types::{ + CompletionRequest, CompletionResponse, CompletionStream, StreamDelta, +}; #[async_trait] pub trait Provider: Send + Sync { fn client(&self) -> &AiClient; - + async fn complete(&self, request: CompletionRequest) -> Result; + + /// 流式补全。默认实现回退到 [`Provider::complete`]， + /// 将完整结果作为单个增量片段发出；支持原生流式的供应商可覆盖此方法。 + async fn complete_stream( + &self, + request: CompletionRequest, + ) -> Result { + let response = self.complete(request).await?; + let delta = StreamDelta { + content: response.content, + finish_reason: response.finish_reason, + }; + Ok(Box::pin(stream::once(async move { Ok(delta) }))) + } } \ No newline at end of file diff --git a/src/core/providers/types.rs b/src/core/providers/types.rs index dc1bc1a..abb0431 100644 --- a/src/core/providers/types.rs +++ b/src/core/providers/types.rs @@ -1,7 +1,24 @@ // core/types.rs +use std::pin::Pin; + +use anyhow::Result; +use futures_util::Stream; use serde::{Deserialize, Serialize}; +/// 流式补全的增量片段。 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StreamDelta { + /// 本次增量的文本内容（可能为空）。 + pub content: String, + /// 若本片段是最后一个，则携带结束原因。 + pub finish_reason: Option, +} + +/// 流式补全返回的事件流类型别名。 +pub type CompletionStream = + Pin> + Send + 'static>>; + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum Role { diff --git a/src/core/runtime/mod.rs b/src/core/runtime/mod.rs index e69de29..8d2970b 100644 --- a/src/core/runtime/mod.rs +++ b/src/core/runtime/mod.rs @@ -0,0 +1,121 @@ +//! 运行时管理 +//! +//! 该模块为 Agent 与工具系统提供统一的「执行上下文」： +//! - 维护当前工作目录（项目根目录） +//! - 提供相对路径解析，并对越界访问进行约束 +//! +//! 注意：JRE / Python 等语言运行时的下载与安装逻辑位于顶层 +//! [`crate::runtime`] 模块，本模块专注于 Agent 侧的工作目录上下文， +//! 二者职责不同，互不重叠。 + +use std::path::{Component, Path, PathBuf}; + +use anyhow::{Context, Result}; + +/// Agent 运行时上下文。 +/// +/// 以项目工作目录为中心，为工具（读写文件、列目录等）提供路径解析能力。 +#[derive(Debug, Clone)] +pub struct Runtime { + working_dir: PathBuf, +} + +impl Runtime { + /// 使用指定工作目录创建运行时。 + pub fn new(working_dir: impl Into) -> Self { + Self { + working_dir: working_dir.into(), + } + } + + /// 基于当前进程工作目录创建运行时。 + pub fn from_current_dir() -> Result { + let dir = std::env::current_dir().context("Failed to get current working directory")?; + Ok(Self::new(dir)) + } + + /// 当前工作目录。 + pub fn working_dir(&self) -> &Path { + &self.working_dir + } + + /// 将用户提供的路径解析为绝对路径。 + /// + /// - 绝对路径原样返回（经规范化处理） + /// - 相对路径基于工作目录拼接 + pub fn resolve(&self, path: impl AsRef) -> PathBuf { + let path = path.as_ref(); + let full = if path.is_absolute() { + path.to_path_buf() + } else { + self.working_dir.join(path) + }; + normalize(&full) + } + + /// 解析路径并校验其位于工作目录之内，防止越界访问。 + /// + /// 对于需要沙箱化工具（如写文件）时使用。 + pub fn resolve_within(&self, path: impl AsRef) -> Result { + let resolved = self.resolve(path); + + let root = normalize(&self.working_dir); + let is_within = resolved + .ancestors() + .any(|ancestor| ancestor == root.as_path()); + + if is_within { + Ok(resolved) + } else { + anyhow::bail!( + "Path '{}' escapes the working directory '{}'", + resolved.display(), + root.display() + ) + } + } +} + +/// 规范化路径：处理 `.` 与 `..` 组件，不访问文件系统。 +fn normalize(path: &Path) -> PathBuf { + let mut result = PathBuf::new(); + + for component in path.components() { + match component { + Component::CurDir => {} + Component::ParentDir => { + // 仅当上一级不是根前缀时才回退，避免越出文件系统根。 + if !matches!(result.components().next_back(), Some(Component::RootDir)) { + result.pop(); + } + } + other => result.push(other.as_os_str()), + } + } + + result +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn resolve_relative_path() { + let rt = Runtime::new("/home/project"); + assert_eq!(rt.resolve("src/main.rs"), PathBuf::from("/home/project/src/main.rs")); + } + + #[test] + fn resolve_absolute_path() { + let rt = Runtime::new("/home/project"); + assert_eq!(rt.resolve("/etc/hosts"), PathBuf::from("/etc/hosts")); + } + + #[test] + fn within_working_dir() { + let rt = Runtime::new("/home/project"); + assert!(rt.resolve_within("src/main.rs").is_ok()); + assert!(rt.resolve_within("../secret").is_err()); + } +} diff --git a/src/core/tools/mod.rs b/src/core/tools/mod.rs index e69de29..316d1f6 100644 --- a/src/core/tools/mod.rs +++ b/src/core/tools/mod.rs @@ -0,0 +1,310 @@ +//! 工具系统 +//! +//! 定义 Agent 可调用的工具抽象 [`Tool`]、工具注册表 [`ToolRegistry`]， +//! 以及若干内置工具（读取文件、写入文件、列出目录）。 +//! +//! 由于当前各供应商的请求体尚未内置原生 function-calling 字段，Agent 采用 +//! 基于提示词的 ReAct 方式调用工具：工具定义会被注入系统提示词，模型以约定的 +//! JSON 代码块表达调用意图，Agent 解析后执行并把结果回灌到对话中。 +//! [`ToolRegistry::definitions_prompt`] 负责生成这段工具说明。 + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use async_trait::async_trait; +use serde_json::{json, Value}; + +use crate::core::runtime::Runtime; + +/// 工具抽象。所有工具都需线程安全，以便在异步任务中共享。 +#[async_trait] +pub trait Tool: Send + Sync { + /// 工具唯一名称（供模型调用时引用）。 + fn name(&self) -> &str; + + /// 工具功能描述。 + fn description(&self) -> &str; + + /// 参数的 JSON Schema 描述。 + fn parameters(&self) -> Value; + + /// 执行工具，返回文本结果（会回灌给模型）。 + async fn execute(&self, args: Value) -> Result; +} + +/// 工具注册表，管理一组可用工具。 +#[derive(Default, Clone)] +pub struct ToolRegistry { + tools: Vec>, +} + +impl ToolRegistry { + pub fn new() -> Self { + Self { tools: Vec::new() } + } + + /// 注册一个工具。 + pub fn register(&mut self, tool: Arc) { + self.tools.push(tool); + } + + /// 按名称查找工具。 + pub fn get(&self, name: &str) -> Option> { + self.tools.iter().find(|t| t.name() == name).cloned() + } + + /// 已注册工具数量。 + pub fn len(&self) -> usize { + self.tools.len() + } + + pub fn is_empty(&self) -> bool { + self.tools.is_empty() + } + + /// 执行指定工具。 + pub async fn execute(&self, name: &str, args: Value) -> Result { + match self.get(name) { + Some(tool) => tool.execute(args).await, + None => anyhow::bail!("Unknown tool: {}", name), + } + } + + /// 生成注入系统提示词的工具说明文本。 + pub fn definitions_prompt(&self) -> String { + if self.tools.is_empty() { + return String::new(); + } + + let mut prompt = String::from( + "# Available Tools\n\n\ + You can use the following tools. To call a tool, respond with ONLY a fenced \ + code block tagged `tool` containing a JSON object with `name` and `arguments`, e.g.:\n\n\ + ```tool\n{\"name\": \"read_file\", \"arguments\": {\"path\": \"src/main.rs\"}}\n```\n\n\ + Wait for the tool result before continuing. Do not fabricate tool outputs.\n\n", + ); + + for tool in &self.tools { + prompt.push_str(&format!( + "- `{}`: {}\n parameters: {}\n", + tool.name(), + tool.description(), + tool.parameters() + )); + } + + prompt + } + + /// 创建带有默认内置工具的注册表。 + pub fn with_defaults(runtime: Arc) -> Self { + let mut registry = Self::new(); + registry.register(Arc::new(ReadFileTool::new(runtime.clone()))); + registry.register(Arc::new(WriteFileTool::new(runtime.clone()))); + registry.register(Arc::new(ListDirTool::new(runtime))); + registry + } +} + +// ============================================================ +// 内置工具：read_file +// ============================================================ + +/// 读取文本文件内容。 +pub struct ReadFileTool { + runtime: Arc, +} + +impl ReadFileTool { + pub fn new(runtime: Arc) -> Self { + Self { runtime } + } +} + +#[async_trait] +impl Tool for ReadFileTool { + fn name(&self) -> &str { + "read_file" + } + + fn description(&self) -> &str { + "Read the text content of a file within the project. Args: {\"path\": string}" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "path": { "type": "string", "description": "File path relative to the project root" } + }, + "required": ["path"] + }) + } + + async fn execute(&self, args: Value) -> Result { + let path = args + .get("path") + .and_then(|v| v.as_str()) + .context("read_file requires a 'path' string argument")?; + + let resolved = self.runtime.resolve_within(path)?; + let content = tokio::fs::read_to_string(&resolved) + .await + .with_context(|| format!("Failed to read file '{}'", resolved.display()))?; + + Ok(content) + } +} + +// ============================================================ +// 内置工具：write_file +// ============================================================ + +/// 写入（覆盖）文本文件，自动创建父目录。 +pub struct WriteFileTool { + runtime: Arc, +} + +impl WriteFileTool { + pub fn new(runtime: Arc) -> Self { + Self { runtime } + } +} + +#[async_trait] +impl Tool for WriteFileTool { + fn name(&self) -> &str { + "write_file" + } + + fn description(&self) -> &str { + "Write text content to a file within the project, creating parent directories as needed. \ + Args: {\"path\": string, \"content\": string}" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "path": { "type": "string", "description": "File path relative to the project root" }, + "content": { "type": "string", "description": "Full text content to write" } + }, + "required": ["path", "content"] + }) + } + + async fn execute(&self, args: Value) -> Result { + let path = args + .get("path") + .and_then(|v| v.as_str()) + .context("write_file requires a 'path' string argument")?; + let content = args + .get("content") + .and_then(|v| v.as_str()) + .context("write_file requires a 'content' string argument")?; + + let resolved = self.runtime.resolve_within(path)?; + if let Some(parent) = resolved.parent() { + tokio::fs::create_dir_all(parent).await.ok(); + } + tokio::fs::write(&resolved, content) + .await + .with_context(|| format!("Failed to write file '{}'", resolved.display()))?; + + Ok(format!("Wrote {} bytes to {}", content.len(), resolved.display())) + } +} + +// ============================================================ +// 内置工具：list_dir +// ============================================================ + +/// 列出目录内容。 +pub struct ListDirTool { + runtime: Arc, +} + +impl ListDirTool { + pub fn new(runtime: Arc) -> Self { + Self { runtime } + } +} + +#[async_trait] +impl Tool for ListDirTool { + fn name(&self) -> &str { + "list_dir" + } + + fn description(&self) -> &str { + "List files and subdirectories of a directory within the project. Args: {\"path\": string}" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "path": { "type": "string", "description": "Directory path relative to the project root" } + }, + "required": ["path"] + }) + } + + async fn execute(&self, args: Value) -> Result { + let path = args + .get("path") + .and_then(|v| v.as_str()) + .unwrap_or("."); + + let resolved = self.runtime.resolve_within(path)?; + let mut entries = tokio::fs::read_dir(&resolved) + .await + .with_context(|| format!("Failed to read directory '{}'", resolved.display()))?; + + let mut lines = Vec::new(); + while let Some(entry) = entries.next_entry().await? { + let file_type = entry.file_type().await?; + let suffix = if file_type.is_dir() { "/" } else { "" }; + lines.push(format!("{}{}", entry.file_name().to_string_lossy(), suffix)); + } + lines.sort(); + + if lines.is_empty() { + Ok(format!("{} is empty", resolved.display())) + } else { + Ok(lines.join("\n")) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn registry_execute_unknown_tool() { + let registry = ToolRegistry::new(); + assert!(registry.execute("nope", json!({})).await.is_err()); + } + + #[tokio::test] + async fn write_then_read_file() { + let dir = std::env::temp_dir().join("omega_tools_test"); + let runtime = Arc::new(Runtime::new(&dir)); + let registry = ToolRegistry::with_defaults(runtime); + + registry + .execute("write_file", json!({ "path": "a.txt", "content": "hello" })) + .await + .unwrap(); + + let content = registry + .execute("read_file", json!({ "path": "a.txt" })) + .await + .unwrap(); + + assert_eq!(content, "hello"); + + tokio::fs::remove_dir_all(&dir).await.ok(); + } +} diff --git a/src/core/transport/http_transoprt.rs b/src/core/transport/http_transoprt.rs index ebd1fa5..67af1f8 100644 --- a/src/core/transport/http_transoprt.rs +++ b/src/core/transport/http_transoprt.rs @@ -127,4 +127,23 @@ impl HttpTransport { Ok(json_stream) } + + /// 执行请求并返回原始 SSE 事件流，由调用方自行解析 `data` 字段。 + /// 相比 [`HttpTransport::execute_stream`]，此方法不会将事件反序列化为 JSON， + /// 因此可以处理诸如 OpenAI 的 `data: [DONE]` 这类非 JSON 结束标记。 + pub async fn execute_sse( + &self, + req: TransportRequest, + ) -> Result { + let request = self.build_request(req)?; + + let response = self.client.execute(request).await?; + + let status = response.status(); + if status.is_client_error() || status.is_server_error() { + return Err(TransportError::Server(status.to_string())); + } + + Ok(SseStream::new(response)) + } } \ No newline at end of file diff --git a/src/pages/chat.rs b/src/pages/chat.rs index a294e85..095bcb0 100644 --- a/src/pages/chat.rs +++ b/src/pages/chat.rs @@ -1,23 +1,297 @@ -use ratatui_kit::{component, element, AnyElement, ElementExt, Hooks}; -use ratatui_kit::components::View; -use ratatui_kit::ratatui::{TerminalOptions, Viewport}; -use ratatui_kit::ratatui::layout::{Direction, Constraint}; +//! 聊天页面：把顶层 Agent 接入 TUI。 +//! +//! 交互模型（Claude Code 风格）： +//! - 已完成的消息通过 `insert_before` 刷入终端原生滚动区（可向上滚动查看历史）； +//! - 当前正在流式生成的回复渲染在页面底部的「实时区」，随 token 到达不断重绘； +//! - 输入框回车提交后，使用 `tokio::spawn` 驱动一轮 Agent 对话，事件通过 mpsc 通道回传。 + +use std::sync::Arc; + +use ratatui_kit::components::{Border, Text, TextArea, TextParagraph, View}; +use ratatui_kit::crossterm::event::{Event, KeyCode, KeyEventKind}; +use ratatui_kit::ratatui::layout::{Constraint, Direction, Position}; use ratatui_kit::ratatui::prelude::Color; -use ratatui_kit::ratatui::style::Style; -use crate::components::input_bar::InputBar; +use ratatui_kit::ratatui::style::{Style, Stylize}; +use ratatui_kit::ratatui::text::Line; +use ratatui_kit::ratatui::widgets::Paragraph; +use ratatui_kit::{ + component, element, AnyElement, Hooks, InsertBeforeHandler, State, SystemContext, + UseContext, UseEvents, UseInsertBefore, UseState, UseTerminalSize, +}; +use tokio::sync::{mpsc, Mutex}; + use crate::components::status_bar::StatusBar; +use crate::core::agent::{Agent, AgentEvent}; +use crate::core::config::AppConfig; + +/// Agent 持有单元：初始化可能失败，失败时保留错误信息用于展示。 +struct AgentCell { + agent: Option, + error: Option, + /// 当前工作目录，用于状态栏展示。 + working_dir: String, +} + +/// 依据环境配置构建 Agent（并尽力开启会话持久化）。 +fn init_agent() -> AgentCell { + let fallback_dir = || { + std::env::current_dir() + .map(|p| p.display().to_string()) + .unwrap_or_default() + }; + + match AppConfig::from_env() { + Ok(cfg) => { + let working_dir = cfg.working_dir.display().to_string(); + match Agent::from_config(&cfg) { + Ok(mut agent) => { + // 持久化为尽力而为：失败仅告警，不阻断对话。 + if let Err(e) = agent.enable_persistence_auto(Some("TUI Chat")) { + log::warn!("Failed to enable chat persistence: {e}"); + } + AgentCell { + agent: Some(agent), + error: None, + working_dir, + } + } + Err(e) => AgentCell { + agent: None, + error: Some(e.to_string()), + working_dir, + }, + } + } + Err(e) => AgentCell { + agent: None, + error: Some(e.to_string()), + working_dir: fallback_dir(), + }, + } +} + +/// 估算文本在给定终端宽度下换行后占用的行数。 +fn wrapped_rows(text: &str, width: u16) -> usize { + let w = if width == 0 { 80 } else { width as usize }; + text.split('\n') + .map(|line| { + let n = line.chars().count(); + if n == 0 { + 1 + } else { + (n + w - 1) / w + } + }) + .sum() +} + +/// 在后台任务中执行一轮 Agent 对话，并把事件映射为 UI 更新。 +#[allow(clippy::too_many_arguments)] +fn spawn_turn( + text: String, + agent_cell: Arc>, + mut streaming: State, + mut status: State, + mut busy: State, + insert_before: InsertBeforeHandler, + term_w: u16, +) { + tokio::spawn(async move { + let (tx, mut rx) = mpsc::unbounded_channel::(); + + // 单独的任务持有 Agent 锁并执行对话，主任务并发消费事件以实现实时刷新。 + let cell = agent_cell.clone(); + let input = text.clone(); + let turn = tokio::spawn(async move { + let mut guard = cell.lock().await; + match guard.agent.as_mut() { + Some(agent) => { + let _ = agent.run_turn(&input, tx).await; + } + None => { + let err = guard + .error + .clone() + .unwrap_or_else(|| "Agent not initialized".to_string()); + let _ = tx.send(AgentEvent::Error(err)); + } + } + // tx 在此处被释放，rx 随之结束。 + }); + + let mut acc = String::new(); + while let Some(event) = rx.recv().await { + match event { + AgentEvent::Token(t) => { + acc.push_str(&t); + streaming.set(acc.clone()); + } + AgentEvent::ToolCall { name, .. } => { + status.set(format!("calling tool: {name} ...")); + } + AgentEvent::ToolResult { name, .. } => { + status.set(format!("tool {name} finished")); + } + AgentEvent::Error(e) => { + status.set(format!("error: {e}")); + } + AgentEvent::Done => break, + } + } + let _ = turn.await; + + // 将完成的回复刷入终端滚动区，并清空实时区。 + let trimmed = acc.trim_end(); + if !trimmed.is_empty() { + let height = wrapped_rows(trimmed, term_w).min(1000) as u16; + insert_before + .render_before(Paragraph::new(trimmed.to_string()), height) + .finish(); + } + + streaming.set(String::new()); + status.set(String::new()); + busy.set(false); + }); +} #[component] pub fn ChatPage(mut hooks: Hooks) -> impl Into> { + // ---- 状态（hook 顺序在所有渲染中保持一致）---- + let mut input = hooks.use_state(String::new); + let streaming = hooks.use_state(String::new); + let mut status = hooks.use_state(String::new); + let mut busy = hooks.use_state(|| false); + let mut should_exit = hooks.use_state(|| false); + let agent_cell = hooks.use_state(|| Arc::new(Mutex::new(init_agent()))); + let (term_w, term_h) = hooks.use_terminal_size(); + let insert_before = hooks.use_insert_before(); + + // 退出处理（非 hook）。 + { + let mut system_ctx = hooks.use_context_mut::(); + if should_exit.get() { + system_ctx.exit(); + } + } + + // 读取初始化错误与工作目录（非阻塞）。 + let (init_error, working_dir) = { + let arc = agent_cell.read().clone(); + match arc.try_lock() { + Ok(cell) => (cell.error.clone(), cell.working_dir.clone()), + Err(_) => (None, String::new()), + } + }; + + // 实时区展示内容：优先流式文本，其次状态提示，最后初始化错误。 + let stream_txt = streaming.read().clone(); + let status_txt = status.read().clone(); + let live_text = if !stream_txt.is_empty() { + stream_txt + } else if busy.get() { + if status_txt.is_empty() { + "thinking...".to_string() + } else { + status_txt + } + } else if let Some(err) = &init_error { + format!("Agent unavailable: {err}") + } else { + String::new() + }; + + // 计算实时区滚动偏移，使最新内容贴底显示。 + let live_height = term_h.saturating_sub(11).max(1) as usize; + let live_rows = wrapped_rows(&live_text, term_w); + let scroll_y = live_rows.saturating_sub(live_height) as u16; + + // ---- 键盘事件 ---- + // State 为 Copy，move 闭包按副本捕获，可在 FnMut 中直接修改（同一底层存储）。 + hooks.use_events(move |event| { + if let Event::Key(key) = event { + if key.kind != KeyEventKind::Press { + return; + } + match key.code { + KeyCode::Esc => { + should_exit.set(true); + } + KeyCode::Enter => { + let text = input.read().trim().to_string(); + if text.is_empty() || busy.get() { + return; + } + + // 回显用户消息到滚动区。 + let echoed = format!("> {text}"); + let height = wrapped_rows(&echoed, term_w).min(1000) as u16; + insert_before + .render_before( + Paragraph::new(Line::styled(echoed, Style::default().fg(Color::Cyan))), + height, + ) + .finish(); + + input.set(String::new()); + busy.set(true); + status.set("thinking...".to_string()); + + spawn_turn( + text, + agent_cell.read().clone(), + streaming, + status, + busy, + insert_before.clone(), + term_w, + ); + } + _ => {} + } + } + }); + + // ---- 渲染 ---- + let mut input_cb = input; + element!( View(flex_direction: Direction::Vertical,) { - View(height: Constraint::Length(10),) { - InputBar() + View(height: Constraint::Min(3),) { + Text( + text: TextParagraph::from(live_text), + style: Style::default(), + wrap: Some(false), + scroll: Position::new(0, scroll_y), + ) + } + View(height: Constraint::Length(6),) { + Border( + height: Constraint::Length(6), + style: Style::default().green(), + top_title: Line::styled("Composer", Style::default().yellow()).centered(), + ) { + TextArea( + value: input.read().to_string(), + is_focus: true, + on_change: move |new_value: String| { + input_cb.set(new_value); + }, + multiline: false, + cursor_style: Style::default().on_green(), + placeholder: Some("Type a message... (Enter to send, Esc to quit)".to_string()), + placeholder_style: Style::default().green(), + ) + } } View(height: Constraint::Length(4),) { StatusBar( - project_path: Some("/home/project/OmegaProject/".into()), - context_length: 123987usize, + project_path: if working_dir.is_empty() { + None + } else { + Some(working_dir.clone().into()) + }, + context_length: live_rows, max_context_length: 1000000usize, style: Style::default(), progress_style: Style::default().fg(Color::Green), @@ -26,4 +300,4 @@ pub fn ChatPage(mut hooks: Hooks) -> impl Into> { } } ) -} \ No newline at end of file +} From 35cbed6fa6a64acef9892ea6a4c84c88b864979a Mon Sep 17 00:00:00 2001 From: fusheng Date: Fri, 16 Oct 2026 22:51:39 +0800 Subject: [PATCH 2/3] =?UTF-8?q?refactor(agent):=20=E6=8C=89=E8=AE=BE?= =?UTF-8?q?=E8=AE=A1=E6=96=87=E6=A1=A3=E9=87=8D=E6=9E=84=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E4=B8=8E=E5=AF=B9=E8=AF=9D=E5=88=9D=E5=A7=8B?= =?UTF-8?q?=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 重构所有数据库模型字段（created_at/updated_at 改为 create_time/update_time） - 统一使用 parse_db_time 函数解析时间字段，移除内联解析逻辑 - 重构 agent_arena 表，新增 chat_uuid、round_uuid、vote_type、comment、status 字段 - 重构 agent_info 表，新增 llm_id、description、temperature、top_p、call_count、avg_score 字段 - 重构 agent_mcp_rel 表，将 mcp_id 更名为 mcp_tool_id - 重构 ai_provider 表，timeout 默认值从 5000 改为 30 秒 - 重构 chat_message 表，新增 round_uuid、input_tokens、output_tokens、duration、model_config 字段 - 重构 chat_session 表，将 status 改为 is_deleted 逻辑删除字段 - 对话页 Agent 初始化从环境变量配置改为数据库驱动加载 - 新增 Ctrl+1..5 快捷键为上一轮回复打星级评分 - 新增记忆系统（MemoryStore）按配置作用域加载与注入 - 新增协调者模式支持（coordinator_agent_id 配置） - 支持通过 OMEGA_CHAT_UUID 恢复历史会话 - clap 启用 derive feature，rmcp 启用 client 与 transport-streamable-http 特性 --- Cargo.lock | 30 + Cargo.toml | 4 +- locales/zh.yml | 16 +- src/cli/commands.rs | 1181 +++++++++++++++++ src/cli/mod.rs | 728 ++++++++++ src/components/welcome.rs | 39 +- src/core/agent/mod.rs | 835 +++++++++++- src/core/arena/mod.rs | 556 ++++++++ src/core/arena/report.rs | 216 +++ src/core/db/db_init.rs | 132 +- src/core/mcp/mod.rs | 289 ++++ src/core/memory/mod.rs | 814 ++++++++++++ src/core/mod.rs | 6 + src/core/project/mod.rs | 193 +++ src/core/prompt/mod.rs | 294 ++++ src/core/providers/ollama/provider.rs | 53 + .../providers/openai_compatible/provider.rs | 55 + src/core/providers/provider.rs | 7 + src/core/scheduler/mod.rs | 958 +++++++++++++ src/db/models/agent_arena.rs | 84 +- src/db/models/agent_info.rs | 112 +- src/db/models/agent_mcp_rel.rs | 55 +- src/db/models/ai_provider.rs | 54 +- src/db/models/chat_message.rs | 104 +- src/db/models/chat_session.rs | 58 +- src/db/models/llm_model.rs | 76 +- src/db/models/mcp_provider.rs | 51 +- src/db/models/mcp_tool.rs | 85 +- src/db/models/memory_info.rs | 85 +- src/db/models/mod.rs | 27 +- src/db/models/project_info.rs | 89 +- src/db/models/prompt_info.rs | 91 +- src/db/models/prompt_placeholder.rs | 94 +- src/db/models/system_api_log.rs | 99 +- src/db/models/system_config.rs | 111 +- src/lib.rs | 43 +- src/main.rs | 7 +- src/pages/chat.rs | 222 +++- src/pages/router.rs | 7 +- src/router/mod.rs | 60 + src/runtime/runtime.rs | 10 + src/services/welcome.rs | 201 +++ 42 files changed, 7552 insertions(+), 679 deletions(-) create mode 100644 src/cli/commands.rs create mode 100644 src/cli/mod.rs create mode 100644 src/core/arena/mod.rs create mode 100644 src/core/arena/report.rs create mode 100644 src/core/mcp/mod.rs create mode 100644 src/core/memory/mod.rs create mode 100644 src/core/project/mod.rs create mode 100644 src/core/prompt/mod.rs create mode 100644 src/core/scheduler/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 1c4721d..a099f08 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -388,6 +388,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" dependencies = [ "clap_builder", + "clap_derive", ] [[package]] @@ -402,6 +403,18 @@ dependencies = [ "strsim", ] +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "clap_lex" version = "1.1.0" @@ -2626,14 +2639,18 @@ dependencies = [ "base64", "chrono", "futures", + "http", "pastey", "pin-project-lite", + "reqwest", "rmcp-macros", "schemars", "serde", "serde_json", + "sse-stream", "thiserror 2.0.18", "tokio", + "tokio-stream", "tokio-util", "tracing", ] @@ -3151,6 +3168,19 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "sse-stream" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4" +dependencies = [ + "bytes", + "futures-util", + "http-body", + "http-body-util", + "pin-project-lite", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" diff --git a/Cargo.toml b/Cargo.toml index 5bef96a..f612fbf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,7 @@ name = "OmegaCode" path = "src/main.rs" [dependencies] -clap = "4.6.1" +clap = { version = "4.6.1", features = ["derive"] } anyhow = "1.0.102" tokio = { version = "1.52.3", features = ["full"] } colored = "3.1.1" @@ -23,7 +23,7 @@ futures-util = "0.3.32" serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.150" serde_urlencoded = "0.7.1" -rmcp = "1.7.0" +rmcp = { version = "1.7.0", features = ["client", "transport-streamable-http-client-reqwest"] } rusqlite = { version = "0.40.0", features = ["bundled"] } ratatui = "0.30.0" crossterm = "0.29.0" diff --git a/locales/zh.yml b/locales/zh.yml index e12d17a..1dbb0d4 100644 --- a/locales/zh.yml +++ b/locales/zh.yml @@ -9,7 +9,15 @@ database_table_create: "数据表 %{table_name} 已创建" welcome.version: "版本: %{version}" welcome.text: "欢迎来到 Omega Code !" welcome.menu: "快速菜单(使用 ↑/↓ 键导航，按 enter 选择项):" -menu_items: - - "打开 Omega Code 控制台面板，配置应用程序" - - "开始新项目" - - "退出" +welcome.menu_items.0: "打开 Omega Code 控制台面板，配置应用程序" +welcome.menu_items.1: "开始新项目" +welcome.menu_items.2: "退出" +welcome.install_message1: "步骤 1：Java 运行时环境……" +welcome.install_message2: "步骤 2：Python 运行时环境……" +welcome.install_message3: "Java 运行时环境安装失败。" +welcome.install_message4: "Python 运行时环境安装失败。" +welcome.install_message5: "正在下载 JRE……" +welcome.install_message6: "正在解压 JRE……" +welcome.install_message7: "正在下载 Python……" +welcome.install_message8: "正在解压 Python……" +welcome.install_message9: "环境已就绪" diff --git a/src/cli/commands.rs b/src/cli/commands.rs new file mode 100644 index 0000000..487a4ff --- /dev/null +++ b/src/cli/commands.rs @@ -0,0 +1,1181 @@ +//! 命令处理 +//! +//! 把 [`crate::cli`] 解析出的子命令落地为具体的管理动作：读写本地 SQLite +//! （经 [`DatabaseManager`]）并把结果打印到标准输出。交互式 TUI 的启动不在本模块， +//! 由 [`crate::run_with`] 决定。 +//! +//! 设计约束（《命令行参数设计》六.开发备注）： +//! - 参数优先级：命令行 > 配置文件 > 系统默认；策略类命令（模型温度、记忆、安全、 +//! 性能）统一以键值形式落地到 `system_config`，作为可被后续启动读取的「配置覆盖」。 +//! - 敏感字段（API Key）仅在需要时写入，不回显明文。 + +use anyhow::{Context, Result}; +use chrono::Utc; +use rusqlite::params; +use std::collections::BTreeMap; +use std::collections::HashMap; +use std::collections::HashSet; + +use crate::core::agent::SessionStore; +use crate::core::arena::report::{self, PromptReportAgent}; +use crate::core::config::parse_provider_kind; +use crate::core::db::DatabaseManager; +use crate::core::arena::{ArenaService, VoteInput}; +use crate::core::prompt::{build_vars, render, validate_prompt}; +use crate::core::providers::{create_provider, ProviderConfig, ProviderKind}; +use crate::db::models::agent_info::AgentInfo; +use crate::db::models::ai_provider::AiProvider; +use crate::db::models::chat_session::ChatSession; +use crate::db::models::llm_model::LLMModel; +use crate::db::models::prompt_info::PromptInfo; +use crate::db::models::system_config::SystemConfig; + +use super::{ + AgentCommand, ArenaCommand, ChatArgs, ChatCommand, Command, MemoryArgs, ModelCommand, + PerformanceArgs, PromptCommand, SecurityArgs, SupplierArgs, +}; + +/// 分发并执行一个子命令。 +pub async fn dispatch(command: Command) -> Result<()> { + match command { + Command::Model { action } => model(action).await, + Command::Supplier(args) => supplier(args), + Command::Agent { action } => agent(action), + Command::Chat(args) => chat(args), + Command::Memory(args) => memory(args), + Command::Security(args) => security(args), + Command::Performance(args) => performance(args), + Command::Arena { action } => arena(action).await, + Command::Prompt { action } => prompt(action), + } +} + +// ============================================================ +// model +// ============================================================ + +async fn model(action: ModelCommand) -> Result<()> { + let db = DatabaseManager::new()?; + match action { + ModelCommand::List { all, .. } => { + let models = db.read_all::()?; + println!("{:<5}{:<28}{:<12}{:<10}", "id", "model_code", "context", "status"); + let mut shown = 0usize; + for m in models { + if !all && m.status != 1 { + continue; + } + println!( + "{:<5}{:<28}{:<12}{:<10}", + m.id, + m.model_code, + m.max_tokens, + if m.status == 1 { "enabled" } else { "disabled" } + ); + shown += 1; + } + if shown == 0 { + println!("(no models matched)"); + } + } + ModelCommand::Status { name, .. } => { + let mut found = db.find_where::("model_code = ?1", &[&name])?; + let m = found + .first() + .with_context(|| format!("model '{}' not found", name))?; + println!("model : {}", m.model_code); + println!("id : {}", m.id); + println!("name : {}", m.name); + println!("provider : #{}", m.provider_id); + println!("context : {}", m.max_tokens); + println!("max_output : {}", m.max_response_tokens); + println!("weight : {}", m.weight); + println!("calls : {}", m.call_count); + println!("status : {}", if m.status == 1 { "enabled" } else { "disabled" }); + // 温度是 Agent 级参数，若曾以命令行覆盖则一并展示。 + if let Some(temp) = read_config(&db, &format!("model.{}.temperature", m.model_code))? { + println!("temperature: {} (override)", temp); + } + let _ = &mut found; + } + ModelCommand::Set { + name, + temperature, + max_tokens, + } => { + let mut found = db.find_where::("model_code = ?1", &[&name])?; + let m = found + .first_mut() + .with_context(|| format!("model '{}' not found", name))?; + + if let Some(mt) = max_tokens { + if mt <= 0 { + anyhow::bail!("--max-tokens must be a positive integer"); + } + m.max_tokens = mt; + db.update(m.clone())?; + println!("set model '{}' max context tokens = {}", name, mt); + } + if let Some(t) = temperature { + if !(0.0..=1.0).contains(&t) { + anyhow::bail!("-t must be within 0.0~1.0"); + } + // 温度属于 Agent 级个性化参数（llm_model 无对应列）， + // 以 system_config 覆盖项持久化，供后续加载时优先采用。 + upsert_config( + &db, + &format!("model.{}.temperature", name), + &t.to_string(), + "float", + &format!("CLI override temperature for model {}", name), + "model", + )?; + println!("set model '{}' temperature override = {}", name, t); + } + if max_tokens.is_none() && temperature.is_none() { + println!("nothing to set (provide -t and/or --max-tokens)"); + } + } + ModelCommand::Sync { supplier } => { + // 解析供应商（先 ID 再名称），与 `supplier` 命令保持一致。 + let mut resolved = match supplier.parse::() { + Ok(id) => db.read_by_id::(id)?, + Err(_) => None, + }; + if resolved.is_none() { + resolved = db + .find_where::("name = ?1", &[&supplier])? + .into_iter() + .next(); + } + let provider = resolved + .with_context(|| format!("supplier '{}' not found", supplier))?; + + // 由供应商行构建配置（kind 取名称，base_url 取 api，key 可空）。 + let config = ProviderConfig { + kind: parse_provider_kind(&provider.name).unwrap_or(ProviderKind::Compatible), + api_key: provider.key.clone().unwrap_or_default(), + base_url: provider.api.clone(), + model: String::new(), + }; + let codes = create_provider(config)?.list_models().await?; + let (added, skipped) = sync_models(&db, provider.id, &codes)?; + println!( + "synced {} model(s) from supplier '{}' (#{}): +{} added, {} existing", + codes.len(), + provider.name, + provider.id, + added, + skipped + ); + } + ModelCommand::Delete { name, yes } => { + // 先做存在性与引用完整性校验（逻辑外键，业务层保障）。 + match check_model_deletable(&db, &name)? { + ModelDeletable::NotFound => { + anyhow::bail!("model '{}' not found", name) + } + ModelDeletable::Blocked { refs } => anyhow::bail!( + "model '{}' is referenced by {} agent(s); reassign them before deleting", + name, + refs + ), + ModelDeletable::Ready(m) => { + if !yes { + println!( + "this will permanently delete model '{}' (#{}); re-run with -y to confirm", + m.model_code, m.id + ); + } else { + db.delete::(m.id)?; + println!("deleted model '{}' (#{})", m.model_code, m.id); + } + } + } + } + } + Ok(()) +} + +// ============================================================ +// supplier +// ============================================================ + +fn supplier(args: SupplierArgs) -> Result<()> { + let db = DatabaseManager::new()?; + + // 先按 ID 匹配，失败再按名称匹配。 + let mut provider = match args.supplier.parse::() { + Ok(id) => db.read_by_id::(id)?, + Err(_) => None, + }; + if provider.is_none() { + provider = db + .find_where::("name = ?1", &[&args.supplier])? + .into_iter() + .next(); + } + let mut provider = + provider.with_context(|| format!("supplier '{}' not found", args.supplier))?; + + if let Some(key) = &args.api_key { + provider.key = Some(key.clone()); + db.update(provider.clone())?; + println!("updated api key for supplier '{}' (#{})", provider.name, provider.id); + } + + println!( + "supplier #{} '{}' -> {} [{}]", + provider.id, + provider.name, + provider.api, + if provider.status == 1 { "enabled" } else { "disabled" } + ); + Ok(()) +} + +// ============================================================ +// agent +// ============================================================ + +fn agent(action: AgentCommand) -> Result<()> { + let db = DatabaseManager::new()?; + match action { + AgentCommand::Use { agent, prompt } => { + let mut a = db + .read_by_id::(agent)? + .with_context(|| format!("agent #{} not found", agent))?; + if let Some(pid) = prompt { + a.prompt_id = Some(pid); + db.update(a.clone())?; + println!("agent #{} prompt override -> #{}", agent, pid); + } + upsert_config( + &db, + "default_agent_id", + &agent.to_string(), + "int", + "CLI selected default agent", + "agent", + )?; + println!("default agent set to #{} ({})", a.id, a.name); + } + AgentCommand::Mcp { agent, mcp } => { + // 校验 Agent 存在。 + db.read_by_id::(agent)? + .with_context(|| format!("agent #{} not found", agent))?; + + // 全量替换该 Agent 的 MCP 授权集合。 + db.conn + .execute("DELETE FROM agent_mcp_rel WHERE agent_id = ?1", params![agent])?; + for tool_id in &mcp { + db.conn.execute( + "INSERT INTO agent_mcp_rel (agent_id, mcp_tool_id) VALUES (?1, ?2)", + params![agent, tool_id], + )?; + } + println!( + "agent #{} now mounts {} MCP tool(s): {:?}", + agent, + mcp.len(), + mcp + ); + } + AgentCommand::Coordinator { agent } => { + // 校验 Agent 存在后，记录为协调者（TUI 加载该 Agent 时升级为调度器）。 + let a = db + .read_by_id::(agent)? + .with_context(|| format!("agent #{} not found", agent))?; + upsert_config( + &db, + "coordinator_agent_id", + &agent.to_string(), + "int", + "CLI selected coordinator agent", + "agent", + )?; + println!("coordinator agent set to #{} ({})", a.id, a.name); + } + } + Ok(()) +} + +// ============================================================ +// chat +// ============================================================ + +fn chat(args: ChatArgs) -> Result<()> { + match args.action { + Some(ChatCommand::Clear) => { + // 上下文缓存为进程内状态，清除不删除数据库历史记录。 + println!("cleared current session context cache (history records kept)"); + Ok(()) + } + Some(ChatCommand::List) => { + let db = DatabaseManager::new()?; + let rows = list_sessions(&db)?; + println!("{:<5}{:<28}{:<8}{:<8}{}", "id", "chat_uuid", "llm", "msgs", "title"); + if rows.is_empty() { + println!("(no sessions)"); + } + for (s, msgs) in rows { + println!( + "{:<5}{:<28}{:<8}{:<8}{}", + s.id, + s.chat_uuid, + s.llm_id, + msgs, + s.title.as_deref().unwrap_or("-"), + ); + } + Ok(()) + } + Some(ChatCommand::Rename { id, title }) => { + let db = DatabaseManager::new()?; + if rename_session(&db, id, &title)? { + println!("renamed session #{} -> '{}'", id, title); + } else { + anyhow::bail!("session #{} not found", id); + } + Ok(()) + } + Some(ChatCommand::Delete { id, yes }) => { + let db = DatabaseManager::new()?; + if !yes { + println!( + "this will permanently delete session #{} and all its messages; re-run with -y to confirm", + id + ); + return Ok(()); + } + let (sessions, messages) = delete_session(&db, id)?; + if sessions == 0 { + anyhow::bail!("session #{} not found", id); + } + println!("deleted session #{} and {} message(s)", id, messages); + Ok(()) + } + None => { + // 仅带 --chat-uuid 的场景由 run_with 拦截以启动 TUI 续跑； + // 若直接分发到此，给出提示。 + match args.chat_uuid { + Some(uuid) => { + println!("to resume session '{}' interactively, run without extra flags", uuid); + Ok(()) + } + None => { + println!("usage: Omega chat --chat-uuid | Omega chat "); + Ok(()) + } + } + } + } +} + +// ============================================================ +// memory / security / performance —— 策略类配置覆盖 +// ============================================================ + +fn memory(args: MemoryArgs) -> Result<()> { + let db = DatabaseManager::new()?; + let value = args.mem_type.as_str(); + upsert_config(&db, "memory_type", value, "string", "CLI memory loading scope", "memory")?; + println!("memory loading scope set to '{}'", value); + Ok(()) +} + +fn security(args: SecurityArgs) -> Result<()> { + let db = DatabaseManager::new()?; + let value = args.mode.as_str(); + upsert_config(&db, "security_mode", value, "string", "CLI MCP security policy", "security")?; + println!("MCP security policy set to '{}'", value); + Ok(()) +} + +fn performance(args: PerformanceArgs) -> Result<()> { + let db = DatabaseManager::new()?; + let value = args.mode.as_str(); + upsert_config(&db, "performance_mode", value, "string", "CLI performance profile", "performance")?; + println!("performance profile set to '{}'", value); + Ok(()) +} + +// ============================================================ +// arena +// ============================================================ + +async fn arena(action: ArenaCommand) -> Result<()> { + let svc = ArenaService::new()?; + match action { + ArenaCommand::Vote { + agent, + score, + comment, + chat_uuid, + round_uuid, + acc, + rel, + flu, + helpful, + } => { + // 收集非空多维度评分（短键 → 1-5）。 + let mut dimensions = BTreeMap::new(); + if let Some(v) = acc { + dimensions.insert("acc".to_string(), v); + } + if let Some(v) = rel { + dimensions.insert("rel".to_string(), v); + } + if let Some(v) = flu { + dimensions.insert("flu".to_string(), v); + } + if let Some(v) = helpful { + dimensions.insert("helpful".to_string(), v); + } + let dims_summary: Vec = dimensions + .iter() + .map(|(k, v)| format!("{}={}", k, v)) + .collect(); + let input = VoteInput { + agent_id: agent, + score, + comment, + chat_uuid: chat_uuid.unwrap_or_default(), + round_uuid, + dimensions, + }; + let id = svc.record_vote(input)?; + println!("recorded vote #{} for agent #{}: {}/5", id, agent, score); + if !dims_summary.is_empty() { + println!("dimensions: {}", dims_summary.join(", ")); + } + // 展示重算后的平均分。 + let avg = svc + .ranking()? + .into_iter() + .find(|r| r.agent_id == agent) + .and_then(|r| r.avg_score); + match avg { + Some(a) => println!("agent #{} avg_score -> {:.1}", agent, a), + None => println!("agent #{} avg_score -> -", agent), + } + } + ArenaCommand::List => { + let ranks = svc.ranking()?; + if ranks.is_empty() { + println!("(no agents)"); + return Ok(()); + } + println!("{:<5}{:<24}{:<12}{:<8}", "id", "name", "avg_score", "votes"); + for r in ranks { + let score = match r.avg_score { + Some(s) => format!("{:.1}", s), + None => "-".to_string(), + }; + println!( + "{:<5}{:<24}{:<12}{:<8}", + r.agent_id, r.name, score, r.vote_count + ); + } + } + ArenaCommand::Report { agent } => { + // 聚合评估数据，复用被评 Agent 自身的供应商/模型配置生成报告。 + let data = svc.report_data(agent)?; + let cfg = SessionStore::new()?.load_agent_config(agent)?; + let provider_config = ProviderConfig { + kind: parse_provider_kind(&cfg.provider_name).unwrap_or(ProviderKind::Compatible), + api_key: cfg.api_key.clone().unwrap_or_default(), + base_url: cfg.base_url.clone(), + model: cfg.model_code.clone(), + }; + let reporter = + PromptReportAgent::new(provider_config, cfg.model_code.clone(), cfg.temperature); + let text = reporter.generate(&data).await?; + println!("{}", text); + let path = report::save_report(agent, &text)?; + println!("report saved to {}", path.display()); + } + } + Ok(()) +} + +// ============================================================ +// prompt (§3.6 版本对比与回滚) +// ============================================================ + +fn prompt(action: PromptCommand) -> Result<()> { + let db = DatabaseManager::new()?; + match action { + PromptCommand::List => { + let mut rows = db.read_all::()?; + // 按名称、再按可解析版本、最后按 id 排序。 + rows.sort_by(|a, b| { + a.name + .cmp(&b.name) + .then_with(|| version_key(a).cmp(&version_key(b))) + .then(a.id.cmp(&b.id)) + }); + println!("{:<5}{:<24}{:<12}{:<8}{}", "id", "name", "version", "role", "category"); + if rows.is_empty() { + println!("(no prompts)"); + } + for p in rows { + println!( + "{:<5}{:<24}{:<12}{:<8}{}", + p.id, + p.name, + p.version.as_deref().unwrap_or("-"), + p.role.as_deref().unwrap_or("-"), + p.category.as_deref().unwrap_or("-"), + ); + } + } + PromptCommand::Show { id } => { + let p = db + .read_by_id::(id)? + .with_context(|| format!("prompt #{} not found", id))?; + println!("id : {}", p.id); + println!("name : {}", p.name); + println!("version : {}", p.version.as_deref().unwrap_or("-")); + println!("role : {}", p.role.as_deref().unwrap_or("-")); + println!("category : {}", p.category.as_deref().unwrap_or("-")); + println!("tags : {}", p.tags.as_deref().unwrap_or("-")); + println!("version_desc: {}", p.version_desc.as_deref().unwrap_or("-")); + println!("---- content ----"); + println!("{}", p.content); + } + PromptCommand::Diff { from, to } => { + let a = db + .read_by_id::(from)? + .with_context(|| format!("prompt #{} not found", from))?; + let b = db + .read_by_id::(to)? + .with_context(|| format!("prompt #{} not found", to))?; + println!( + "diff: #{} {} ({}) -> #{} {} ({})", + a.id, + a.name, + a.version.as_deref().unwrap_or("-"), + b.id, + b.name, + b.version.as_deref().unwrap_or("-"), + ); + for line in diff_lines(&a.content, &b.content) { + println!("{}", line); + } + } + PromptCommand::Rollback { id, desc } => { + let target = db + .read_by_id::(id)? + .with_context(|| format!("prompt #{} not found", id))?; + let (new_id, new_version) = rollback_prompt(&db, &target, desc)?; + println!( + "rolled back '{}' to content of #{} as new version {} (id #{})", + target.name, target.id, new_version, new_id + ); + } + PromptCommand::Validate { id } => { + let p = db + .read_by_id::(id)? + .with_context(|| format!("prompt #{} not found", id))?; + // 已知键 = 该提示词的占位符 + 内置 {AgentList}/{MCPList}。 + let store = SessionStore::new()?; + let placeholders = store.load_prompt_placeholders(id)?; + let mut known: HashSet = + placeholders.iter().map(|(k, _, _)| k.clone()).collect(); + known.insert("AgentList".to_string()); + known.insert("MCPList".to_string()); + + let report = validate_prompt(&p.content, &known); + println!( + "prompt #{} '{}' ({}) validation:", + p.id, + p.name, + p.version.as_deref().unwrap_or("-") + ); + println!( + " placeholders: {}", + if report.keys.is_empty() { + "(none)".to_string() + } else { + report.keys.join(", ") + } + ); + if report.is_ok() { + println!(" result: OK"); + } else { + if !report.unresolved.is_empty() { + println!(" unresolved: {}", report.unresolved.join(", ")); + } + if !report.malformed.is_empty() { + println!(" malformed : {}", report.malformed.join(" ")); + } + if report.unclosed { + println!(" unclosed : yes"); + } + println!(" result: FAILED"); + } + } + PromptCommand::Render { id, agent } => { + let p = db + .read_by_id::(id)? + .with_context(|| format!("prompt #{} not found", id))?; + let store = SessionStore::new()?; + let placeholders = store.load_prompt_placeholders(id)?; + // {AgentList}/{MCPList} 仅在指定 Agent 上下文时填充，否则为空。 + let (agents, tools) = match agent { + Some(aid) => { + let agents = store.load_callable_agents(aid)?; + let tools = store + .load_agent_mcp_tools(aid)? + .into_iter() + .map(|t| t.name) + .collect(); + (agents, tools) + } + None => (Vec::new(), Vec::new()), + }; + let vars = build_vars(&agents, &tools, &placeholders); + println!("{}", render(&p.content, &vars)); + } + } + Ok(()) +} + +/// 以 `target` 的内容生成同名新版本（追加，保留全部历史），返回（新 id，新版本号）。 +/// +/// 版本号取同名现有版本集合的最大可解析版本的补丁位 +1；不可解析则回退 `v1.0.0`。 +fn rollback_prompt( + db: &DatabaseManager, + target: &PromptInfo, + desc: Option, +) -> Result<(i64, String)> { + let siblings = db.find_where::("name = ?1", &[&target.name])?; + let versions: Vec = siblings.iter().filter_map(|p| p.version.clone()).collect(); + let new_version = next_version(&versions); + let note = desc.unwrap_or_else(|| { + format!( + "rollback to #{} ({})", + target.id, + target.version.as_deref().unwrap_or("-") + ) + }); + let new_id = db.create(PromptInfo { + id: 0, + name: target.name.clone(), + content: target.content.clone(), + description: target.description.clone(), + category: target.category.clone(), + role: target.role.clone(), + version: Some(new_version.clone()), + version_desc: Some(note), + tags: target.tags.clone(), + create_time: Utc::now(), + update_time: Utc::now(), + remark: Some("auto-rollback".to_string()), + })?; + Ok((new_id, new_version)) +} + +/// 解析形如 `v1.2.3` / `1.2.3` 的版本号为可比较三元组；无法解析返回 `None`。 +fn parse_version(v: &str) -> Option<(u32, u32, u32)> { + let s = v.trim().trim_start_matches(['v', 'V']); + let mut parts = s.split('.'); + let major = parts.next()?.trim().parse::().ok()?; + let minor = parts.next()?.trim().parse::().ok()?; + let patch = parts.next()?.trim().parse::().ok()?; + Some((major, minor, patch)) +} + +/// 排序键：可解析版本用其三元组，不可解析退化为 `(0,0,0)`（排在最前）。 +fn version_key(p: &PromptInfo) -> (u32, u32, u32) { + p.version.as_deref().and_then(parse_version).unwrap_or((0, 0, 0)) +} + +/// 基于现有版本集合计算下一个补丁版本号；集合为空或均不可解析时返回 `v1.0.0`。 +fn next_version(existing: &[String]) -> String { + match existing.iter().filter_map(|s| parse_version(s)).max() { + None => "v1.0.0".to_string(), + Some((ma, mi, pa)) => format!("v{}.{}.{}", ma, mi, pa + 1), + } +} + +/// 行级 diff（基于 LCS），返回带 ` `（未变）/`- `（删除）/`+ `（新增）前缀的行。 +fn diff_lines(old: &str, new: &str) -> Vec { + let a: Vec<&str> = old.lines().collect(); + let b: Vec<&str> = new.lines().collect(); + let (n, m) = (a.len(), b.len()); + + // dp[i][j] = a[i..] 与 b[j..] 的 LCS 长度。 + let mut dp = vec![vec![0u32; m + 1]; n + 1]; + for i in (0..n).rev() { + for j in (0..m).rev() { + dp[i][j] = if a[i] == b[j] { + dp[i + 1][j + 1] + 1 + } else { + dp[i + 1][j].max(dp[i][j + 1]) + }; + } + } + + let mut out = Vec::new(); + let (mut i, mut j) = (0, 0); + while i < n && j < m { + if a[i] == b[j] { + out.push(format!(" {}", a[i])); + i += 1; + j += 1; + } else if dp[i + 1][j] >= dp[i][j + 1] { + out.push(format!("- {}", a[i])); + i += 1; + } else { + out.push(format!("+ {}", b[j])); + j += 1; + } + } + while i < n { + out.push(format!("- {}", a[i])); + i += 1; + } + while j < m { + out.push(format!("+ {}", b[j])); + j += 1; + } + out +} + +// ============================================================ +// system_config helpers +// ============================================================ + +/// 写入（或更新）一个 `system_config` 键值。 +fn upsert_config( + db: &DatabaseManager, + key: &str, + value: &str, + key_type: &str, + description: &str, + tag: &str, +) -> Result<()> { + let mut existing = db.find_where::("key = ?1", &[&key])?; + if let Some(cfg) = existing.first_mut() { + cfg.value = value.to_string(); + db.update(cfg.clone())?; + } else { + let cfg = SystemConfig { + key: key.to_string(), + value: value.to_string(), + key_type: Some(key_type.to_string()), + description: Some(description.to_string()), + tag: Some(tag.to_string()), + sort: 0, + create_time: Utc::now(), + update_time: Utc::now(), + remark: None, + }; + db.create(cfg)?; + } + Ok(()) +} + +/// 读取一个 `system_config` 键值（不存在返回 `None`）。 +fn read_config(db: &DatabaseManager, key: &str) -> Result> { + let found = db.find_where::("key = ?1", &[&key])?; + Ok(found.into_iter().next().map(|c| c.value)) +} + +// ============================================================ +// model auto-sync (§3.4) +// ============================================================ + +/// 将供应商返回的模型标识同步入 `llm_model`（按 `provider_id + model_code` 幂等）。 +/// +/// 仅新增缺失项，已存在的模型保持不变（模型基础参数不支持自动覆写）。 +/// 返回（新增数，已存在跳过数）；空白标识被忽略。 +fn sync_models(db: &DatabaseManager, provider_id: i32, codes: &[String]) -> Result<(usize, usize)> { + let mut added = 0usize; + let mut skipped = 0usize; + for raw in codes { + let code = raw.trim(); + if code.is_empty() { + continue; + } + let code = code.to_string(); + let existing = db.find_where::( + "provider_id = ?1 AND model_code = ?2", + &[&provider_id, &code], + )?; + if existing.first().is_some() { + skipped += 1; + continue; + } + db.create(LLMModel { + id: 0, + provider_id, + avatar: None, + name: code.clone(), + model_code: code, + max_tokens: 0, + max_response_tokens: 0, + weight: 0, + call_count: 0, + status: 1, + create_time: Utc::now(), + update_time: Utc::now(), + remark: Some("auto-synced".to_string()), + })?; + added += 1; + } + Ok((added, skipped)) +} + +/// `model delete` 的前置校验结果。 +enum ModelDeletable { + /// 模型不存在。 + NotFound, + /// 被 `refs` 个 Agent 引用，不可删除。 + Blocked { refs: usize }, + /// 可删除，附带待删模型行。 + Ready(LLMModel), +} + +/// 校验模型是否可删除：不存在 / 被 Agent 引用 / 可删除。 +fn check_model_deletable(db: &DatabaseManager, code: &str) -> Result { + let found = db.find_where::("model_code = ?1", &[&code])?; + let Some(m) = found.into_iter().next() else { + return Ok(ModelDeletable::NotFound); + }; + let refs = db.find_where::("llm_id = ?1", &[&m.id])?.len(); + if refs > 0 { + Ok(ModelDeletable::Blocked { refs }) + } else { + Ok(ModelDeletable::Ready(m)) + } +} + +// ============================================================ +// chat session helpers (§3.9) +// ============================================================ + +/// 列出所有会话及各自消息数（按 id 升序）。 +fn list_sessions(db: &DatabaseManager) -> Result> { + let mut sessions = db.read_all::()?; + sessions.sort_by_key(|s| s.id); + + // 一次性统计每个 chat_uuid 的消息数。 + let mut stmt = db + .conn + .prepare("SELECT chat_uuid, COUNT(*) FROM chat_message GROUP BY chat_uuid")?; + let counts: HashMap = stmt + .query_map(params![], |row| { + Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)? as usize)) + })? + .filter_map(|r| r.ok()) + .collect(); + + Ok(sessions + .into_iter() + .map(|s| { + let msgs = counts.get(&s.chat_uuid).copied().unwrap_or(0); + (s, msgs) + }) + .collect()) +} + +/// 重命名会话标题；返回是否找到并更新。 +fn rename_session(db: &DatabaseManager, id: i32, title: &str) -> Result { + let Some(mut s) = db.read_by_id::(id)? else { + return Ok(false); + }; + s.title = Some(title.to_string()); + db.update(s)?; + Ok(true) +} + +/// 物理删除会话及其全部消息；返回（删除的会话数，删除的消息数）。 +fn delete_session(db: &DatabaseManager, id: i32) -> Result<(usize, usize)> { + let Some(s) = db.read_by_id::(id)? else { + return Ok((0, 0)); + }; + let messages = db + .conn + .execute("DELETE FROM chat_message WHERE chat_uuid = ?1", params![s.chat_uuid])?; + let sessions = db + .conn + .execute("DELETE FROM chat_session WHERE id = ?1", params![id])?; + Ok((sessions, messages)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::db::Model; + use crate::db::models::chat_message::ChatMessage; + + /// 构造仅含所需表的内存库，验证 upsert/read 配置往返。 + fn temp_db() -> DatabaseManager { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch(SystemConfig::create_table_sql()).unwrap(); + // DatabaseManager::new 会打开磁盘库；测试直接用已有连接构造。 + DatabaseManager { conn } + } + + #[test] + fn upsert_then_read_config_roundtrip() { + let db = temp_db(); + assert_eq!(read_config(&db, "security_mode").unwrap(), None); + + upsert_config(&db, "security_mode", "strict", "string", "d", "security").unwrap(); + assert_eq!(read_config(&db, "security_mode").unwrap().as_deref(), Some("strict")); + + // 再次写入应更新而非重复插入。 + upsert_config(&db, "security_mode", "bypass", "string", "d", "security").unwrap(); + assert_eq!(read_config(&db, "security_mode").unwrap().as_deref(), Some("bypass")); + let all = db.read_all::().unwrap(); + assert_eq!(all.len(), 1); + } + + #[test] + fn sync_models_inserts_new_and_is_idempotent() { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch(LLMModel::create_table_sql()).unwrap(); + let db = DatabaseManager { conn }; + + let codes = vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string()]; + assert_eq!(sync_models(&db, 1, &codes).unwrap(), (2, 0)); + // 再次同步应全部跳过（幂等）。 + assert_eq!(sync_models(&db, 1, &codes).unwrap(), (0, 2)); + // 不同 provider_id 视为不同模型。 + assert_eq!(sync_models(&db, 2, &codes).unwrap(), (2, 0)); + // 空白项被忽略。 + assert_eq!(sync_models(&db, 1, &[" ".to_string()]).unwrap(), (0, 0)); + + let all = db.read_all::().unwrap(); + assert_eq!(all.len(), 4); + let synced = all.iter().find(|m| m.provider_id == 1 && m.model_code == "gpt-4o").unwrap(); + assert_eq!(synced.name, "gpt-4o"); + assert_eq!(synced.status, 1); + assert_eq!(synced.remark.as_deref(), Some("auto-synced")); + } + + fn make_prompt(name: &str, version: &str, content: &str) -> PromptInfo { + PromptInfo { + id: 0, + name: name.to_string(), + content: content.to_string(), + description: None, + category: None, + role: Some("System".to_string()), + version: Some(version.to_string()), + version_desc: None, + tags: None, + create_time: Utc::now(), + update_time: Utc::now(), + remark: None, + } + } + + #[test] + fn parse_version_handles_prefix_and_garbage() { + assert_eq!(parse_version("v1.2.3"), Some((1, 2, 3))); + assert_eq!(parse_version("1.0.0"), Some((1, 0, 0))); + assert_eq!(parse_version("v0.9"), None); + assert_eq!(parse_version("nightly"), None); + } + + #[test] + fn next_version_bumps_patch_of_max() { + assert_eq!(next_version(&[]), "v1.0.0"); + assert_eq!(next_version(&["v1.0.0".to_string()]), "v1.0.1"); + assert_eq!( + next_version(&[ + "v1.0.0".to_string(), + "v1.2.9".to_string(), + "v1.1.0".to_string() + ]), + "v1.2.10" + ); + // 均不可解析时回退。 + assert_eq!(next_version(&["nightly".to_string()]), "v1.0.0"); + } + + #[test] + fn diff_lines_marks_context_add_remove() { + let out = diff_lines("a\nb\nc", "a\nB\nc\nd"); + assert!(out.contains(&" a".to_string())); + assert!(out.contains(&"- b".to_string())); + assert!(out.contains(&"+ B".to_string())); + assert!(out.contains(&" c".to_string())); + assert!(out.contains(&"+ d".to_string())); + // 无差异时全为上下文行。 + let same = diff_lines("x\ny", "x\ny"); + assert_eq!(same, vec![" x".to_string(), " y".to_string()]); + } + + #[test] + fn rollback_prompt_appends_new_version_with_target_content() { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch(PromptInfo::create_table_sql()).unwrap(); + let db = DatabaseManager { conn }; + db.create(make_prompt("coder", "v1.0.0", "old content")).unwrap(); + db.create(make_prompt("coder", "v1.1.0", "new content")).unwrap(); + + // 回滚到 v1.0.0（id #1）。 + let target = db.read_by_id::(1).unwrap().unwrap(); + let (new_id, new_version) = rollback_prompt(&db, &target, None).unwrap(); + assert_eq!(new_version, "v1.1.1"); + + let created = db.read_by_id::(new_id as i32).unwrap().unwrap(); + assert_eq!(created.content, "old content"); + assert_eq!(created.name, "coder"); + assert!(created.version_desc.unwrap().contains("rollback")); + assert_eq!(created.remark.as_deref(), Some("auto-rollback")); + // 历史保留：共 3 行。 + assert_eq!(db.read_all::().unwrap().len(), 3); + } + + #[test] + fn rollback_prompt_uses_custom_desc() { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch(PromptInfo::create_table_sql()).unwrap(); + let db = DatabaseManager { conn }; + db.create(make_prompt("writer", "v2.0.0", "body")).unwrap(); + let target = db.read_by_id::(1).unwrap().unwrap(); + let (_, ver) = rollback_prompt(&db, &target, Some("revert bug".to_string())).unwrap(); + assert_eq!(ver, "v2.0.1"); + let created = db.read_by_id::(2).unwrap().unwrap(); + assert_eq!(created.version_desc.as_deref(), Some("revert bug")); + } + + fn make_session(uuid: &str, title: Option<&str>, llm_id: i32) -> ChatSession { + ChatSession { + id: 0, + chat_uuid: uuid.to_string(), + title: title.map(|t| t.to_string()), + llm_id, + is_deleted: 0, + create_time: Utc::now(), + update_time: Utc::now(), + remark: None, + } + } + + fn make_message(uuid: &str) -> ChatMessage { + ChatMessage { + id: 0, + chat_uuid: uuid.to_string(), + round_uuid: None, + role: "user".to_string(), + content: "hi".to_string(), + input_tokens: 0, + output_tokens: 0, + duration: 0, + model_config: None, + create_time: Utc::now(), + update_time: Utc::now(), + remark: None, + } + } + + fn make_model(code: &str) -> LLMModel { + LLMModel { + id: 0, + provider_id: 1, + avatar: None, + name: code.to_string(), + model_code: code.to_string(), + max_tokens: 0, + max_response_tokens: 0, + weight: 0, + call_count: 0, + status: 1, + create_time: Utc::now(), + update_time: Utc::now(), + remark: None, + } + } + + fn session_db() -> DatabaseManager { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch(ChatSession::create_table_sql()).unwrap(); + conn.execute_batch(ChatMessage::create_table_sql()).unwrap(); + DatabaseManager { conn } + } + + #[test] + fn list_sessions_counts_messages_per_uuid() { + let db = session_db(); + db.create(make_session("u1", Some("first"), 1)).unwrap(); + db.create(make_session("u2", None, 2)).unwrap(); + // u1 两条消息，u2 一条。 + db.create(make_message("u1")).unwrap(); + db.create(make_message("u1")).unwrap(); + db.create(make_message("u2")).unwrap(); + + let rows = list_sessions(&db).unwrap(); + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].0.chat_uuid, "u1"); + assert_eq!(rows[0].1, 2); + assert_eq!(rows[1].0.chat_uuid, "u2"); + assert_eq!(rows[1].1, 1); + } + + #[test] + fn rename_session_updates_title_or_reports_missing() { + let db = session_db(); + db.create(make_session("u1", None, 1)).unwrap(); + assert!(rename_session(&db, 1, "planning").unwrap()); + assert_eq!( + db.read_by_id::(1).unwrap().unwrap().title.as_deref(), + Some("planning") + ); + // 不存在的会话。 + assert!(!rename_session(&db, 99, "x").unwrap()); + } + + #[test] + fn delete_session_removes_rows_and_messages() { + let db = session_db(); + db.create(make_session("u1", None, 1)).unwrap(); + db.create(make_session("u2", None, 1)).unwrap(); + db.create(make_message("u1")).unwrap(); + db.create(make_message("u1")).unwrap(); + + assert_eq!(delete_session(&db, 1).unwrap(), (1, 2)); + // u1 已删，u2 保留。 + assert!(db.read_by_id::(1).unwrap().is_none()); + assert!(db.read_by_id::(2).unwrap().is_some()); + // 再删不存在的会话返回 (0,0)。 + assert_eq!(delete_session(&db, 1).unwrap(), (0, 0)); + } + + #[test] + fn check_model_deletable_guards_references() { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch(LLMModel::create_table_sql()).unwrap(); + conn.execute_batch(AgentInfo::create_table_sql()).unwrap(); + let db = DatabaseManager { conn }; + // #1 used 被 Agent 引用，#2 free 无引用。 + db.create(make_model("used")).unwrap(); + db.create(make_model("free")).unwrap(); + db.conn + .execute( + "INSERT INTO agent_info (id, llm_id, name, status) VALUES (1, 1, 'A', 1)", + params![], + ) + .unwrap(); + + assert!(matches!( + check_model_deletable(&db, "missing").unwrap(), + ModelDeletable::NotFound + )); + assert!(matches!( + check_model_deletable(&db, "used").unwrap(), + ModelDeletable::Blocked { refs: 1 } + )); + match check_model_deletable(&db, "free").unwrap() { + ModelDeletable::Ready(m) => assert_eq!(m.model_code, "free"), + _ => panic!("expected Ready"), + } + } +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs new file mode 100644 index 0000000..f0a6fab --- /dev/null +++ b/src/cli/mod.rs @@ -0,0 +1,728 @@ +//! 命令行参数体系 +//! +//! 依据《命令行参数设计》定义 Omega CLI 的完整参数结构：全局启动参数 +//! （Workspace 路径、`--mode` 启动模式）与六大类子命令（模型 / 供应商 / Agent / +//! 会话 / 记忆 / 安全 / 性能）。 +//! +//! 参数优先级：**命令行参数 > 配置文件 > 系统默认配置**（见设计文档六.开发备注）。 +//! 本模块只负责「解析与校验」，具体命令的业务落地在 [`commands`] 中实现。 + +pub mod commands; + +use std::path::PathBuf; + +use clap::{Args, Parser, Subcommand, ValueEnum}; + +/// Omega CLI 顶层参数。 +/// +/// 无子命令时以 [`Cli::workspace`] 解析出的路径作为 Workspace 启动交互式 TUI； +/// 带子命令时执行对应的管理动作（见 [`commands::dispatch`]）。 +#[derive(Parser, Debug)] +#[command( + name = "Omega", + version, + about = "Omega Code —— CLI-First 本地 AI Agent 全栈运行时", + long_about = None, + arg_required_else_help = false +)] +pub struct Cli { + /// 本地路径（绝对或相对），作为 Omega 的 Workspace，用于工程配置与文件索引构建。 + #[arg(value_name = "LocalPath", global = false)] + pub path: Option, + + /// 启动模式：single（单模式）/ workspace（工作空间模式），决定 Agent、MCP 的加载范围与权限管控。 + #[arg(long, value_enum, default_value_t = Mode::Single)] + pub mode: Mode, + + /// 子命令（缺省时启动交互式 TUI）。 + #[command(subcommand)] + pub command: Option, +} + +/// 启动模式。 +#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)] +pub enum Mode { + /// 单模式：仅当前 Workspace，加载范围最小。 + Single, + /// 工作空间模式：完整加载工作空间的 Agent、MCP 能力。 + Workspace, +} + +impl Mode { + /// 解析后的模式字符串（用于写入运行配置）。 + pub fn as_str(&self) -> &'static str { + match self { + Mode::Single => "single", + Mode::Workspace => "workspace", + } + } +} + +impl Cli { + /// 解析 Workspace 路径：显式 `LocalPath` 优先，否则回退到当前工作目录。 + /// + /// 相对路径基于当前进程工作目录解析为绝对路径。 + pub fn workspace(&self) -> anyhow::Result { + let base = match &self.path { + Some(p) => p.clone(), + None => std::env::current_dir()?, + }; + Ok(if base.is_absolute() { + base + } else { + std::env::current_dir()?.join(base) + }) + } +} + +/// 全部子命令。 +#[derive(Subcommand, Debug)] +pub enum Command { + /// LLM 模型管理：list / status / set / sync / delete。 + Model { + #[command(subcommand)] + action: ModelCommand, + }, + /// AI 供应商配置：加载指定供应商，可选覆盖 API 密钥。 + Supplier(SupplierArgs), + /// Agent 管理：use / mcp。 + Agent { + #[command(subcommand)] + action: AgentCommand, + }, + /// 会话管理：加载会话、清除上下文，或 list / rename / delete。 + Chat(ChatArgs), + /// 记忆加载：指定加载的记忆类型。 + Memory(MemoryArgs), + /// 安全策略：指定 MCP 安全策略。 + Security(SecurityArgs), + /// 性能模式：指定 LLM 使用策略与 Token 消耗档位。 + Performance(PerformanceArgs), + /// 竞技场：对 Agent 单轮输出星级打分与查看排行榜。 + Arena { + #[command(subcommand)] + action: ArenaCommand, + }, + /// 提示词管理：list / show / diff / rollback / validate / render（§3.6 版本对比、回滚与校验预览）。 + Prompt { + #[command(subcommand)] + action: PromptCommand, + }, +} + +/// `model` 子命令。 +#[derive(Subcommand, Debug)] +pub enum ModelCommand { + /// 列出模型（默认仅启用；`-a` 含未启用）。 + List { + /// 列出所有可用模型（含未启用）。 + #[arg(short = 'a')] + all: bool, + /// 跳过确认步骤。 + #[arg(short = 'y')] + yes: bool, + }, + /// 查看指定模型的状态与参数配置。 + Status { + /// 模型标识（`model_code`），如 deepseek-v4。 + name: String, + /// 跳过确认步骤。 + #[arg(short = 'y')] + yes: bool, + }, + /// 配置模型参数（覆盖模型默认配置）。 + Set { + /// 模型标识（`model_code`）。 + name: String, + /// 采样温度，取值 0.0~1.0（Agent 级参数）。 + #[arg(short = 't')] + temperature: Option, + /// 最大上下文 Token（正整数）。 + #[arg(long = "max-tokens")] + max_tokens: Option, + }, + /// 从供应商 API 自动同步模型列表到本地库（§3.4 模型自动同步）。 + Sync { + /// 供应商 ID（或名称）。 + #[arg(short = 's')] + supplier: String, + }, + /// 删除模型（物理删除；被 Agent 引用时拒绝，§3.4 增删查）。 + Delete { + /// 模型标识（`model_code`）。 + name: String, + /// 跳过二次确认。 + #[arg(short = 'y')] + yes: bool, + }, +} + +/// `supplier` 子命令参数。 +#[derive(Args, Debug)] +pub struct SupplierArgs { + /// 供应商 ID（或名称）。 + #[arg(short = 's')] + pub supplier: String, + /// API 密钥（优先级高于配置文件）。 + #[arg(long = "api-key")] + pub api_key: Option, +} + +/// `agent` 子命令。 +#[derive(Subcommand, Debug)] +pub enum AgentCommand { + /// 指定默认使用的 Agent。 + Use { + /// Agent ID。 + #[arg(short = 'a')] + agent: i32, + /// 覆盖 Agent 默认绑定的 Prompt ID。 + #[arg(long)] + prompt: Option, + }, + /// 指定 Agent 挂载的 MCP 工具集合。 + Mcp { + /// Agent ID。 + #[arg(short = 'a')] + agent: i32, + /// MCP 工具 ID 列表，逗号分隔（如 1,2,3）。 + #[arg(long, value_delimiter = ',', required = true)] + mcp: Vec, + }, + /// 指定作为多 Agent 调度器（Coordinator）的 Agent。 + Coordinator { + /// Agent ID。 + #[arg(short = 'a')] + agent: i32, + }, +} + +/// `chat` 子命令参数（既可加载会话，也可执行 `clear`）。 +#[derive(Args, Debug)] +pub struct ChatArgs { + /// 会话 UUID（与数据库 `chat_uuid` 一致），用于加载历史、继续对话。 + #[arg(long = "chat-uuid")] + pub chat_uuid: Option, + /// 会话操作子命令。 + #[command(subcommand)] + pub action: Option, +} + +/// `chat` 的操作子命令。 +#[derive(Subcommand, Debug)] +pub enum ChatCommand { + /// 清除当前会话上下文缓存，重置会话状态（不删除历史记录）。 + Clear, + /// 列出会话（id、标题、UUID、绑定模型、消息数）。 + List, + /// 重命名会话标题（§3.9）。 + Rename { + /// 会话 ID。 + id: i32, + /// 新标题。 + title: String, + }, + /// 物理删除会话及其全部消息（需 `-y` 二次确认，§3.9）。 + Delete { + /// 会话 ID。 + id: i32, + /// 跳过二次确认。 + #[arg(short = 'y')] + yes: bool, + }, +} + +/// `memory` 子命令参数。 +#[derive(Args, Debug)] +pub struct MemoryArgs { + /// 记忆类型：user / project / all / none。 + #[arg(long = "type", value_enum)] + pub mem_type: MemoryType, +} + +/// 记忆类型。 +#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)] +pub enum MemoryType { + /// 用户记忆。 + User, + /// 项目记忆。 + Project, + /// 全部记忆。 + All, + /// 不加载记忆。 + None, +} + +impl MemoryType { + pub fn as_str(&self) -> &'static str { + match self { + MemoryType::User => "user", + MemoryType::Project => "project", + MemoryType::All => "all", + MemoryType::None => "none", + } + } +} + +/// `security` 子命令参数。 +#[derive(Args, Debug)] +pub struct SecurityArgs { + /// 安全策略：default / strict / bypass。 + #[arg(long, value_enum)] + pub mode: SecurityMode, +} + +/// MCP 安全策略。 +#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)] +pub enum SecurityMode { + /// 默认策略。 + Default, + /// 严格策略。 + Strict, + /// 宽松策略。 + Bypass, +} + +impl SecurityMode { + pub fn as_str(&self) -> &'static str { + match self { + SecurityMode::Default => "default", + SecurityMode::Strict => "strict", + SecurityMode::Bypass => "bypass", + } + } +} + +/// `performance` 子命令参数。 +#[derive(Args, Debug)] +pub struct PerformanceArgs { + /// 性能模式：performance / balance / quality。 + #[arg(long, value_enum)] + pub mode: PerformanceMode, +} + +/// 性能模式。 +#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)] +pub enum PerformanceMode { + /// 性能模式（低延迟、低 Token 消耗）。 + Performance, + /// 平衡模式。 + Balance, + /// 质量模式（高质量、高 Token 消耗）。 + Quality, +} + +impl PerformanceMode { + pub fn as_str(&self) -> &'static str { + match self { + PerformanceMode::Performance => "performance", + PerformanceMode::Balance => "balance", + PerformanceMode::Quality => "quality", + } + } +} + +/// `arena` 子命令。 +#[derive(Subcommand, Debug)] +pub enum ArenaCommand { + /// 对指定 Agent 打分（1-5 星）并附可选评价。 + Vote { + /// Agent ID。 + #[arg(short = 'a')] + agent: i32, + /// 星级评分，取值 1-5。 + #[arg(long)] + score: i8, + /// 文字评价（可选）。 + #[arg(long)] + comment: Option, + /// 关联会话 UUID（可选）。 + #[arg(long = "chat-uuid")] + chat_uuid: Option, + /// 关联单轮消息 UUID（可选）。 + #[arg(long = "round-uuid")] + round_uuid: Option, + /// 准确性维度评分（1-5，可选）。 + #[arg(long)] + acc: Option, + /// 相关性维度评分（1-5，可选）。 + #[arg(long)] + rel: Option, + /// 流畅性维度评分（1-5，可选）。 + #[arg(long)] + flu: Option, + /// 有用性维度评分（1-5，可选）。 + #[arg(long)] + helpful: Option, + }, + /// 按平均评分展示 Agent 排行榜。 + List, + /// 生成指定 Agent 的 Prompt 优化报告（LLM 驱动，打印并落盘）。 + Report { + /// Agent ID。 + #[arg(short = 'a')] + agent: i32, + }, +} + +/// `prompt` 子命令（§3.6 提示词管理与版本能力）。 +#[derive(Subcommand, Debug)] +pub enum PromptCommand { + /// 列出全部提示词（按名称、版本排序）。 + List, + /// 查看指定提示词版本的完整内容与元信息。 + Show { + /// 提示词 ID。 + id: i32, + }, + /// 对比两个提示词版本的内容差异（行级 diff）。 + Diff { + /// 旧版本提示词 ID。 + from: i32, + /// 新版本提示词 ID。 + to: i32, + }, + /// 回滚：以指定历史版本内容生成同名新版本（追加，不覆盖历史）。 + Rollback { + /// 目标历史版本提示词 ID。 + id: i32, + /// 新版本备注（可选）。 + #[arg(long)] + desc: Option, + }, + /// 校验提示词语法与占位符完整性（§3.6.5）。 + Validate { + /// 提示词 ID。 + id: i32, + }, + /// 渲染预览：用占位符值填充并打印最终提示词（§3.6.5）。 + Render { + /// 提示词 ID。 + id: i32, + /// 以指定 Agent 上下文填充 {AgentList}/{MCPList}（可选）。 + #[arg(long)] + agent: Option, + }, +} + +#[cfg(test)] +mod tests { + use super::*; + use clap::CommandFactory; + + /// CLI 定义自身必须合法（ clap derive 结构体校验）。 + #[test] + fn cli_definition_is_valid() { + Cli::command().debug_assert(); + } + + #[test] + fn parses_global_path_and_mode() { + let cli = Cli::try_parse_from(["Omega", "./proj", "--mode", "workspace"]).unwrap(); + assert_eq!(cli.path, Some(PathBuf::from("./proj"))); + assert_eq!(cli.mode, Mode::Workspace); + assert!(cli.command.is_none()); + } + + #[test] + fn mode_defaults_to_single() { + let cli = Cli::try_parse_from(["Omega"]).unwrap(); + assert_eq!(cli.mode, Mode::Single); + } + + #[test] + fn parses_model_subcommands() { + let cli = Cli::try_parse_from(["Omega", "model", "list", "-a"]).unwrap(); + match cli.command { + Some(Command::Model { + action: ModelCommand::List { all, .. }, + }) => assert!(all), + other => panic!("unexpected: {other:?}"), + } + + let cli = Cli::try_parse_from(["Omega", "model", "set", "gpt-4o", "-t", "0.3", "--max-tokens", "8192"]).unwrap(); + match cli.command { + Some(Command::Model { + action: + ModelCommand::Set { + name, + temperature, + max_tokens, + }, + }) => { + assert_eq!(name, "gpt-4o"); + assert_eq!(temperature, Some(0.3)); + assert_eq!(max_tokens, Some(8192)); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_model_sync() { + let cli = Cli::try_parse_from(["Omega", "model", "sync", "-s", "openai"]).unwrap(); + match cli.command { + Some(Command::Model { + action: ModelCommand::Sync { supplier }, + }) => assert_eq!(supplier, "openai"), + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_model_delete() { + let cli = Cli::try_parse_from(["Omega", "model", "delete", "gpt-4o", "-y"]).unwrap(); + match cli.command { + Some(Command::Model { + action: ModelCommand::Delete { name, yes }, + }) => { + assert_eq!(name, "gpt-4o"); + assert!(yes); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_chat_session_management() { + assert!(matches!( + Cli::try_parse_from(["Omega", "chat", "list"]).unwrap().command, + Some(Command::Chat(ChatArgs { + action: Some(ChatCommand::List), + .. + })) + )); + + let cli = Cli::try_parse_from(["Omega", "chat", "rename", "4", "Sprint planning"]).unwrap(); + match cli.command { + Some(Command::Chat(ChatArgs { + action: Some(ChatCommand::Rename { id, title }), + .. + })) => { + assert_eq!(id, 4); + assert_eq!(title, "Sprint planning"); + } + other => panic!("unexpected: {other:?}"), + } + + let cli = Cli::try_parse_from(["Omega", "chat", "delete", "4", "-y"]).unwrap(); + match cli.command { + Some(Command::Chat(ChatArgs { + action: Some(ChatCommand::Delete { id, yes }), + .. + })) => { + assert_eq!(id, 4); + assert!(yes); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_agent_mcp_tool_ids() { + let cli = Cli::try_parse_from(["Omega", "agent", "mcp", "-a", "5", "--mcp", "1,2,3"]).unwrap(); + match cli.command { + Some(Command::Agent { + action: AgentCommand::Mcp { agent, mcp }, + }) => { + assert_eq!(agent, 5); + assert_eq!(mcp, vec![1, 2, 3]); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_agent_coordinator() { + let cli = Cli::try_parse_from(["Omega", "agent", "coordinator", "-a", "3"]).unwrap(); + match cli.command { + Some(Command::Agent { + action: AgentCommand::Coordinator { agent }, + }) => assert_eq!(agent, 3), + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_arena_vote() { + let cli = Cli::try_parse_from( + ["Omega", "arena", "vote", "-a", "2", "--score", "5", "--comment", "good"], + ) + .unwrap(); + match cli.command { + Some(Command::Arena { + action: + ArenaCommand::Vote { + agent, + score, + comment, + chat_uuid, + round_uuid, + .. + }, + }) => { + assert_eq!(agent, 2); + assert_eq!(score, 5); + assert_eq!(comment.as_deref(), Some("good")); + assert!(chat_uuid.is_none()); + assert!(round_uuid.is_none()); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_arena_vote_with_dims() { + let cli = Cli::try_parse_from( + ["Omega", "arena", "vote", "-a", "2", "--score", "4", "--acc", "5", "--rel", "3"], + ) + .unwrap(); + match cli.command { + Some(Command::Arena { + action: ArenaCommand::Vote { acc, rel, flu, helpful, .. }, + }) => { + assert_eq!(acc, Some(5)); + assert_eq!(rel, Some(3)); + assert_eq!(flu, None); + assert_eq!(helpful, None); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_arena_report() { + let cli = Cli::try_parse_from(["Omega", "arena", "report", "-a", "7"]).unwrap(); + assert!(matches!( + cli.command, + Some(Command::Arena { + action: ArenaCommand::Report { agent: 7 } + }) + )); + } + + #[test] + fn parses_arena_list() { + let cli = Cli::try_parse_from(["Omega", "arena", "list"]).unwrap(); + assert!(matches!( + cli.command, + Some(Command::Arena { + action: ArenaCommand::List + }) + )); + } + + #[test] + fn parses_prompt_subcommands() { + let cli = Cli::try_parse_from(["Omega", "prompt", "list"]).unwrap(); + assert!(matches!( + cli.command, + Some(Command::Prompt { + action: PromptCommand::List + }) + )); + + let cli = Cli::try_parse_from(["Omega", "prompt", "diff", "3", "7"]).unwrap(); + match cli.command { + Some(Command::Prompt { + action: PromptCommand::Diff { from, to }, + }) => { + assert_eq!(from, 3); + assert_eq!(to, 7); + } + other => panic!("unexpected: {other:?}"), + } + + let cli = + Cli::try_parse_from(["Omega", "prompt", "rollback", "5", "--desc", "revert"]).unwrap(); + match cli.command { + Some(Command::Prompt { + action: PromptCommand::Rollback { id, desc }, + }) => { + assert_eq!(id, 5); + assert_eq!(desc.as_deref(), Some("revert")); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_prompt_validate_and_render() { + assert!(matches!( + Cli::try_parse_from(["Omega", "prompt", "validate", "8"]).unwrap().command, + Some(Command::Prompt { + action: PromptCommand::Validate { id: 8 } + }) + )); + + let cli = Cli::try_parse_from(["Omega", "prompt", "render", "8", "--agent", "3"]).unwrap(); + match cli.command { + Some(Command::Prompt { + action: PromptCommand::Render { id, agent }, + }) => { + assert_eq!(id, 8); + assert_eq!(agent, Some(3)); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_chat_clear_and_uuid() { + let cli = Cli::try_parse_from(["Omega", "chat", "clear"]).unwrap(); + match cli.command { + Some(Command::Chat(ChatArgs { + action: Some(ChatCommand::Clear), + .. + })) => {} + other => panic!("unexpected: {other:?}"), + } + + let cli = Cli::try_parse_from(["Omega", "chat", "--chat-uuid", "abc-123"]).unwrap(); + match cli.command { + Some(Command::Chat(ChatArgs { chat_uuid, action } )) => { + assert_eq!(chat_uuid.as_deref(), Some("abc-123")); + assert!(action.is_none()); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parses_memory_security_performance() { + let cli = Cli::try_parse_from(["Omega", "memory", "--type", "project"]).unwrap(); + assert!(matches!( + cli.command, + Some(Command::Memory(MemoryArgs { + mem_type: MemoryType::Project + })) + )); + + let cli = Cli::try_parse_from(["Omega", "security", "--mode", "strict"]).unwrap(); + assert!(matches!( + cli.command, + Some(Command::Security(SecurityArgs { + mode: SecurityMode::Strict + })) + )); + + let cli = Cli::try_parse_from(["Omega", "performance", "--mode", "quality"]).unwrap(); + assert!(matches!( + cli.command, + Some(Command::Performance(PerformanceArgs { + mode: PerformanceMode::Quality + })) + )); + } + + #[test] + fn workspace_resolves_relative_and_absolute() { + let abs = Cli::try_parse_from(["Omega", "C:/some/proj"]).unwrap(); + assert_eq!(abs.workspace().unwrap(), PathBuf::from("C:/some/proj")); + + let rel = Cli::try_parse_from(["Omega", "proj"]).unwrap(); + let resolved = rel.workspace().unwrap(); + assert!(resolved.is_absolute()); + assert!(resolved.ends_with("proj")); + } +} diff --git a/src/components/welcome.rs b/src/components/welcome.rs index ca572b9..653f164 100644 --- a/src/components/welcome.rs +++ b/src/components/welcome.rs @@ -1,4 +1,4 @@ -use ratatui_kit::{Component, ComponentDrawer, ComponentUpdater, Hooks, Props, UseEvents, UseFuture, UseState}; +use ratatui_kit::{Component, ComponentDrawer, ComponentUpdater, Hooks, Props, UseEvents, UseExit, UseFuture, UseRouter, UseState}; use ratatui_kit::crossterm::event::{Event, KeyCode, KeyEventKind}; @@ -10,6 +10,7 @@ use ratatui_kit::ratatui::{ use crate::runtime::download::DownloadProgress; use crate::runtime::extract::ExtractProgress; use crate::runtime::runtime::{ensure_jre, ensure_python}; +use crate::services::welcome as panel_service; #[derive(Props)] pub struct WelcomeProps { @@ -37,6 +38,8 @@ pub struct Welcome { menu_style: Style, version: Option, show_menu: bool, + /// 菜单模式下的状态提示（如打开控制台面板的结果）。 + menu_message: String, jre_ready: bool, python_ready: bool, @@ -61,6 +64,7 @@ impl Component for Welcome { menu_style: props.menu_style, version: props.version.clone(), show_menu: props.show_menu, + menu_message: String::new(), jre_ready: false, python_ready: false, @@ -98,6 +102,12 @@ impl Component for Welcome { let mut speed = hooks.use_state(|| self.speed); let mut state_index = hooks.use_state(|| self.selected_index); + let menu_message = hooks.use_state(|| self.menu_message.clone()); + + // 导航与退出句柄：须在条件性 use_future 之前无条件获取，以保持钩子调用顺序稳定。 + let mut navigate = hooks.use_navigate(); + let mut exit = hooks.use_exit(); + hooks.use_events(move |event| { if jre_ready.get() && python_ready.get() { if let Event::Key(key) = event { @@ -119,9 +129,21 @@ impl Component for Welcome { } KeyCode::Enter => match state_index.get() { - 0 => todo!("打开配置"), - 1 => todo!("开启新项目"), - 2 => todo!("退出程序"), + // 打开控制台面板：后台启动 java -jar 并在浏览器打开（缺 jar/JRE 时提示）。 + 0 => { + let mut msg = menu_message.clone(); + std::thread::spawn(move || { + let text = match panel_service::open_console_panel() { + Ok(outcome) => outcome.message(), + Err(e) => format!("cannot open console panel: {e}"), + }; + msg.set(text); + }); + } + // 开启新项目：进入聊天页。 + 1 => navigate.push("/chat"), + // 退出程序。 + 2 => exit(), _ => {} }, @@ -229,6 +251,7 @@ impl Component for Welcome { self.show_menu = jre_ready.get() && python_ready.get(); self.selected_index = state_index.get(); + self.menu_message = menu_message.read().clone(); self.stage = stage.read().clone(); self.percent = percent.get(); @@ -343,6 +366,14 @@ impl Component for Welcome { style, ))); } + + if !self.menu_message.is_empty() { + lines.push(Line::from("")); + lines.push(Line::from(Span::styled( + self.menu_message.clone(), + self.style.add_modifier(Modifier::DIM), + ))); + } } Paragraph::new(lines).render(area, buf); diff --git a/src/core/agent/mod.rs b/src/core/agent/mod.rs index 22a13f3..b2586d9 100644 --- a/src/core/agent/mod.rs +++ b/src/core/agent/mod.rs @@ -16,7 +16,9 @@ //! while let Some(event) = rx.recv().await { /* ... */ } //! ``` +use std::path::PathBuf; use std::sync::{Arc, Mutex}; +use std::time::Instant; use anyhow::{Context, Result}; use futures_util::{Stream, StreamExt}; @@ -25,9 +27,12 @@ use rusqlite::{params, Connection, OptionalExtension}; use serde_json::{json, Value}; use tokio::sync::mpsc; -use crate::core::config::AppConfig; +use crate::core::config::{parse_provider_kind, AppConfig, DEFAULT_SYSTEM_PROMPT}; +use crate::core::mcp::{self, McpToolSpec}; +use crate::core::memory::{MemoryExtractor, MemoryScope, MemoryStore}; +use crate::core::prompt; use crate::core::providers::{ - create_provider, CompletionRequest, Message, Provider, Role, + create_provider, CompletionRequest, Message, Provider, ProviderConfig, ProviderKind, Role, }; use crate::core::runtime::Runtime; use crate::core::tools::ToolRegistry; @@ -100,6 +105,57 @@ pub struct ModelRef { pub model_name: Option, } +/// 从数据库加载的 Agent 完整配置（agent_info + llm_model + ai_provider + prompt_info + agent_mcp_rel）。 +#[derive(Debug, Clone)] +pub struct DbAgentConfig { + /// Agent 主键。 + pub agent_id: i32, + /// Agent 名称。 + pub agent_name: String, + /// 绑定的 LLM 模型 ID（`agent_info.llm_id`）。 + pub llm_id: i32, + /// 绑定的 Prompt ID（`agent_info.prompt_id`，用于加载自定义占位符）。 + pub prompt_id: Option, + /// 供应商名称（用于推断 [`ProviderKind`]）。 + pub provider_name: String, + /// 供应商 API Base URL。 + pub base_url: String, + /// 供应商 API 密钥。 + pub api_key: Option, + /// 模型调用标识（`llm_model.model_code`）。 + pub model_code: String, + /// 单次输出上限（`llm_model.max_response_tokens`，0 表示不限）。 + pub max_response_tokens: i32, + /// 系统提示词（`prompt_info.content`，可能为空）。 + pub system_prompt: String, + /// 采样温度。 + pub temperature: Option, + /// 核采样 top_p。 + pub top_p: Option, + /// 授权的 MCP 工具 ID 列表（`agent_mcp_rel`）。 + pub mcp_tool_ids: Vec, +} + +/// 待持久化的一条消息记录（对应 `chat_message` 新 schema）。 +pub struct MessageRecord<'a> { + /// 会话 UUID，关联 `chat_session.chat_uuid`。 + pub chat_uuid: &'a str, + /// 单轮消息 UUID，同一轮 user/assistant/tool 共享。 + pub round_uuid: Option<&'a str>, + /// 消息角色：user/assistant/system/tool。 + pub role: &'a str, + /// 消息正文。 + pub content: &'a str, + /// 单轮输入 Token 消耗。 + pub input_tokens: i32, + /// 单轮输出 Token 消耗。 + pub output_tokens: i32, + /// 调用耗时（毫秒）。 + pub duration: i32, + /// 模型参数快照（JSON）。 + pub model_config: Option<&'a str>, +} + /// 会话存储：把对话消息持久化到 SQLite。 /// /// 内部使用 [`Mutex`] 保证 `Send + Sync`，可安全地被 Agent 在异步任务中共享。 @@ -124,6 +180,14 @@ impl SessionStore { .map_err(|_| anyhow::anyhow!("SessionStore mutex poisoned")) } + /// 测试专用：用已有连接构造（例如内存库）。 + #[cfg(test)] + pub fn from_connection(conn: Connection) -> Self { + Self { + conn: Mutex::new(conn), + } + } + /// 确保 `ai_provider`→`llm_model` 记录存在，返回 `llm_model` 的 id。 /// /// 供应商按 `name` 匹配、模型按 `(provider_id, model_code)` 匹配；缺失时以真实配置创建， @@ -144,7 +208,7 @@ impl SessionStore { None => { conn.execute( "INSERT INTO ai_provider (name, avatar, api, key, timeout, status) \ - VALUES (?1, '', ?2, ?3, 5000, 1)", + VALUES (?1, NULL, ?2, ?3, 30, 1)", params![m.provider_name, m.provider_api, m.provider_key], )?; conn.last_insert_rowid() as i32 @@ -166,7 +230,7 @@ impl SessionStore { None => { conn.execute( "INSERT INTO llm_model (provider_id, name, avatar, model_code, max_tokens, max_response_tokens, weight, call_count, status) \ - VALUES (?1, ?2, '', ?3, 0, 0, 0, 0, 1)", + VALUES (?1, ?2, NULL, ?3, 0, 0, 0, 0, 1)", params![provider_id, model_name, m.model_code], )?; conn.last_insert_rowid() as i32 @@ -180,24 +244,29 @@ impl SessionStore { pub fn create_session(&self, uuid: &str, title: Option<&str>, llm_id: i32) -> Result<()> { let conn = self.lock()?; conn.execute( - "INSERT INTO chat_session (chat_uuid, title, llm_id, status) VALUES (?1, ?2, ?3, 1)", + "INSERT INTO chat_session (chat_uuid, title, llm_id, is_deleted) VALUES (?1, ?2, ?3, 0)", params![uuid, title, llm_id], )?; Ok(()) } - /// 保存一条消息。 - pub fn save_message( - &self, - session_uuid: &str, - role: &str, - content: &str, - model: Option<&str>, - ) -> Result<()> { + /// 保存一条消息（`chat_message` 新 schema：按 `round_uuid` 分轮，记录 Token/耗时/模型快照）。 + pub fn save_message(&self, rec: &MessageRecord) -> Result<()> { let conn = self.lock()?; conn.execute( - "INSERT INTO chat_message (session_id, role, content, model, status) VALUES (?1, ?2, ?3, ?4, 1)", - params![session_uuid, role, content, model], + "INSERT INTO chat_message \ + (chat_uuid, round_uuid, role, content, input_tokens, output_tokens, duration, model_config) \ + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + params![ + rec.chat_uuid, + rec.round_uuid, + rec.role, + rec.content, + rec.input_tokens, + rec.output_tokens, + rec.duration, + rec.model_config, + ], )?; Ok(()) } @@ -206,7 +275,7 @@ impl SessionStore { pub fn load_history(&self, session_uuid: &str) -> Result> { let conn = self.lock()?; let mut stmt = conn.prepare( - "SELECT role, content FROM chat_message WHERE session_id = ?1 ORDER BY id ASC", + "SELECT role, content FROM chat_message WHERE chat_uuid = ?1 ORDER BY id ASC", )?; let rows = stmt.query_map(params![session_uuid], |row| { @@ -228,6 +297,185 @@ impl SessionStore { Ok(history) } + + /// 从数据库加载某个 Agent 的完整配置。 + /// + /// 关联查询 `agent_info`→`llm_model`→`ai_provider` 与 `prompt_info`， + /// 并读取 `agent_mcp_rel` 得到授权的 MCP 工具 ID 列表。任一主链缺失时报错。 + pub fn load_agent_config(&self, agent_id: i32) -> Result { + let conn = self.lock()?; + + let row = conn + .query_row( + "SELECT a.name, a.llm_id, a.prompt_id, a.temperature, a.top_p, \ + m.model_code, m.max_response_tokens, \ + p.name, p.api, p.key \ + FROM agent_info a \ + JOIN llm_model m ON m.id = a.llm_id \ + JOIN ai_provider p ON p.id = m.provider_id \ + WHERE a.id = ?1", + params![agent_id], + |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, i32>(1)?, + row.get::<_, Option>(2)?, + row.get::<_, Option>(3)?, + row.get::<_, Option>(4)?, + row.get::<_, String>(5)?, + row.get::<_, i32>(6)?, + row.get::<_, String>(7)?, + row.get::<_, String>(8)?, + row.get::<_, Option>(9)?, + )) + }, + ) + .optional()? + .with_context(|| { + format!( + "agent_info id={} not found or its llm_model/ai_provider is missing", + agent_id + ) + })?; + + let ( + agent_name, + llm_id, + prompt_id, + temperature, + top_p, + model_code, + max_response_tokens, + provider_name, + base_url, + api_key, + ) = row; + + // 提示词内容（可选）。 + let system_prompt = match prompt_id { + Some(pid) => conn + .query_row( + "SELECT content FROM prompt_info WHERE id = ?1", + params![pid], + |row| row.get::<_, String>(0), + ) + .optional()? + .unwrap_or_default(), + None => String::new(), + }; + + // 授权的 MCP 工具 ID。 + let mut stmt = + conn.prepare("SELECT mcp_tool_id FROM agent_mcp_rel WHERE agent_id = ?1")?; + let mcp_tool_ids: Vec = stmt + .query_map(params![agent_id], |row| row.get::<_, i32>(0))? + .filter_map(|r| r.ok()) + .collect(); + + Ok(DbAgentConfig { + agent_id, + agent_name, + llm_id, + prompt_id, + provider_name, + base_url, + api_key, + model_code, + max_response_tokens, + system_prompt, + temperature, + top_p, + mcp_tool_ids, + }) + } + + /// 加载某个 Agent 已授权且启用的 MCP 工具规格（用于工具调度）。 + /// + /// 关联 `agent_mcp_rel`→`mcp_tool`→`mcp_provider`，仅返回工具与提供方均处于 + /// 启用状态（`status = 1`）的记录，从而在执行前完成权限与可用性过滤。 + pub fn load_agent_mcp_tools(&self, agent_id: i32) -> Result> { + let conn = self.lock()?; + let mut stmt = conn.prepare( + "SELECT t.id, t.tool_code, t.name, t.description, t.parameters, p.url \ + FROM agent_mcp_rel r \ + JOIN mcp_tool t ON t.id = r.mcp_tool_id \ + JOIN mcp_provider p ON p.id = t.provider_id \ + WHERE r.agent_id = ?1 AND t.status = 1 AND p.status = 1", + )?; + + let specs = stmt + .query_map(params![agent_id], |row| { + Ok(McpToolSpec { + tool_id: row.get(0)?, + tool_code: row.get(1)?, + name: row.get(2)?, + description: row.get(3)?, + parameters: row.get(4)?, + provider_url: row.get(5)?, + }) + })? + .filter_map(|r| r.ok()) + .collect(); + + Ok(specs) + } + + /// 加载可被当前 Agent 调用的其它启用 Agent 名称（用于 `{AgentList}` 占位符）。 + /// + /// 系统未建立 agent-agent 关联表，故取「除自身外所有 `status = 1` 的 Agent」作为合理近似。 + pub fn load_callable_agents(&self, exclude_id: i32) -> Result> { + let conn = self.lock()?; + let mut stmt = conn + .prepare("SELECT name FROM agent_info WHERE status = 1 AND id != ?1 ORDER BY id ASC")?; + let names = stmt + .query_map(params![exclude_id], |row| row.get::<_, String>(0))? + .filter_map(|r| r.ok()) + .collect(); + Ok(names) + } + + /// 加载某个 Prompt 的启用占位符（`key`, `value`, `default_value`），用于自定义占位符解析。 + pub fn load_prompt_placeholders( + &self, + prompt_id: i32, + ) -> Result, Option)>> { + let conn = self.lock()?; + let mut stmt = conn.prepare( + "SELECT key, value, default_value FROM prompt_placeholder \ + WHERE prompt_id = ?1 AND status = 1", + )?; + let rows = stmt + .query_map(params![prompt_id], |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, Option>(1)?, + row.get::<_, Option>(2)?, + )) + })? + .filter_map(|r| r.ok()) + .collect(); + Ok(rows) + } + + /// 记录一次外部 API（大模型）调用日志到 `system_api_log`。 + /// + /// `rel_uuid` 关联业务实体（如 `chat_message.round_uuid`），`status` 为 1 成功 / 0 失败。 + pub fn log_api( + &self, + rel_uuid: Option<&str>, + status: i8, + token_cost: i32, + duration: i32, + content: Option<&str>, + ) -> Result<()> { + let conn = self.lock()?; + conn.execute( + "INSERT INTO system_api_log (rel_uuid, status, token_cost, duration, content) \ + VALUES (?1, ?2, ?3, ?4, ?5)", + params![rel_uuid, status, token_cost, duration, content], + )?; + Ok(()) + } } // ============================================================ @@ -248,6 +496,22 @@ pub struct Agent { provider_name: String, provider_api: String, provider_key: Option, + /// 采样温度（来自 Agent 配置，作用于每次请求）。 + temperature: Option, + /// 核采样 top_p（记录于模型快照；当前请求体未透传，留待 provider 扩展）。 + top_p: Option, + /// 单次输出上限。 + max_tokens: Option, + /// 数据库驱动的 Agent 主键（`from_database` 时设置）。 + agent_id: Option, + /// 数据库驱动的 LLM 模型主键（`start_session` 建会话时使用）。 + llm_id: Option, + /// 已渲染的记忆上下文（会话初始化时一次性注入系统提示词）。 + memory_context: String, + /// 记忆抽取器（作用域非 None 时启用，每轮对话后异步抽取记忆）。 + extractor: Option>, + /// 最近一轮对话的 round_uuid（供 TUI 内投票关联本轮回复）。 + last_round_uuid: Option, } impl Agent { @@ -270,6 +534,14 @@ impl Agent { provider_name: String::new(), provider_api: String::new(), provider_key: None, + temperature: None, + top_p: None, + max_tokens: None, + agent_id: None, + llm_id: None, + memory_context: String::new(), + extractor: None, + last_round_uuid: None, } } @@ -291,7 +563,8 @@ impl Agent { config.system_prompt.clone(), tools, ) - .with_provider_meta(config.provider_kind.as_str(), config.base_url.clone(), key)) + .with_provider_meta(config.provider_kind.as_str(), config.base_url.clone(), key) + .with_sampling(config.temperature, None, config.max_tokens)) } /// 设置供应商元数据（用于持久化时建立 provider→model 关联）。 @@ -307,6 +580,158 @@ impl Agent { self } + /// 设置采样参数（温度 / top_p / 单次输出上限）。 + pub fn with_sampling( + mut self, + temperature: Option, + top_p: Option, + max_tokens: Option, + ) -> Self { + self.temperature = temperature; + self.top_p = top_p; + self.max_tokens = max_tokens; + self + } + + /// 基于数据库配置构造 Agent（数据库驱动）。 + /// + /// 从 `agent_info`+`llm_model`+`ai_provider`+`prompt_info`+`agent_mcp_rel` 加载完整配置， + /// 据此创建供应商、系统提示词与采样参数，并绑定持久化存储（会话由 [`Agent::start_session`] 创建）。 + /// `mcp_tool_ids` 供后续 MCP 工具调度做权限校验（见 `crate::core::tools`）。 + pub fn from_database(agent_id: i32, working_dir: PathBuf) -> Result { + let store = SessionStore::new()?; + let cfg = store.load_agent_config(agent_id)?; + // 已授权且启用的 MCP 工具（权限校验：仅 agent_mcp_rel 中登记且状态启用者可用）。 + let mcp_specs = store.load_agent_mcp_tools(agent_id)?; + // 先收集工具名称用于 `{MCPList}` 占位符，再交给 build_tools 消费 mcp_specs。 + let mcp_names: Vec = mcp_specs.iter().map(|s| s.name.clone()).collect(); + + let provider_kind = + parse_provider_kind(&cfg.provider_name).unwrap_or(ProviderKind::Compatible); + let provider = create_provider(ProviderConfig { + kind: provider_kind, + api_key: cfg.api_key.clone().unwrap_or_default(), + base_url: cfg.base_url.clone(), + model: cfg.model_code.clone(), + })?; + + let runtime = Arc::new(Runtime::new(working_dir)); + let mut tools = ToolRegistry::with_defaults(runtime); + for tool in mcp::build_tools(mcp_specs) { + tools.register(tool); + } + + let raw_prompt = if cfg.system_prompt.trim().is_empty() { + DEFAULT_SYSTEM_PROMPT.to_string() + } else { + cfg.system_prompt.clone() + }; + + // 解析系统提示词占位符：`{AgentList}`=可调用 Agent，`{MCPList}`=已授权工具， + // 其余自定义占位符来自 prompt_placeholder（加载失败不阻断，退化为空）。 + let callable = store.load_callable_agents(cfg.agent_id).unwrap_or_default(); + let placeholders = cfg + .prompt_id + .map(|pid| store.load_prompt_placeholders(pid).unwrap_or_default()) + .unwrap_or_default(); + let vars = prompt::build_vars(&callable, &mcp_names, &placeholders); + let system_prompt = prompt::render(&raw_prompt, &vars); + + let max_tokens = if cfg.max_response_tokens > 0 { + Some(cfg.max_response_tokens as u32) + } else { + None + }; + + let mut agent = Self::new(provider, cfg.model_code.clone(), system_prompt, tools) + .with_provider_meta(&cfg.provider_name, &cfg.base_url, cfg.api_key.clone()) + .with_sampling(cfg.temperature, cfg.top_p, max_tokens); + agent.agent_id = Some(cfg.agent_id); + agent.llm_id = Some(cfg.llm_id); + agent.store = Some(Arc::new(store)); + Ok(agent) + } + + /// 将该 Agent 升级为多 Agent 调度器（Coordinator）。 + /// + /// 复用自身供应商/模型配置构造 [`crate::core::scheduler::AgentWorkerFactory`]，注册四个 + /// 调度工具（create/resume/stop/await worker），追加协调者提示词，并放宽工具迭代上限 + /// （单工具/轮 + 默认 8 次不足以协调多 Worker）。仅用于被标记为 `coordinator_agent_id` + /// 的 Agent（见 `src/pages/chat.rs`）。 + pub fn enable_coordinator(&mut self, working_dir: PathBuf) -> Result<()> { + use crate::core::scheduler::{self, AgentWorkerFactory, WorkerPool}; + + // `provider_api` 即 base_url（见 `with_provider_meta`）。 + let provider_config = ProviderConfig { + kind: parse_provider_kind(&self.provider_name).unwrap_or(ProviderKind::Compatible), + api_key: self.provider_key.clone().unwrap_or_default(), + base_url: self.provider_api.clone(), + model: self.model.clone(), + }; + let factory = Arc::new(AgentWorkerFactory { + provider_config, + working_dir, + temperature: self.temperature, + top_p: self.top_p, + max_tokens: self.max_tokens, + }); + let pool = Arc::new(WorkerPool::default()); + for tool in scheduler::build_tools(pool, factory) { + self.tools.register(tool); + } + self.system_prompt.push_str(scheduler::COORDINATOR_PROMPT); + self.max_iterations = self.max_iterations.max(24); + Ok(()) + } + + /// 为该 Agent 启用记忆系统。 + /// + /// 立即渲染当前作用域下的记忆上下文（一次性注入系统提示词，见 [`Agent::build_system_prompt`]）； + /// 若作用域非 [`MemoryScope::None`]，则复用自身供应商/模型配置构造 [`MemoryExtractor`]， + /// 在每轮对话结束后异步抽取并落库记忆（见 `run_turn`）。记忆上下文不随本轮抽取刷新， + /// 新记忆下次会话生效（符合「黑盒自动管理」，避免每轮读文件）。 + pub fn enable_memory(&mut self, store: Arc, scope: MemoryScope) -> Result<()> { + self.memory_context = store.render_context(scope)?; + if scope != MemoryScope::None { + // `provider_api` 即 base_url（见 `with_provider_meta`）。 + let provider_config = ProviderConfig { + kind: parse_provider_kind(&self.provider_name).unwrap_or(ProviderKind::Compatible), + api_key: self.provider_key.clone().unwrap_or_default(), + base_url: self.provider_api.clone(), + model: self.model.clone(), + }; + self.extractor = Some(Arc::new(MemoryExtractor::new( + provider_config, + self.model.clone(), + self.temperature, + store, + scope, + ))); + } + Ok(()) + } + + /// 数据库驱动 Agent 的数据库配置主键（若非数据库驱动则为 `None`）。 + pub fn agent_id(&self) -> Option { + self.agent_id + } + + /// 使用已绑定的存储创建新会话（数据库驱动场景，`llm_id` 已知）。 + /// + /// 若尚未绑定存储则自动打开一个。会话标题缺省时用 Agent 名称或 "Chat"。 + pub fn start_session(&mut self, title: Option<&str>) -> Result<()> { + let store = match &self.store { + Some(s) => s.clone(), + None => Arc::new(SessionStore::new()?), + }; + let llm_id = self.llm_id.unwrap_or(0); + let uuid = generate_uuid(); + store.create_session(&uuid, title, llm_id)?; + self.store = Some(store); + self.session_uuid = uuid; + Ok(()) + } + /// 设置最大工具调用迭代次数。 pub fn with_max_iterations(mut self, max_iterations: usize) -> Self { self.max_iterations = max_iterations.max(1); @@ -324,6 +749,11 @@ impl Agent { &self.session_uuid } + /// 最近一轮对话的 round_uuid（尚无完成轮次时为 `None`）。 + pub fn last_round_uuid(&self) -> Option<&str> { + self.last_round_uuid.as_deref() + } + /// 当前对话历史。 pub fn history(&self) -> &[Message] { &self.history @@ -382,7 +812,12 @@ impl Agent { role: Role::User, content: user_input.to_string(), }); - self.persist("user", user_input, None); + // 本轮对话共享同一个 round_uuid，用于在 chat_message 中聚合单轮消息。 + let round_uuid = generate_uuid(); + self.last_round_uuid = Some(round_uuid.clone()); + let round = round_uuid.as_str(); + let input_tokens = estimate_tokens(user_input); + self.persist(Some(round), "user", user_input, input_tokens, 0, 0); let mut iteration = 0; loop { @@ -395,18 +830,25 @@ impl Agent { }); messages.extend(self.history.iter().cloned()); + // 请求侧 Token 估算（流式无真实 usage，用字符数近似）。 + let req_tokens: i32 = messages.iter().map(|m| estimate_tokens(&m.content)).sum(); + let request = CompletionRequest { model: self.model.clone(), messages, - temperature: None, - max_tokens: None, + temperature: self.temperature, + max_tokens: self.max_tokens, stream: true, }; + let started = Instant::now(); let mut stream = match self.provider.complete_stream(request).await { Ok(stream) => stream, Err(e) => { - let _ = tx.send(AgentEvent::Error(e.to_string())); + let elapsed = started.elapsed().as_millis().min(i32::MAX as u128) as i32; + let msg = e.to_string(); + self.persist_api_log(Some(round), 0, 0, elapsed, Some(&msg)); + let _ = tx.send(AgentEvent::Error(msg)); return Err(e); } }; @@ -421,17 +863,23 @@ impl Agent { } } Err(e) => { - let _ = tx.send(AgentEvent::Error(e.to_string())); + let elapsed = started.elapsed().as_millis().min(i32::MAX as u128) as i32; + let msg = e.to_string(); + self.persist_api_log(Some(round), 0, req_tokens, elapsed, Some(&msg)); + let _ = tx.send(AgentEvent::Error(msg)); return Err(e); } } } + let duration = started.elapsed().as_millis().min(i32::MAX as u128) as i32; + let output_tokens = estimate_tokens(&accumulated); + self.persist_api_log(Some(round), 1, req_tokens + output_tokens, duration, None); self.history.push(Message { role: Role::Assistant, content: accumulated.clone(), }); - self.persist("assistant", &accumulated, Some(self.model.as_str())); + self.persist(Some(round), "assistant", &accumulated, 0, output_tokens, duration); // 若模型请求调用工具且尚未超出迭代上限，则执行并把结果回灌。 if iteration < self.max_iterations { @@ -457,11 +905,33 @@ impl Agent { role: Role::User, content: tool_message.clone(), }); - self.persist("tool", &tool_message, None); + self.persist(Some(round), "tool", &tool_message, 0, 0, 0); continue; } } + // 一轮结束后异步抽取记忆（作用域非 None 且有助手输出时）。 + if !accumulated.trim().is_empty() { + if let Some(extractor) = &self.extractor { + let ex = extractor.clone(); + let round_messages = vec![ + Message { + role: Role::User, + content: user_input.to_string(), + }, + Message { + role: Role::Assistant, + content: accumulated.clone(), + }, + ]; + tokio::spawn(async move { + if let Err(e) = ex.extract_from_round(&round_messages).await { + warn!("memory extraction failed: {}", e); + } + }); + } + } + let _ = tx.send(AgentEvent::Done); break; } @@ -482,24 +952,96 @@ impl Agent { tokio_stream::wrappers::UnboundedReceiverStream::new(rx) } - /// 组合系统提示词与工具说明。 + /// 组合系统提示词、记忆上下文与工具说明。 fn build_system_prompt(&self) -> String { let tools = self.tools.definitions_prompt(); - if tools.is_empty() { - self.system_prompt.clone() - } else { - format!("{}\n\n{}", self.system_prompt, tools) + let mut prompt = self.system_prompt.clone(); + if !self.memory_context.trim().is_empty() { + prompt.push_str("\n\n"); + prompt.push_str(&self.memory_context); } + if !tools.is_empty() { + prompt.push_str("\n\n"); + prompt.push_str(&tools); + } + prompt } /// 持久化一条消息（若已开启持久化）。 - fn persist(&self, role: &str, content: &str, model: Option<&str>) { + fn persist( + &self, + round_uuid: Option<&str>, + role: &str, + content: &str, + input_tokens: i32, + output_tokens: i32, + duration: i32, + ) { if let Some(store) = &self.store { - if let Err(e) = store.save_message(&self.session_uuid, role, content, model) { + // 仅助手消息记录模型参数快照。 + let snapshot = if role == "assistant" { + Some(self.model_config_snapshot()) + } else { + None + }; + let rec = MessageRecord { + chat_uuid: &self.session_uuid, + round_uuid, + role, + content, + input_tokens, + output_tokens, + duration, + model_config: snapshot.as_deref(), + }; + if let Err(e) = store.save_message(&rec) { warn!("Failed to persist '{}' message: {}", role, e); } } } + + /// 记录一次主对话 LLM 调用日志到 `system_api_log`（若已开启持久化）。 + /// + /// 日志失败仅告警，不影响对话主流程。仅记录主对话调用（记忆抽取、MCP 工具调用不计）。 + fn persist_api_log( + &self, + rel_uuid: Option<&str>, + status: i8, + token_cost: i32, + duration: i32, + content: Option<&str>, + ) { + if let Some(store) = &self.store { + if let Err(e) = store.log_api(rel_uuid, status, token_cost, duration, content) { + warn!("Failed to log api call: {}", e); + } + } + } + + /// 构建模型参数快照（JSON），记录本轮实际使用的模型与关键参数。 + fn model_config_snapshot(&self) -> String { + json!({ + "model": self.model, + "provider": self.provider_name, + "temperature": self.temperature, + "top_p": self.top_p, + "max_tokens": self.max_tokens, + "max_iterations": self.max_iterations, + }) + .to_string() + } +} + +/// 粗略估算文本 Token 数（约每 4 个字符 1 个 token）。 +/// +/// 流式增量不携带 usage 统计，此处用字符数近似，保证 `chat_message` +/// 的 Token 列有可用值；若后续 provider 暴露真实 usage 可直接替换。 +fn estimate_tokens(text: &str) -> i32 { + let chars = text.chars().count(); + if chars == 0 { + return 0; + } + ((chars + 3) / 4) as i32 } /// 生成一个会话唯一标识（无外部 uuid 依赖）。 @@ -535,4 +1077,235 @@ mod tests { fn uuid_is_unique_enough() { assert_ne!(generate_uuid(), generate_uuid()); } + + #[test] + fn estimate_tokens_approximates_by_chars() { + assert_eq!(estimate_tokens(""), 0); + assert_eq!(estimate_tokens("abcd"), 1); + assert_eq!(estimate_tokens("abcde"), 2); + } + + /// 构造含完整关联链的内存库，验证 `load_agent_config` 的 JOIN 查询。 + #[test] + fn load_agent_config_joins_provider_model_prompt() { + use crate::core::db::Model; + use crate::db::models::agent_info::AgentInfo; + use crate::db::models::agent_mcp_rel::AgentMcpRel; + use crate::db::models::ai_provider::AiProvider; + use crate::db::models::llm_model::LLMModel; + use crate::db::models::prompt_info::PromptInfo; + + let conn = Connection::open_in_memory().unwrap(); + conn.execute_batch(AiProvider::create_table_sql()).unwrap(); + conn.execute_batch(LLMModel::create_table_sql()).unwrap(); + conn.execute_batch(PromptInfo::create_table_sql()).unwrap(); + conn.execute_batch(AgentInfo::create_table_sql()).unwrap(); + conn.execute_batch(AgentMcpRel::create_table_sql()).unwrap(); + + conn.execute( + "INSERT INTO ai_provider (id, name, api, key, timeout, status) VALUES (1, 'openai', 'https://api.openai.com/v1', 'sk-x', 30, 1)", + [], + ).unwrap(); + conn.execute( + "INSERT INTO llm_model (id, provider_id, name, model_code, max_response_tokens, status) VALUES (10, 1, 'GPT-4o', 'gpt-4o', 4096, 1)", + [], + ).unwrap(); + conn.execute( + "INSERT INTO prompt_info (id, name, content) VALUES (100, 'sys', 'You are helpful.')", + [], + ).unwrap(); + conn.execute( + "INSERT INTO agent_info (id, llm_id, name, prompt_id, temperature, top_p, status) VALUES (5, 10, 'Coder', 100, 0.3, 0.9, 1)", + [], + ).unwrap(); + conn.execute( + "INSERT INTO agent_mcp_rel (agent_id, mcp_tool_id) VALUES (5, 77)", + [], + ).unwrap(); + conn.execute( + "INSERT INTO agent_mcp_rel (agent_id, mcp_tool_id) VALUES (5, 88)", + [], + ).unwrap(); + + let store = SessionStore::from_connection(conn); + let cfg = store.load_agent_config(5).unwrap(); + assert_eq!(cfg.agent_name, "Coder"); + assert_eq!(cfg.llm_id, 10); + assert_eq!(cfg.provider_name, "openai"); + assert_eq!(cfg.base_url, "https://api.openai.com/v1"); + assert_eq!(cfg.api_key.as_deref(), Some("sk-x")); + assert_eq!(cfg.model_code, "gpt-4o"); + assert_eq!(cfg.max_response_tokens, 4096); + assert_eq!(cfg.system_prompt, "You are helpful."); + assert_eq!(cfg.temperature, Some(0.3)); + assert_eq!(cfg.top_p, Some(0.9)); + assert_eq!(cfg.mcp_tool_ids, vec![77, 88]); + + // 不存在的 agent 应报错。 + assert!(store.load_agent_config(999).is_err()); + } + + /// 验证 `load_agent_mcp_tools` 的三表 JOIN 与启用状态过滤。 + #[test] + fn load_agent_mcp_tools_filters_disabled() { + use crate::core::db::Model; + use crate::db::models::agent_mcp_rel::AgentMcpRel; + use crate::db::models::mcp_provider::McpProvider; + use crate::db::models::mcp_tool::McpTool; + + let conn = Connection::open_in_memory().unwrap(); + conn.execute_batch(McpProvider::create_table_sql()).unwrap(); + conn.execute_batch(McpTool::create_table_sql()).unwrap(); + conn.execute_batch(AgentMcpRel::create_table_sql()).unwrap(); + + // 两个 provider：一个启用、一个停用。 + conn.execute( + "INSERT INTO mcp_provider (id, name, url, timeout, status) VALUES (1, 'py', 'http://localhost:8000/mcp', 30, 1)", + [], + ).unwrap(); + conn.execute( + "INSERT INTO mcp_provider (id, name, url, timeout, status) VALUES (2, 'off', 'http://localhost:9000/mcp', 30, 0)", + [], + ).unwrap(); + // 工具：add（启用）、echo（停用）、calc（启用但归属停用 provider）。 + conn.execute( + "INSERT INTO mcp_tool (id, provider_id, name, tool_code, description, parameters, status) VALUES (77, 1, 'Add', 'add', 'add two', '{\"type\":\"object\"}', 1)", + [], + ).unwrap(); + conn.execute( + "INSERT INTO mcp_tool (id, provider_id, name, tool_code, status) VALUES (78, 1, 'Echo', 'echo', 0)", + [], + ).unwrap(); + conn.execute( + "INSERT INTO mcp_tool (id, provider_id, name, tool_code, status) VALUES (79, 2, 'Calc', 'calc', 1)", + [], + ).unwrap(); + // agent 5 授权全部三个工具。 + for tool_id in [77, 78, 79] { + conn.execute( + "INSERT INTO agent_mcp_rel (agent_id, mcp_tool_id) VALUES (5, ?1)", + params![tool_id], + ).unwrap(); + } + + let store = SessionStore::from_connection(conn); + let specs = store.load_agent_mcp_tools(5).unwrap(); + // 仅 add 满足「工具启用 且 provider 启用」。 + assert_eq!(specs.len(), 1); + assert_eq!(specs[0].tool_id, 77); + assert_eq!(specs[0].tool_code, "add"); + assert_eq!(specs[0].name, "Add"); + assert_eq!(specs[0].description.as_deref(), Some("add two")); + assert_eq!(specs[0].parameters.as_deref(), Some("{\"type\":\"object\"}")); + assert_eq!(specs[0].provider_url, "http://localhost:8000/mcp"); + + // 未授权的 agent 得到空集合。 + assert!(store.load_agent_mcp_tools(999).unwrap().is_empty()); + } + + /// 验证 `log_api` 写入 `system_api_log` 后可查得各列。 + #[test] + fn log_api_persists_row() { + use crate::core::db::Model; + use crate::db::models::system_api_log::SystemApiLog; + + let conn = Connection::open_in_memory().unwrap(); + conn.execute_batch(SystemApiLog::create_table_sql()).unwrap(); + let store = SessionStore::from_connection(conn); + + store + .log_api(Some("round-1"), 1, 42, 123, None) + .unwrap(); + store + .log_api(Some("round-2"), 0, 0, 5, Some("boom")) + .unwrap(); + + let conn = store.conn.lock().unwrap(); + let (status, token_cost, duration, rel_uuid): (i8, i32, i32, Option) = conn + .query_row( + "SELECT status, token_cost, duration, rel_uuid FROM system_api_log WHERE rel_uuid = ?1", + params!["round-1"], + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)), + ) + .unwrap(); + assert_eq!(status, 1); + assert_eq!(token_cost, 42); + assert_eq!(duration, 123); + assert_eq!(rel_uuid.as_deref(), Some("round-1")); + + let (status, content): (i8, Option) = conn + .query_row( + "SELECT status, content FROM system_api_log WHERE rel_uuid = ?1", + params!["round-2"], + |row| Ok((row.get(0)?, row.get(1)?)), + ) + .unwrap(); + assert_eq!(status, 0); + assert_eq!(content.as_deref(), Some("boom")); + } + + /// 验证 `load_callable_agents` 排除自身且仅返回 status=1。 + #[test] + fn load_callable_agents_excludes_self_and_disabled() { + use crate::core::db::Model; + use crate::db::models::agent_info::AgentInfo; + + let conn = Connection::open_in_memory().unwrap(); + conn.execute_batch(AgentInfo::create_table_sql()).unwrap(); + conn.execute( + "INSERT INTO agent_info (id, llm_id, name, status) VALUES (1, 1, 'Self', 1)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO agent_info (id, llm_id, name, status) VALUES (2, 1, 'Enabled', 1)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO agent_info (id, llm_id, name, status) VALUES (3, 1, 'Disabled', 0)", + [], + ) + .unwrap(); + + let store = SessionStore::from_connection(conn); + let names = store.load_callable_agents(1).unwrap(); + assert_eq!(names, vec!["Enabled".to_string()]); + } + + /// 验证 `load_prompt_placeholders` 仅返回 status=1 的行，且携带 value/default_value。 + #[test] + fn load_prompt_placeholders_filters_disabled() { + use crate::core::db::Model; + use crate::db::models::prompt_placeholder::PromptPlaceholder; + + let conn = Connection::open_in_memory().unwrap(); + conn.execute_batch(PromptPlaceholder::create_table_sql()) + .unwrap(); + conn.execute( + "INSERT INTO prompt_placeholder (prompt_id, key, value, default_value, status) \ + VALUES (100, 'City', 'Beijing', 'N/A', 1)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO prompt_placeholder (prompt_id, key, value, default_value, status) \ + VALUES (100, 'Off', 'x', 'y', 0)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO prompt_placeholder (prompt_id, key, value, default_value, status) \ + VALUES (200, 'Other', 'z', NULL, 1)", + [], + ) + .unwrap(); + + let store = SessionStore::from_connection(conn); + let rows = store.load_prompt_placeholders(100).unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].0, "City"); + assert_eq!(rows[0].1.as_deref(), Some("Beijing")); + assert_eq!(rows[0].2.as_deref(), Some("N/A")); + } } diff --git a/src/core/arena/mod.rs b/src/core/arena/mod.rs new file mode 100644 index 0000000..1eeed48 --- /dev/null +++ b/src/core/arena/mod.rs @@ -0,0 +1,556 @@ +//! 竞技场投票 +//! +//! 对应《系统功能设计与数据库设计文档》3.8 Agent 竞技场模块：用户可对单轮会话、 +//! 单 Agent 输出进行**星级打分（1-5）+ 文字评价**，评价数据落 `agent_arena`，并反向 +//! 聚合更新 `agent_info.avg_score`（0-5 均值）与 `call_count`。 +//! +//! 设计说明： +//! - 投票经 CLI（`Omega arena vote` / `arena list`）与 TUI（Ctrl+1..5）落库； +//! - 多维度评估（准确性/相关性/流畅性/有用性，各 1-5）序列化为 JSON 存入 +//! `agent_arena.remark`，`vote_type` 仍为综合星级，不改表结构； +//! - [`report`] 子模块提供 LLM 驱动的 Prompt 优化报告智能体（3.8.3 数据闭环）， +//! 经 `Omega arena report` 触发，聚合数据见 [`ArenaService::report_data`]。 + +pub mod report; + +use std::collections::BTreeMap; + +use anyhow::{bail, Context, Result}; +use chrono::Utc; +use rusqlite::params; + +use crate::core::db::DatabaseManager; +use crate::db::models::agent_arena::AgentArena; +use crate::db::models::agent_info::AgentInfo; + +/// 多维度评估定义：(短键, 中文标签)。序列化入 `agent_arena.remark` JSON。 +pub const DIMENSIONS: &[(&str, &str)] = &[ + ("acc", "准确性"), + ("rel", "相关性"), + ("flu", "流畅性"), + ("helpful", "有用性"), +]; + +/// 一次投票的输入。 +#[derive(Debug, Clone)] +pub struct VoteInput { + /// 被投票的 Agent ID（`agent_info.id`）。 + pub agent_id: i32, + /// 星级评分，取值 1-5。 + pub score: i8, + /// 文字评价（可选）。 + pub comment: Option, + /// 关联会话 UUID（`agent_arena.chat_uuid` NOT NULL，未提供时为空串）。 + pub chat_uuid: String, + /// 关联单轮消息 UUID（可选）。 + pub round_uuid: Option, + /// 多维度评分（短键 → 1-5），序列化存入 `agent_arena.remark`；空表示未评多维度。 + pub dimensions: BTreeMap, +} + +/// 排行榜中一个 Agent 的聚合评分。 +#[derive(Debug, Clone)] +pub struct AgentRank { + pub agent_id: i32, + pub name: String, + /// 平均星级评分（0-5，无有效投票时为 `None`）。 + pub avg_score: Option, + /// 有效投票数。 + pub vote_count: i64, +} + +/// 某 Agent 的评估聚合数据，供 Prompt 优化报告智能体分析。 +#[derive(Debug, Clone)] +pub struct ReportData { + pub agent_id: i32, + pub agent_name: String, + /// 当前系统提示词（`prompt_info.content`）。 + pub system_prompt: String, + /// 综合平均星级（0-5）。 + pub avg_score: Option, + /// 有效投票总数。 + pub vote_count: i64, + /// 星级分布，下标 0..4 对应 1..5 星的票数。 + pub distribution: [i64; 5], + /// 多维度均值：(维度键, 均值, 样本数)。 + pub dimensions: Vec<(String, f32, i64)>, + /// 近期文字评价（最多 50 条，按时间倒序）。 + pub comments: Vec, +} + +/// 竞技场服务：投票落库、评分聚合、排行榜。 +pub struct ArenaService { + db: DatabaseManager, +} + +impl ArenaService { + /// 打开应用数据库构造服务。 + pub fn new() -> Result { + Ok(Self { + db: DatabaseManager::new()?, + }) + } + + /// 测试专用：用已有 [`DatabaseManager`]（例如内存库）构造。 + #[cfg(test)] + fn from_db(db: DatabaseManager) -> Self { + Self { db } + } + + /// 记录一条星级投票，并同步重算该 Agent 的 `avg_score`、递增 `call_count`。 + /// + /// 返回新写入的 `agent_arena` 行 id。评分越界或 Agent 不存在时报错。 + pub fn record_vote(&self, v: VoteInput) -> Result { + if !(1..=5).contains(&v.score) { + bail!("score must be within 1-5"); + } + // 多维度评分（可选）：逐维校验后序列化为 remark JSON；空则 remark = None。 + let remark = if v.dimensions.is_empty() { + None + } else { + for (k, val) in &v.dimensions { + if !(1..=5).contains(val) { + bail!("dimension '{}' must be within 1-5", k); + } + } + Some(serde_json::to_string(&v.dimensions)?) + }; + let mut agent = self + .db + .read_by_id::(v.agent_id)? + .with_context(|| format!("agent #{} not found", v.agent_id))?; + + let row = AgentArena { + id: 0, + chat_uuid: v.chat_uuid, + round_uuid: v.round_uuid, + agent_id: v.agent_id, + vote_type: v.score, + comment: v.comment, + status: 1, + create_time: Utc::now(), + update_time: Utc::now(), + remark, + }; + let id = self.db.create(row)?; + + // 聚合重算平均分并累计调用次数（每次投票视作一次被评价的调用）。 + let (avg, _count) = self.score_stats(v.agent_id)?; + agent.avg_score = avg; + agent.call_count += 1; + self.db.update(agent)?; + + Ok(id) + } + + /// 计算某 Agent 的有效投票平均分（保留 1 位小数）与票数；无票返回 `(None, 0)`。 + fn score_stats(&self, agent_id: i32) -> Result<(Option, i64)> { + Ok(self.db.conn.query_row( + "SELECT AVG(vote_type), COUNT(*) FROM agent_arena WHERE agent_id = ?1 AND status = 1", + params![agent_id], + |row| { + let avg: Option = row.get(0)?; + let count: i64 = row.get(1)?; + Ok((avg.map(round1), count)) + }, + )?) + } + + /// 幂等重算并写回某 Agent 的 `avg_score`（供手动纠正）。返回重算后的均值。 + pub fn recompute(&self, agent_id: i32) -> Result> { + let mut agent = self + .db + .read_by_id::(agent_id)? + .with_context(|| format!("agent #{} not found", agent_id))?; + let (avg, _count) = self.score_stats(agent_id)?; + agent.avg_score = avg; + self.db.update(agent)?; + Ok(avg) + } + + /// 按平均评分降序返回全部 Agent 的排行榜（未评分者置底）。 + pub fn ranking(&self) -> Result> { + let mut stmt = self.db.conn.prepare( + "SELECT a.id, a.name, a.avg_score, \ + (SELECT COUNT(*) FROM agent_arena r WHERE r.agent_id = a.id AND r.status = 1) \ + FROM agent_info a \ + ORDER BY a.avg_score DESC, a.id ASC", + )?; + let rows = stmt.query_map([], |row| { + Ok(AgentRank { + agent_id: row.get(0)?, + name: row.get(1)?, + avg_score: row.get(2)?, + vote_count: row.get(3)?, + }) + })?; + + let mut out = Vec::new(); + for r in rows { + out.push(r?); + } + Ok(out) + } + + /// 汇总某 Agent 的评估数据（星级分布、多维度均值、文字评价、当前提示词）， + /// 供 Prompt 优化报告智能体分析。Agent 不存在时报错。 + pub fn report_data(&self, agent_id: i32) -> Result { + let agent = self + .db + .read_by_id::(agent_id)? + .with_context(|| format!("agent #{} not found", agent_id))?; + + // 当前系统提示词（prompt_id → prompt_info.content）。 + let system_prompt = match agent.prompt_id { + Some(pid) => self + .db + .conn + .query_row( + "SELECT content FROM prompt_info WHERE id = ?1", + params![pid], + |row| row.get::<_, String>(0), + ) + .unwrap_or_default(), + None => String::new(), + }; + + // 星级分布（1..5），忽略越界值。 + let mut distribution = [0i64; 5]; + { + let mut stmt = self.db.conn.prepare( + "SELECT vote_type, COUNT(*) FROM agent_arena \ + WHERE agent_id = ?1 AND status = 1 GROUP BY vote_type", + )?; + let rows = stmt.query_map(params![agent_id], |row| { + Ok((row.get::<_, i32>(0)?, row.get::<_, i64>(1)?)) + })?; + for r in rows { + let (vt, cnt) = r?; + if (1..=5).contains(&vt) { + distribution[(vt - 1) as usize] = cnt; + } + } + } + let vote_count: i64 = distribution.iter().sum(); + + let dimensions = self.dimension_stats(agent_id)?; + + // 近期非空文字评价（最多 50 条，按时间倒序）。 + let mut comments = Vec::new(); + { + let mut stmt = self.db.conn.prepare( + "SELECT comment FROM agent_arena \ + WHERE agent_id = ?1 AND status = 1 AND comment IS NOT NULL AND comment != '' \ + ORDER BY id DESC LIMIT 50", + )?; + let rows = stmt.query_map(params![agent_id], |row| row.get::<_, String>(0))?; + for r in rows { + comments.push(r?); + } + } + + Ok(ReportData { + agent_id, + agent_name: agent.name, + system_prompt, + avg_score: agent.avg_score, + vote_count, + distribution, + dimensions, + comments, + }) + } + + /// 解析各行 `remark` 中的多维度评分，按 [`DIMENSIONS`] 顺序返回 (键, 均值, 样本数)。 + /// + /// 非法 JSON 行跳过；未知维度键按字典序追加在已声明维度之后。 + fn dimension_stats(&self, agent_id: i32) -> Result> { + let mut acc: BTreeMap = BTreeMap::new(); // key -> (sum, count) + { + let mut stmt = self.db.conn.prepare( + "SELECT remark FROM agent_arena \ + WHERE agent_id = ?1 AND status = 1 AND remark IS NOT NULL", + )?; + let rows = stmt.query_map(params![agent_id], |row| row.get::<_, String>(0))?; + for r in rows { + let raw = r?; + if let Ok(map) = serde_json::from_str::>(&raw) { + for (k, val) in map { + let e = acc.entry(k).or_insert((0, 0)); + e.0 += val as i64; + e.1 += 1; + } + } + } + } + + let mut out = Vec::new(); + let mut seen = std::collections::HashSet::new(); + for (key, _label) in DIMENSIONS { + if let Some((s, c)) = acc.get(*key) { + out.push((key.to_string(), round1(*s as f64 / *c as f64), *c)); + seen.insert(*key); + } + } + for (k, (s, c)) in &acc { + if !seen.contains(k.as_str()) { + out.push((k.clone(), round1(*s as f64 / *c as f64), *c)); + } + } + Ok(out) + } +} + +/// 四舍五入到 1 位小数。 +fn round1(x: f64) -> f32 { + ((x * 10.0).round() / 10.0) as f32 +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::db::Model; + + /// 构造含 `agent_info`+`agent_arena` 表的内存库，并插入 `n` 个 Agent（id 从 1 起）。 + fn test_service(agents: usize) -> ArenaService { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch(AgentInfo::create_table_sql()).unwrap(); + conn.execute_batch(AgentArena::create_table_sql()).unwrap(); + for i in 1..=agents { + conn.execute( + "INSERT INTO agent_info (id, llm_id, name, status) VALUES (?1, 1, ?2, 1)", + params![i as i32, format!("Agent-{}", i)], + ) + .unwrap(); + } + ArenaService::from_db(DatabaseManager { conn }) + } + + fn vote(agent_id: i32, score: i8) -> VoteInput { + VoteInput { + agent_id, + score, + comment: None, + chat_uuid: String::new(), + round_uuid: None, + dimensions: BTreeMap::new(), + } + } + + fn dims(pairs: &[(&str, i8)]) -> BTreeMap { + pairs.iter().map(|(k, v)| (k.to_string(), *v)).collect() + } + + fn avg_of(svc: &ArenaService, agent_id: i32) -> Option { + svc.db + .read_by_id::(agent_id) + .unwrap() + .unwrap() + .avg_score + } + + fn call_count_of(svc: &ArenaService, agent_id: i32) -> i64 { + svc.db + .read_by_id::(agent_id) + .unwrap() + .unwrap() + .call_count + } + + #[test] + fn record_vote_inserts_and_updates_avg() { + let svc = test_service(1); + let id = svc.record_vote(vote(1, 5)).unwrap(); + assert!(id > 0); + + let rows = svc + .db + .find_where::("agent_id = ?1", &[&1i32]) + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].vote_type, 5); + assert_eq!(rows[0].status, 1); + + assert_eq!(avg_of(&svc, 1), Some(5.0)); + assert_eq!(call_count_of(&svc, 1), 1); + } + + #[test] + fn avg_rounds_to_one_decimal() { + let svc = test_service(1); + svc.record_vote(vote(1, 5)).unwrap(); + svc.record_vote(vote(1, 4)).unwrap(); + assert_eq!(avg_of(&svc, 1), Some(4.5)); + assert_eq!(call_count_of(&svc, 1), 2); + } + + #[test] + fn invalid_score_rejected() { + let svc = test_service(1); + assert!(svc.record_vote(vote(1, 0)).is_err()); + assert!(svc.record_vote(vote(1, 6)).is_err()); + // 越界投票不应写入任何行。 + let rows = svc + .db + .find_where::("agent_id = ?1", &[&1i32]) + .unwrap(); + assert!(rows.is_empty()); + } + + #[test] + fn nonexistent_agent_rejected() { + let svc = test_service(1); + let err = svc.record_vote(vote(99, 5)).unwrap_err(); + assert!(err.to_string().contains("agent #99 not found")); + } + + #[test] + fn comment_and_round_uuid_persisted() { + let svc = test_service(1); + svc.record_vote(VoteInput { + agent_id: 1, + score: 3, + comment: Some("ok".to_string()), + chat_uuid: "c-1".to_string(), + round_uuid: Some("r-1".to_string()), + dimensions: BTreeMap::new(), + }) + .unwrap(); + let rows = svc + .db + .find_where::("agent_id = ?1", &[&1i32]) + .unwrap(); + assert_eq!(rows[0].comment.as_deref(), Some("ok")); + assert_eq!(rows[0].chat_uuid, "c-1"); + assert_eq!(rows[0].round_uuid.as_deref(), Some("r-1")); + } + + #[test] + fn ranking_orders_by_score_desc() { + let svc = test_service(3); + // Agent-1 均分 3，Agent-2 均分 5，Agent-3 无投票。 + svc.record_vote(vote(1, 3)).unwrap(); + svc.record_vote(vote(2, 5)).unwrap(); + svc.record_vote(vote(2, 5)).unwrap(); + + let ranks = svc.ranking().unwrap(); + assert_eq!(ranks.len(), 3); + // 最高分在前。 + assert_eq!(ranks[0].agent_id, 2); + assert_eq!(ranks[0].avg_score, Some(5.0)); + assert_eq!(ranks[0].vote_count, 2); + // 次高。 + assert_eq!(ranks[1].agent_id, 1); + assert_eq!(ranks[1].avg_score, Some(3.0)); + assert_eq!(ranks[1].vote_count, 1); + // 未评分置底。 + assert_eq!(ranks[2].agent_id, 3); + assert_eq!(ranks[2].avg_score, None); + assert_eq!(ranks[2].vote_count, 0); + } + + #[test] + fn recompute_is_idempotent() { + let svc = test_service(1); + svc.record_vote(vote(1, 4)).unwrap(); + // 手动置空后重算应恢复。 + let mut a = svc.db.read_by_id::(1).unwrap().unwrap(); + a.avg_score = None; + svc.db.update(a).unwrap(); + assert_eq!(avg_of(&svc, 1), None); + assert_eq!(svc.recompute(1).unwrap(), Some(4.0)); + assert_eq!(avg_of(&svc, 1), Some(4.0)); + } + + #[test] + fn record_vote_serializes_dimensions_to_remark() { + let svc = test_service(1); + svc.record_vote(VoteInput { + agent_id: 1, + score: 4, + comment: None, + chat_uuid: "c-1".to_string(), + round_uuid: None, + dimensions: dims(&[("acc", 5), ("rel", 4)]), + }) + .unwrap(); + let rows = svc + .db + .find_where::("agent_id = ?1", &[&1i32]) + .unwrap(); + assert_eq!(rows[0].vote_type, 4); + let remark = rows[0].remark.clone().unwrap(); + let parsed: BTreeMap = serde_json::from_str(&remark).unwrap(); + assert_eq!(parsed.get("acc"), Some(&5)); + assert_eq!(parsed.get("rel"), Some(&4)); + } + + #[test] + fn record_vote_rejects_out_of_range_dimension() { + let svc = test_service(1); + let r = svc.record_vote(VoteInput { + agent_id: 1, + score: 4, + comment: None, + chat_uuid: String::new(), + round_uuid: None, + dimensions: dims(&[("acc", 6)]), + }); + assert!(r.is_err()); + // 越界维度不应写入任何行。 + let rows = svc + .db + .find_where::("agent_id = ?1", &[&1i32]) + .unwrap(); + assert!(rows.is_empty()); + } + + #[test] + fn report_data_aggregates_distribution_dims_comments() { + let svc = test_service(1); + svc.record_vote(VoteInput { + agent_id: 1, + score: 5, + comment: Some("great".to_string()), + chat_uuid: String::new(), + round_uuid: None, + dimensions: dims(&[("acc", 5), ("rel", 5)]), + }) + .unwrap(); + svc.record_vote(VoteInput { + agent_id: 1, + score: 3, + comment: Some("ok".to_string()), + chat_uuid: String::new(), + round_uuid: None, + dimensions: dims(&[("acc", 3), ("rel", 5)]), + }) + .unwrap(); + svc.record_vote(VoteInput { + agent_id: 1, + score: 3, + comment: None, + chat_uuid: String::new(), + round_uuid: None, + dimensions: BTreeMap::new(), + }) + .unwrap(); + + let data = svc.report_data(1).unwrap(); + assert_eq!(data.agent_name, "Agent-1"); + assert_eq!(data.vote_count, 3); + // 分布：5 星 1 票，3 星 2 票。 + assert_eq!(data.distribution[4], 1); + assert_eq!(data.distribution[2], 2); + // acc 均值 (5+3)/2=4.0（样本 2）；rel 均值 5.0。 + let acc = data.dimensions.iter().find(|d| d.0 == "acc").unwrap(); + assert_eq!(acc.1, 4.0); + assert_eq!(acc.2, 2); + let rel = data.dimensions.iter().find(|d| d.0 == "rel").unwrap(); + assert_eq!(rel.1, 5.0); + // 评论按 id DESC，仅 2 条非空。 + assert_eq!(data.comments.len(), 2); + assert_eq!(data.comments[0], "ok"); + assert_eq!(data.comments[1], "great"); + } +} diff --git a/src/core/arena/report.rs b/src/core/arena/report.rs new file mode 100644 index 0000000..fe843e9 --- /dev/null +++ b/src/core/arena/report.rs @@ -0,0 +1,216 @@ +//! Prompt 优化报告智能体（设计 3.8.3 数据闭环） +//! +//! 复用竞技场评估数据（星级分布、多维度均值、用户文字评价、当前系统提示词），以独立 +//! provider 调 LLM 产出结构化的 Prompt 优化建议报告，并可落盘留档。 +//! +//! 与 [`crate::core::memory::MemoryExtractor`] 同款模式：构造**独立** provider（不借用 +//! Agent 的 `Box`），以便在 CLI 命令中一次性调用。LLM 调用不在单测覆盖内， +//! 单测仅验证纯函数 [`render_report_input`] 与落盘 [`save_report_to`]。 + +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result}; +use chrono::Utc; + +use crate::core::providers::{create_provider, CompletionRequest, Message, ProviderConfig, Role}; +use crate::platform::get_platform_app_dir; + +use super::{ReportData, DIMENSIONS}; + +/// 报告智能体系统提示词。 +pub const REPORT_PROMPT: &str = "\ +你是一名 Prompt 优化专家。下面给出某个 Agent 的当前系统提示词、竞技场星级评分分布、\ +多维度评估均值与用户文字评价。请据此产出一份结构化的 Prompt 优化建议报告，使用 Markdown，包含： +1. 问题诊断：结合低分档位与差评，指出当前提示词可能导致的问题； +2. 改进建议：针对每个问题给出可操作的改进方向； +3. 建议的新提示词要点：给出优化后提示词的关键条目（可直接采用的措辞）。 +请保持客观、具体，避免空泛。"; + +/// Prompt 优化报告智能体。 +pub struct PromptReportAgent { + provider_config: ProviderConfig, + model: String, + temperature: Option, +} + +impl PromptReportAgent { + /// 构造报告智能体（复用被评 Agent 的供应商/模型配置）。 + pub fn new(provider_config: ProviderConfig, model: String, temperature: Option) -> Self { + Self { + provider_config, + model, + temperature, + } + } + + /// 调用 LLM 生成优化报告正文（Markdown）。 + pub async fn generate(&self, data: &ReportData) -> Result { + let provider = create_provider(self.provider_config.clone())?; + let request = CompletionRequest { + model: self.model.clone(), + messages: vec![ + Message { + role: Role::System, + content: REPORT_PROMPT.to_string(), + }, + Message { + role: Role::User, + content: render_report_input(data), + }, + ], + temperature: self.temperature, + max_tokens: None, + stream: false, + }; + let response = provider.complete(request).await?; + Ok(response.content) + } +} + +/// 将评估数据拼装为供 LLM 分析的 Markdown 输入（纯函数，可测）。 +pub fn render_report_input(data: &ReportData) -> String { + let mut out = String::new(); + out.push_str(&format!( + "# Agent 评估数据：{} (#{})\n\n", + data.agent_name, data.agent_id + )); + + out.push_str("## 当前系统提示词\n"); + if data.system_prompt.trim().is_empty() { + out.push_str("(未配置)\n"); + } else { + out.push_str(&data.system_prompt); + out.push('\n'); + } + out.push('\n'); + + let avg = match data.avg_score { + Some(a) => format!("{:.1}", a), + None => "-".to_string(), + }; + out.push_str(&format!( + "## 综合评分\n平均分：{} / 5，有效投票：{}\n\n", + avg, data.vote_count + )); + + out.push_str("## 星级分布\n"); + for star in 1..=5 { + out.push_str(&format!( + "- {} 星：{}\n", + star, data.distribution[(star - 1) as usize] + )); + } + out.push('\n'); + + out.push_str("## 多维度评估均值\n"); + if data.dimensions.is_empty() { + out.push_str("(无多维度数据)\n"); + } else { + for (key, avg, count) in &data.dimensions { + let label = DIMENSIONS + .iter() + .find(|(k, _)| *k == key.as_str()) + .map(|(_, l)| *l) + .unwrap_or(key.as_str()); + out.push_str(&format!( + "- {} ({})：{:.1} / 5（样本 {}）\n", + label, key, avg, count + )); + } + } + out.push('\n'); + + out.push_str("## 用户评价\n"); + if data.comments.is_empty() { + out.push_str("(无文字评价)\n"); + } else { + for c in &data.comments { + out.push_str(&format!("- {}\n", c)); + } + } + + out +} + +/// 将报告写入 `~/.omega/reports/agent__.md`，返回文件路径。 +pub fn save_report(agent_id: i32, content: &str) -> Result { + let dir = get_platform_app_dir()?.join("reports"); + save_report_to(&dir, agent_id, content) +} + +/// 写入指定目录（供测试传 tempdir）。 +fn save_report_to(dir: &Path, agent_id: i32, content: &str) -> Result { + std::fs::create_dir_all(dir) + .with_context(|| format!("failed to create report dir: {}", dir.display()))?; + let ts = Utc::now().format("%Y%m%d-%H%M%S").to_string(); + let path = dir.join(format!("agent_{}_{}.md", agent_id, ts)); + std::fs::write(&path, content) + .with_context(|| format!("failed to write report: {}", path.display()))?; + Ok(path) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn sample_data() -> ReportData { + ReportData { + agent_id: 7, + agent_name: "Coder".to_string(), + system_prompt: "You are a coder.".to_string(), + avg_score: Some(3.5), + vote_count: 4, + distribution: [0, 1, 1, 1, 1], + dimensions: vec![ + ("acc".to_string(), 4.0, 2), + ("rel".to_string(), 3.0, 2), + ], + comments: vec!["too verbose".to_string(), "helpful".to_string()], + } + } + + #[test] + fn render_report_input_contains_sections() { + let text = render_report_input(&sample_data()); + assert!(text.contains("Coder")); + assert!(text.contains("准确性")); + assert!(text.contains("相关性")); + assert!(text.contains("too verbose")); + assert!(text.contains("You are a coder.")); + // 星级分布计数（distribution[4]=1 → “5 星：1”）。 + assert!(text.contains("5 星：1")); + } + + #[test] + fn render_report_input_handles_empty() { + let data = ReportData { + agent_id: 1, + agent_name: "A".to_string(), + system_prompt: String::new(), + avg_score: None, + vote_count: 0, + distribution: [0; 5], + dimensions: vec![], + comments: vec![], + }; + let text = render_report_input(&data); + assert!(text.contains("(未配置)")); + assert!(text.contains("(无多维度数据)")); + assert!(text.contains("(无文字评价)")); + } + + #[test] + fn save_report_to_writes_file() { + let dir = std::env::temp_dir().join(format!( + "omega-report-{}", + Utc::now().timestamp_nanos_opt().unwrap_or_default() + )); + let path = save_report_to(&dir, 7, "# hello").unwrap(); + assert!(path.exists()); + let name = path.file_name().unwrap().to_string_lossy().to_string(); + assert!(name.starts_with("agent_7_")); + assert!(name.ends_with(".md")); + assert_eq!(std::fs::read_to_string(&path).unwrap(), "# hello"); + let _ = std::fs::remove_dir_all(&dir); + } +} diff --git a/src/core/db/db_init.rs b/src/core/db/db_init.rs index 115f1bd..b763222 100644 --- a/src/core/db/db_init.rs +++ b/src/core/db/db_init.rs @@ -266,20 +266,142 @@ fn create_table_system_config(conn: &Connection) -> Result<()> { fn create_indexes(conn: &Connection) -> Result<()> { conn.execute_batch( " + -- 逻辑外键普通索引（关联字段） CREATE INDEX IF NOT EXISTS idx_chat_session_chat_uuid ON chat_session(chat_uuid); - CREATE INDEX IF NOT EXISTS idx_chat_message_session_id ON chat_message(session_id); + CREATE INDEX IF NOT EXISTS idx_chat_message_chat_uuid ON chat_message(chat_uuid); + CREATE INDEX IF NOT EXISTS idx_chat_message_round_uuid ON chat_message(round_uuid); CREATE INDEX IF NOT EXISTS idx_agent_arena_agent_id ON agent_arena(agent_id); - CREATE INDEX IF NOT EXISTS idx_agent_arena_arena_id ON agent_arena(arena_id); - CREATE INDEX IF NOT EXISTS idx_memory_info_session_id ON memory_info(session_id); + CREATE INDEX IF NOT EXISTS idx_agent_arena_chat_uuid ON agent_arena(chat_uuid); + CREATE INDEX IF NOT EXISTS idx_memory_info_memory_uuid ON memory_info(memory_uuid); CREATE INDEX IF NOT EXISTS idx_agent_mcp_rel_agent_id ON agent_mcp_rel(agent_id); - CREATE INDEX IF NOT EXISTS idx_agent_mcp_rel_mcp_id ON agent_mcp_rel(mcp_id); + CREATE INDEX IF NOT EXISTS idx_agent_mcp_rel_mcp_tool_id ON agent_mcp_rel(mcp_tool_id); CREATE INDEX IF NOT EXISTS idx_prompt_placeholder_prompt_id ON prompt_placeholder(prompt_id); CREATE INDEX IF NOT EXISTS idx_llm_model_provider_id ON llm_model(provider_id); CREATE INDEX IF NOT EXISTS idx_mcp_tool_provider_id ON mcp_tool(provider_id); - CREATE INDEX IF NOT EXISTS idx_agent_info_model_id ON agent_info(model_id); + CREATE INDEX IF NOT EXISTS idx_agent_info_llm_id ON agent_info(llm_id); CREATE INDEX IF NOT EXISTS idx_agent_info_prompt_id ON agent_info(prompt_id); + CREATE INDEX IF NOT EXISTS idx_system_api_log_rel_uuid ON system_api_log(rel_uuid); + + -- 唯一业务字段索引 + CREATE UNIQUE INDEX IF NOT EXISTS uk_llm_model_model_code ON llm_model(model_code); + CREATE UNIQUE INDEX IF NOT EXISTS uk_project_info_path ON project_info(path); + CREATE UNIQUE INDEX IF NOT EXISTS uk_prompt_placeholder_key ON prompt_placeholder(prompt_id, key); ", )?; info!("{}", t!("database_table_create", table_name = "table_indexes")); Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::db::DatabaseManager; + use crate::db::models::chat_message::ChatMessage; + use crate::db::models::chat_session::ChatSession; + use crate::db::models::project_info::ProjectInfo; + + /// 在内存库上执行全量建表 + 建索引，验证所有 CREATE TABLE / CREATE INDEX + /// 语句在运行时均为合法 SQL（列名、唯一索引与设计文档 schema 一致）。 + fn mem_db() -> DatabaseManager { + let conn = Connection::open_in_memory().expect("open in-memory db"); + init_all_tables(&conn).expect("init_all_tables should succeed"); + DatabaseManager { conn } + } + + #[test] + fn init_all_tables_creates_schema_and_indexes() { + let db = mem_db(); + // 15 张业务表均应存在 + let tables = [ + "ai_provider", "llm_model", "prompt_info", "agent_info", + "prompt_placeholder", "agent_arena", "chat_session", "chat_message", + "memory_info", "mcp_provider", "mcp_tool", "agent_mcp_rel", + "project_info", "system_api_log", "system_config", + ]; + for name in tables { + let exists: bool = db + .conn + .query_row( + "SELECT COUNT(*) > 0 FROM sqlite_master WHERE type='table' AND name=?1", + [name], + |row| row.get(0), + ) + .unwrap(); + assert!(exists, "table {} should exist", name); + } + // 唯一索引应已创建 + let idx: bool = db + .conn + .query_row( + "SELECT COUNT(*) > 0 FROM sqlite_master WHERE type='index' AND name='uk_llm_model_model_code'", + [], + |row| row.get(0), + ) + .unwrap(); + assert!(idx, "unique index on llm_model(model_code) should exist"); + } + + #[test] + fn chat_session_and_message_roundtrip() { + let db = mem_db(); + let session = ChatSession { + id: 0, + chat_uuid: "chat-1".to_string(), + title: Some("TUI Chat".to_string()), + llm_id: 1, + is_deleted: 0, + create_time: chrono::Utc::now(), + update_time: chrono::Utc::now(), + remark: None, + }; + let sid = db.create(session).unwrap(); + assert!(sid > 0); + + let msg = ChatMessage { + id: 0, + chat_uuid: "chat-1".to_string(), + round_uuid: Some("round-1".to_string()), + role: "user".to_string(), + content: "hello".to_string(), + input_tokens: 3, + output_tokens: 5, + duration: 42, + model_config: Some("{\"temperature\":0.7}".to_string()), + create_time: chrono::Utc::now(), + update_time: chrono::Utc::now(), + remark: None, + }; + db.create(msg).unwrap(); + + let loaded: Vec = db + .find_where("chat_uuid = ?1", &[&"chat-1".to_string()]) + .unwrap(); + assert_eq!(loaded.len(), 1); + assert_eq!(loaded[0].content, "hello"); + assert_eq!(loaded[0].input_tokens, 3); + assert_eq!(loaded[0].output_tokens, 5); + assert_eq!(loaded[0].duration, 42); + assert_eq!(loaded[0].round_uuid.as_deref(), Some("round-1")); + } + + #[test] + fn project_info_path_unique_index_enforced() { + let db = mem_db(); + let mk = |path: &str| ProjectInfo { + id: 0, + name: "proj".to_string(), + path: path.to_string(), + create_at: None, + updata_at: None, + is_deleted: 0, + create_time: chrono::Utc::now(), + update_time: chrono::Utc::now(), + remark: None, + }; + db.create(mk("/a/b")).unwrap(); + // 相同 path 应触发唯一索引冲突 + assert!(db.create(mk("/a/b")).is_err()); + // 不同 path 允许 + db.create(mk("/c/d")).unwrap(); + } +} diff --git a/src/core/mcp/mod.rs b/src/core/mcp/mod.rs new file mode 100644 index 0000000..27bb0cd --- /dev/null +++ b/src/core/mcp/mod.rs @@ -0,0 +1,289 @@ +//! MCP 工具调度 +//! +//! 依据《系统功能设计与数据库设计文档》四.10/四.11 与《后端开发手册》：Agent 通过 +//! `agent_mcp_rel` 获得授权的 MCP 工具集合（`mcp_tool`），每个工具归属于某个 +//! MCP Server（`mcp_provider`）。本模块负责把这些「数据库登记的远程工具」适配成本地 +//! [`Tool`]，并在调用时经由 [`rmcp`] 客户端（Streamable HTTP 传输）路由到对应的 +//! Python MCP Server 执行。 +//! +//! 设计要点： +//! - **权限校验**：只有出现在授权表中的 `tool_code` 才能被路由，未授权调用直接拒绝。 +//! - **路由**：以 `tool_code → provider_url` 建立映射，调用时按工具的归属地址转发。 +//! - **适配**：[`McpTool`] 把远程工具包装为本地 [`Tool`]，其名称即 `tool_code`， +//! 参数 Schema 取自 `mcp_tool.parameters`（缺省为空 object）。 +//! +//! [`Tool`]: crate::core::tools::Tool +//! [`rmcp`]: https://docs.rs/rmcp + +use std::collections::HashMap; +use std::sync::Arc; + +use anyhow::{Context, Result}; +use async_trait::async_trait; +use rmcp::model::{CallToolRequestParams, CallToolResult, ClientInfo}; +use rmcp::transport::StreamableHttpClientTransport; +use rmcp::ServiceExt; +use serde_json::{json, Map, Value}; + +use crate::core::tools::Tool; + +/// 从数据库加载的一个已授权 MCP 工具规格。 +#[derive(Debug, Clone)] +pub struct McpToolSpec { + /// `mcp_tool.id`。 + pub tool_id: i32, + /// 工具调用标识（`mcp_tool.tool_code`），也是本地 [`Tool::name`]。 + pub tool_code: String, + /// 工具展示名（`mcp_tool.name`）。 + pub name: String, + /// 工具描述（`mcp_tool.description`，可空）。 + pub description: Option, + /// 入参 JSON Schema（`mcp_tool.parameters`，可空）。 + pub parameters: Option, + /// 归属 MCP Server 地址（`mcp_provider.url`）。 + pub provider_url: String, +} + +/// MCP 工具路由器：维护 `tool_code → provider_url` 的授权映射并转发调用。 +/// +/// 由 [`build_tools`] 在多个 [`McpTool`] 之间共享（`Arc`），保证「未登记即不可调用」 +/// 的权限边界集中在一处。 +#[derive(Debug, Default, Clone)] +pub struct McpToolRouter { + routes: HashMap, +} + +impl McpToolRouter { + /// 依据一组授权工具规格构建路由表。 + pub fn from_specs(specs: &[McpToolSpec]) -> Self { + let mut routes = HashMap::new(); + for spec in specs { + routes.insert(spec.tool_code.clone(), spec.provider_url.clone()); + } + Self { routes } + } + + /// 判断某个 `tool_code` 是否已授权（存在于路由表中）。 + pub fn is_authorized(&self, tool_code: &str) -> bool { + self.routes.contains_key(tool_code) + } + + /// 取得某个 `tool_code` 归属的 MCP Server 地址。 + pub fn provider_url(&self, tool_code: &str) -> Option<&str> { + self.routes.get(tool_code).map(|s| s.as_str()) + } + + /// 已授权工具数量。 + pub fn len(&self) -> usize { + self.routes.len() + } + + pub fn is_empty(&self) -> bool { + self.routes.is_empty() + } + + /// 路由调用一个远程 MCP 工具。 + /// + /// 未授权的 `tool_code` 会被直接拒绝（权限校验），否则转发到其归属 Server。 + pub async fn call(&self, tool_code: &str, args: Map) -> Result { + let url = self + .provider_url(tool_code) + .with_context(|| format!("MCP tool '{}' is not authorized", tool_code))?; + call_mcp_tool(url, tool_code, args).await + } +} + +/// 经由 Streamable HTTP 传输连接 MCP Server，调用单个工具并抽取文本结果。 +/// +/// 连接为一次性短连接：调用完成后主动 `cancel` 释放后台任务。若服务端返回 +/// `is_error == true`，则以错误形式上抛（附服务端返回的文本）。 +pub async fn call_mcp_tool( + url: &str, + tool_code: &str, + args: Map, +) -> Result { + let transport = StreamableHttpClientTransport::from_uri(url.to_string()); + let client = ClientInfo::default() + .serve(transport) + .await + .with_context(|| format!("failed to connect to MCP server at {}", url))?; + + let result = client + .call_tool(CallToolRequestParams::new(tool_code.to_string()).with_arguments(args)) + .await + .with_context(|| format!("MCP tool '{}' call failed", tool_code)); + + // 无论成功与否都尽量释放连接（忽略 cancel 自身的错误）。 + let _ = client.cancel().await; + + let result = result?; + if result.is_error == Some(true) { + anyhow::bail!( + "MCP tool '{}' returned error: {}", + tool_code, + extract_text(&result) + ); + } + Ok(extract_text(&result)) +} + +/// 从 [`CallToolResult`] 中抽取文本：优先拼接所有文本内容块， +/// 若无文本则回退到 `structured_content` 的 JSON 字符串。 +fn extract_text(result: &CallToolResult) -> String { + let mut text = String::new(); + for content in &result.content { + if let Some(t) = content.as_text() { + if !text.is_empty() { + text.push('\n'); + } + text.push_str(&t.text); + } + } + if text.is_empty() { + if let Some(structured) = &result.structured_content { + return structured.to_string(); + } + } + text +} + +/// 把一个远程 MCP 工具适配为本地 [`Tool`]。 +/// +/// 名称即 `tool_code`，参数 Schema 取自 `mcp_tool.parameters`，执行时经共享的 +/// [`McpToolRouter`] 路由到归属 Server。 +pub struct McpTool { + code: String, + description: String, + parameters: Value, + router: Arc, +} + +impl McpTool { + /// 依据规格与共享路由器构造。 + pub fn new(spec: &McpToolSpec, router: Arc) -> Self { + let description = spec + .description + .clone() + .filter(|d| !d.trim().is_empty()) + .unwrap_or_else(|| spec.name.clone()); + Self { + code: spec.tool_code.clone(), + description, + parameters: parse_parameters(spec.parameters.as_deref()), + router, + } + } +} + +/// 解析 `mcp_tool.parameters`（JSON Schema 字符串）。 +/// +/// 缺省或非法时回退为一个空 object schema，保证提示词中始终有可用的参数描述。 +fn parse_parameters(raw: Option<&str>) -> Value { + let default = json!({ "type": "object", "properties": {} }); + match raw { + Some(s) if !s.trim().is_empty() => serde_json::from_str::(s).unwrap_or(default), + _ => default, + } +} + +#[async_trait] +impl Tool for McpTool { + fn name(&self) -> &str { + &self.code + } + + fn description(&self) -> &str { + &self.description + } + + fn parameters(&self) -> Value { + self.parameters.clone() + } + + async fn execute(&self, args: Value) -> Result { + let map = match args { + Value::Object(m) => m, + Value::Null => Map::new(), + other => { + let mut m = Map::new(); + m.insert("value".to_string(), other); + m + } + }; + self.router.call(&self.code, map).await + } +} + +/// 依据一组已授权工具规格构建本地 [`Tool`] 集合（共享同一个 [`McpToolRouter`]）。 +/// +/// 返回的工具可直接注册进 [`ToolRegistry`](crate::core::tools::ToolRegistry)。 +pub fn build_tools(specs: Vec) -> Vec> { + let router = Arc::new(McpToolRouter::from_specs(&specs)); + specs + .iter() + .map(|spec| Arc::new(McpTool::new(spec, router.clone())) as Arc) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn spec(tool_code: &str, url: &str) -> McpToolSpec { + McpToolSpec { + tool_id: 1, + tool_code: tool_code.to_string(), + name: tool_code.to_string(), + description: None, + parameters: None, + provider_url: url.to_string(), + } + } + + #[test] + fn router_authorizes_only_registered_tools() { + let router = McpToolRouter::from_specs(&[ + spec("add", "http://localhost:8000/mcp"), + spec("echo", "http://localhost:8001/mcp"), + ]); + assert!(router.is_authorized("add")); + assert!(router.is_authorized("echo")); + assert!(!router.is_authorized("drop_table")); + assert_eq!(router.provider_url("add"), Some("http://localhost:8000/mcp")); + assert_eq!(router.provider_url("nope"), None); + assert_eq!(router.len(), 2); + } + + #[tokio::test] + async fn router_rejects_unauthorized_call() { + let router = McpToolRouter::from_specs(&[spec("add", "http://localhost:8000/mcp")]); + let err = router.call("secret", Map::new()).await.unwrap_err(); + assert!(err.to_string().contains("not authorized")); + } + + #[test] + fn parse_parameters_falls_back_to_empty_object() { + let fallback = json!({ "type": "object", "properties": {} }); + assert_eq!(parse_parameters(None), fallback); + assert_eq!(parse_parameters(Some("")), fallback); + assert_eq!(parse_parameters(Some("not json")), fallback); + + let schema = json!({ "type": "object", "properties": { "a": { "type": "integer" } } }); + assert_eq!(parse_parameters(Some(&schema.to_string())), schema); + } + + #[test] + fn build_tools_maps_specs_with_shared_router() { + let tools = build_tools(vec![ + spec("add", "http://localhost:8000/mcp"), + spec("echo", "http://localhost:8001/mcp"), + ]); + assert_eq!(tools.len(), 2); + assert_eq!(tools[0].name(), "add"); + assert_eq!(tools[1].name(), "echo"); + // 描述缺省时回退到 name。 + assert_eq!(tools[0].description(), "add"); + // 参数缺省时为空 object schema。 + assert_eq!(tools[0].parameters()["type"], "object"); + } +} diff --git a/src/core/memory/mod.rs b/src/core/memory/mod.rs new file mode 100644 index 0000000..7153a42 --- /dev/null +++ b/src/core/memory/mod.rs @@ -0,0 +1,814 @@ +//! 记忆系统 +//! +//! 对应《系统功能设计与数据库设计文档》3.10 记忆模块。分两类记忆： +//! - **用户记忆**（`user`）：长期编码习惯、对话偏好、使用行为，全局生效；正文存于 +//! `~/.omega/memory/.md`（[`crate::platform::get_platform_app_dir`]）。 +//! - **项目记忆**（`project`）：绑定单一项目上下文，抑制幻觉、固定业务边界；正文存于 +//! `/.omega/memory/.md`，仅在 Workspace 模式（`project_dir` 存在）可用。 +//! +//! 存储采用「数据库 + 本地文件」双存储：`memory_info` 表仅存**索引** +//! （memory_uuid / memory_type / keywords / is_vector=0 / vector_id=NULL / status）， +//! **正文写文件**。FAISS/LanceDB 向量检索、上下文压缩、过期清理属设计明确的 +//! 「前期不实现」范围，本模块不涉及。 +//! +//! 加载作用域来自 `system_config.memory_type`（由 `Omega memory --type` 写入），映射为 +//! [`MemoryScope`]；未设置时为 [`MemoryScope::None`]（记忆关闭）。 +//! +//! [`MemoryExtractor`] 提供自动抽取：一轮对话结束后用独立 provider 调 LLM 总结可长期 +//! 沉淀的信息，去重后经 [`MemoryStore::save`] 写入。 + +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +use anyhow::{anyhow, Context, Result}; +use log::warn; +use rand::RngExt; +use rusqlite::{params, Connection, OptionalExtension}; +use serde_json::Value; + +use crate::core::db::Model; +use crate::core::providers::{ + create_provider, CompletionRequest, Message, ProviderConfig, Role, +}; +use crate::db::models::memory_info::MemoryInfo; +use crate::platform::{get_database_path, get_platform_app_dir}; + +/// 单次注入提示词的最大记忆条数。 +pub const MAX_INJECT: usize = 20; +/// 注入提示词段的最大字符数（防止提示词膨胀）。 +pub const MAX_INJECT_CHARS: usize = 4000; + +// ============================================================ +// 作用域与领域对象 +// ============================================================ + +/// 记忆加载作用域（对应 `system_config.memory_type`）。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MemoryScope { + /// 仅加载用户记忆。 + User, + /// 仅加载项目记忆。 + Project, + /// 加载用户 + 项目记忆。 + All, + /// 关闭记忆功能。 + None, +} + +impl MemoryScope { + /// 从字符串解析作用域；无法识别的值一律视为 [`MemoryScope::None`]（安全关闭）。 + pub fn parse(s: &str) -> Self { + match s.trim().to_lowercase().as_str() { + "user" => MemoryScope::User, + "project" => MemoryScope::Project, + "all" => MemoryScope::All, + _ => MemoryScope::None, + } + } + + /// 稳定小写标识。 + pub fn as_str(&self) -> &'static str { + match self { + MemoryScope::User => "user", + MemoryScope::Project => "project", + MemoryScope::All => "all", + MemoryScope::None => "none", + } + } + + /// 该作用域是否允许加载/写入指定 `memory_type`（"user" / "project"）。 + pub fn allows_type(&self, memory_type: &str) -> bool { + match self { + MemoryScope::All => memory_type == "user" || memory_type == "project", + MemoryScope::User => memory_type == "user", + MemoryScope::Project => memory_type == "project", + MemoryScope::None => false, + } + } +} + +/// 一条记忆的领域对象：索引字段 + 文件正文。 +#[derive(Debug, Clone)] +pub struct MemoryRecord { + /// 记忆唯一标识（同时作为正文文件名 `.md`）。 + pub memory_uuid: String, + /// 记忆类型："user" / "project"。 + pub memory_type: String, + /// 检索关键词。 + pub keywords: Vec, + /// 记忆正文（来自文件）。 + pub content: String, +} + +/// 归一化关键词：去空白、去空项、排序、去重后以逗号连接（用于存储与去重比较）。 +fn normalize_keywords(keywords: &[String]) -> String { + let mut v: Vec = keywords + .iter() + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect(); + v.sort(); + v.dedup(); + v.join(",") +} + +/// 拆分存储的关键词串为向量。 +fn split_keywords(s: Option) -> Vec { + match s { + Some(x) => x + .split(',') + .map(|p| p.trim().to_string()) + .filter(|p| !p.is_empty()) + .collect(), + None => Vec::new(), + } +} + +/// 生成记忆唯一标识（无外部 uuid 依赖）。 +fn new_memory_uuid() -> String { + let ts = chrono::Utc::now().timestamp(); + let rand_part: u32 = rand::rng().random_range(0..u32::MAX); + format!("m-{:x}-{:08x}", ts, rand_part) +} + +// ============================================================ +// MemoryStore +// ============================================================ + +/// 记忆存储：`memory_info` 索引（SQLite）+ 正文文件（本地目录）。 +/// +/// 内部用 [`Mutex`] 保证 `Send + Sync`，可被 Agent 在异步任务中共享。 +pub struct MemoryStore { + conn: Mutex, + /// 用户记忆正文目录（`~/.omega/memory`）。 + user_dir: PathBuf, + /// 项目记忆正文目录（`/.omega/memory`）；`None` 表示非 Workspace 模式。 + project_dir: Option, +} + +impl MemoryStore { + /// 打开记忆存储。 + /// + /// `project_root` 为 `Some` 时启用项目记忆（目录 `/.omega/memory`）， + /// 传入 Workspace 根路径；`None`（single 模式）时仅支持用户记忆。 + pub fn new(project_root: Option) -> Result { + let path = get_database_path()?; + let conn = Connection::open(&path) + .with_context(|| format!("Failed to open database: {}", path.display()))?; + Self::ensure_table(&conn)?; + let user_dir = get_platform_app_dir()?.join("memory"); + let project_dir = project_root.map(|p| p.join(".omega").join("memory")); + Ok(Self { + conn: Mutex::new(conn), + user_dir, + project_dir, + }) + } + + /// 测试专用：用已有连接与目录构造。 + #[cfg(test)] + fn from_parts(conn: Connection, user_dir: PathBuf, project_dir: Option) -> Self { + Self::ensure_table(&conn).unwrap(); + Self { + conn: Mutex::new(conn), + user_dir, + project_dir, + } + } + + /// 确保 `memory_info` 表存在（启动时已建，此处幂等兜底以支持独立使用）。 + fn ensure_table(conn: &Connection) -> Result<()> { + conn.execute_batch(MemoryInfo::create_table_sql())?; + Ok(()) + } + + fn lock(&self) -> Result> { + self.conn + .lock() + .map_err(|_| anyhow!("MemoryStore mutex poisoned")) + } + + /// 指定记忆类型的正文目录。项目记忆在 `project_dir` 缺失时报错。 + fn dir_for(&self, memory_type: &str) -> Result { + match memory_type { + "project" => self + .project_dir + .clone() + .ok_or_else(|| anyhow!("project memory unavailable in single mode")), + _ => Ok(self.user_dir.clone()), + } + } + + /// 读取某条记忆的正文文件。 + fn read_content(&self, memory_type: &str, memory_uuid: &str) -> Result { + let path = self + .dir_for(memory_type)? + .join(format!("{}.md", memory_uuid)); + std::fs::read_to_string(&path) + .with_context(|| format!("failed to read memory file: {}", path.display())) + } + + /// 保存（或更新）一条记忆：写正文文件 + upsert `memory_info` 索引行。 + /// + /// 以 `memory_uuid` 为幂等键：已存在则更新 keywords/status/update_time，否则插入。 + pub fn save(&self, rec: &MemoryRecord) -> Result<()> { + let dir = self.dir_for(&rec.memory_type)?; + std::fs::create_dir_all(&dir) + .with_context(|| format!("failed to create memory dir: {}", dir.display()))?; + let path = dir.join(format!("{}.md", rec.memory_uuid)); + std::fs::write(&path, &rec.content) + .with_context(|| format!("failed to write memory file: {}", path.display()))?; + + let keywords = normalize_keywords(&rec.keywords); + let conn = self.lock()?; + let existing: Option = conn + .query_row( + "SELECT id FROM memory_info WHERE memory_uuid = ?1", + params![rec.memory_uuid], + |row| row.get(0), + ) + .optional()?; + match existing { + Some(id) => { + conn.execute( + "UPDATE memory_info SET keywords = ?1, status = 1, \ + update_time = CURRENT_TIMESTAMP WHERE id = ?2", + params![keywords, id], + )?; + } + None => { + conn.execute( + "INSERT INTO memory_info \ + (memory_uuid, memory_type, is_vector, vector_id, keywords, status) \ + VALUES (?1, ?2, 0, NULL, ?3, 1)", + params![rec.memory_uuid, rec.memory_type, keywords], + )?; + } + } + Ok(()) + } + + /// 按作用域加载记忆（正文来自文件），按更新时间倒序，上限 [`MAX_INJECT`] 条。 + /// + /// 正文文件缺失的记忆跳过并告警（不阻断加载）。[`MemoryScope::None`] 返回空。 + pub fn load(&self, scope: MemoryScope) -> Result> { + if scope == MemoryScope::None { + return Ok(Vec::new()); + } + let where_sql = match scope { + MemoryScope::User => "status = 1 AND memory_type = 'user'", + MemoryScope::Project => "status = 1 AND memory_type = 'project'", + _ => "status = 1", + }; + let sql = format!( + "SELECT memory_uuid, memory_type, keywords FROM memory_info \ + WHERE {} ORDER BY update_time DESC, id DESC LIMIT ?1", + where_sql + ); + let conn = self.lock()?; + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map(params![MAX_INJECT as i64], |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, String>(1)?, + row.get::<_, Option>(2)?, + )) + })?; + + let mut out = Vec::new(); + for row in rows { + let (memory_uuid, memory_type, kw) = row?; + match self.read_content(&memory_type, &memory_uuid) { + Ok(content) => out.push(MemoryRecord { + memory_uuid, + memory_type, + keywords: split_keywords(kw), + content, + }), + Err(e) => warn!("skip memory ({}): {}", memory_uuid, e), + } + } + Ok(out) + } + + /// 列出记忆（等价于 [`MemoryStore::load`]，含正文，供展示/调试）。 + pub fn list(&self, scope: MemoryScope) -> Result> { + self.load(scope) + } + + /// 物理删除一条记忆（文件 + 索引行）。返回是否删除了索引行。 + pub fn delete(&self, memory_uuid: &str) -> Result { + let conn = self.lock()?; + let memory_type: Option = conn + .query_row( + "SELECT memory_type FROM memory_info WHERE memory_uuid = ?1", + params![memory_uuid], + |row| row.get(0), + ) + .optional()?; + let removed = conn.execute( + "DELETE FROM memory_info WHERE memory_uuid = ?1", + params![memory_uuid], + )?; + if let Some(mt) = memory_type { + if let Ok(dir) = self.dir_for(&mt) { + let path = dir.join(format!("{}.md", memory_uuid)); + let _ = std::fs::remove_file(path); + } + } + Ok(removed > 0) + } + + /// 按关键词模糊检索记忆（`keywords LIKE %kw%`），受作用域约束。 + pub fn search(&self, keyword: &str, scope: MemoryScope) -> Result> { + if scope == MemoryScope::None || keyword.trim().is_empty() { + return Ok(Vec::new()); + } + let type_sql = match scope { + MemoryScope::User => " AND memory_type = 'user'", + MemoryScope::Project => " AND memory_type = 'project'", + _ => "", + }; + let sql = format!( + "SELECT memory_uuid, memory_type, keywords FROM memory_info \ + WHERE status = 1 AND keywords LIKE ?1{} ORDER BY update_time DESC, id DESC", + type_sql + ); + let like = format!("%{}%", keyword.trim()); + let conn = self.lock()?; + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map(params![like], |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, String>(1)?, + row.get::<_, Option>(2)?, + )) + })?; + + let mut out = Vec::new(); + for row in rows { + let (memory_uuid, memory_type, kw) = row?; + match self.read_content(&memory_type, &memory_uuid) { + Ok(content) => out.push(MemoryRecord { + memory_uuid, + memory_type, + keywords: split_keywords(kw), + content, + }), + Err(e) => warn!("skip memory ({}): {}", memory_uuid, e), + } + } + Ok(out) + } + + /// 去重启发式：是否已存在同类型且关键词集合完全相同的有效记忆。 + pub fn exists_same(&self, memory_type: &str, keywords: &[String]) -> bool { + let kw = normalize_keywords(keywords); + if kw.is_empty() { + return false; + } + let conn = match self.lock() { + Ok(c) => c, + Err(_) => return false, + }; + conn.query_row( + "SELECT 1 FROM memory_info WHERE status = 1 AND memory_type = ?1 AND keywords = ?2 LIMIT 1", + params![memory_type, kw], + |_| Ok(()), + ) + .optional() + .ok() + .flatten() + .is_some() + } + + /// 渲染注入 Agent 系统提示词的记忆段；无记忆或作用域为 None 时返回空串。 + pub fn render_context(&self, scope: MemoryScope) -> Result { + let records = self.load(scope)?; + if records.is_empty() { + return Ok(String::new()); + } + let mut user_lines: Vec = Vec::new(); + let mut project_lines: Vec = Vec::new(); + for r in records { + let line = format!("- [{}] {}", r.keywords.join(", "), r.content.trim()); + if r.memory_type == "project" { + project_lines.push(line); + } else { + user_lines.push(line); + } + } + + let mut s = format!("# Memory (scope={})\n", scope.as_str()); + if !user_lines.is_empty() { + s.push_str("## User Memory\n"); + for l in &user_lines { + s.push_str(l); + s.push('\n'); + } + } + if !project_lines.is_empty() { + s.push_str("## Project Memory\n"); + for l in &project_lines { + s.push_str(l); + s.push('\n'); + } + } + + if s.chars().count() > MAX_INJECT_CHARS { + s = s.chars().take(MAX_INJECT_CHARS).collect(); + } + Ok(s) + } +} + +// ============================================================ +// 自动抽取 +// ============================================================ + +/// 从对话中抽取到的一条候选记忆（解析自 LLM 输出）。 +#[derive(Debug, Clone)] +pub struct ExtractedMemory { + /// 记忆类型："user" / "project"。 + pub memory_type: String, + /// 记忆正文。 + pub content: String, + /// 关键词。 + pub keywords: Vec, +} + +/// 记忆抽取指令。 +pub const EXTRACTION_PROMPT: &str = "\ +You are a memory extraction assistant for a local coding agent. +Review the conversation round below and extract ONLY information worth remembering LONG-TERM. + +Extract two kinds of memories: +- \"user\": durable user coding habits, tooling/style preferences, or communication preferences. +- \"project\": durable project facts, constraints, conventions, or business boundaries. + +Rules: +- Ignore one-off task details, transient state, secrets/credentials, and anything sensitive. +- Keep each memory concise and self-contained (one fact per entry). +- If nothing is worth remembering, output an empty array. + +Output ONLY a fenced code block tagged ```memory containing a JSON array. Each element: +{\"type\": \"user\" | \"project\", \"content\": \"\", \"keywords\": [\"\", ...]} + +Example: +```memory +[{\"type\": \"user\", \"content\": \"Prefers Rust with anyhow for errors\", \"keywords\": [\"rust\", \"error-handling\"]}] +``` +"; + +/// 把角色转为文本标识。 +fn role_str(role: &Role) -> &'static str { + match role { + Role::System => "system", + Role::User => "user", + Role::Assistant => "assistant", + } +} + +/// 把一组消息渲染为抽取用的对话文本。 +fn render_transcript(messages: &[Message]) -> String { + let mut s = String::new(); + for m in messages { + s.push_str(role_str(&m.role)); + s.push_str(": "); + s.push_str(m.content.trim()); + s.push('\n'); + } + s +} + +/// 解析 LLM 输出中的 ```memory 围栏 JSON 数组为候选记忆列表。 +/// +/// 取**最后一个** ```memory 围栏块；无围栏或 JSON 非法时返回空 Vec（不抛错）。 +/// 仅接受 `type` 为 "user"/"project" 且 `content` 非空的元素。 +pub fn parse_extracted_memories(text: &str) -> Vec { + const MARKER: &str = "```memory"; + let Some(start) = text.rfind(MARKER) else { + return Vec::new(); + }; + let after = &text[start + MARKER.len()..]; + let Some(end) = after.find("```") else { + return Vec::new(); + }; + let json_str = after[..end].trim(); + let value: Value = match serde_json::from_str(json_str) { + Ok(v) => v, + Err(_) => return Vec::new(), + }; + let Some(arr) = value.as_array() else { + return Vec::new(); + }; + + let mut out = Vec::new(); + for item in arr { + let memory_type = item + .get("type") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if memory_type != "user" && memory_type != "project" { + continue; + } + let content = item + .get("content") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if content.trim().is_empty() { + continue; + } + let keywords = item + .get("keywords") + .and_then(|v| v.as_array()) + .map(|a| { + a.iter() + .filter_map(|k| k.as_str().map(|s| s.trim().to_string())) + .filter(|s| !s.is_empty()) + .collect() + }) + .unwrap_or_default(); + out.push(ExtractedMemory { + memory_type, + content, + keywords, + }); + } + out +} + +/// 记忆抽取器：一轮对话结束后用独立 provider 调 LLM 总结记忆并落库。 +/// +/// 使用与 Agent 相同的供应商/模型配置，但构造**独立** provider（不借用 Agent 的 +/// `Box`），以便在 `tokio::spawn` 的独立任务中运行、不阻塞对话。 +pub struct MemoryExtractor { + provider_config: ProviderConfig, + model: String, + temperature: Option, + store: Arc, + scope: MemoryScope, +} + +impl MemoryExtractor { + /// 构造抽取器。 + pub fn new( + provider_config: ProviderConfig, + model: String, + temperature: Option, + store: Arc, + scope: MemoryScope, + ) -> Self { + Self { + provider_config, + model, + temperature, + store, + scope, + } + } + + /// 从一轮对话中抽取并写入记忆，返回实际写入条数。 + /// + /// 作用域为 None 或无消息时直接返回 0。抽取到的候选按作用域过滤、跳过空内容与 + /// 已存在（同类型 + 同关键词集合）者，其余生成新 uuid 落库。 + pub async fn extract_from_round(&self, messages: &[Message]) -> Result { + if self.scope == MemoryScope::None || messages.is_empty() { + return Ok(0); + } + let provider = create_provider(self.provider_config.clone())?; + let request = CompletionRequest { + model: self.model.clone(), + messages: vec![ + Message { + role: Role::System, + content: EXTRACTION_PROMPT.to_string(), + }, + Message { + role: Role::User, + content: render_transcript(messages), + }, + ], + temperature: self.temperature, + max_tokens: None, + stream: false, + }; + let response = provider.complete(request).await?; + let candidates = parse_extracted_memories(&response.content); + + let mut written = 0usize; + for c in candidates { + if !self.scope.allows_type(&c.memory_type) { + continue; + } + if self.store.exists_same(&c.memory_type, &c.keywords) { + continue; + } + let rec = MemoryRecord { + memory_uuid: new_memory_uuid(), + memory_type: c.memory_type, + keywords: c.keywords, + content: c.content, + }; + self.store.save(&rec)?; + written += 1; + } + Ok(written) + } +} + +// ============================================================ +// 测试 +// ============================================================ + +#[cfg(test)] +mod tests { + use super::*; + + fn temp_dir(tag: &str) -> PathBuf { + let p = std::env::temp_dir().join(format!("omega-mem-{}-{}", tag, new_memory_uuid())); + std::fs::create_dir_all(&p).unwrap(); + p + } + + fn test_store(with_project: bool) -> MemoryStore { + let conn = Connection::open_in_memory().unwrap(); + let project_dir = if with_project { Some(temp_dir("proj")) } else { None }; + MemoryStore::from_parts(conn, temp_dir("user"), project_dir) + } + + fn rec(uuid: &str, ty: &str, kws: &[&str], content: &str) -> MemoryRecord { + MemoryRecord { + memory_uuid: uuid.to_string(), + memory_type: ty.to_string(), + keywords: kws.iter().map(|s| s.to_string()).collect(), + content: content.to_string(), + } + } + + #[test] + fn scope_parse_and_allows() { + assert_eq!(MemoryScope::parse("USER"), MemoryScope::User); + assert_eq!(MemoryScope::parse(" project "), MemoryScope::Project); + assert_eq!(MemoryScope::parse("all"), MemoryScope::All); + assert_eq!(MemoryScope::parse("bogus"), MemoryScope::None); + assert_eq!(MemoryScope::parse(""), MemoryScope::None); + + assert!(MemoryScope::All.allows_type("user")); + assert!(MemoryScope::All.allows_type("project")); + assert!(MemoryScope::User.allows_type("user")); + assert!(!MemoryScope::User.allows_type("project")); + assert!(MemoryScope::Project.allows_type("project")); + assert!(!MemoryScope::None.allows_type("user")); + assert_eq!(MemoryScope::All.as_str(), "all"); + } + + #[test] + fn save_load_roundtrip_user() { + let store = test_store(false); + store + .save(&rec("u1", "user", &["rust", "errors"], "Prefers anyhow")) + .unwrap(); + // 正文文件写入用户目录。 + assert!(store.user_dir.join("u1.md").exists()); + + let loaded = store.load(MemoryScope::User).unwrap(); + assert_eq!(loaded.len(), 1); + assert_eq!(loaded[0].memory_uuid, "u1"); + assert_eq!(loaded[0].content, "Prefers anyhow"); + let mut kws = loaded[0].keywords.clone(); + kws.sort(); + assert_eq!(kws, vec!["errors", "rust"]); + } + + #[test] + fn scope_filters_records() { + let store = test_store(true); + store.save(&rec("u1", "user", &["a"], "user mem")).unwrap(); + store + .save(&rec("p1", "project", &["b"], "project mem")) + .unwrap(); + + assert_eq!(store.load(MemoryScope::User).unwrap().len(), 1); + assert_eq!(store.load(MemoryScope::Project).unwrap().len(), 1); + assert_eq!(store.load(MemoryScope::All).unwrap().len(), 2); + assert_eq!(store.load(MemoryScope::None).unwrap().len(), 0); + } + + #[test] + fn project_save_without_dir_errors() { + let store = test_store(false); + let err = store + .save(&rec("p1", "project", &["x"], "no project dir")) + .unwrap_err(); + assert!(err.to_string().contains("project memory unavailable")); + } + + #[test] + fn delete_removes_file_and_row() { + let store = test_store(false); + store.save(&rec("u1", "user", &["a"], "mem")).unwrap(); + assert!(store.user_dir.join("u1.md").exists()); + + assert!(store.delete("u1").unwrap()); + assert!(!store.user_dir.join("u1.md").exists()); + assert_eq!(store.load(MemoryScope::User).unwrap().len(), 0); + // 再次删除返回 false（无索引行）。 + assert!(!store.delete("u1").unwrap()); + } + + #[test] + fn save_is_upsert_by_uuid() { + let store = test_store(false); + store.save(&rec("u1", "user", &["a"], "v1")).unwrap(); + store.save(&rec("u1", "user", &["b"], "v2")).unwrap(); + let loaded = store.load(MemoryScope::User).unwrap(); + assert_eq!(loaded.len(), 1, "upsert should not duplicate row"); + assert_eq!(loaded[0].content, "v2"); + assert_eq!(loaded[0].keywords, vec!["b"]); + } + + #[test] + fn search_by_keyword() { + let store = test_store(false); + store + .save(&rec("u1", "user", &["rust", "tokio"], "async")) + .unwrap(); + store.save(&rec("u2", "user", &["python"], "sync")).unwrap(); + + let hits = store.search("tokio", MemoryScope::All).unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].memory_uuid, "u1"); + + assert_eq!(store.search("golang", MemoryScope::All).unwrap().len(), 0); + // 作用域约束：project 作用域下搜不到 user 记忆。 + assert_eq!(store.search("tokio", MemoryScope::Project).unwrap().len(), 0); + } + + #[test] + fn exists_same_dedups_by_keyword_set() { + let store = test_store(false); + store.save(&rec("u1", "user", &["a", "b"], "mem")).unwrap(); + // 顺序不同的同集合应判为已存在。 + assert!(store.exists_same( + "user", + &["b".to_string(), "a".to_string()] + )); + assert!(!store.exists_same("user", &["a".to_string()])); + assert!(!store.exists_same("project", &["a".to_string(), "b".to_string()])); + assert!(!store.exists_same("user", &[])); + } + + #[test] + fn render_context_formats_and_guards() { + let store = test_store(true); + assert_eq!(store.render_context(MemoryScope::None).unwrap(), ""); + assert_eq!(store.render_context(MemoryScope::All).unwrap(), ""); + + store.save(&rec("u1", "user", &["a"], "user mem")).unwrap(); + store + .save(&rec("p1", "project", &["b"], "project mem")) + .unwrap(); + let ctx = store.render_context(MemoryScope::All).unwrap(); + assert!(ctx.starts_with("# Memory (scope=all)")); + assert!(ctx.contains("## User Memory")); + assert!(ctx.contains("## Project Memory")); + assert!(ctx.contains("user mem")); + assert!(ctx.contains("project mem")); + } + + #[test] + fn parse_extracted_memories_valid() { + let text = "Here you go:\n```memory\n[\ + {\"type\":\"user\",\"content\":\"likes rust\",\"keywords\":[\"rust\"]},\ + {\"type\":\"project\",\"content\":\"uses sqlite\",\"keywords\":[\"db\"]}\ + ]\n```"; + let out = parse_extracted_memories(text); + assert_eq!(out.len(), 2); + assert_eq!(out[0].memory_type, "user"); + assert_eq!(out[0].content, "likes rust"); + assert_eq!(out[0].keywords, vec!["rust"]); + assert_eq!(out[1].memory_type, "project"); + } + + #[test] + fn parse_extracted_memories_filters_invalid_entries() { + let text = "```memory\n[\ + {\"type\":\"other\",\"content\":\"x\"},\ + {\"type\":\"user\",\"content\":\" \"},\ + {\"type\":\"user\",\"content\":\"ok\",\"keywords\":[\"k\"]}\ + ]\n```"; + let out = parse_extracted_memories(text); + assert_eq!(out.len(), 1); + assert_eq!(out[0].content, "ok"); + } + + #[test] + fn parse_extracted_memories_falls_back_empty() { + assert!(parse_extracted_memories("no fence here").is_empty()); + assert!(parse_extracted_memories("```memory\nnot json\n```").is_empty()); + assert!(parse_extracted_memories("```memory\n{\"type\":\"user\"}\n```").is_empty()); + } +} diff --git a/src/core/mod.rs b/src/core/mod.rs index cec949c..b396836 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -1,9 +1,15 @@ pub mod transport; pub mod tools; +pub mod mcp; pub mod config; pub mod runtime; pub mod db; pub mod agent; +pub mod arena; +pub mod memory; +pub mod project; +pub mod prompt; +pub mod scheduler; pub mod providers; pub mod client; mod errors; \ No newline at end of file diff --git a/src/core/project/mod.rs b/src/core/project/mod.rs new file mode 100644 index 0000000..e695296 --- /dev/null +++ b/src/core/project/mod.rs @@ -0,0 +1,193 @@ +//! 项目信息自动登记 +//! +//! 对应《系统功能设计与数据库设计文档》四.13 `project_info`：以工作区绝对路径 `path` +//! 唯一标识一个项目。CLI-First 下不额外提供命令，改由运行时在 **workspace 模式** 启动 +//! TUI 时自动登记当前打开的项目（见 `src/lib.rs::launch_tui`）： +//! - 首次打开：写入 `name`（目录名）+ `path` + `create_at`/`updata_at`； +//! - 重复打开：仅刷新 `updata_at` 并复活 `is_deleted`，`id` 保持稳定。 +//! +//! `create_at`/`updata_at` 为项目业务时间，统一用 `%Y-%m-%d %H:%M:%S`（`parse_db_time` 可解析）。 + +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result}; +use chrono::Utc; + +use crate::core::db::DatabaseManager; +use crate::db::models::project_info::ProjectInfo; + +/// 生成业务时间字符串（`%Y-%m-%d %H:%M:%S`，UTC）。 +fn now_str() -> String { + Utc::now().format("%Y-%m-%d %H:%M:%S").to_string() +} + +/// 登记（或刷新）一个项目，返回其 `project_info.id`。 +/// +/// 以规整后的绝对路径为业务唯一键：命中既有记录则刷新 `updata_at` 并复活 `is_deleted`； +/// 否则新建记录（`name` 取路径末段目录名，缺失时回退整路径串）。 +pub fn register(db: &DatabaseManager, path: &Path) -> Result { + let abs = normalize(path)?; + let path_str = abs.to_string_lossy().to_string(); + let name = abs + .file_name() + .map(|n| n.to_string_lossy().to_string()) + .filter(|n| !n.is_empty()) + .unwrap_or_else(|| path_str.clone()); + + let now = now_str(); + + let existing = db.find_where::("path = ?1", &[&path_str])?; + if let Some(mut row) = existing.into_iter().next() { + // 重复打开：刷新业务更新时间并复活逻辑删除标记。 + let id = row.id; + row.updata_at = Some(now); + row.is_deleted = 0; + db.update(row)?; + return Ok(id); + } + + let row = ProjectInfo { + id: 0, + name, + path: path_str, + create_at: Some(now.clone()), + updata_at: Some(now), + is_deleted: 0, + create_time: Utc::now(), + update_time: Utc::now(), + remark: None, + }; + let id = db.create(row)?; + Ok(id as i32) +} + +/// 当运行于 workspace 模式时，登记当前打开的工作区。 +/// +/// 仅在 `OMEGA_MODE == "workspace"` 时实际写库，工作区路径取自 `OMEGA_WORKING_DIR` +/// （回退进程当前目录）；非 workspace 模式返回 `Ok(None)`，不做任何登记。 +pub fn register_opened_from_env() -> Result> { + let mode = std::env::var("OMEGA_MODE").unwrap_or_default(); + if mode != "workspace" { + return Ok(None); + } + + let path = match std::env::var("OMEGA_WORKING_DIR") { + Ok(dir) if !dir.trim().is_empty() => PathBuf::from(dir), + _ => std::env::current_dir().context("failed to resolve working dir")?, + }; + + let db = DatabaseManager::new()?; + register(&db, &path).map(Some) +} + +/// 列出全部项目记录。 +pub fn list(db: &DatabaseManager) -> Result> { + db.read_all::() +} + +/// 物理删除一个项目记录，返回受影响行数。 +pub fn remove(db: &DatabaseManager, id: i32) -> Result { + db.delete::(id) +} + +/// 规整为绝对路径（不要求路径存在，仅做规范化）。 +fn normalize(path: &Path) -> Result { + if path.is_absolute() { + Ok(path.to_path_buf()) + } else { + let base = std::env::current_dir().context("failed to resolve current dir")?; + Ok(base.join(path)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::db::Model; + + /// 构造含 `project_info` 表的内存库。 + fn test_db() -> DatabaseManager { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + conn.execute_batch(ProjectInfo::create_table_sql()).unwrap(); + DatabaseManager { conn } + } + + #[test] + fn register_creates_with_dir_name() { + let db = test_db(); + let path = if cfg!(windows) { + Path::new("C:\\workspace\\demo") + } else { + Path::new("/workspace/demo") + }; + + let id = register(&db, path).unwrap(); + assert!(id > 0); + + let row = db.read_by_id::(id).unwrap().unwrap(); + assert_eq!(row.name, "demo"); + assert!(row.create_at.is_some()); + assert!(row.updata_at.is_some()); + assert_eq!(row.is_deleted, 0); + } + + #[test] + fn register_same_path_is_stable() { + let db = test_db(); + let path = if cfg!(windows) { + Path::new("C:\\workspace\\demo") + } else { + Path::new("/workspace/demo") + }; + + let first = register(&db, path).unwrap(); + let second = register(&db, path).unwrap(); + assert_eq!(first, second); + + let rows = db.read_all::().unwrap(); + assert_eq!(rows.len(), 1); + } + + #[test] + fn register_revives_deleted() { + let db = test_db(); + let path = if cfg!(windows) { + Path::new("C:\\workspace\\demo") + } else { + Path::new("/workspace/demo") + }; + + let id = register(&db, path).unwrap(); + // 手动置为已删除，再登记应复活。 + let mut row = db.read_by_id::(id).unwrap().unwrap(); + row.is_deleted = 1; + db.update(row).unwrap(); + + register(&db, path).unwrap(); + let row = db.read_by_id::(id).unwrap().unwrap(); + assert_eq!(row.is_deleted, 0); + } + + #[test] + fn list_counts_and_remove_deletes() { + let db = test_db(); + let p1 = if cfg!(windows) { + Path::new("C:\\workspace\\a") + } else { + Path::new("/workspace/a") + }; + let p2 = if cfg!(windows) { + Path::new("C:\\workspace\\b") + } else { + Path::new("/workspace/b") + }; + + let id1 = register(&db, p1).unwrap(); + register(&db, p2).unwrap(); + assert_eq!(list(&db).unwrap().len(), 2); + + assert_eq!(remove(&db, id1).unwrap(), 1); + assert_eq!(list(&db).unwrap().len(), 1); + assert!(db.read_by_id::(id1).unwrap().is_none()); + } +} diff --git a/src/core/prompt/mod.rs b/src/core/prompt/mod.rs new file mode 100644 index 0000000..bd66d35 --- /dev/null +++ b/src/core/prompt/mod.rs @@ -0,0 +1,294 @@ +//! Prompt 占位符解析 +//! +//! 对应《系统功能设计与数据库设计文档》四.5 `prompt_placeholder`：系统提示词中的 `{key}` +//! 占位符在加载 Agent 时被替换为实际值。内置两个约定占位符： +//! - `{AgentList}`：当前可调用的其它 Agent 名称列表； +//! - `{MCPList}`：当前 Agent 已授权的 MCP 工具名称列表； +//! +//! 其余占位符来自 `prompt_placeholder` 表（`key`→`value`，缺省回退 `default_value`）。 +//! 未知/空/跨行/嵌套花括号的占位符**原样保留**，避免因模板书写差异破坏正文。 + +use std::collections::HashMap; +use std::collections::HashSet; + +/// 将 `template` 中的 `{key}` 占位符按 `vars` 替换。 +/// +/// 扫描按 ASCII 花括号字节位切片（多字节 UTF-8 续字节均 `>= 0x80`，故不会误伤中文正文）： +/// - 命中 `vars` 的键：替换为其值； +/// - 未知键、空 `{}`、含换行或嵌套花括号：原样保留； +/// - 无闭合 `}`：原样输出剩余部分。 +pub fn render(template: &str, vars: &HashMap) -> String { + let bytes = template.as_bytes(); + let mut out = String::new(); + let mut i = 0usize; + + while i < bytes.len() { + // 定位下一个 `{`（其前的正文原样拷贝）。 + let Some(rel) = bytes[i..].iter().position(|&b| b == b'{') else { + out.push_str(&template[i..]); + break; + }; + let brace = i + rel; + out.push_str(&template[i..brace]); + + let after = brace + 1; + // 无闭合 `}`：原样输出剩余并结束。 + let Some(rel_end) = bytes[after..].iter().position(|&b| b == b'}') else { + out.push_str(&template[brace..]); + break; + }; + let end = after + rel_end; // `}` 的字节下标 + let key = &template[after..end]; + + let placeholder_valid = + !key.is_empty() && !key.contains('\n') && !key.contains('{') && !key.contains('}'); + + if placeholder_valid { + match vars.get(key) { + Some(val) => out.push_str(val), + // 未知键：原样保留 `{key}`。 + None => out.push_str(&template[brace..=end]), + } + } else { + // 非法占位符（空/跨行/嵌套）：原样保留。 + out.push_str(&template[brace..=end]); + } + + i = end + 1; + } + + out +} + +/// 构建占位符变量表。 +/// +/// - `AgentList`/`MCPList` 渲染为 `- item` 多行列表（空集合则为空串）； +/// - 自定义占位符 `key → value.or(default_value).unwrap_or("")`（`value` 优先于 `default_value`）。 +pub fn build_vars( + callable_agents: &[String], + mcp_tools: &[String], + placeholders: &[(String, Option, Option)], +) -> HashMap { + let mut vars = HashMap::new(); + vars.insert("AgentList".to_string(), render_list(callable_agents)); + vars.insert("MCPList".to_string(), render_list(mcp_tools)); + + for (key, value, default_value) in placeholders { + let resolved = value.clone().or_else(|| default_value.clone()).unwrap_or_default(); + vars.insert(key.clone(), resolved); + } + + vars +} + +/// 将条目渲染为 `- item` 多行列表；空集合返回空串。 +fn render_list(items: &[String]) -> String { + items + .iter() + .map(|item| format!("- {}", item)) + .collect::>() + .join("\n") +} + +/// 占位符扫描结果（用于§3.6.5 语法合法性校验）。 +#[derive(Debug, Default, PartialEq, Eq)] +pub struct PlaceholderScan { + /// 合法占位符键（按出现顺序去重）。 + pub keys: Vec, + /// 非法花括号片段（空 `{}`、跨行、嵌套），原样保留。 + pub malformed: Vec, + /// 是否存在未闭合的 `{`。 + pub unclosed: bool, +} + +/// 扫描模板中的占位符，规则与 [`render`] 一致。 +/// +/// - `{key}` 且 key 非空、不含换行/花括号 → 合法键（去重）； +/// - 空 `{}`、跨行、嵌套花括号 → 记为非法片段； +/// - 出现 `{` 但无后续 `}` → 记为未闭合。 +pub fn scan_placeholders(template: &str) -> PlaceholderScan { + let bytes = template.as_bytes(); + let mut scan = PlaceholderScan::default(); + let mut i = 0usize; + + while i < bytes.len() { + let Some(rel) = bytes[i..].iter().position(|&b| b == b'{') else { + break; + }; + let brace = i + rel; + let after = brace + 1; + let Some(rel_end) = bytes[after..].iter().position(|&b| b == b'}') else { + scan.unclosed = true; + break; + }; + let end = after + rel_end; + let key = &template[after..end]; + let valid = + !key.is_empty() && !key.contains('\n') && !key.contains('{') && !key.contains('}'); + if valid { + if !scan.keys.iter().any(|k| k == key) { + scan.keys.push(key.to_string()); + } + } else { + scan.malformed.push(template[brace..=end].to_string()); + } + i = end + 1; + } + + scan +} + +/// 提示词校验报告（§3.6.5 语法 + 占位符完整性）。 +#[derive(Debug, PartialEq, Eq)] +pub struct ValidationReport { + /// 模板中出现的合法占位符键（去重）。 + pub keys: Vec, + /// 未在 `known` 中声明的键（完整性缺失）。 + pub unresolved: Vec, + /// 非法花括号片段。 + pub malformed: Vec, + /// 是否存在未闭合花括号。 + pub unclosed: bool, +} + +impl ValidationReport { + /// 无任何语法与完整性问题时为真。 + pub fn is_ok(&self) -> bool { + self.unresolved.is_empty() && self.malformed.is_empty() && !self.unclosed + } +} + +/// 校验模板：扫描占位符，并以 `known`（已声明的占位符键集）判定完整性。 +pub fn validate_prompt(template: &str, known: &HashSet) -> ValidationReport { + let scan = scan_placeholders(template); + let unresolved = scan + .keys + .iter() + .filter(|k| !known.contains(*k)) + .cloned() + .collect(); + ValidationReport { + keys: scan.keys, + unresolved, + malformed: scan.malformed, + unclosed: scan.unclosed, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn vars_of(pairs: &[(&str, &str)]) -> HashMap { + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + } + + #[test] + fn render_replaces_known_key() { + let vars = vars_of(&[("Name", "Omega")]); + assert_eq!(render("Hi {Name}!", &vars), "Hi Omega!"); + } + + #[test] + fn render_keeps_unknown_key() { + let vars = vars_of(&[("Name", "Omega")]); + assert_eq!(render("{Name} and {Other}", &vars), "Omega and {Other}"); + } + + #[test] + fn render_without_braces_is_unchanged() { + let vars = vars_of(&[]); + assert_eq!(render("plain text 中文正文", &vars), "plain text 中文正文"); + } + + #[test] + fn render_handles_unclosed_brace() { + let vars = vars_of(&[("Name", "Omega")]); + assert_eq!(render("Hi {Name", &vars), "Hi {Name"); + } + + #[test] + fn render_keeps_empty_and_nested_and_multiline() { + let vars = vars_of(&[("a", "A")]); + assert_eq!(render("{}", &vars), "{}"); + assert_eq!(render("{a{b}}", &vars), "{a{b}}"); + assert_eq!(render("{a\nb}", &vars), "{a\nb}"); + } + + #[test] + fn render_preserves_chinese_around_placeholder() { + let vars = vars_of(&[("名称", "欧米茄")]); + assert_eq!(render("你好 {名称} 世界", &vars), "你好 欧米茄 世界"); + } + + #[test] + fn build_vars_lists_and_custom() { + let agents = vec!["Coder".to_string(), "Writer".to_string()]; + let tools = vec!["add".to_string()]; + let placeholders = vec![ + ("City".to_string(), Some("Beijing".to_string()), Some("N/A".to_string())), + ("Lang".to_string(), None, Some("Rust".to_string())), + ("Empty".to_string(), None, None), + ]; + + let vars = build_vars(&agents, &tools, &placeholders); + assert_eq!(vars["AgentList"], "- Coder\n- Writer"); + assert_eq!(vars["MCPList"], "- add"); + // value 优先于 default_value。 + assert_eq!(vars["City"], "Beijing"); + // 无 value 回退 default_value。 + assert_eq!(vars["Lang"], "Rust"); + // 两者皆无为空串。 + assert_eq!(vars["Empty"], ""); + } + + #[test] + fn build_vars_empty_lists_are_empty_string() { + let vars = build_vars(&[], &[], &[]); + assert_eq!(vars["AgentList"], ""); + assert_eq!(vars["MCPList"], ""); + } + + #[test] + fn scan_collects_keys_dedup_and_malformed() { + let scan = scan_placeholders("{A} {B} {A} {} {c{d}}"); + assert_eq!(scan.keys, vec!["A".to_string(), "B".to_string()]); + assert_eq!(scan.malformed, vec!["{}".to_string(), "{c{d}".to_string()]); + assert!(!scan.unclosed); + } + + #[test] + fn scan_flags_unclosed_brace() { + let scan = scan_placeholders("ok {A} then {B"); + assert_eq!(scan.keys, vec!["A".to_string()]); + assert!(scan.unclosed); + } + + #[test] + fn scan_plain_text_has_nothing() { + let scan = scan_placeholders("no placeholders 中文"); + assert!(scan.keys.is_empty()); + assert!(scan.malformed.is_empty()); + assert!(!scan.unclosed); + } + + #[test] + fn validate_prompt_reports_unresolved_and_syntax() { + let mut known = HashSet::new(); + known.insert("AgentList".to_string()); + known.insert("City".to_string()); + + // City/AgentList 已声明，Lang 未声明，另有一个非法片段。 + let report = validate_prompt("{City} {AgentList} {Lang} {}", &known); + assert_eq!(report.unresolved, vec!["Lang".to_string()]); + assert_eq!(report.malformed, vec!["{}".to_string()]); + assert!(!report.is_ok()); + + // 全部已声明且无非法片段时通过。 + let ok = validate_prompt("{City} {AgentList}", &known); + assert!(ok.is_ok()); + } +} diff --git a/src/core/providers/ollama/provider.rs b/src/core/providers/ollama/provider.rs index 9032181..1c0aa69 100644 --- a/src/core/providers/ollama/provider.rs +++ b/src/core/providers/ollama/provider.rs @@ -59,4 +59,57 @@ impl Provider for OllamaProvider { usage: None, }) } + + async fn list_models(&self) -> Result> { + let mut headers = HashMap::new(); + headers.insert("Content-Type".to_string(), "application/json".to_string()); + + let url = format!("{}/api/tags", self.base_url); + + let value: serde_json::Value = self + .client + .get_json(&url, headers, &HashMap::::new()) + .await + .map_err(|e| anyhow::anyhow!("Ollama tags API error: {}", e))?; + + Ok(parse_model_names(&value)) + } +} + +/// 从 Ollama `/api/tags` 响应（`{"models":[{"name":...}]}`）中提取模型名。 +fn parse_model_names(value: &serde_json::Value) -> Vec { + value + .get("models") + .and_then(|m| m.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|item| item.get("name").and_then(|n| n.as_str())) + .map(|s| s.to_string()) + .collect() + }) + .unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use super::parse_model_names; + use serde_json::json; + + #[test] + fn extracts_names_from_models_array() { + let value = json!({ + "models": [ + {"name": "llama3:latest", "size": 1}, + {"name": "qwen2.5:7b", "size": 2} + ] + }); + assert_eq!(parse_model_names(&value), vec!["llama3:latest", "qwen2.5:7b"]); + } + + #[test] + fn missing_or_malformed_models_yields_empty() { + assert!(parse_model_names(&json!({})).is_empty()); + assert!(parse_model_names(&json!({"models": "nope"})).is_empty()); + assert!(parse_model_names(&json!({"models": [{"size": 1}]})).is_empty()); + } } \ No newline at end of file diff --git a/src/core/providers/openai_compatible/provider.rs b/src/core/providers/openai_compatible/provider.rs index e90b25a..814a1b3 100644 --- a/src/core/providers/openai_compatible/provider.rs +++ b/src/core/providers/openai_compatible/provider.rs @@ -54,6 +54,21 @@ impl Provider for OpenAiCompatibleProvider { Ok(CompletionResponse::from(response)) } + async fn list_models(&self) -> Result> { + let mut headers = HashMap::new(); + headers.insert("Authorization".to_string(), format!("Bearer {}", self.api_key)); + + let url = format!("{}/models", self.base_url); + + let value: serde_json::Value = self + .client + .get_json(&url, headers, &HashMap::::new()) + .await + .map_err(|e| anyhow::anyhow!("OpenAI models API error: {}", e))?; + + Ok(parse_model_ids(&value)) + } + async fn complete_stream( &self, mut request: CompletionRequest, @@ -117,4 +132,44 @@ impl Provider for OpenAiCompatibleProvider { Ok(Box::pin(stream)) } +} + +/// 从 OpenAI 兼容 `/models` 响应（`{"data":[{"id":...}]}`）中提取模型标识。 +fn parse_model_ids(value: &serde_json::Value) -> Vec { + value + .get("data") + .and_then(|d| d.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|item| item.get("id").and_then(|i| i.as_str())) + .map(|s| s.to_string()) + .collect() + }) + .unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use super::parse_model_ids; + use serde_json::json; + + #[test] + fn extracts_ids_from_data_array() { + let value = json!({ + "object": "list", + "data": [ + {"id": "gpt-4o", "object": "model"}, + {"id": "gpt-4o-mini", "object": "model"} + ] + }); + assert_eq!(parse_model_ids(&value), vec!["gpt-4o", "gpt-4o-mini"]); + } + + #[test] + fn missing_or_malformed_data_yields_empty() { + assert!(parse_model_ids(&json!({})).is_empty()); + assert!(parse_model_ids(&json!({"data": "nope"})).is_empty()); + // 缺少 id 字段的项被跳过。 + assert!(parse_model_ids(&json!({"data": [{"object": "model"}]})).is_empty()); + } } \ No newline at end of file diff --git a/src/core/providers/provider.rs b/src/core/providers/provider.rs index 169b300..3835683 100644 --- a/src/core/providers/provider.rs +++ b/src/core/providers/provider.rs @@ -12,6 +12,13 @@ pub trait Provider: Send + Sync { async fn complete(&self, request: CompletionRequest) -> Result; + /// 拉取该供应商可用的模型标识列表（用于《设计文档》§3.4 模型自动同步）。 + /// + /// 默认不支持；由具备模型清单接口的供应商（如 OpenAI 兼容、Ollama）覆盖实现。 + async fn list_models(&self) -> Result> { + anyhow::bail!("this provider does not support listing models") + } + /// 流式补全。默认实现回退到 [`Provider::complete`]， /// 将完整结果作为单个增量片段发出；支持原生流式的供应商可覆盖此方法。 async fn complete_stream( diff --git a/src/core/scheduler/mod.rs b/src/core/scheduler/mod.rs new file mode 100644 index 0000000..7b5bd1a --- /dev/null +++ b/src/core/scheduler/mod.rs @@ -0,0 +1,958 @@ +//! 多 Agent 调度器（Coordinator / Tool / Worker 三层架构） +//! +//! 对应《调度器设计与实现》：Coordinator 作为「顶层协调者」负责拆分任务、调度 Worker +//! 子代理并行执行、汇总结果；本模块提供支撑调度层的四个工具（创建 / 续跑 / 停止 / 等待） +//! 以及 Worker 生命周期管理（进程内内存池）。 +//! +//! 设计要点： +//! - **解耦**：调度层不直接依赖 [`Provider`](crate::core::providers::Provider) / `AiClient`， +//! 而是经 [`WorkerContext`] / [`WorkerFactory`] 抽象驱动 Worker。生产环境用真实 [`Agent`] +//! 实现（[`AgentWorkerFactory`]），测试用 echo 桩实现，从而调度逻辑可独立单测。 +//! - **并发**：`create_worker` 后台异步起 Worker 立即返回 `worker_id`，`await_workers` +//! 阻塞收集标准化结果，实现多 Worker 并行。 +//! - **可追溯**：Worker 各自建立独立会话（经 `chat_session`/`chat_message` 持久化）， +//! 进程内池保留上下文以支持续跑（`resume_worker`）。 + +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use anyhow::{Context, Result}; +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use tokio::sync::mpsc; +use tokio::task::AbortHandle; + +use crate::core::agent::{Agent, AgentEvent}; +use crate::core::providers::{create_provider, ProviderConfig}; +use crate::core::runtime::Runtime; +use crate::core::tools::{Tool, ToolRegistry}; + +/// `await_workers` 轮询间隔。 +const POLL_INTERVAL: Duration = Duration::from_millis(50); +/// `await_workers` 总超时（防止个别 Worker 卡死导致协调者永久阻塞）。 +pub const AWAIT_TIMEOUT: Duration = Duration::from_secs(600); + +// ============================================================ +// 类型 +// ============================================================ + +/// Worker 执行状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum WorkerStatus { + /// 运行中。 + Running, + /// 已成功完成。 + Completed, + /// 执行失败。 + Failed, + /// 被协调者中止。 + Stopped, +} + +impl WorkerStatus { + pub fn as_str(&self) -> &'static str { + match self { + WorkerStatus::Running => "running", + WorkerStatus::Completed => "completed", + WorkerStatus::Failed => "failed", + WorkerStatus::Stopped => "stopped", + } + } + + /// 是否处于终止态（不再变化，`await_workers` 可停止等待）。 + pub fn is_terminal(&self) -> bool { + !matches!(self, WorkerStatus::Running) + } +} + +/// Worker 上报的标准化结果。 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct WorkerReport { + /// Worker 唯一标识。 + pub worker_id: String, + /// 任务标识（取自 Worker 上报，缺省回退 `worker_id`）。 + pub task_id: String, + /// 业务状态："success" / "failed"。 + pub status: String, + /// 结果摘要。 + pub summary: String, + /// 具体产出（文件路径 / commit 哈希 / 关键发现等）。 + pub artifacts: Vec, +} + +/// Worker 执行上下文抽象：执行一条自包含指令并返回最终产出文本。 +/// +/// 实现内部保留历史，以支持 [`ResumeWorkerTool`] 复用上下文续跑。 +#[async_trait] +pub trait WorkerContext: Send + Sync { + /// 执行一条指令，返回 Worker 的最终产出文本（应含标准化 `result` 块）。 + async fn run_instruction(&self, instruction: String) -> Result; +} + +/// Worker 工厂：按子任务类型创建一个独立的 [`WorkerContext`]。 +pub trait WorkerFactory: Send + Sync { + fn create(&self, task_type: &str) -> Arc; +} + +/// 池中一个 Worker 的运行态句柄。 +struct WorkerHandle { + status: WorkerStatus, + ctx: Arc, + report: Option, + abort: Option, +} + +/// Worker 池：进程内内存态，管理所有 Worker 的生命周期与结果。 +/// +/// 内部用 [`std::sync::Mutex`] 保护，临界区短小、**不跨 `await` 持有**。 +#[derive(Default)] +pub struct WorkerPool { + inner: std::sync::Mutex>, +} + +impl WorkerPool { + fn guard(&self) -> std::sync::MutexGuard<'_, HashMap> { + // 中毒时恢复内部数据，避免因单次 panic 使整个池不可用。 + self.inner.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// 登记一个新 Worker（初始为 Running）。 + pub fn insert(&self, id: String, ctx: Arc) { + self.guard().insert( + id, + WorkerHandle { + status: WorkerStatus::Running, + ctx, + report: None, + abort: None, + }, + ); + } + + /// 回填后台任务的 abort 句柄（供 `stop_worker` 中止）。 + pub fn set_abort(&self, id: &str, abort: AbortHandle) { + if let Some(h) = self.guard().get_mut(id) { + h.abort = Some(abort); + } + } + + /// 标记 Worker 终止并写入报告。 + pub fn finish(&self, id: &str, status: WorkerStatus, report: Option) { + if let Some(h) = self.guard().get_mut(id) { + h.status = status; + h.report = report; + h.abort = None; + } + } + + /// 置回 Running（续跑前调用），清空旧报告。 + pub fn set_running(&self, id: &str) { + if let Some(h) = self.guard().get_mut(id) { + h.status = WorkerStatus::Running; + h.report = None; + } + } + + /// 中止 Worker：触发 abort 并置为 Stopped；保留 ctx 以便后续续跑。 + /// 返回该 Worker 是否存在。 + pub fn stop(&self, id: &str) -> bool { + let mut g = self.guard(); + match g.get_mut(id) { + Some(h) => { + if let Some(a) = h.abort.take() { + a.abort(); + } + h.status = WorkerStatus::Stopped; + true + } + None => false, + } + } + + /// 取出 Worker 的执行上下文（用于续跑，复用历史）。 + pub fn get_ctx(&self, id: &str) -> Option> { + self.guard().get(id).map(|h| h.ctx.clone()) + } + + /// 当前所有 Worker 的 id。 + pub fn ids(&self) -> Vec { + self.guard().keys().cloned().collect() + } + + /// 指定的所有 id 是否均已终止（不存在的 id 视为无需等待）。 + pub fn all_terminal(&self, ids: &[String]) -> bool { + let g = self.guard(); + ids.iter().all(|id| match g.get(id) { + Some(h) => h.status.is_terminal(), + None => true, + }) + } + + /// 收集指定 id 的状态与报告快照（跳过不存在的 id）。 + pub fn snapshot(&self, ids: &[String]) -> Vec<(String, WorkerStatus, Option)> { + let g = self.guard(); + ids.iter() + .filter_map(|id| { + g.get(id) + .map(|h| (id.clone(), h.status.clone(), h.report.clone())) + }) + .collect() + } +} + +// ============================================================ +// 生产实现：基于真实 Agent 的 Worker +// ============================================================ + +/// 生产环境的 Worker 工厂：用协调者的供应商/模型配置构造独立 [`Agent`]。 +pub struct AgentWorkerFactory { + /// 供应商配置（复用协调者 Agent）。 + pub provider_config: ProviderConfig, + /// 工作目录（Workspace 路径）。 + pub working_dir: PathBuf, + /// 采样温度。 + pub temperature: Option, + /// 核采样 top_p。 + pub top_p: Option, + /// 单次输出上限。 + pub max_tokens: Option, +} + +impl AgentWorkerFactory { + fn build_agent(&self, task_type: &str) -> Result { + let provider = create_provider(self.provider_config.clone())?; + let runtime = Arc::new(Runtime::new(self.working_dir.clone())); + let tools = ToolRegistry::with_defaults(runtime); + + let key = if self.provider_config.api_key.is_empty() { + None + } else { + Some(self.provider_config.api_key.clone()) + }; + + let mut agent = Agent::new( + provider, + self.provider_config.model.clone(), + WORKER_PROMPT.to_string(), + tools, + ) + .with_provider_meta( + self.provider_config.kind.as_str(), + self.provider_config.base_url.clone(), + key, + ) + .with_sampling(self.temperature, self.top_p, self.max_tokens); + + // 每个 Worker 建立独立会话，便于复盘（尽力而为，失败仅告警）。 + if let Err(e) = agent.enable_persistence_auto(Some(task_type)) { + log::warn!("worker persistence init failed: {e}"); + } + Ok(agent) + } +} + +impl WorkerFactory for AgentWorkerFactory { + fn create(&self, task_type: &str) -> Arc { + match self.build_agent(task_type) { + Ok(agent) => Arc::new(AgentWorkerContext { + agent: tokio::sync::Mutex::new(agent), + }), + // 构造失败不 panic：返回一个运行期报错的上下文，交由调度层标记 Failed。 + Err(e) => Arc::new(FailedWorkerContext { + error: e.to_string(), + }), + } + } +} + +/// 基于 [`Agent`] 的 Worker 上下文。 +pub struct AgentWorkerContext { + agent: tokio::sync::Mutex, +} + +#[async_trait] +impl WorkerContext for AgentWorkerContext { + async fn run_instruction(&self, instruction: String) -> Result { + let (tx, mut rx) = mpsc::unbounded_channel::(); + let mut agent = self.agent.lock().await; + // 无界通道不会阻塞发送：先驱动整轮，再收集累积的助手文本。 + agent.run_turn(&instruction, tx).await?; + let mut acc = String::new(); + while let Some(ev) = rx.recv().await { + if let AgentEvent::Token(t) = ev { + acc.push_str(&t); + } + } + Ok(acc) + } +} + +/// 构造失败时的兜底上下文：运行即报错。 +struct FailedWorkerContext { + error: String, +} + +#[async_trait] +impl WorkerContext for FailedWorkerContext { + async fn run_instruction(&self, _instruction: String) -> Result { + anyhow::bail!("worker unavailable: {}", self.error) + } +} + +// ============================================================ +// 提示词 +// ============================================================ + +/// 追加到协调者 Agent 系统提示词后的调度器说明。 +pub const COORDINATOR_PROMPT: &str = "\n\n# Coordinator Role (Multi-Agent Scheduler)\n\ +You are the top-level Coordinator. You DO NOT execute concrete tasks yourself; you decompose the \ +user's software-engineering request into self-contained subtasks, orchestrate Worker sub-agents \ +through the scheduler tools, and aggregate their results.\n\n\ +Pipeline: research -> implement -> verify.\n\n\ +Scheduler tools:\n\ +- `create_worker` {instruction, task_type?}: start a Worker asynchronously; returns worker_id \ +immediately. Use it for parallelizable work.\n\ +- `await_workers` {ids?}: block until the given (or all) workers finish; returns their standardized \ +reports. ALWAYS await before aggregating.\n\ +- `resume_worker` {worker_id, instruction}: give an existing worker a new instruction, reusing its \ +context (use for fixes after a failure).\n\ +- `stop_worker` {worker_id}: terminate a misguided worker (it can be resumed later).\n\n\ +Rules:\n\ +- Instructions MUST be self-contained: include file paths, exact operations, and completion \ +criteria. Workers cannot see this conversation.\n\ +- Parallelism: read-only tasks (research/query) run in parallel; write tasks are serialized per \ +file (never let two workers edit the same file concurrently).\n\ +- Aggregation is YOUR job: never ask a worker to inspect another worker's output. Wait for all \ +parallel subtasks via `await_workers`, then integrate the reports yourself.\n\ +- On worker failure, prefer `resume_worker` with a corrective instruction reusing its error \ +context; after repeated failures, stop and report to the user.\n\ +- Trivial actions (reading one file, running one command) do not need a worker.\n\n\ +Each worker ends with a fenced ```result JSON block: \ +{\"task_id\",\"status\":\"success|failed\",\"summary\",\"artifacts\":[...]}. `await_workers` returns \ +these reports for you to summarize.\n"; + +/// Worker 子代理的系统提示词。 +const WORKER_PROMPT: &str = "You are a Worker sub-agent in a multi-agent scheduler. You receive a \ +single self-contained instruction and must complete it independently using the available file \ +tools (read_file/write_file/list_dir) and any mounted MCP tools. You have NO access to the \ +coordinator's or the user's conversation; rely only on the instruction.\n\n\ +Do the work requested, then END your reply with exactly one fenced code block tagged `result` \ +containing JSON:\n\ +```result\n\ +{\"task_id\": \"\", \"status\": \ +\"success\", \"summary\": \"\", \"artifacts\": [\"\"]}\n\ +```\n\ +If you cannot complete the instruction, set status to \"failed\" and explain why in summary. Do \ +not fabricate results."; + +// ============================================================ +// 结果解析 +// ============================================================ + +/// 从 Worker 产出文本解析标准化报告（取最后一个 ```result 围栏块）。 +/// +/// 无有效围栏块时回退：`status="success"`、`summary` 为文本截断（前 500 字）。 +pub fn parse_worker_report(worker_id: &str, text: &str) -> WorkerReport { + const MARKER: &str = "```result"; + if let Some(start) = text.rfind(MARKER) { + let after = &text[start + MARKER.len()..]; + if let Some(end) = after.find("```") { + let json_str = after[..end].trim(); + if let Ok(v) = serde_json::from_str::(json_str) { + let raw_status = v + .get("status") + .and_then(|s| s.as_str()) + .unwrap_or("success"); + let status = if raw_status == "failed" { + "failed".to_string() + } else { + "success".to_string() + }; + return WorkerReport { + worker_id: worker_id.to_string(), + task_id: v + .get("task_id") + .and_then(|s| s.as_str()) + .unwrap_or(worker_id) + .to_string(), + status, + summary: v + .get("summary") + .and_then(|s| s.as_str()) + .unwrap_or("") + .to_string(), + artifacts: v + .get("artifacts") + .and_then(|a| a.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|x| x.as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(), + }; + } + } + } + + let summary: String = text.chars().take(500).collect(); + WorkerReport { + worker_id: worker_id.to_string(), + task_id: worker_id.to_string(), + status: "success".to_string(), + summary, + artifacts: vec![], + } +} + +/// 生成一个 Worker 唯一标识。 +fn new_worker_id() -> String { + use rand::RngExt; + let n: u32 = rand::rng().random_range(0..u32::MAX); + format!("w-{:08x}", n) +} + +/// 后台运行一个 Worker 指令并把结果写回池（`create_worker`/`resume_worker` 共用）。 +async fn run_and_finish( + pool: Arc, + worker_id: String, + ctx: Arc, + instruction: String, +) { + let outcome = ctx.run_instruction(instruction).await; + let (status, report) = match outcome { + Ok(text) => { + let report = parse_worker_report(&worker_id, &text); + let status = if report.status == "failed" { + WorkerStatus::Failed + } else { + WorkerStatus::Completed + }; + (status, Some(report)) + } + Err(e) => ( + WorkerStatus::Failed, + Some(WorkerReport { + worker_id: worker_id.clone(), + task_id: worker_id.clone(), + status: "failed".to_string(), + summary: format!("worker execution error: {e}"), + artifacts: vec![], + }), + ), + }; + pool.finish(&worker_id, status, report); +} + +// ============================================================ +// 调度工具 +// ============================================================ + +/// `create_worker`：异步启动一个新 Worker，立即返回 `worker_id`。 +pub struct CreateWorkerTool { + pool: Arc, + factory: Arc, +} + +#[async_trait] +impl Tool for CreateWorkerTool { + fn name(&self) -> &str { + "create_worker" + } + + fn description(&self) -> &str { + "Start a new Worker sub-agent asynchronously to execute a self-contained instruction. \ + Returns a worker_id immediately (does not block). Args: {\"instruction\": string, \ + \"task_type\"?: string}" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "instruction": { "type": "string", "description": "Self-contained subtask instruction (file paths, operations, completion criteria)" }, + "task_type": { "type": "string", "description": "Subtask type: research / implement / verify / generic" } + }, + "required": ["instruction"] + }) + } + + async fn execute(&self, args: Value) -> Result { + let instruction = args + .get("instruction") + .and_then(|v| v.as_str()) + .context("create_worker requires an 'instruction' string argument")? + .to_string(); + let task_type = args + .get("task_type") + .and_then(|v| v.as_str()) + .unwrap_or("generic"); + + let worker_id = new_worker_id(); + let ctx = self.factory.create(task_type); + self.pool.insert(worker_id.clone(), ctx.clone()); + + let join = tokio::spawn(run_and_finish( + self.pool.clone(), + worker_id.clone(), + ctx, + instruction, + )); + self.pool.set_abort(&worker_id, join.abort_handle()); + + Ok(json!({ "worker_id": worker_id, "status": "running" }).to_string()) + } +} + +/// `await_workers`：阻塞等待指定（或全部）Worker 终止并汇总报告。 +pub struct AwaitWorkersTool { + pool: Arc, +} + +#[async_trait] +impl Tool for AwaitWorkersTool { + fn name(&self) -> &str { + "await_workers" + } + + fn description(&self) -> &str { + "Block until the given workers (or all workers when 'ids' is omitted) reach a terminal \ + state, then return their aggregated reports. Args: {\"ids\"?: string[]}" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "ids": { "type": "array", "items": { "type": "string" }, "description": "Worker ids to await; omit to await all" } + } + }) + } + + async fn execute(&self, args: Value) -> Result { + let requested: Vec = args + .get("ids") + .and_then(|v| v.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|x| x.as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(); + + let targets = if requested.is_empty() { + self.pool.ids() + } else { + requested + }; + + if targets.is_empty() { + return Ok(json!({ "workers": [], "complete": true, "note": "no workers to await" }) + .to_string()); + } + + let deadline = Instant::now() + AWAIT_TIMEOUT; + loop { + if self.pool.all_terminal(&targets) || Instant::now() >= deadline { + break; + } + tokio::time::sleep(POLL_INTERVAL).await; + } + + let complete = self.pool.all_terminal(&targets); + let workers: Vec = self + .pool + .snapshot(&targets) + .into_iter() + .map(|(id, status, report)| { + json!({ + "worker_id": id, + "status": status.as_str(), + "report": report.map(|r| serde_json::to_value(r).unwrap_or(Value::Null)).unwrap_or(Value::Null), + }) + }) + .collect(); + + Ok(json!({ "workers": workers, "complete": complete }).to_string()) + } +} + +/// `resume_worker`：给已有 Worker 下发新指令，复用其上下文续跑。 +pub struct ResumeWorkerTool { + pool: Arc, +} + +#[async_trait] +impl Tool for ResumeWorkerTool { + fn name(&self) -> &str { + "resume_worker" + } + + fn description(&self) -> &str { + "Send a new instruction to an existing worker, reusing its prior context (history/errors). \ + Runs asynchronously. Args: {\"worker_id\": string, \"instruction\": string}" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "worker_id": { "type": "string", "description": "Target worker id" }, + "instruction": { "type": "string", "description": "New self-contained instruction (e.g. a corrective fix)" } + }, + "required": ["worker_id", "instruction"] + }) + } + + async fn execute(&self, args: Value) -> Result { + let worker_id = args + .get("worker_id") + .and_then(|v| v.as_str()) + .context("resume_worker requires a 'worker_id' string argument")? + .to_string(); + let instruction = args + .get("instruction") + .and_then(|v| v.as_str()) + .context("resume_worker requires an 'instruction' string argument")? + .to_string(); + + let ctx = self + .pool + .get_ctx(&worker_id) + .with_context(|| format!("worker '{}' not found", worker_id))?; + + self.pool.set_running(&worker_id); + let join = tokio::spawn(run_and_finish( + self.pool.clone(), + worker_id.clone(), + ctx, + instruction, + )); + self.pool.set_abort(&worker_id, join.abort_handle()); + + Ok(json!({ "worker_id": worker_id, "status": "resumed" }).to_string()) + } +} + +/// `stop_worker`：中止一个 Worker（保留上下文，可后续续跑）。 +pub struct StopWorkerTool { + pool: Arc, +} + +#[async_trait] +impl Tool for StopWorkerTool { + fn name(&self) -> &str { + "stop_worker" + } + + fn description(&self) -> &str { + "Terminate a running worker (e.g. wrong direction). The worker can be resumed later. \ + Args: {\"worker_id\": string}" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "worker_id": { "type": "string", "description": "Target worker id" } + }, + "required": ["worker_id"] + }) + } + + async fn execute(&self, args: Value) -> Result { + let worker_id = args + .get("worker_id") + .and_then(|v| v.as_str()) + .context("stop_worker requires a 'worker_id' string argument")? + .to_string(); + + if !self.pool.stop(&worker_id) { + anyhow::bail!("worker '{}' not found", worker_id); + } + Ok(json!({ "worker_id": worker_id, "status": "stopped" }).to_string()) + } +} + +/// 构建四个调度工具（共享同一个 [`WorkerPool`] 与 [`WorkerFactory`]）。 +pub fn build_tools( + pool: Arc, + factory: Arc, +) -> Vec> { + vec![ + Arc::new(CreateWorkerTool { + pool: pool.clone(), + factory, + }), + Arc::new(AwaitWorkersTool { pool: pool.clone() }), + Arc::new(ResumeWorkerTool { pool: pool.clone() }), + Arc::new(StopWorkerTool { pool }), + ] +} + +#[cfg(test)] +mod tests { + use super::*; + + /// echo 桩上下文：记录收到的指令，按配置返回成功/失败的标准化结果。 + struct EchoWorkerContext { + fail: bool, + delay: Duration, + calls: Arc>>, + } + + #[async_trait] + impl WorkerContext for EchoWorkerContext { + async fn run_instruction(&self, instruction: String) -> Result { + self.calls.lock().unwrap().push(instruction.clone()); + tokio::time::sleep(self.delay).await; + if self.fail { + Ok("```result\n{\"status\":\"failed\",\"summary\":\"boom\"}\n```".to_string()) + } else { + Ok(format!( + "```result\n{{\"task_id\":\"t1\",\"status\":\"success\",\"summary\":\"done: {}\",\"artifacts\":[\"a.txt\"]}}\n```", + instruction + )) + } + } + } + + struct EchoFactory { + fail: bool, + delay: Duration, + calls: Arc>>, + } + + impl WorkerFactory for EchoFactory { + fn create(&self, _task_type: &str) -> Arc { + Arc::new(EchoWorkerContext { + fail: self.fail, + delay: self.delay, + calls: self.calls.clone(), + }) + } + } + + fn setup(fail: bool, delay: Duration) -> (Arc, Vec>, Arc>>) { + let pool = Arc::new(WorkerPool::default()); + let calls = Arc::new(std::sync::Mutex::new(Vec::new())); + let factory = Arc::new(EchoFactory { + fail, + delay, + calls: calls.clone(), + }); + let tools = build_tools(pool.clone(), factory); + (pool, tools, calls) + } + + fn tool<'a>(tools: &'a [Arc], name: &str) -> &'a Arc { + tools.iter().find(|t| t.name() == name).expect("tool exists") + } + + fn wid_of(create_out: &str) -> String { + serde_json::from_str::(create_out) + .unwrap() + .get("worker_id") + .unwrap() + .as_str() + .unwrap() + .to_string() + } + + #[test] + fn build_tools_registers_four_schedulers() { + let (_pool, tools, _calls) = setup(false, Duration::ZERO); + let mut names: Vec<&str> = tools.iter().map(|t| t.name()).collect(); + names.sort(); + assert_eq!( + names, + vec!["await_workers", "create_worker", "resume_worker", "stop_worker"] + ); + } + + #[test] + fn parse_worker_report_valid_block() { + let text = "prefix\n```result\n{\"task_id\":\"T9\",\"status\":\"failed\",\"summary\":\"nope\",\"artifacts\":[\"x.rs\",\"y.rs\"]}\n```\ntrail"; + let rep = parse_worker_report("w-1", text); + assert_eq!(rep.worker_id, "w-1"); + assert_eq!(rep.task_id, "T9"); + assert_eq!(rep.status, "failed"); + assert_eq!(rep.summary, "nope"); + assert_eq!(rep.artifacts, vec!["x.rs".to_string(), "y.rs".to_string()]); + } + + #[test] + fn parse_worker_report_falls_back_without_fence() { + let rep = parse_worker_report("w-2", "just some plain text with no result block"); + assert_eq!(rep.worker_id, "w-2"); + assert_eq!(rep.task_id, "w-2"); + assert_eq!(rep.status, "success"); + assert!(rep.summary.contains("plain text")); + assert!(rep.artifacts.is_empty()); + } + + #[tokio::test] + async fn create_then_await_completes() { + let (_pool, tools, calls) = setup(false, Duration::from_millis(5)); + let out = tool(&tools, "create_worker") + .execute(json!({ "instruction": "research auth", "task_type": "research" })) + .await + .unwrap(); + let wid = wid_of(&out); + + let awaited = tool(&tools, "await_workers") + .execute(json!({ "ids": [wid.clone()] })) + .await + .unwrap(); + let v: Value = serde_json::from_str(&awaited).unwrap(); + assert_eq!(v.get("complete").unwrap().as_bool(), Some(true)); + let w0 = &v.get("workers").unwrap().as_array().unwrap()[0]; + assert_eq!(w0.get("status").unwrap().as_str(), Some("completed")); + assert_eq!( + w0.pointer("/report/status").and_then(|x| x.as_str()), + Some("success") + ); + assert_eq!(calls.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn failed_worker_is_marked_failed() { + let (_pool, tools, _calls) = setup(true, Duration::from_millis(5)); + let out = tool(&tools, "create_worker") + .execute(json!({ "instruction": "do thing" })) + .await + .unwrap(); + let wid = wid_of(&out); + + let awaited = tool(&tools, "await_workers") + .execute(json!({})) + .await + .unwrap(); + let v: Value = serde_json::from_str(&awaited).unwrap(); + let w0 = &v.get("workers").unwrap().as_array().unwrap()[0]; + assert_eq!(w0.get("worker_id").unwrap().as_str(), Some(wid.as_str())); + assert_eq!(w0.get("status").unwrap().as_str(), Some("failed")); + } + + #[tokio::test] + async fn stop_worker_marks_stopped() { + // 长延时确保在 Worker 完成前中止。 + let (pool, tools, _calls) = setup(false, Duration::from_millis(500)); + let out = tool(&tools, "create_worker") + .execute(json!({ "instruction": "long task" })) + .await + .unwrap(); + let wid = wid_of(&out); + + let stopped = tool(&tools, "stop_worker") + .execute(json!({ "worker_id": wid.clone() })) + .await + .unwrap(); + assert!(stopped.contains("stopped")); + + let awaited = tool(&tools, "await_workers") + .execute(json!({ "ids": [wid.clone()] })) + .await + .unwrap(); + let v: Value = serde_json::from_str(&awaited).unwrap(); + let w0 = &v.get("workers").unwrap().as_array().unwrap()[0]; + assert_eq!(w0.get("status").unwrap().as_str(), Some("stopped")); + + // 未知 id 停止应报错。 + assert!(tool(&tools, "stop_worker") + .execute(json!({ "worker_id": "nope" })) + .await + .is_err()); + let _ = pool; + } + + #[tokio::test] + async fn resume_worker_reuses_context() { + let (_pool, tools, calls) = setup(false, Duration::from_millis(5)); + let out = tool(&tools, "create_worker") + .execute(json!({ "instruction": "first" })) + .await + .unwrap(); + let wid = wid_of(&out); + tool(&tools, "await_workers") + .execute(json!({ "ids": [wid.clone()] })) + .await + .unwrap(); + + // 续跑：复用同一 ctx，指令累积到共享 calls。 + tool(&tools, "resume_worker") + .execute(json!({ "worker_id": wid.clone(), "instruction": "second" })) + .await + .unwrap(); + let awaited = tool(&tools, "await_workers") + .execute(json!({ "ids": [wid.clone()] })) + .await + .unwrap(); + let v: Value = serde_json::from_str(&awaited).unwrap(); + assert_eq!(v.get("complete").unwrap().as_bool(), Some(true)); + + let log = calls.lock().unwrap(); + assert_eq!(log.len(), 2); + assert_eq!(log[0], "first"); + assert_eq!(log[1], "second"); + + // 未知 worker 续跑应报错。 + assert!(tool(&tools, "resume_worker") + .execute(json!({ "worker_id": "ghost", "instruction": "x" })) + .await + .is_err()); + } + + #[tokio::test] + async fn await_empty_pool_returns_no_workers() { + let (_pool, tools, _calls) = setup(false, Duration::ZERO); + let out = tool(&tools, "await_workers").execute(json!({})).await.unwrap(); + let v: Value = serde_json::from_str(&out).unwrap(); + assert_eq!(v.get("workers").unwrap().as_array().unwrap().len(), 0); + assert_eq!(v.get("complete").unwrap().as_bool(), Some(true)); + } + + #[tokio::test] + async fn create_worker_requires_instruction() { + let (_pool, tools, _calls) = setup(false, Duration::ZERO); + assert!(tool(&tools, "create_worker") + .execute(json!({ "task_type": "research" })) + .await + .is_err()); + } + + #[test] + fn pool_all_terminal_and_snapshot_skip_missing() { + let pool = WorkerPool::default(); + let ctx: Arc = Arc::new(FailedWorkerContext { + error: "x".into(), + }); + pool.insert("w-a".into(), ctx.clone()); + pool.insert("w-b".into(), ctx.clone()); + pool.finish( + "w-a", + WorkerStatus::Completed, + Some(WorkerReport { + worker_id: "w-a".into(), + task_id: "w-a".into(), + status: "success".into(), + summary: "ok".into(), + artifacts: vec![], + }), + ); + + // w-a 终止、w-b 仍运行 -> 整体未终止。 + assert!(!pool.all_terminal(&["w-a".into(), "w-b".into()])); + // 仅 w-a -> 终止；不存在的 id 视为终止。 + assert!(pool.all_terminal(&["w-a".into(), "ghost".into()])); + + let snap = pool.snapshot(&["w-a".into(), "ghost".into()]); + assert_eq!(snap.len(), 1); + assert_eq!(snap[0].0, "w-a"); + assert_eq!(snap[0].1, WorkerStatus::Completed); + } +} diff --git a/src/db/models/agent_arena.rs b/src/db/models/agent_arena.rs index 7ba42ac..e692cc2 100644 --- a/src/db/models/agent_arena.rs +++ b/src/db/models/agent_arena.rs @@ -5,15 +5,31 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// 竞技场投票表（`agent_arena`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.6 / 3.8：记录用户对单轮回答的 +/// **星级评分（1-5）与文字评价**，通过 `chat_uuid`/`round_uuid` 关联具体消息； +/// 评分聚合为 `agent_info.avg_score`（0-5 均值）。 #[derive(Clone, Debug)] pub struct AgentArena { pub id: i32, + /// 会话 UUID，关联 `chat_session.chat_uuid`。 + pub chat_uuid: String, + /// 单轮消息 UUID，关联 `chat_message.round_uuid`。 + pub round_uuid: Option, + /// 被投票的 Agent ID，关联 `agent_info.id`。 pub agent_id: i32, - pub arena_id: i32, - pub created_at: DateTime, - pub updated_at: DateTime, + /// 投票类型：星级评分 1-5（竞技场打分，聚合为 `agent_info.avg_score`）。 + pub vote_type: i8, + /// 用户评价文本。 + pub comment: Option, + /// 状态：1 有效 / 0 无效。 + pub status: i8, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -26,62 +42,62 @@ impl Model for AgentArena { r#" CREATE TABLE IF NOT EXISTS agent_arena ( id INTEGER PRIMARY KEY AUTOINCREMENT, + chat_uuid TEXT NOT NULL, + round_uuid TEXT, agent_id INTEGER NOT NULL, - arena_id INTEGER NOT NULL, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + vote_type INTEGER NOT NULL DEFAULT 0, + comment TEXT, + status INTEGER NOT NULL DEFAULT 1, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, + chat_uuid: row.get("chat_uuid")?, + round_uuid: row.get("round_uuid")?, agent_id: row.get("agent_id")?, - arena_id: row.get("arena_id")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + vote_type: row.get("vote_type")?, + comment: row.get("comment")?, + status: row.get("status")?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "agent_id, arena_id, remark" + "chat_uuid, round_uuid, agent_id, vote_type, comment, status, remark" } fn insert_values(&self) -> Vec> { vec![ + Box::new(self.chat_uuid.clone()), + Box::new(self.round_uuid.clone()), Box::new(self.agent_id), - Box::new(self.arena_id), + Box::new(self.vote_type), + Box::new(self.comment.clone()), + Box::new(self.status), Box::new(self.remark.clone()), ] } fn update_set_clause(&self) -> String { r#" + chat_uuid = ?, + round_uuid = ?, agent_id = ?, - arena_id = ?, - updated_at = CURRENT_TIMESTAMP, + vote_type = ?, + comment = ?, + status = ?, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -90,8 +106,12 @@ impl Model for AgentArena { fn update_values(&self) -> Vec> { vec![ + Box::new(self.chat_uuid.clone()), + Box::new(self.round_uuid.clone()), Box::new(self.agent_id), - Box::new(self.arena_id), + Box::new(self.vote_type), + Box::new(self.comment.clone()), + Box::new(self.status), Box::new(self.remark.clone()), ] } diff --git a/src/db/models/agent_info.rs b/src/db/models/agent_info.rs index 29c8e22..4e04423 100644 --- a/src/db/models/agent_info.rs +++ b/src/db/models/agent_info.rs @@ -5,18 +5,39 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// 智能体表（`agent_info`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.3：Agent 以 LLM + 专属 Prompt 为核心， +/// 绑定 MCP 工具（经 `agent_mcp_rel` 中间表）、支持跨 Agent 能力复用。 +/// 提示词 1 对 1 绑定，模型个性化参数（温度、top_p）在此配置生效。 #[derive(Clone, Debug)] pub struct AgentInfo { pub id: i32, + /// LLM 模型 ID，关联 `llm_model.id` 逻辑外键。 + pub llm_id: i32, + /// 智能体名称，非空。 pub name: String, - pub avatar: String, - pub model_id: i32, - pub prompt_id: i32, - pub status: i8, // default 1 - pub created_at: DateTime, - pub updated_at: DateTime, + /// 智能体描述（简介说明）。 + pub description: Option, + /// 提示词 ID，关联 `prompt_info.id`，一对一逻辑唯一绑定。 + pub prompt_id: Option, + /// 头像地址（图片 URL / 本地路径）。 + pub avatar: Option, + /// 温度，decimal(2,1)。 + pub temperature: Option, + /// 核心采样 top_p，decimal(2,1)。 + pub top_p: Option, + /// 累计调用次数。 + pub call_count: i64, + /// 平均评分，0-5 分，竞技场评价聚合均值。 + pub avg_score: Option, + /// 启用状态：0-禁用 1-启用。 + pub status: i8, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -29,61 +50,60 @@ impl Model for AgentInfo { r#" CREATE TABLE IF NOT EXISTS agent_info ( id INTEGER PRIMARY KEY AUTOINCREMENT, + llm_id INTEGER NOT NULL, name TEXT NOT NULL, - avatar TEXT NOT NULL, - model_id INTEGER NOT NULL, - prompt_id INTEGER NOT NULL, + description TEXT, + prompt_id INTEGER, + avatar TEXT, + temperature REAL, + top_p REAL, + call_count INTEGER NOT NULL DEFAULT 0, + avg_score REAL, status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, + llm_id: row.get("llm_id")?, name: row.get("name")?, - avatar: row.get("avatar")?, - model_id: row.get("model_id")?, + description: row.get("description")?, prompt_id: row.get("prompt_id")?, + avatar: row.get("avatar")?, + temperature: row.get("temperature")?, + top_p: row.get("top_p")?, + call_count: row.get("call_count")?, + avg_score: row.get("avg_score")?, status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "name, avatar, model_id, prompt_id, status, remark" + "llm_id, name, description, prompt_id, avatar, temperature, top_p, call_count, avg_score, status, remark" } fn insert_values(&self) -> Vec> { vec![ + Box::new(self.llm_id), Box::new(self.name.clone()), - Box::new(self.avatar.clone()), - Box::new(self.model_id), + Box::new(self.description.clone()), Box::new(self.prompt_id), + Box::new(self.avatar.clone()), + Box::new(self.temperature), + Box::new(self.top_p), + Box::new(self.call_count), + Box::new(self.avg_score), Box::new(self.status), Box::new(self.remark.clone()), ] @@ -91,12 +111,17 @@ impl Model for AgentInfo { fn update_set_clause(&self) -> String { r#" + llm_id = ?, name = ?, - avatar = ?, - model_id = ?, + description = ?, prompt_id = ?, + avatar = ?, + temperature = ?, + top_p = ?, + call_count = ?, + avg_score = ?, status = ?, - updated_at = CURRENT_TIMESTAMP, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -105,10 +130,15 @@ impl Model for AgentInfo { fn update_values(&self) -> Vec> { vec![ + Box::new(self.llm_id), Box::new(self.name.clone()), - Box::new(self.avatar.clone()), - Box::new(self.model_id), + Box::new(self.description.clone()), Box::new(self.prompt_id), + Box::new(self.avatar.clone()), + Box::new(self.temperature), + Box::new(self.top_p), + Box::new(self.call_count), + Box::new(self.avg_score), Box::new(self.status), Box::new(self.remark.clone()), ] diff --git a/src/db/models/agent_mcp_rel.rs b/src/db/models/agent_mcp_rel.rs index 0b6c1a6..da6404a 100644 --- a/src/db/models/agent_mcp_rel.rs +++ b/src/db/models/agent_mcp_rel.rs @@ -5,15 +5,22 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// Agent 与 MCP 工具关联表（`agent_mcp_rel`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.12：多对多关联，声明某个 Agent +/// 被授权可调用哪些 MCP 工具，用于工具调度时的权限校验。 #[derive(Clone, Debug)] pub struct AgentMcpRel { pub id: i32, + /// Agent ID，关联 `agent_info.id`。 pub agent_id: i32, - pub mcp_id: i32, - pub created_at: DateTime, - pub updated_at: DateTime, + /// MCP 工具 ID，关联 `mcp_tool.id`。 + pub mcp_tool_id: i32, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -27,52 +34,36 @@ impl Model for AgentMcpRel { CREATE TABLE IF NOT EXISTS agent_mcp_rel ( id INTEGER PRIMARY KEY AUTOINCREMENT, agent_id INTEGER NOT NULL, - mcp_id INTEGER NOT NULL, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + mcp_tool_id INTEGER NOT NULL, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, agent_id: row.get("agent_id")?, - mcp_id: row.get("mcp_id")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + mcp_tool_id: row.get("mcp_tool_id")?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "agent_id, mcp_id, remark" + "agent_id, mcp_tool_id, remark" } fn insert_values(&self) -> Vec> { vec![ Box::new(self.agent_id), - Box::new(self.mcp_id), + Box::new(self.mcp_tool_id), Box::new(self.remark.clone()), ] } @@ -80,8 +71,8 @@ impl Model for AgentMcpRel { fn update_set_clause(&self) -> String { r#" agent_id = ?, - mcp_id = ?, - updated_at = CURRENT_TIMESTAMP, + mcp_tool_id = ?, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -91,7 +82,7 @@ impl Model for AgentMcpRel { fn update_values(&self) -> Vec> { vec![ Box::new(self.agent_id), - Box::new(self.mcp_id), + Box::new(self.mcp_tool_id), Box::new(self.remark.clone()), ] } diff --git a/src/db/models/ai_provider.rs b/src/db/models/ai_provider.rs index b49f9d1..27f4d21 100644 --- a/src/db/models/ai_provider.rs +++ b/src/db/models/ai_provider.rs @@ -5,19 +5,30 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// AI 供应商表（`ai_provider`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.1：统一管理接入的 AI 厂商服务， +/// 除 Ollama 外 API 地址硬编码预置，密钥隔离存储。 #[derive(Clone, Debug)] pub struct AiProvider { pub id: i32, + /// 供应商名称，非空。 pub name: String, - pub avatar: String, + /// 头像地址（图片 URL / 本地路径）。 + pub avatar: Option, + /// API 地址，非空。 pub api: String, + /// API 密钥，可为空（Ollama 非强制）。 pub key: Option, + /// 请求超时时间，单位：秒，默认 30。 pub timeout: i32, + /// 启用状态：0-禁用 1-启用。 pub status: i8, - pub created_at: DateTime, - pub updated_at: DateTime, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -31,22 +42,21 @@ impl Model for AiProvider { CREATE TABLE IF NOT EXISTS ai_provider ( id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, - avatar TEXT NOT NULL, + avatar TEXT, api TEXT NOT NULL, key TEXT, - timeout INTEGER NOT NULL DEFAULT 5000, + timeout INTEGER NOT NULL DEFAULT 30, status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, @@ -56,24 +66,8 @@ impl Model for AiProvider { key: row.get("key")?, timeout: row.get("timeout")?, status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + create_time, + update_time, remark: row.get("remark")?, }) } @@ -102,7 +96,7 @@ impl Model for AiProvider { key = ?, timeout = ?, status = ?, - updated_at = CURRENT_TIMESTAMP, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -124,4 +118,4 @@ impl Model for AiProvider { fn primary_key_value(&self) -> Box { Box::new(self.id) } -} \ No newline at end of file +} diff --git a/src/db/models/chat_message.rs b/src/db/models/chat_message.rs index d5d6627..d75577b 100644 --- a/src/db/models/chat_message.rs +++ b/src/db/models/chat_message.rs @@ -5,18 +5,35 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// 消息表（`chat_message`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.8：拆分每轮对话独立存储， +/// 通过 `chat_uuid` 关联会话、`round_uuid` 标识单轮消息，并记录 Token 消耗、 +/// 调用耗时与模型参数快照。 #[derive(Clone, Debug)] pub struct ChatMessage { pub id: i32, - pub session_id: String, + /// 会话 UUID，关联 `chat_session.chat_uuid`。 + pub chat_uuid: String, + /// 单轮消息 UUID，单轮唯一标识。 + pub round_uuid: Option, + /// 消息角色：user/assistant/system。 pub role: String, + /// 消息内容（对话正文）。 pub content: String, - pub model: Option, - pub status: i8, // default 1 - pub created_at: DateTime, - pub updated_at: DateTime, + /// 单轮输入消耗 Token 数。 + pub input_tokens: i32, + /// 单轮输出消耗 Token 数。 + pub output_tokens: i32, + /// 调用耗时，单位：毫秒。 + pub duration: i32, + /// 模型参数快照（JSON 存储当时模型配置）。 + pub model_config: Option, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -29,74 +46,70 @@ impl Model for ChatMessage { r#" CREATE TABLE IF NOT EXISTS chat_message ( id INTEGER PRIMARY KEY AUTOINCREMENT, - session_id TEXT NOT NULL, + chat_uuid TEXT NOT NULL, + round_uuid TEXT, role TEXT NOT NULL, content TEXT NOT NULL, - model TEXT, - status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + duration INTEGER NOT NULL DEFAULT 0, + model_config TEXT, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, - session_id: row.get("session_id")?, + chat_uuid: row.get("chat_uuid")?, + round_uuid: row.get("round_uuid")?, role: row.get("role")?, content: row.get("content")?, - model: row.get("model")?, - status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + input_tokens: row.get("input_tokens")?, + output_tokens: row.get("output_tokens")?, + duration: row.get("duration")?, + model_config: row.get("model_config")?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "session_id, role, content, model, status, remark" + "chat_uuid, round_uuid, role, content, input_tokens, output_tokens, duration, model_config, remark" } fn insert_values(&self) -> Vec> { vec![ - Box::new(self.session_id.clone()), + Box::new(self.chat_uuid.clone()), + Box::new(self.round_uuid.clone()), Box::new(self.role.clone()), Box::new(self.content.clone()), - Box::new(self.model.clone()), - Box::new(self.status), + Box::new(self.input_tokens), + Box::new(self.output_tokens), + Box::new(self.duration), + Box::new(self.model_config.clone()), Box::new(self.remark.clone()), ] } fn update_set_clause(&self) -> String { r#" - session_id = ?, + chat_uuid = ?, + round_uuid = ?, role = ?, content = ?, - model = ?, - status = ?, - updated_at = CURRENT_TIMESTAMP, + input_tokens = ?, + output_tokens = ?, + duration = ?, + model_config = ?, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -105,11 +118,14 @@ impl Model for ChatMessage { fn update_values(&self) -> Vec> { vec![ - Box::new(self.session_id.clone()), + Box::new(self.chat_uuid.clone()), + Box::new(self.round_uuid.clone()), Box::new(self.role.clone()), Box::new(self.content.clone()), - Box::new(self.model.clone()), - Box::new(self.status), + Box::new(self.input_tokens), + Box::new(self.output_tokens), + Box::new(self.duration), + Box::new(self.model_config.clone()), Box::new(self.remark.clone()), ] } diff --git a/src/db/models/chat_session.rs b/src/db/models/chat_session.rs index dcefd5a..4701ec0 100644 --- a/src/db/models/chat_session.rs +++ b/src/db/models/chat_session.rs @@ -5,17 +5,27 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// 会话表（`chat_session`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.7：一个项目可包含多个会话， +/// `chat_uuid` 全局唯一并对外暴露。按《后端开发手册》3.2，本表保留 `is_deleted` +/// 逻辑删除字段（仅业务层过滤，无数据恢复）。 #[derive(Clone, Debug)] pub struct ChatSession { pub id: i32, + /// 会话唯一标识，全局唯一、对外暴露。 pub chat_uuid: String, + /// 会话名称。 pub title: Option, + /// 绑定 LLM ID，关联 `llm_model.id` 逻辑外键。 pub llm_id: i32, - pub status: i8, // default 1 - pub created_at: DateTime, - pub updated_at: DateTime, + /// 逻辑删除：0-否 1-是。 + pub is_deleted: i8, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -31,48 +41,32 @@ impl Model for ChatSession { chat_uuid TEXT NOT NULL, title TEXT, llm_id INTEGER NOT NULL, - status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + is_deleted INTEGER NOT NULL DEFAULT 0, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, chat_uuid: row.get("chat_uuid")?, title: row.get("title")?, llm_id: row.get("llm_id")?, - status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + is_deleted: row.get("is_deleted")?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "chat_uuid, title, llm_id, status, remark" + "chat_uuid, title, llm_id, is_deleted, remark" } fn insert_values(&self) -> Vec> { @@ -80,7 +74,7 @@ impl Model for ChatSession { Box::new(self.chat_uuid.clone()), Box::new(self.title.clone()), Box::new(self.llm_id), - Box::new(self.status), + Box::new(self.is_deleted), Box::new(self.remark.clone()), ] } @@ -90,8 +84,8 @@ impl Model for ChatSession { chat_uuid = ?, title = ?, llm_id = ?, - status = ?, - updated_at = CURRENT_TIMESTAMP, + is_deleted = ?, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -103,7 +97,7 @@ impl Model for ChatSession { Box::new(self.chat_uuid.clone()), Box::new(self.title.clone()), Box::new(self.llm_id), - Box::new(self.status), + Box::new(self.is_deleted), Box::new(self.remark.clone()), ] } diff --git a/src/db/models/llm_model.rs b/src/db/models/llm_model.rs index 1be9e4b..7b8958c 100644 --- a/src/db/models/llm_model.rs +++ b/src/db/models/llm_model.rs @@ -5,22 +5,36 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// LLM 模型表（`llm_model`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.2：模型为公共基座能力，无独立运行参数， +/// 个性化配置（温度、上下文限制等）由 Agent 独立配置。`model_code` 逻辑唯一。 #[derive(Clone, Debug)] pub struct LLMModel { pub id: i32, + /// 供应商 ID，关联 `ai_provider.id` 逻辑外键。 pub provider_id: i32, - pub name: Option, - pub avatar: String, + /// 头像地址（图片 URL / 本地路径）。 + pub avatar: Option, + /// 模型名称，非空。 + pub name: String, + /// 模型标识，接口调用唯一标识，非空、逻辑唯一。 pub model_code: String, - pub max_tokens: i32, // default 0 - pub max_response_tokens: i32, // default 0 - pub weight: i32, // default 0 - pub call_count: i64, // default 0 - pub status: i8, // default 1 - pub created_at: DateTime, - pub updated_at: DateTime, + /// 最大上下文（最大输入 token）。 + pub max_tokens: i32, + /// 单次输出上限（单次生成最大 token）。 + pub max_response_tokens: i32, + /// 负载分发权重。 + pub weight: i32, + /// 累计调用次数。 + pub call_count: i64, + /// 启用状态：0-禁用 1-启用。 + pub status: i8, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -34,67 +48,51 @@ impl Model for LLMModel { CREATE TABLE IF NOT EXISTS llm_model ( id INTEGER PRIMARY KEY AUTOINCREMENT, provider_id INTEGER NOT NULL, - name TEXT, - avatar TEXT NOT NULL, + avatar TEXT, + name TEXT NOT NULL, model_code TEXT NOT NULL, max_tokens INTEGER NOT NULL DEFAULT 0, max_response_tokens INTEGER NOT NULL DEFAULT 0, weight INTEGER NOT NULL DEFAULT 0, call_count INTEGER NOT NULL DEFAULT 0, status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, provider_id: row.get("provider_id")?, - name: row.get("name")?, avatar: row.get("avatar")?, + name: row.get("name")?, model_code: row.get("model_code")?, max_tokens: row.get("max_tokens")?, max_response_tokens: row.get("max_response_tokens")?, weight: row.get("weight")?, call_count: row.get("call_count")?, status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "provider_id, name, avatar, model_code, max_tokens, max_response_tokens, weight, call_count, status, remark" + "provider_id, avatar, name, model_code, max_tokens, max_response_tokens, weight, call_count, status, remark" } fn insert_values(&self) -> Vec> { vec![ Box::new(self.provider_id), - Box::new(self.name.clone()), Box::new(self.avatar.clone()), + Box::new(self.name.clone()), Box::new(self.model_code.clone()), Box::new(self.max_tokens), Box::new(self.max_response_tokens), @@ -108,15 +106,15 @@ impl Model for LLMModel { fn update_set_clause(&self) -> String { r#" provider_id = ?, - name = ?, avatar = ?, + name = ?, model_code = ?, max_tokens = ?, max_response_tokens = ?, weight = ?, call_count = ?, status = ?, - updated_at = CURRENT_TIMESTAMP, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -126,8 +124,8 @@ impl Model for LLMModel { fn update_values(&self) -> Vec> { vec![ Box::new(self.provider_id), - Box::new(self.name.clone()), Box::new(self.avatar.clone()), + Box::new(self.name.clone()), Box::new(self.model_code.clone()), Box::new(self.max_tokens), Box::new(self.max_response_tokens), diff --git a/src/db/models/mcp_provider.rs b/src/db/models/mcp_provider.rs index 60719e4..5a346a8 100644 --- a/src/db/models/mcp_provider.rs +++ b/src/db/models/mcp_provider.rs @@ -5,19 +5,30 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// MCP 服务提供方表（`mcp_provider`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.10：登记一个 MCP Server（如 Python 工具端）， +/// 记录其连接地址、鉴权、超时与健康检查接口。`timeout` 单位为秒。 #[derive(Clone, Debug)] pub struct McpProvider { pub id: i32, + /// 提供方名称。 pub name: String, + /// MCP Server 连接地址。 pub url: String, + /// 鉴权密钥（可空）。 pub auth_key: Option, - pub timeout: i32, // ms default 5000 + /// 请求超时，单位：秒。 + pub timeout: i32, + /// 健康检查接口地址（可空）。 pub health_api: Option, - pub status: i8, // default 1 - pub created_at: DateTime, - pub updated_at: DateTime, + /// 状态：1 启用 / 0 停用。 + pub status: i8, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -33,19 +44,19 @@ impl Model for McpProvider { name TEXT NOT NULL, url TEXT NOT NULL, auth_key TEXT, - timeout INTEGER NOT NULL DEFAULT 5000, + timeout INTEGER NOT NULL DEFAULT 30, health_api TEXT, status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, @@ -55,24 +66,8 @@ impl Model for McpProvider { timeout: row.get("timeout")?, health_api: row.get("health_api")?, status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + create_time, + update_time, remark: row.get("remark")?, }) } @@ -101,7 +96,7 @@ impl Model for McpProvider { timeout = ?, health_api = ?, status = ?, - updated_at = CURRENT_TIMESTAMP, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() diff --git a/src/db/models/mcp_tool.rs b/src/db/models/mcp_tool.rs index 2f98812..ee66d4a 100644 --- a/src/db/models/mcp_tool.rs +++ b/src/db/models/mcp_tool.rs @@ -5,20 +5,34 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// MCP 工具表（`mcp_tool`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.11：登记某个 MCP Server 暴露的工具， +/// `tool_code` 为调用标识，`parameters` 描述入参 Schema，`response_format` 描述返回结构。 #[derive(Clone, Debug)] pub struct McpTool { pub id: i32, + /// 所属 MCP 提供方 ID，关联 `mcp_provider.id`。 pub provider_id: i32, + /// 工具展示名称。 pub name: String, + /// 工具调用标识（唯一编码）。 pub tool_code: String, - pub tool_type: String, - pub description: String, - pub input_schema: String, - pub status: i8, // default 1 - pub created_at: DateTime, - pub updated_at: DateTime, + /// 工具描述（可空）。 + pub description: Option, + /// 工具标签/分类（可空）。 + pub tag: Option, + /// 入参 Schema（JSON，可空）。 + pub parameters: Option, + /// 返回结构说明（JSON，可空）。 + pub response_format: Option, + /// 状态：1 启用 / 0 停用。 + pub status: i8, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -34,54 +48,40 @@ impl Model for McpTool { provider_id INTEGER NOT NULL, name TEXT NOT NULL, tool_code TEXT NOT NULL, - tool_type TEXT NOT NULL, - description TEXT NOT NULL, - input_schema TEXT NOT NULL, + description TEXT, + tag TEXT, + parameters TEXT, + response_format TEXT, status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, provider_id: row.get("provider_id")?, name: row.get("name")?, tool_code: row.get("tool_code")?, - tool_type: row.get("tool_type")?, description: row.get("description")?, - input_schema: row.get("input_schema")?, + tag: row.get("tag")?, + parameters: row.get("parameters")?, + response_format: row.get("response_format")?, status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "provider_id, name, tool_code, tool_type, description, input_schema, status, remark" + "provider_id, name, tool_code, description, tag, parameters, response_format, status, remark" } fn insert_values(&self) -> Vec> { @@ -89,9 +89,10 @@ impl Model for McpTool { Box::new(self.provider_id), Box::new(self.name.clone()), Box::new(self.tool_code.clone()), - Box::new(self.tool_type.clone()), Box::new(self.description.clone()), - Box::new(self.input_schema.clone()), + Box::new(self.tag.clone()), + Box::new(self.parameters.clone()), + Box::new(self.response_format.clone()), Box::new(self.status), Box::new(self.remark.clone()), ] @@ -102,11 +103,12 @@ impl Model for McpTool { provider_id = ?, name = ?, tool_code = ?, - tool_type = ?, description = ?, - input_schema = ?, + tag = ?, + parameters = ?, + response_format = ?, status = ?, - updated_at = CURRENT_TIMESTAMP, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -118,9 +120,10 @@ impl Model for McpTool { Box::new(self.provider_id), Box::new(self.name.clone()), Box::new(self.tool_code.clone()), - Box::new(self.tool_type.clone()), Box::new(self.description.clone()), - Box::new(self.input_schema.clone()), + Box::new(self.tag.clone()), + Box::new(self.parameters.clone()), + Box::new(self.response_format.clone()), Box::new(self.status), Box::new(self.remark.clone()), ] diff --git a/src/db/models/memory_info.rs b/src/db/models/memory_info.rs index 9882f21..2529230 100644 --- a/src/db/models/memory_info.rs +++ b/src/db/models/memory_info.rs @@ -5,17 +5,30 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// 记忆信息表（`memory_info`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.9：以 `memory_uuid` 唯一标识一条记忆， +/// 区分用户级/项目级记忆，并支持向量化索引（`is_vector`/`vector_id`）与关键词检索。 #[derive(Clone, Debug)] pub struct MemoryInfo { pub id: i32, - pub session_id: String, - pub content: String, + /// 记忆唯一标识。 + pub memory_uuid: String, + /// 记忆类型：user（用户级）/ project（项目级）。 pub memory_type: String, - pub status: i8, // default 1 - pub created_at: DateTime, - pub updated_at: DateTime, + /// 是否已向量化：1 是 / 0 否。 + pub is_vector: i8, + /// 向量索引 ID（外部向量库引用）。 + pub vector_id: Option, + /// 检索关键词（逗号分隔或 JSON）。 + pub keywords: Option, + /// 状态：1 有效 / 0 失效。 + pub status: i8, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -28,58 +41,48 @@ impl Model for MemoryInfo { r#" CREATE TABLE IF NOT EXISTS memory_info ( id INTEGER PRIMARY KEY AUTOINCREMENT, - session_id TEXT NOT NULL, - content TEXT NOT NULL, + memory_uuid TEXT NOT NULL, memory_type TEXT NOT NULL, + is_vector INTEGER NOT NULL DEFAULT 0, + vector_id TEXT, + keywords TEXT, status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, - session_id: row.get("session_id")?, - content: row.get("content")?, + memory_uuid: row.get("memory_uuid")?, memory_type: row.get("memory_type")?, + is_vector: row.get("is_vector")?, + vector_id: row.get("vector_id")?, + keywords: row.get("keywords")?, status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "session_id, content, memory_type, status, remark" + "memory_uuid, memory_type, is_vector, vector_id, keywords, status, remark" } fn insert_values(&self) -> Vec> { vec![ - Box::new(self.session_id.clone()), - Box::new(self.content.clone()), + Box::new(self.memory_uuid.clone()), Box::new(self.memory_type.clone()), + Box::new(self.is_vector), + Box::new(self.vector_id.clone()), + Box::new(self.keywords.clone()), Box::new(self.status), Box::new(self.remark.clone()), ] @@ -87,11 +90,13 @@ impl Model for MemoryInfo { fn update_set_clause(&self) -> String { r#" - session_id = ?, - content = ?, + memory_uuid = ?, memory_type = ?, + is_vector = ?, + vector_id = ?, + keywords = ?, status = ?, - updated_at = CURRENT_TIMESTAMP, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -100,9 +105,11 @@ impl Model for MemoryInfo { fn update_values(&self) -> Vec> { vec![ - Box::new(self.session_id.clone()), - Box::new(self.content.clone()), + Box::new(self.memory_uuid.clone()), Box::new(self.memory_type.clone()), + Box::new(self.is_vector), + Box::new(self.vector_id.clone()), + Box::new(self.keywords.clone()), Box::new(self.status), Box::new(self.remark.clone()), ] diff --git a/src/db/models/mod.rs b/src/db/models/mod.rs index 34128e6..150d746 100644 --- a/src/db/models/mod.rs +++ b/src/db/models/mod.rs @@ -1,3 +1,6 @@ +use chrono::{DateTime, NaiveDateTime, Utc}; +use rusqlite::Result as RusqliteResult; + pub mod ai_provider; pub mod llm_model; pub mod agent_info; @@ -12,4 +15,26 @@ pub mod mcp_tool; pub mod agent_mcp_rel; pub mod project_info; pub mod system_api_log; -pub mod system_config; \ No newline at end of file +pub mod system_config; + +/// 解析数据库时间字段。 +/// +/// 兼容两种格式： +/// - RFC3339（如 `2026-01-01T12:00:00Z`）； +/// - SQLite `CURRENT_TIMESTAMP` 产出的 `YYYY-MM-DD HH:MM:SS`。 +/// +/// 统一返回 UTC 时间，避免各 model 重复实现且格式不一致导致的解析失败。 +pub fn parse_db_time(raw: &str) -> RusqliteResult> { + if let Ok(dt) = DateTime::parse_from_rfc3339(raw) { + return Ok(dt.with_timezone(&Utc)); + } + + let naive = NaiveDateTime::parse_from_str(raw, "%Y-%m-%d %H:%M:%S") + .or_else(|_| NaiveDateTime::parse_from_str(raw, "%Y-%m-%dT%H:%M:%S")) + .or_else(|_| NaiveDateTime::parse_from_str(raw, "%Y-%m-%d %H:%M:%S%.f")) + .map_err(|e| { + rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e)) + })?; + + Ok(DateTime::::from_naive_utc_and_offset(naive, Utc)) +} diff --git a/src/db/models/project_info.rs b/src/db/models/project_info.rs index 9987e2e..099f85f 100644 --- a/src/db/models/project_info.rs +++ b/src/db/models/project_info.rs @@ -5,17 +5,29 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// 项目信息表（`project_info`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.13：以工作区绝对路径 `path` 唯一标识一个项目， +/// 记录项目自身的时间信息（`create_at`/`updata_at`，字段名沿用设计文档拼写）与逻辑删除标记。 +/// 注意：`create_at`/`updata_at` 为项目业务时间，与审计字段 `create_time`/`update_time` 并存。 #[derive(Clone, Debug)] pub struct ProjectInfo { pub id: i32, + /// 项目名称。 pub name: String, - pub icon: String, - pub project_type: String, - pub status: i8, // default 1 - pub created_at: DateTime, - pub updated_at: DateTime, + /// 项目工作区绝对路径（唯一）。 + pub path: String, + /// 项目创建时间（业务字段，字符串原样存储）。 + pub create_at: Option, + /// 项目更新时间（业务字段，字符串原样存储；字段名沿用设计文档拼写）。 + pub updata_at: Option, + /// 逻辑删除标记：0 正常 / 1 已删除。 + pub is_deleted: i8, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -29,58 +41,45 @@ impl Model for ProjectInfo { CREATE TABLE IF NOT EXISTS project_info ( id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, - icon TEXT NOT NULL, - project_type TEXT NOT NULL, - status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + path TEXT NOT NULL, + create_at TEXT, + updata_at TEXT, + is_deleted INTEGER NOT NULL DEFAULT 0, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, name: row.get("name")?, - icon: row.get("icon")?, - project_type: row.get("project_type")?, - status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + path: row.get("path")?, + create_at: row.get("create_at")?, + updata_at: row.get("updata_at")?, + is_deleted: row.get("is_deleted")?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "name, icon, project_type, status, remark" + "name, path, create_at, updata_at, is_deleted, remark" } fn insert_values(&self) -> Vec> { vec![ Box::new(self.name.clone()), - Box::new(self.icon.clone()), - Box::new(self.project_type.clone()), - Box::new(self.status), + Box::new(self.path.clone()), + Box::new(self.create_at.clone()), + Box::new(self.updata_at.clone()), + Box::new(self.is_deleted), Box::new(self.remark.clone()), ] } @@ -88,10 +87,11 @@ impl Model for ProjectInfo { fn update_set_clause(&self) -> String { r#" name = ?, - icon = ?, - project_type = ?, - status = ?, - updated_at = CURRENT_TIMESTAMP, + path = ?, + create_at = ?, + updata_at = ?, + is_deleted = ?, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -101,9 +101,10 @@ impl Model for ProjectInfo { fn update_values(&self) -> Vec> { vec![ Box::new(self.name.clone()), - Box::new(self.icon.clone()), - Box::new(self.project_type.clone()), - Box::new(self.status), + Box::new(self.path.clone()), + Box::new(self.create_at.clone()), + Box::new(self.updata_at.clone()), + Box::new(self.is_deleted), Box::new(self.remark.clone()), ] } diff --git a/src/db/models/prompt_info.rs b/src/db/models/prompt_info.rs index 4ac944b..a267f98 100644 --- a/src/db/models/prompt_info.rs +++ b/src/db/models/prompt_info.rs @@ -5,16 +5,34 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// 提示词表（`prompt_info`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.4：支持模板化配置、占位符动态渲染、 +/// 版本迭代（同名称下 `version` 唯一）与质量校验。 #[derive(Clone, Debug)] pub struct PromptInfo { pub id: i32, + /// 提示词名称，非空。 pub name: String, + /// 提示词内容，支持占位符模板。 pub content: String, - pub status: i8, // default 1 - pub created_at: DateTime, - pub updated_at: DateTime, + /// 提示词简介。 + pub description: Option, + /// 业务分类标签。 + pub category: Option, + /// 角色类型，枚举约束：System/User。 + pub role: Option, + /// 版本号，如 v1.0.0，同名称下版本唯一。 + pub version: Option, + /// 版本备注（迭代说明）。 + pub version_desc: Option, + /// 标签信息。 + pub tags: Option, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -29,54 +47,53 @@ impl Model for PromptInfo { id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, content TEXT NOT NULL, - status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + description TEXT, + category TEXT, + role TEXT, + version TEXT, + version_desc TEXT, + tags TEXT, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, name: row.get("name")?, content: row.get("content")?, - status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + description: row.get("description")?, + category: row.get("category")?, + role: row.get("role")?, + version: row.get("version")?, + version_desc: row.get("version_desc")?, + tags: row.get("tags")?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "name, content, status, remark" + "name, content, description, category, role, version, version_desc, tags, remark" } fn insert_values(&self) -> Vec> { vec![ Box::new(self.name.clone()), Box::new(self.content.clone()), - Box::new(self.status), + Box::new(self.description.clone()), + Box::new(self.category.clone()), + Box::new(self.role.clone()), + Box::new(self.version.clone()), + Box::new(self.version_desc.clone()), + Box::new(self.tags.clone()), Box::new(self.remark.clone()), ] } @@ -85,8 +102,13 @@ impl Model for PromptInfo { r#" name = ?, content = ?, - status = ?, - updated_at = CURRENT_TIMESTAMP, + description = ?, + category = ?, + role = ?, + version = ?, + version_desc = ?, + tags = ?, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -97,7 +119,12 @@ impl Model for PromptInfo { vec![ Box::new(self.name.clone()), Box::new(self.content.clone()), - Box::new(self.status), + Box::new(self.description.clone()), + Box::new(self.category.clone()), + Box::new(self.role.clone()), + Box::new(self.version.clone()), + Box::new(self.version_desc.clone()), + Box::new(self.tags.clone()), Box::new(self.remark.clone()), ] } diff --git a/src/db/models/prompt_placeholder.rs b/src/db/models/prompt_placeholder.rs index fee5b73..22252f4 100644 --- a/src/db/models/prompt_placeholder.rs +++ b/src/db/models/prompt_placeholder.rs @@ -5,17 +5,37 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// 占位符表（`prompt_placeholder`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.5：Prompt 占位符定义， +/// `(prompt_id, key)` 联合唯一。 +/// +/// 注：设计文档字段名为 `default`，因其为 SQL 保留字，落地列名统一用 `default_value`， +/// 语义一致（占位符兜底默认值）。 #[derive(Clone, Debug)] pub struct PromptPlaceholder { pub id: i32, + /// 提示词 ID，关联 `prompt_info.id` 逻辑外键。 pub prompt_id: i32, - pub placeholder: String, - pub value: String, - pub status: i8, // default 1 - pub created_at: DateTime, - pub updated_at: DateTime, + /// 占位符键，唯一标识（联合唯一约束 prompt_id + key）。 + pub key: String, + /// 占位符标签（前端展示名称）。 + pub label: Option, + /// 占位符值（填充值）。 + pub value: Option, + /// 默认值（兜底），对应设计文档 `default` 字段。 + pub default_value: Option, + /// 输入类型：输入 / 下拉 / 文本域。 + pub input_type: Option, + /// 是否必填：0-否 1-是。 + pub required: i8, + /// 启用状态：0-禁用 1-启用。 + pub status: i8, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -29,57 +49,53 @@ impl Model for PromptPlaceholder { CREATE TABLE IF NOT EXISTS prompt_placeholder ( id INTEGER PRIMARY KEY AUTOINCREMENT, prompt_id INTEGER NOT NULL, - placeholder TEXT NOT NULL, - value TEXT NOT NULL, + key TEXT NOT NULL, + label TEXT, + value TEXT, + default_value TEXT, + input_type TEXT, + required INTEGER NOT NULL DEFAULT 0, status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; - let updated_at: String = row.get("updated_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, prompt_id: row.get("prompt_id")?, - placeholder: row.get("placeholder")?, + key: row.get("key")?, + label: row.get("label")?, value: row.get("value")?, + default_value: row.get("default_value")?, + input_type: row.get("input_type")?, + required: row.get("required")?, status: row.get("status")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, - updated_at: DateTime::parse_from_rfc3339(&updated_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "prompt_id, placeholder, value, status, remark" + "prompt_id, key, label, value, default_value, input_type, required, status, remark" } fn insert_values(&self) -> Vec> { vec![ Box::new(self.prompt_id), - Box::new(self.placeholder.clone()), + Box::new(self.key.clone()), + Box::new(self.label.clone()), Box::new(self.value.clone()), + Box::new(self.default_value.clone()), + Box::new(self.input_type.clone()), + Box::new(self.required), Box::new(self.status), Box::new(self.remark.clone()), ] @@ -88,10 +104,14 @@ impl Model for PromptPlaceholder { fn update_set_clause(&self) -> String { r#" prompt_id = ?, - placeholder = ?, + key = ?, + label = ?, value = ?, + default_value = ?, + input_type = ?, + required = ?, status = ?, - updated_at = CURRENT_TIMESTAMP, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -101,8 +121,12 @@ impl Model for PromptPlaceholder { fn update_values(&self) -> Vec> { vec![ Box::new(self.prompt_id), - Box::new(self.placeholder.clone()), + Box::new(self.key.clone()), + Box::new(self.label.clone()), Box::new(self.value.clone()), + Box::new(self.default_value.clone()), + Box::new(self.input_type.clone()), + Box::new(self.required), Box::new(self.status), Box::new(self.remark.clone()), ] diff --git a/src/db/models/system_api_log.rs b/src/db/models/system_api_log.rs index f6c7650..5c485ab 100644 --- a/src/db/models/system_api_log.rs +++ b/src/db/models/system_api_log.rs @@ -5,18 +5,28 @@ use rusqlite::{ ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// 系统 API 调用日志表（`system_api_log`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.14：记录每次外部 API（大模型/工具）调用， +/// 通过 `rel_uuid` 关联业务实体（如 `chat_message.round_uuid`），并统计 Token 消耗与耗时。 #[derive(Clone, Debug)] pub struct SystemApiLog { pub id: i32, - pub api_path: String, - pub method: String, - pub status_code: i32, - pub response_time: i64, - pub request_data: String, - pub response_data: String, - pub created_at: DateTime, + /// 关联业务 UUID（可空）。 + pub rel_uuid: Option, + /// 调用状态：1 成功 / 0 失败。 + pub status: i8, + /// 本次调用 Token 消耗。 + pub token_cost: i32, + /// 调用耗时，单位：毫秒。 + pub duration: i32, + /// 调用内容/详情（可空，JSON 或文本）。 + pub content: Option, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -29,66 +39,58 @@ impl Model for SystemApiLog { r#" CREATE TABLE IF NOT EXISTS system_api_log ( id INTEGER PRIMARY KEY AUTOINCREMENT, - api_path TEXT NOT NULL, - method TEXT NOT NULL, - status_code INTEGER NOT NULL, - response_time INTEGER NOT NULL, - request_data TEXT NOT NULL, - response_data TEXT NOT NULL, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + rel_uuid TEXT, + status INTEGER NOT NULL DEFAULT 1, + token_cost INTEGER NOT NULL DEFAULT 0, + duration INTEGER NOT NULL DEFAULT 0, + content TEXT, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at: String = row.get("created_at")?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { id: row.get("id")?, - api_path: row.get("api_path")?, - method: row.get("method")?, - status_code: row.get("status_code")?, - response_time: row.get("response_time")?, - request_data: row.get("request_data")?, - response_data: row.get("response_data")?, - created_at: DateTime::parse_from_rfc3339(&created_at) - .map(|v| v.with_timezone(&Utc)) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?, + rel_uuid: row.get("rel_uuid")?, + status: row.get("status")?, + token_cost: row.get("token_cost")?, + duration: row.get("duration")?, + content: row.get("content")?, + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "api_path, method, status_code, response_time, request_data, response_data, remark" + "rel_uuid, status, token_cost, duration, content, remark" } fn insert_values(&self) -> Vec> { vec![ - Box::new(self.api_path.clone()), - Box::new(self.method.clone()), - Box::new(self.status_code), - Box::new(self.response_time), - Box::new(self.request_data.clone()), - Box::new(self.response_data.clone()), + Box::new(self.rel_uuid.clone()), + Box::new(self.status), + Box::new(self.token_cost), + Box::new(self.duration), + Box::new(self.content.clone()), Box::new(self.remark.clone()), ] } fn update_set_clause(&self) -> String { r#" - api_path = ?, - method = ?, - status_code = ?, - response_time = ?, - request_data = ?, - response_data = ?, + rel_uuid = ?, + status = ?, + token_cost = ?, + duration = ?, + content = ?, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() @@ -97,12 +99,11 @@ impl Model for SystemApiLog { fn update_values(&self) -> Vec> { vec![ - Box::new(self.api_path.clone()), - Box::new(self.method.clone()), - Box::new(self.status_code), - Box::new(self.response_time), - Box::new(self.request_data.clone()), - Box::new(self.response_data.clone()), + Box::new(self.rel_uuid.clone()), + Box::new(self.status), + Box::new(self.token_cost), + Box::new(self.duration), + Box::new(self.content.clone()), Box::new(self.remark.clone()), ] } diff --git a/src/db/models/system_config.rs b/src/db/models/system_config.rs index a824db9..a9bc6ba 100644 --- a/src/db/models/system_config.rs +++ b/src/db/models/system_config.rs @@ -1,23 +1,33 @@ -use chrono::{DateTime, NaiveDateTime, Utc}; +use chrono::{DateTime, Utc}; use rusqlite::{ Result as RusqliteResult, Row, ToSql, }; +use super::parse_db_time; use crate::core::db::Model; +/// 系统配置表（`system_config`）。 +/// +/// 对应《系统功能设计与数据库设计文档》四.15：以字符串 `key` 为主键的键值配置， +/// 支持类型标注、分组标签与排序，供 CLI/Control Plane 读取系统级默认配置。 #[derive(Clone, Debug)] pub struct SystemConfig { + /// 配置键（主键）。 pub key: String, + /// 配置值。 pub value: String, - pub key_type: String, - pub description: String, - pub tag: String, + /// 值类型标注（string/int/bool/json 等，可空）。 + pub key_type: Option, + /// 配置说明（可空）。 + pub description: Option, + /// 分组标签（可空）。 + pub tag: Option, + /// 排序序号。 pub sort: i32, - pub status: i8, - pub created_at: DateTime, - pub updated_at: DateTime, + pub create_time: DateTime, + pub update_time: DateTime, pub remark: Option, } @@ -31,89 +41,39 @@ impl Model for SystemConfig { CREATE TABLE IF NOT EXISTS system_config ( key TEXT PRIMARY KEY, value TEXT NOT NULL, - key_type TEXT NOT NULL, - description TEXT NOT NULL, - tag TEXT NOT NULL, + key_type TEXT, + description TEXT, + tag TEXT, sort INTEGER NOT NULL DEFAULT 0, - status INTEGER NOT NULL DEFAULT 1, - created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + create_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + update_time TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, remark TEXT ) "# } fn from_row(row: &Row) -> RusqliteResult { - let created_at_str: String = - row.get("created_at")?; - - let updated_at_str: String = - row.get("updated_at")?; - - let created_at = - NaiveDateTime::parse_from_str( - &created_at_str, - "%Y-%m-%d %H:%M:%S", - ) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?; - - let updated_at = - NaiveDateTime::parse_from_str( - &updated_at_str, - "%Y-%m-%d %H:%M:%S", - ) - .map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 0, - rusqlite::types::Type::Text, - Box::new(e), - ) - })?; + let create_time = parse_db_time(&row.get::<_, String>("create_time")?)?; + let update_time = parse_db_time(&row.get::<_, String>("update_time")?)?; Ok(Self { key: row.get("key")?, - value: row.get("value")?, - key_type: row.get("key_type")?, - description: row.get("description")?, - tag: row.get("tag")?, - sort: row.get("sort")?, - - status: row.get("status")?, - - created_at: - DateTime::::from_naive_utc_and_offset( - created_at, - Utc, - ), - - updated_at: - DateTime::::from_naive_utc_and_offset( - updated_at, - Utc, - ), - + create_time, + update_time, remark: row.get("remark")?, }) } fn insert_columns() -> &'static str { - "key, value, key_type, description, tag, sort, status, remark" + "key, value, key_type, description, tag, sort, remark" } - fn insert_values( - &self, - ) -> Vec> { + fn insert_values(&self) -> Vec> { vec![ Box::new(self.key.clone()), Box::new(self.value.clone()), @@ -121,7 +81,6 @@ impl Model for SystemConfig { Box::new(self.description.clone()), Box::new(self.tag.clone()), Box::new(self.sort), - Box::new(self.status), Box::new(self.remark.clone()), ] } @@ -133,24 +92,20 @@ impl Model for SystemConfig { description = ?, tag = ?, sort = ?, - status = ?, - updated_at = CURRENT_TIMESTAMP, + update_time = CURRENT_TIMESTAMP, remark = ? "# .trim() .to_string() } - fn update_values( - &self, - ) -> Vec> { + fn update_values(&self) -> Vec> { vec![ Box::new(self.value.clone()), Box::new(self.key_type.clone()), Box::new(self.description.clone()), Box::new(self.tag.clone()), Box::new(self.sort), - Box::new(self.status), Box::new(self.remark.clone()), ] } @@ -159,9 +114,7 @@ impl Model for SystemConfig { "key" } - fn primary_key_value( - &self, - ) -> Box { + fn primary_key_value(&self) -> Box { Box::new(self.key.clone()) } -} \ No newline at end of file +} diff --git a/src/lib.rs b/src/lib.rs index 44dc6d5..c0a4bc0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,6 +11,7 @@ use std::io; pub mod components; pub mod core; +pub mod cli; pub mod db; pub mod health; pub mod pages; @@ -27,7 +28,42 @@ i18n!("locales", fallback = "en"); #[tokio::main] pub async fn run() -> anyhow::Result<()> { + use clap::Parser; + + use crate::cli::{Cli, Command}; + info!("{}", t!("logger_is_initialized")); + + let mut cli = Cli::parse(); + + // 解析 Workspace 与启动模式，并注入运行环境（优先级：命令行 > 配置文件 > 系统默认）。 + let workspace = cli.workspace()?; + // Rust 2024：`set_var` 为 unsafe；此处单线程启动阶段修改，安全可控。 + unsafe { + std::env::set_var("OMEGA_WORKING_DIR", &workspace); + std::env::set_var("OMEGA_MODE", cli.mode.as_str()); + } + info!("workspace = {} (mode = {})", workspace.display(), cli.mode.as_str()); + + match cli.command.take() { + // 无子命令：以解析出的 Workspace 启动交互式 TUI。 + None => launch_tui().await, + // `chat --chat-uuid `（无 clear）：续跑指定会话，注入后进入 TUI。 + Some(Command::Chat(chat)) if chat.chat_uuid.is_some() && chat.action.is_none() => { + let uuid = chat.chat_uuid.clone().unwrap_or_default(); + unsafe { + std::env::set_var("OMEGA_CHAT_UUID", &uuid); + } + info!("resuming chat session = {}", uuid); + launch_tui().await + } + // 其余管理型子命令：执行后直接返回，不进入 TUI。 + Some(cmd) => cli::commands::dispatch(cmd).await, + } +} + +/// 健康检查通过后启动交互式 TUI。 +async fn launch_tui() -> anyhow::Result<()> { info!("{}", t!("test_message", name = "OmegaCode")); info!("{}", t!("current_locale", locale_name = "en")); @@ -36,6 +72,11 @@ pub async fn run() -> anyhow::Result<()> { return Err(anyhow::anyhow!("Health check failed")); } + // workspace 模式下自动登记当前打开的项目（失败不阻断 TUI 启动）。 + if let Err(e) = crate::core::project::register_opened_from_env() { + log::warn!("project register failed: {e}"); + } + element!(Root) .into_any() .fullscreen() @@ -50,7 +91,7 @@ pub fn Root(_hooks: Hooks) -> impl Into> { element!( RouterProvider( routes: app_routes(), - index_path: "/router", + index_path: crate::router::initial_route(), ) ) } diff --git a/src/main.rs b/src/main.rs index d67e0bf..9d8149d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,21 +2,22 @@ extern crate OmegaCode; #[macro_use] extern crate log; -#[macro_use] -extern crate flexi_logger; use flexi_logger::{ - Cleanup, Criterion, Duplicate, FileSpec, Logger, Naming, WriteMode, detailed_format, + Cleanup, Criterion, FileSpec, Logger, Naming, WriteMode, detailed_format, }; fn main() { init_logger(); if let Err(e) = OmegaCode::run() { error!("Application error: {}", e); + // CLI 子命令的错误需反馈到控制台（日志仅写文件），否则表现为静默 exit 1。 + eprintln!("Error: {}", e); // Print full error chain let mut current = e.source(); while let Some(cause) = current { error!("Caused by: {}", cause); + eprintln!(" caused by: {}", cause); current = cause.source(); } std::process::exit(1); diff --git a/src/pages/chat.rs b/src/pages/chat.rs index 095bcb0..02e4a7b 100644 --- a/src/pages/chat.rs +++ b/src/pages/chat.rs @@ -5,10 +5,13 @@ //! - 当前正在流式生成的回复渲染在页面底部的「实时区」，随 token 到达不断重绘； //! - 输入框回车提交后，使用 `tokio::spawn` 驱动一轮 Agent 对话，事件通过 mpsc 通道回传。 +use std::collections::BTreeMap; +use std::path::PathBuf; use std::sync::Arc; +use anyhow::Result; use ratatui_kit::components::{Border, Text, TextArea, TextParagraph, View}; -use ratatui_kit::crossterm::event::{Event, KeyCode, KeyEventKind}; +use ratatui_kit::crossterm::event::{Event, KeyCode, KeyEventKind, KeyModifiers}; use ratatui_kit::ratatui::layout::{Constraint, Direction, Position}; use ratatui_kit::ratatui::prelude::Color; use ratatui_kit::ratatui::style::{Style, Stylize}; @@ -22,7 +25,10 @@ use tokio::sync::{mpsc, Mutex}; use crate::components::status_bar::StatusBar; use crate::core::agent::{Agent, AgentEvent}; -use crate::core::config::AppConfig; +use crate::core::arena::{ArenaService, VoteInput}; +use crate::core::db::DatabaseManager; +use crate::core::memory::{MemoryScope, MemoryStore}; +use crate::db::models::system_config::SystemConfig; /// Agent 持有单元：初始化可能失败，失败时保留错误信息用于展示。 struct AgentCell { @@ -32,40 +38,125 @@ struct AgentCell { working_dir: String, } -/// 依据环境配置构建 Agent（并尽力开启会话持久化）。 +/// 从 `system_config` 读取一个整数型 Agent 主键配置（如 `default_agent_id` / +/// `coordinator_agent_id`，均由 CLI 写入）。 +/// +/// 返回 `Ok(None)` 表示尚未配置或值非法；`Err` 表示数据库不可用。 +fn resolve_agent_id_config(key: &str) -> Result> { + let db = DatabaseManager::new()?; + let found = db.find_where::("key = ?1", &[&key])?; + Ok(found + .into_iter() + .next() + .and_then(|c| c.value.trim().parse::().ok())) +} + +/// 从 `system_config` 读取一个字符串型配置（如 `memory_type`）。 +/// +/// 返回 `Ok(None)` 表示尚未配置；`Err` 表示数据库不可用。 +fn resolve_string_config(key: &str) -> Result> { + let db = DatabaseManager::new()?; + let found = db.find_where::("key = ?1", &[&key])?; + Ok(found + .into_iter() + .next() + .map(|c| c.value.trim().to_string()) + .filter(|v| !v.is_empty())) +} + +/// 依据数据库配置构建 Agent（数据库驱动，并尽力开启会话持久化）。 +/// +/// 配置来源严格为数据库：从 `system_config.default_agent_id` 解析默认 Agent，再经 +/// [`Agent::from_database`] 加载 `agent_info`+`llm_model`+`ai_provider`+`prompt_info` +/// 及已授权的 MCP 工具。若未配置默认 Agent 或加载失败，则返回带错误信息的空单元， +/// 由页面展示引导（不回退到环境变量配置）。 fn init_agent() -> AgentCell { - let fallback_dir = || { - std::env::current_dir() - .map(|p| p.display().to_string()) - .unwrap_or_default() + // 工作目录来自 CLI 注入的 Workspace 路径（OMEGA_WORKING_DIR），与 Agent 配置来源无关。 + let working_dir = std::env::var("OMEGA_WORKING_DIR") + .map(PathBuf::from) + .or_else(|_| std::env::current_dir()) + .unwrap_or_else(|_| PathBuf::from(".")); + let working_dir_str = working_dir.display().to_string(); + + let agent_id = match resolve_agent_id_config("default_agent_id") { + Ok(Some(id)) => id, + Ok(None) => { + return AgentCell { + agent: None, + error: Some( + "no default agent configured; run `Omega agent use -a ` first".to_string(), + ), + working_dir: working_dir_str, + }; + } + Err(e) => { + return AgentCell { + agent: None, + error: Some(format!("failed to resolve default agent: {e}")), + working_dir: working_dir_str, + }; + } }; - match AppConfig::from_env() { - Ok(cfg) => { - let working_dir = cfg.working_dir.display().to_string(); - match Agent::from_config(&cfg) { - Ok(mut agent) => { - // 持久化为尽力而为：失败仅告警，不阻断对话。 - if let Err(e) = agent.enable_persistence_auto(Some("TUI Chat")) { - log::warn!("Failed to enable chat persistence: {e}"); - } - AgentCell { - agent: Some(agent), - error: None, - working_dir, + // 若该 Agent 被标记为协调者（`coordinator_agent_id`），则升级为多 Agent 调度器。 + let coordinator_id = resolve_agent_id_config("coordinator_agent_id").unwrap_or(None); + + match Agent::from_database(agent_id, working_dir.clone()) { + Ok(mut agent) => { + if Some(agent_id) == coordinator_id { + if let Err(e) = agent.enable_coordinator(working_dir.clone()) { + log::warn!("Failed to enable coordinator: {e}"); + } + } + // 记忆系统：按 `system_config.memory_type` 作用域加载并注入记忆上下文， + // 作用域非 none 时启用每轮对话后的异步记忆抽取。失败仅告警，不阻断对话。 + let memory_scope = MemoryScope::parse( + &resolve_string_config("memory_type") + .ok() + .flatten() + .unwrap_or_default(), + ); + if memory_scope != MemoryScope::None { + // Workspace 模式启用项目记忆（正文存 `/.omega/memory`）； + // single 模式仅用户记忆（project_dir = None）。 + let project_root = + if std::env::var("OMEGA_MODE").as_deref() == Ok("workspace") { + Some(working_dir.clone()) + } else { + None + }; + match MemoryStore::new(project_root) { + Ok(store) => { + if let Err(e) = agent.enable_memory(Arc::new(store), memory_scope) { + log::warn!("Failed to enable memory: {e}"); + } } + Err(e) => log::warn!("Failed to open memory store: {e}"), } - Err(e) => AgentCell { - agent: None, - error: Some(e.to_string()), - working_dir, - }, + } + // 会话初始化为尽力而为：失败仅告警，不阻断对话。 + // 若 CLI 指定了 --chat-uuid（经环境变量注入），则加载该历史会话续跑； + // 否则新建一个持久化会话（llm_id 已由数据库配置绑定）。 + let resume = std::env::var("OMEGA_CHAT_UUID") + .ok() + .filter(|s| !s.trim().is_empty()); + let init = match resume { + Some(uuid) => agent.load_session(uuid), + None => agent.start_session(Some("TUI Chat")), + }; + if let Err(e) = init { + log::warn!("Failed to initialize chat session: {e}"); + } + AgentCell { + agent: Some(agent), + error: None, + working_dir: working_dir_str, } } Err(e) => AgentCell { agent: None, - error: Some(e.to_string()), - working_dir: fallback_dir(), + error: Some(format!("failed to load agent #{agent_id}: {e}")), + working_dir: working_dir_str, }, } } @@ -247,6 +338,81 @@ pub fn ChatPage(mut hooks: Hooks) -> impl Into> { term_w, ); } + // Ctrl+1..5：给上一轮助手回复打综合星级（不填多维度，多维度录入以 CLI 为准）。 + KeyCode::Char(c @ '1'..='5') + if key.modifiers.contains(KeyModifiers::CONTROL) => + { + if busy.get() { + return; + } + let score = c.to_digit(10).unwrap_or(1) as i8; + let cell = agent_cell.read().clone(); + let mut vote_status = status; + let insert = insert_before.clone(); + let width = term_w; + tokio::spawn(async move { + // 先取锁读取标识，随即释放，避免跨阻塞 DB 调用持锁。 + let (agent_id, chat_uuid, round_uuid) = { + let guard = cell.lock().await; + match guard.agent.as_ref() { + Some(a) => ( + a.agent_id(), + a.session_uuid().to_string(), + a.last_round_uuid().map(|s| s.to_string()), + ), + None => (None, String::new(), None), + } + }; + let (Some(aid), Some(round)) = (agent_id, round_uuid) else { + vote_status.set("nothing to rate yet".to_string()); + return; + }; + let svc = match ArenaService::new() { + Ok(s) => s, + Err(e) => { + vote_status.set(format!("vote failed: {e}")); + return; + } + }; + let input = VoteInput { + agent_id: aid, + score, + comment: None, + chat_uuid, + round_uuid: Some(round), + dimensions: BTreeMap::new(), + }; + match svc.record_vote(input) { + Ok(_) => { + let avg = svc + .ranking() + .ok() + .and_then(|rs| { + rs.into_iter() + .find(|r| r.agent_id == aid) + .and_then(|r| r.avg_score) + }); + let line = match avg { + Some(a) => { + format!("rated last response: {}/5 (avg -> {:.1})", score, a) + } + None => format!("rated last response: {}/5", score), + }; + let height = wrapped_rows(&line, width).min(1000) as u16; + insert + .render_before( + Paragraph::new(Line::styled( + line, + Style::default().fg(Color::Green), + )), + height, + ) + .finish(); + } + Err(e) => vote_status.set(format!("vote failed: {e}")), + } + }); + } _ => {} } } @@ -279,7 +445,7 @@ pub fn ChatPage(mut hooks: Hooks) -> impl Into> { }, multiline: false, cursor_style: Style::default().on_green(), - placeholder: Some("Type a message... (Enter to send, Esc to quit)".to_string()), + placeholder: Some("Type a message... (Enter to send, Ctrl+1..5 rate last reply, Esc to quit)".to_string()), placeholder_style: Style::default().green(), ) } diff --git a/src/pages/router.rs b/src/pages/router.rs index 07c8fbb..713af07 100644 --- a/src/pages/router.rs +++ b/src/pages/router.rs @@ -23,9 +23,10 @@ pub fn RouterPage(mut hooks: Hooks) -> impl Into> { } }); - //todo 这里的逻辑是通过数据库查询是否初始化，从而确定要路由到 Welcome 还是 Chat - //todo 为了开发方便，这里的逻辑还需要判断是否是开发环境，如果是开发环境，默认路由到路由页面，路由页面在生产环境不对外暴露 - //todo 同时路由函数禁止在生成环境中调用，避免在生产环境中路由到路由页面 + // 本页为「开发调试路由页」：仅在开发环境（debug 构建或设置 OMEGA_DEV）作为入口， + // 由 `crate::router::initial_route()` 决定；生产环境入口直接为 /welcome 或 /chat， + // 永不进入本页，故下方 1/2/3 手动路由不对生产环境暴露 + // （见《系统功能设计》路由约束：路由调试函数禁止在生产环境调用）。 element!( View(height: Constraint::Length(10),) { diff --git a/src/router/mod.rs b/src/router/mod.rs index e4e7313..cab05f2 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -10,4 +10,64 @@ pub fn app_routes() -> Vec { "/router" => RouterPage, }; routes +} + +/// 是否为开发环境：debug 构建，或显式设置 `OMEGA_DEV` 环境变量。 +/// +/// 开发环境下入口路由指向 `/router`（手动 1/2/3 调试页）；生产环境永不进入 +/// `/router`，满足《系统功能设计》「路由调试页不对生产环境暴露」的约束。 +pub fn is_dev() -> bool { + cfg!(debug_assertions) || std::env::var("OMEGA_DEV").is_ok() +} + +/// 应用是否已完成初始化配置：以 `system_config.default_agent_id` 是否已写入为准 +/// （由 CLI `agent use` 设置）。数据库不可用或未配置均视为「未初始化」。 +/// +/// 与 chat 页 `init_agent` 的纯数据库驱动配置来源保持一致：无默认 Agent 时进入 +/// 欢迎页引导，已配置时直接进入聊天页。 +pub fn is_configured() -> bool { + use crate::core::db::DatabaseManager; + use crate::db::models::system_config::SystemConfig; + + let Ok(db) = DatabaseManager::new() else { + return false; + }; + db.find_where::("key = ?1", &[&"default_agent_id"]) + .ok() + .and_then(|rows| rows.into_iter().next()) + .map(|c| c.value.trim().parse::().is_ok()) + .unwrap_or(false) +} + +/// 依据「是否开发环境 / 是否已配置」决定入口路由（纯函数，可测）。 +pub fn choose_initial_route(dev: bool, configured: bool) -> &'static str { + if dev { + "/router" + } else if configured { + "/chat" + } else { + "/welcome" + } +} + +/// 计算本次启动的入口路由路径（供 `Root` 的 `RouterProvider.index_path` 使用）。 +pub fn initial_route() -> &'static str { + choose_initial_route(is_dev(), is_configured()) +} + +#[cfg(test)] +mod tests { + use super::choose_initial_route; + + #[test] + fn dev_env_always_enters_router() { + assert_eq!(choose_initial_route(true, false), "/router"); + assert_eq!(choose_initial_route(true, true), "/router"); + } + + #[test] + fn prod_routes_by_configuration() { + assert_eq!(choose_initial_route(false, true), "/chat"); + assert_eq!(choose_initial_route(false, false), "/welcome"); + } } \ No newline at end of file diff --git a/src/runtime/runtime.rs b/src/runtime/runtime.rs index 0d544df..e69a531 100644 --- a/src/runtime/runtime.rs +++ b/src/runtime/runtime.rs @@ -21,6 +21,16 @@ pub struct RuntimePaths { pub python: PathBuf, } +/// 同步探测已安装的 JRE `java` 可执行文件路径（不触发下载）。 +/// +/// 仅扫描 `~/.omega/runtime/jre/current`；未安装或找不到时返回 `None`。 +/// 供「打开控制台面板」等需要就地启动 Java 进程的场景使用。 +pub fn installed_java() -> Option { + let runtime_root = get_app_runtime_dir().ok()?; + let current_dir = runtime_root.join("jre").join("current"); + find_java_binary(¤t_dir, current_platform().os).ok() +} + pub async fn ensure_jre< DF, EF, diff --git a/src/services/welcome.rs b/src/services/welcome.rs index e69de29..53557d6 100644 --- a/src/services/welcome.rs +++ b/src/services/welcome.rs @@ -0,0 +1,201 @@ +//! 欢迎页服务：启动本地「控制台面板」并在浏览器中打开。 +//! +//! 控制台面板（Spring Boot 后端 + Vue 前端）由 `build.sh` 打包为与 `OmegaCode` +//! 可执行文件同目录的 `console_panel.jar`，运行于 。本模块 +//! 负责：探测面板是否已运行、定位 jar 与已安装的 JRE、以 `java -jar` 后台启动面板， +//! 并调用系统默认浏览器打开面板地址。 +//! +//! 设计约束：所有探测/启动均为「尽力而为」——缺 jar 或缺 JRE 时返回带清晰原因的错误， +//! 由 TUI 以状态行提示用户，绝不 panic、不阻断欢迎页其余交互。 + +use std::net::{SocketAddr, TcpStream}; +use std::path::PathBuf; +use std::process::Command; +use std::time::Duration; + +use anyhow::{Context, Result}; + +use crate::platform::detect::{current_platform, Os}; +use crate::runtime::runtime::installed_java; + +/// 控制台面板服务端口。 +pub const CONSOLE_PANEL_PORT: u16 = 5200; +/// 控制台面板访问地址。 +pub const CONSOLE_PANEL_URL: &str = "http://localhost:5200"; +/// 打包后的面板 jar 文件名（与可执行文件同目录）。 +pub const CONSOLE_PANEL_JAR: &str = "console_panel.jar"; +/// 等待面板端口就绪的轮询次数（每次 500ms，合计约 30s）。 +const PANEL_READY_POLLS: u32 = 60; + +/// 面板启动结果，用于 TUI 状态提示。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PanelOutcome { + /// 面板已在运行，直接打开浏览器。 + AlreadyRunning, + /// 面板未运行，已后台启动 `java -jar`，正等待端口就绪后自动打开浏览器。 + Starting, +} + +impl PanelOutcome { + /// 面向用户的一句话状态描述。 + pub fn message(&self) -> String { + match self { + PanelOutcome::AlreadyRunning => { + format!("console panel already running — opened {CONSOLE_PANEL_URL}") + } + PanelOutcome::Starting => { + format!("starting console panel — will open {CONSOLE_PANEL_URL} when ready") + } + } + } +} + +/// 定位 `console_panel.jar`：优先可执行文件同目录，其次当前工作目录。 +/// +/// 开发环境（`cargo run`）通常两处都没有该 jar，返回 `None`。 +pub fn locate_console_jar() -> Option { + if let Ok(exe) = std::env::current_exe() { + if let Some(dir) = exe.parent() { + let candidate = dir.join(CONSOLE_PANEL_JAR); + if candidate.is_file() { + return Some(candidate); + } + } + } + let cwd = PathBuf::from(CONSOLE_PANEL_JAR); + if cwd.is_file() { + return Some(cwd); + } + None +} + +/// 探测指定端口是否已可连接（用于判断面板是否已运行）。 +fn port_open(port: u16) -> bool { + let addr = SocketAddr::from(([127, 0, 0, 1], port)); + TcpStream::connect_timeout(&addr, Duration::from_millis(300)).is_ok() +} + +/// 探测控制台面板是否已在运行。 +pub fn panel_running() -> bool { + port_open(CONSOLE_PANEL_PORT) +} + +/// 依据操作系统构造「打开浏览器」所需的可执行程序与参数（纯函数，可测）。 +/// +/// Windows 使用 `cmd /C start "" `（空标题占位），macOS 使用 `open`， +/// Linux 使用 `xdg-open`。 +pub fn browser_command(os: Os, url: &str) -> (&'static str, Vec) { + match os { + Os::Windows => ( + "cmd", + vec![ + "/C".to_string(), + "start".to_string(), + String::new(), + url.to_string(), + ], + ), + Os::MacOS => ("open", vec![url.to_string()]), + Os::Linux => ("xdg-open", vec![url.to_string()]), + } +} + +/// 用系统默认浏览器打开给定 URL（非阻塞 `spawn`）。 +pub fn open_in_browser(url: &str) -> Result<()> { + let os = current_platform().os; + let (program, args) = browser_command(os, url); + Command::new(program) + .args(&args) + .spawn() + .with_context(|| format!("failed to launch browser via '{program}'"))?; + Ok(()) +} + +/// 打开控制台面板：若未运行则后台启动 `java -jar`，随后打开浏览器。 +/// +/// 返回 [`PanelOutcome`] 描述所采取的动作。缺 jar 或缺 JRE 时返回错误（附清晰原因）， +/// 调用方（欢迎页）据此提示用户，不影响其余菜单项。 +pub fn open_console_panel() -> Result { + if panel_running() { + open_in_browser(CONSOLE_PANEL_URL)?; + return Ok(PanelOutcome::AlreadyRunning); + } + + let jar = locate_console_jar().context( + "console_panel.jar not found next to the executable (build the bundle with ./build.sh first)", + )?; + let java = installed_java().context( + "JRE not installed; finish the runtime installation on the welcome screen first", + )?; + + Command::new(&java) + .arg("-jar") + .arg(&jar) + .spawn() + .with_context(|| format!("failed to start console panel via '{}'", java.display()))?; + + // 后台轮询端口就绪再打开浏览器，避免阻塞 TUI；超时也尝试打开一次（面板可能启动较慢）。 + std::thread::spawn(|| { + for _ in 0..PANEL_READY_POLLS { + if panel_running() { + let _ = open_in_browser(CONSOLE_PANEL_URL); + return; + } + std::thread::sleep(Duration::from_millis(500)); + } + let _ = open_in_browser(CONSOLE_PANEL_URL); + }); + + Ok(PanelOutcome::Starting) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::TcpListener; + + #[test] + fn browser_command_matches_platform_conventions() { + let (program, args) = browser_command(Os::Windows, "http://localhost:5200"); + assert_eq!(program, "cmd"); + assert_eq!( + args, + vec![ + "/C".to_string(), + "start".to_string(), + String::new(), + "http://localhost:5200".to_string() + ] + ); + + let (program, args) = browser_command(Os::MacOS, "http://x"); + assert_eq!(program, "open"); + assert_eq!(args, vec!["http://x".to_string()]); + + let (program, args) = browser_command(Os::Linux, "http://x"); + assert_eq!(program, "xdg-open"); + assert_eq!(args, vec!["http://x".to_string()]); + } + + #[test] + fn port_open_detects_active_listener() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let port = listener.local_addr().unwrap().port(); + // 监听中：应探测为开启。 + assert!(port_open(port)); + // 显式持有 listener 至断言之后，避免端口在断言前被释放/复用。 + drop(listener); + } + + #[test] + fn panel_outcome_messages_mention_url() { + assert!(PanelOutcome::AlreadyRunning.message().contains(CONSOLE_PANEL_URL)); + assert!(PanelOutcome::Starting.message().contains(CONSOLE_PANEL_URL)); + } + + #[test] + fn locate_console_jar_does_not_panic() { + // 开发环境下通常无 jar；仅验证调用安全（返回 Option，不 panic）。 + let _ = locate_console_jar(); + } +} From 4c93973f3ed1e97cb28564fb910fd72d4361c97a Mon Sep 17 00:00:00 2001 From: fusheng Date: Sat, 17 Oct 2026 22:17:50 +0800 Subject: [PATCH 3/3] =?UTF-8?q?fix(app):=20=E4=BF=AE=E5=A4=8D=E8=B7=AF?= =?UTF-8?q?=E7=94=B1=E8=B6=8A=E7=95=8C=20panic=20=E5=92=8C=E7=BB=84?= =?UTF-8?q?=E4=BB=B6=E4=B8=8A=E4=B8=8B=E6=96=87=E7=BC=BA=E5=A4=B1=E5=B4=A9?= =?UTF-8?q?=E6=BA=83?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 将路由跳转从 push 改为 replace，规避 RouterHistory 越界 panic - 在 Welcome 组件 update 中手动挂载 context_stack，修复 "context not available" 崩溃 - 移除 layout.rs 中多余的 main 函数 - 清理 lib.rs、status_bar.rs 等文件中未使用的导入和变量 - 为 progress_bar 和 stream 模块添加 #[allow(dead_code)] 标记 - 修复 status_bar 中 buf.get_mut 越界风险，改用 cell_mut - 补充 status_bar 返回类型生命周期标注 - 移除 providers 模块中未使用的 pub use 导出 - 在 Cargo.toml 中配置 non_snake_case = allow 以支持大写驼峰 bin 名 --- Cargo.toml | 5 ++++ src/components/progress_bar.rs | 4 ++- src/components/status_bar.rs | 14 ++++------ src/components/welcome.rs | 31 +++++++++++++-------- src/core/client/mod.rs | 2 +- src/core/providers/claude_messages/mod.rs | 4 +-- src/core/providers/openai_compatible/mod.rs | 4 +-- src/core/transport/stream.rs | 2 ++ src/lib.rs | 5 +--- src/pages/layout.rs | 11 +------- src/pages/router.rs | 10 +++++-- 11 files changed, 47 insertions(+), 45 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index f612fbf..84f6544 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,11 @@ edition = "2024" [lib] path = "src/lib.rs" +# crate/bin 名沿用产品名 `OmegaCode`（大写驼峰），与可执行文件名保持一致； +# 显式关闭 non_snake_case 以抑制“crate 名应为 snake case”警告。 +[lints.rust] +non_snake_case = "allow" + [[bin]] name = "OmegaCode" path = "src/main.rs" diff --git a/src/components/progress_bar.rs b/src/components/progress_bar.rs index af3a07c..76cbeb3 100644 --- a/src/components/progress_bar.rs +++ b/src/components/progress_bar.rs @@ -1,10 +1,12 @@ use ratatui_kit::prelude::*; use ratatui_kit::ratatui::{ style::{Modifier, Style}, - text::{Line, Span}, + text::Span, }; use ratatui_kit::ratatui::layout::Direction; +// 预留：进度条组件属性。当前 Welcome 组件内联渲染安装进度条，ProgressBar 尚未接入，故允许未构造。 +#[allow(dead_code)] #[derive(Debug, Clone, Props)] pub struct ProgressBarProps { pub style: Style, diff --git a/src/components/status_bar.rs b/src/components/status_bar.rs index 02dd386..634f90c 100644 --- a/src/components/status_bar.rs +++ b/src/components/status_bar.rs @@ -4,13 +4,12 @@ use ratatui_kit::{ ComponentUpdater, Hooks, Props, - UseState, }; use ratatui_kit::ratatui::{ style::{Color, Style, Stylize}, text::{Line, Span}, - widgets::{Block, Borders, Paragraph, Widget}, + widgets::{Paragraph, Widget}, layout::{Alignment, Rect}, buffer::Buffer, }; @@ -100,9 +99,9 @@ impl StatusBar { fn render_border_top(&self, buf: &mut Buffer, area: Rect) { if area.height > 0 { for x in area.left()..area.right() { - buf.get_mut(x, area.top()) - .set_char('─') - .set_style(self.border_style); + if let Some(cell) = buf.cell_mut((x, area.top())) { + cell.set_char('─').set_style(self.border_style); + } } } } @@ -123,7 +122,7 @@ impl StatusBar { } /// 创建内容行，包含左侧路径和右侧进度条 - fn create_content_line(&self, total_width: usize) -> Line { + fn create_content_line(&self, total_width: usize) -> Line<'_> { let mut spans = Vec::new(); // 计算左右两部分的最大宽度 @@ -149,7 +148,6 @@ impl StatusBar { // 为右侧进度条预留空间 let progress_bar_width = std::cmp::min(right_max_width, 20); // 最大20字符宽的进度条 - let required_space = progress_bar_width + 3; // 加上一些间隔 let available_middle_space = total_width.saturating_sub(current_path_width + progress_bar_width); // 添加中间空白 @@ -187,7 +185,7 @@ impl StatusBar { } /// 创建上下文长度进度条 - fn create_progress_bar(&self, max_width: usize) -> Vec { + fn create_progress_bar(&self, max_width: usize) -> Vec> { let mut spans = Vec::new(); if self.max_context_length == 0 { diff --git a/src/components/welcome.rs b/src/components/welcome.rs index 653f164..9b6e941 100644 --- a/src/components/welcome.rs +++ b/src/components/welcome.rs @@ -83,23 +83,30 @@ impl Component for Welcome { &mut self, props: &mut Self::Props<'_>, mut hooks: Hooks, - _updater: &mut ComponentUpdater, + updater: &mut ComponentUpdater, ) { + // 手写 `impl Component`（非 `#[component]` 宏）不会自动注入上下文栈： + // 宏生成的 `update` 会先执行 `hooks.with_context_stack(updater.component_context_stack())`， + // 而此处 `hooks` 由框架构造时 `context` 为 `None`。若不手动挂载即调用 + // `use_navigate`/`use_exit`/`use_context*` 等依赖上下文的钩子，会触发 + // `use_context.rs` 的 "context not available" 崩溃。故在任何钩子调用前先挂载上下文栈。 + let mut hooks = hooks.with_context_stack(updater.component_context_stack()); + self.style = props.style; self.menu_style = props.menu_style; self.version = props.version.clone(); - let mut jre_ready = hooks.use_state(|| self.jre_ready); - let mut python_ready = hooks.use_state(|| self.python_ready); - let mut is_installing = hooks.use_state(|| self.is_installing); - let mut install_message = hooks.use_state(|| self.install_message.clone()); + let jre_ready = hooks.use_state(|| self.jre_ready); + let python_ready = hooks.use_state(|| self.python_ready); + let is_installing = hooks.use_state(|| self.is_installing); + let install_message = hooks.use_state(|| self.install_message.clone()); let mut install_started = hooks.use_state(|| false); - let mut stage = hooks.use_state(|| self.stage.clone()); - let mut percent = hooks.use_state(|| self.percent); - let mut downloaded = hooks.use_state(|| self.downloaded); - let mut total = hooks.use_state(|| self.total); - let mut speed = hooks.use_state(|| self.speed); + let stage = hooks.use_state(|| self.stage.clone()); + let percent = hooks.use_state(|| self.percent); + let downloaded = hooks.use_state(|| self.downloaded); + let total = hooks.use_state(|| self.total); + let speed = hooks.use_state(|| self.speed); let mut state_index = hooks.use_state(|| self.selected_index); let menu_message = hooks.use_state(|| self.menu_message.clone()); @@ -140,8 +147,8 @@ impl Component for Welcome { msg.set(text); }); } - // 开启新项目：进入聊天页。 - 1 => navigate.push("/chat"), + // 开启新项目：进入聊天页（replace 不增长历史栈，规避 RouterHistory::push 越界缺陷）。 + 1 => navigate.replace("/chat"), // 退出程序。 2 => exit(), _ => {} diff --git a/src/core/client/mod.rs b/src/core/client/mod.rs index 98e8f40..93a3d88 100644 --- a/src/core/client/mod.rs +++ b/src/core/client/mod.rs @@ -9,7 +9,7 @@ pub use ai_client::AiClient; use anyhow::{anyhow, Result}; use bytes::Bytes; use futures_util::stream::Stream; -use rand::{Rng, RngExt}; +use rand::RngExt; use reqwest::{ multipart, Client, diff --git a/src/core/providers/claude_messages/mod.rs b/src/core/providers/claude_messages/mod.rs index 0654bea..bbfde7e 100644 --- a/src/core/providers/claude_messages/mod.rs +++ b/src/core/providers/claude_messages/mod.rs @@ -2,6 +2,4 @@ mod request; mod response; mod provider; -pub use provider::ClaudeProvider; -pub use request::ClaudeRequest; -pub use response::ClaudeResponse; \ No newline at end of file +pub use provider::ClaudeProvider; \ No newline at end of file diff --git a/src/core/providers/openai_compatible/mod.rs b/src/core/providers/openai_compatible/mod.rs index 71baa27..63c5650 100644 --- a/src/core/providers/openai_compatible/mod.rs +++ b/src/core/providers/openai_compatible/mod.rs @@ -2,6 +2,4 @@ mod request; mod response; mod provider; -pub use provider::OpenAiCompatibleProvider; -pub use request::OpenAiRequest; -pub use response::OpenAiResponse; \ No newline at end of file +pub use provider::OpenAiCompatibleProvider; \ No newline at end of file diff --git a/src/core/transport/stream.rs b/src/core/transport/stream.rs index 694ba11..e7b8cd1 100644 --- a/src/core/transport/stream.rs +++ b/src/core/transport/stream.rs @@ -1,5 +1,7 @@ use futures_util::StreamExt; +// 预留：SSE 原始字节流调试打印。当前流式响应由 AiClient::post_sse 处理，此函数尚未接入。 +#[allow(dead_code)] pub async fn read_sse( response: reqwest::Response, ) -> anyhow::Result<()> { diff --git a/src/lib.rs b/src/lib.rs index c0a4bc0..d4ba0b2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,13 +1,10 @@ #[macro_use] extern crate rust_i18n; -use crate::pages::layout::App; use log::{error, info}; -use ratatui_kit::crossterm::event::{Event, KeyCode, KeyEventKind}; use ratatui_kit::prelude::RouterProvider; -use ratatui_kit::{AnyElement, ElementExt, Hooks, UseEvents, UseRouter, component, element}; +use ratatui_kit::{AnyElement, ElementExt, Hooks, component, element}; use rust_i18n::t; -use std::io; pub mod components; pub mod core; diff --git a/src/pages/layout.rs b/src/pages/layout.rs index b228ccd..94f1fc8 100644 --- a/src/pages/layout.rs +++ b/src/pages/layout.rs @@ -2,7 +2,7 @@ use crate::utils; use ratatui::{ layout::{Constraint, Direction, Layout, Rect}, - style::{Color, Modifier, Style}, + style::{Color, Style}, text::{Line, Span}, widgets::{Block, Borders, Paragraph}, Frame, DefaultTerminal, @@ -165,13 +165,4 @@ impl App { )); frame.render_widget(prompt, area); } -} - -// 标准主函数（极简、规范） -fn main() -> anyhow::Result<()> { - let mut terminal = ratatui::init(); - let mut app = App::default(); - let result = app.run(&mut terminal); - ratatui::restore(); - result } \ No newline at end of file diff --git a/src/pages/router.rs b/src/pages/router.rs index 713af07..bad5d09 100644 --- a/src/pages/router.rs +++ b/src/pages/router.rs @@ -14,9 +14,13 @@ pub fn RouterPage(mut hooks: Hooks) -> impl Into> { if let Event::Key(key_event) = event { if key_event.kind == KeyEventKind::Press { match key_event.code { - KeyCode::Char('1') => navigate.push("/welcome"), - KeyCode::Char('2') => navigate.push("/chat"), - KeyCode::Char('3') => navigate.push("/router"), + // 用 replace 而非 push：ratatui-kit 0.5.9 的 RouterHistory::push 在 + // 历史长度达到 max_length 后会 pop_front 却不同步回退 current， + // 导致 insert(current) 越界 panic（index out of bounds）。 + // 调试路由页只需切换当前页，replace 不增长历史栈，天然规避该缺陷。 + KeyCode::Char('1') => navigate.replace("/welcome"), + KeyCode::Char('2') => navigate.replace("/chat"), + KeyCode::Char('3') => navigate.replace("/router"), _ => {} } }