添加qiming-rcoder模块
This commit is contained in:
84
qiming-rcoder/crates/agent_abstraction/Cargo.toml
Normal file
84
qiming-rcoder/crates/agent_abstraction/Cargo.toml
Normal file
@@ -0,0 +1,84 @@
|
||||
[package]
|
||||
name = "agent_abstraction"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
authors = ["Your Name <your.email@example.com>"]
|
||||
license = "MIT OR Apache-2.0"
|
||||
description = "Agent abstraction layer for RCoder"
|
||||
publish = false
|
||||
|
||||
[features]
|
||||
default = []
|
||||
# Pingora 代理模式:API Key 和 Base URL 使用占位符,由 Pingora 代理注入真实值
|
||||
# Windows 不支持 Pingora,禁用此 feature 时直接使用真实的 API Key 和 Base URL
|
||||
proxy = []
|
||||
|
||||
[dependencies]
|
||||
# 依赖其他 crate
|
||||
agent_config = { path = "../agent_config" }
|
||||
shared_types = { path = "../shared_types" }
|
||||
|
||||
# 异步运行时
|
||||
tokio = { workspace = true }
|
||||
|
||||
# 异步 trait
|
||||
async-trait = { workspace = true }
|
||||
|
||||
# 错误处理
|
||||
thiserror = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
|
||||
# 并发集合
|
||||
dashmap = { workspace = true }
|
||||
|
||||
# 缓存
|
||||
moka = { workspace = true }
|
||||
|
||||
# 目录操作
|
||||
dirs = { workspace = true }
|
||||
|
||||
# 文件系统监听
|
||||
notify = { workspace = true }
|
||||
|
||||
# 日期时间
|
||||
chrono = { workspace = true, features = ["serde"] }
|
||||
|
||||
# UUID 生成
|
||||
uuid = { workspace = true, features = ["v4"] }
|
||||
|
||||
# 序列化
|
||||
serde_json = { workspace = true }
|
||||
|
||||
# 日志
|
||||
tracing = { workspace = true }
|
||||
|
||||
# 兼容性
|
||||
tokio-util = { workspace = true }
|
||||
|
||||
# ACP 协议 - 官方 SDK(完全 Send-safe,无需 LocalSet)
|
||||
agent-client-protocol = { workspace = true }
|
||||
|
||||
# 跨平台命令查找
|
||||
which = { workspace = true }
|
||||
|
||||
# MCP 协议 - rmcp 官方库: https://github.com/modelcontextprotocol/rust-sdk
|
||||
rmcp = { workspace = true }
|
||||
|
||||
# 序列化
|
||||
serde = { workspace = true }
|
||||
|
||||
# 进程组管理(用于清理整个进程树)
|
||||
# Unix: process-group (setpgid), Windows: job-object
|
||||
# process-wrap 内部通过 #[cfg] 控制各平台类型的可用性
|
||||
process-wrap = { version = "9.0", features = ["tokio1", "process-group", "job-object"] }
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
# Unix 信号处理
|
||||
nix = { workspace = true, features = ["signal"] }
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
# process-wrap CreationFlags 需要 PROCESS_CREATION_FLAGS 类型
|
||||
windows = { version = "0.62", features = ["Win32_System_Threading"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3.8"
|
||||
191
qiming-rcoder/crates/agent_abstraction/src/acp/connection.rs
Normal file
191
qiming-rcoder/crates/agent_abstraction/src/acp/connection.rs
Normal file
@@ -0,0 +1,191 @@
|
||||
//! Agent connection wrapper for ACP protocol.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use agent_client_protocol::schema::{CancelNotification, PromptRequest, SessionId};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
// 重新导出 shared_types 中的统一类型
|
||||
pub use shared_types::{CancelNotificationRequestWrapper, CancelResult};
|
||||
|
||||
/// Agent connection wrapper
|
||||
#[derive(Debug)]
|
||||
pub struct AgentConnection {
|
||||
/// Project ID (业务主键)
|
||||
/// project_id 是业务的唯一标识符,用于关联项目和对应的 agent 实例
|
||||
/// 可以根据 project_id 找到对应的 session_id
|
||||
pub project_id: String,
|
||||
|
||||
/// Service Type (服务类型)
|
||||
/// 用于区分不同类型的服务,对应不同的 Docker 镜像和运行环境
|
||||
/// 当前所有业务都使用 ServiceType::RCoder
|
||||
pub service_type: shared_types::ServiceType,
|
||||
|
||||
/// Session ID (可空)
|
||||
/// session_id 是 ACP 协议中的会话标识符,在 agent 的 newSession 成功后由 agent 返回
|
||||
/// 创建 agent 时可能没有 session_id,需要在会话建立后更新
|
||||
pub session_id: Option<SessionId>,
|
||||
|
||||
/// Prompt sender channel
|
||||
pub prompt_tx: Arc<mpsc::UnboundedSender<PromptRequest>>,
|
||||
/// Cancel sender channel - wrapped in Arc to avoid Debug requirement
|
||||
pub cancel_tx: Arc<mpsc::UnboundedSender<CancelNotificationRequestWrapper>>,
|
||||
}
|
||||
|
||||
impl AgentConnection {
|
||||
/// Create a new agent connection
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `project_id` - 业务主键,唯一标识项目
|
||||
/// * `service_type` - 服务类型,对应不同的 Docker 镜像和运行环境
|
||||
/// * `session_id` - ACP 会话 ID(可空,创建时可能没有)
|
||||
/// * `prompt_tx` - 提示消息发送通道
|
||||
/// * `cancel_tx` - 取消消息发送通道
|
||||
pub fn new(
|
||||
project_id: String,
|
||||
service_type: shared_types::ServiceType,
|
||||
session_id: Option<SessionId>,
|
||||
prompt_tx: Arc<mpsc::UnboundedSender<PromptRequest>>,
|
||||
cancel_tx: Arc<mpsc::UnboundedSender<CancelNotificationRequestWrapper>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
project_id,
|
||||
service_type,
|
||||
session_id,
|
||||
prompt_tx,
|
||||
cancel_tx,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get project ID (业务主键)
|
||||
pub fn project_id(&self) -> &str {
|
||||
&self.project_id
|
||||
}
|
||||
|
||||
/// Get session ID (可空)
|
||||
pub fn session_id(&self) -> Option<&SessionId> {
|
||||
self.session_id.as_ref()
|
||||
}
|
||||
|
||||
/// Update session ID after successful newSession
|
||||
/// 在 agent 的 newSession 成功后调用,更新 session_id
|
||||
pub fn set_session_id(&mut self, session_id: SessionId) {
|
||||
self.session_id = Some(session_id);
|
||||
}
|
||||
|
||||
/// Check if session is ready (has session_id)
|
||||
pub fn has_session(&self) -> bool {
|
||||
self.session_id.is_some()
|
||||
}
|
||||
|
||||
/// Send prompt
|
||||
///
|
||||
/// 通过 channel 发送 prompt 到 LocalSet 中运行的 agent
|
||||
/// session_id 由服务端自动管理,调用方无需关心
|
||||
///
|
||||
/// 设计说明:
|
||||
/// - SACP 版本支持 Send trait,可以在 tokio::spawn 中运行
|
||||
/// - 使用 MPSC channel 解耦调用方和 agent 运行环境
|
||||
/// - newSession 在 agent 启动后自动执行
|
||||
/// - prompt handler 会自动处理 session_id 的覆盖
|
||||
pub async fn send_prompt(
|
||||
&self,
|
||||
prompt: PromptRequest,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
self.prompt_tx.send(prompt).map_err(|e| {
|
||||
Box::new(std::io::Error::new(std::io::ErrorKind::Other, e))
|
||||
as Box<dyn std::error::Error + Send + Sync>
|
||||
})
|
||||
}
|
||||
|
||||
/// Send cancel
|
||||
///
|
||||
/// 通过 channel 发送取消请求到 LocalSet 中运行的 agent
|
||||
/// 返回一个 receiver,调用方可以通过它等待取消结果
|
||||
///
|
||||
/// # Returns
|
||||
/// - `Ok(receiver)`: 请求已发送,通过 receiver.await 获取取消结果
|
||||
/// - `Err`: 发送请求失败(channel 已关闭)
|
||||
pub async fn send_cancel(
|
||||
&self,
|
||||
cancel_notification: CancelNotification,
|
||||
) -> Result<oneshot::Receiver<CancelResult>, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let (result_tx, result_rx) = oneshot::channel();
|
||||
|
||||
self.cancel_tx
|
||||
.send(CancelNotificationRequestWrapper {
|
||||
cancel_notification,
|
||||
result_tx,
|
||||
})
|
||||
.map_err(|e| {
|
||||
Box::new(std::io::Error::new(std::io::ErrorKind::Other, e))
|
||||
as Box<dyn std::error::Error + Send + Sync>
|
||||
})?;
|
||||
|
||||
Ok(result_rx)
|
||||
}
|
||||
|
||||
/// Send cancel and wait for result
|
||||
///
|
||||
/// 发送取消请求并等待结果
|
||||
pub async fn send_cancel_and_wait(
|
||||
&self,
|
||||
cancel_notification: CancelNotification,
|
||||
) -> Result<CancelResult, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let result_rx = self.send_cancel(cancel_notification).await?;
|
||||
|
||||
result_rx.await.map_err(|e| {
|
||||
Box::new(std::io::Error::new(
|
||||
std::io::ErrorKind::Other,
|
||||
format!("Failed to receive cancel result: {}", e),
|
||||
)) as Box<dyn std::error::Error + Send + Sync>
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Agent status type (re-exported from shared_types)
|
||||
pub type AgentStatus = shared_types::AgentStatus;
|
||||
|
||||
/// Connection status enum
|
||||
#[repr(u8)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum ConnectionStatus {
|
||||
/// Connection is being established
|
||||
Connecting = 1,
|
||||
/// Connection is active
|
||||
Connected = 2,
|
||||
/// Connection is idle
|
||||
Idle = 3,
|
||||
/// Connection has an error
|
||||
Error = 4,
|
||||
/// Connection is closed
|
||||
Closed = 5,
|
||||
}
|
||||
|
||||
impl ConnectionStatus {
|
||||
/// Convert from u8
|
||||
pub fn from_u8(value: u8) -> Self {
|
||||
match value {
|
||||
1 => ConnectionStatus::Connecting,
|
||||
2 => ConnectionStatus::Connected,
|
||||
3 => ConnectionStatus::Idle,
|
||||
4 => ConnectionStatus::Error,
|
||||
5 => ConnectionStatus::Closed,
|
||||
_ => ConnectionStatus::Closed,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert to u8
|
||||
pub fn to_u8(self) -> u8 {
|
||||
self as u8
|
||||
}
|
||||
}
|
||||
|
||||
/// Connection statistics
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ConnectionStats {
|
||||
pub total_connections: u64,
|
||||
pub active_connections: u64,
|
||||
pub idle_connections: u64,
|
||||
pub error_connections: u64,
|
||||
}
|
||||
20
qiming-rcoder/crates/agent_abstraction/src/acp/mod.rs
Normal file
20
qiming-rcoder/crates/agent_abstraction/src/acp/mod.rs
Normal file
@@ -0,0 +1,20 @@
|
||||
//! ACP connection management module.
|
||||
|
||||
mod connection;
|
||||
|
||||
pub use connection::{
|
||||
AgentConnection, AgentStatus, CancelNotificationRequestWrapper, CancelResult, ConnectionStats,
|
||||
ConnectionStatus,
|
||||
};
|
||||
|
||||
/// Legacy type alias for backward compatibility
|
||||
pub type Connection = AgentConnection;
|
||||
|
||||
/// Placeholder error type
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum AcpError {
|
||||
#[error("Connection error: {0}")]
|
||||
Connection(String),
|
||||
#[error("Other error: {0}")]
|
||||
Other(String),
|
||||
}
|
||||
31
qiming-rcoder/crates/agent_abstraction/src/error.rs
Normal file
31
qiming-rcoder/crates/agent_abstraction/src/error.rs
Normal file
@@ -0,0 +1,31 @@
|
||||
//! Agent abstraction layer error types
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
/// Agent abstraction layer error
|
||||
#[derive(Error, Debug)]
|
||||
pub enum AgentAbstractionError {
|
||||
#[error("connection error: {0}")]
|
||||
Connection(String),
|
||||
|
||||
#[error("registry error: {0}")]
|
||||
Registry(String),
|
||||
|
||||
#[error("process error: {0}")]
|
||||
Process(String),
|
||||
|
||||
#[error("other error: {0}")]
|
||||
Other(String),
|
||||
|
||||
#[error("serialization error: {0}")]
|
||||
Serialization(#[from] serde_json::Error),
|
||||
|
||||
#[error("I/O error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
|
||||
#[error("type cast error")]
|
||||
Cast,
|
||||
|
||||
#[error("not found: {0}")]
|
||||
NotFound(String),
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
618
qiming-rcoder/crates/agent_abstraction/src/launcher/lifecycle.rs
Normal file
618
qiming-rcoder/crates/agent_abstraction/src/launcher/lifecycle.rs
Normal file
@@ -0,0 +1,618 @@
|
||||
//! Agent生命周期管理
|
||||
//!
|
||||
//! 基于RAII原则的简洁生命周期管理设计
|
||||
//!
|
||||
//! ## 僵尸进程问题解决方案
|
||||
//!
|
||||
//! 核心问题:Drop trait 是同步的,无法 await child.wait()
|
||||
//!
|
||||
//! 解决方案:
|
||||
//! 1. **后台回收任务**:立即启动后台任务 wait() 子进程
|
||||
//! 2. **进程组终止**:使用 nix::kill 发送信号到进程组
|
||||
//! 3. **三重保障**:PID 1 的 process_reaper 模块兜底
|
||||
//!
|
||||
//! ## 进程组说明
|
||||
//!
|
||||
//! 使用 `process-wrap` crate 创建真正的进程组:
|
||||
//! - 启动时使用 `ProcessGroup::leader()` 创建进程组
|
||||
//! - 终止时发送 `kill(-pgid, SIGKILL)` 到整个进程组
|
||||
//! - 能够正确清理子进程及其所有孙进程
|
||||
|
||||
use anyhow::Result;
|
||||
use dashmap::DashMap;
|
||||
use process_wrap::tokio::ChildWrapper;
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
};
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use agent_client_protocol::schema::SessionId;
|
||||
use shared_types::{AgentLifecycle, ModelProviderConfig};
|
||||
|
||||
/// Agent生命周期守卫
|
||||
///
|
||||
/// 遵循RAII原则,当守卫被drop时自动清理agent资源
|
||||
///
|
||||
/// ## 僵尸进程避免机制
|
||||
///
|
||||
/// 1. **后台回收任务**:构造时立即启动 tokio::spawn 等待子进程
|
||||
/// 2. **进程组终止**:Drop 时发送信号到进程组(使用 nix::kill)
|
||||
/// 3. **PID 1 兜底**:process_reaper 模块自动回收所有孤儿进程
|
||||
///
|
||||
/// ## 进程组信号
|
||||
///
|
||||
/// 在 Unix 上,使用负的进程组 ID 发送信号:
|
||||
/// - `kill(-pgid, SIGKILL)` 杀死整个进程组
|
||||
/// - 这会终止子进程及其所有后代(如果子进程创建了真正的进程组)
|
||||
pub struct AgentLifecycleGuard {
|
||||
inner: Arc<AgentLifecycleInner>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AgentLifecycleGuard {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("AgentLifecycleGuard")
|
||||
.field("project_id", &self.inner.project_id)
|
||||
.field("session_id", &self.inner.session_id)
|
||||
.field("pgid", &self.inner.pgid)
|
||||
.field("stopped", &self.inner.stopped.load(Ordering::SeqCst))
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
struct AgentLifecycleInner {
|
||||
project_id: String,
|
||||
session_id: SessionId,
|
||||
/// 🔥 进程组 ID(当前实现:使用 child.pid 作为伪进程组)
|
||||
///
|
||||
/// 注意:当前实现使用子进程的 PID 作为 PGID。
|
||||
/// - 如果子进程通过 setsid() 创建了真正的进程组,kill(-pgid) 会杀死整个进程树
|
||||
/// - 如果子进程没有创建进程组,kill(-pgid) 只会杀死子进程本身
|
||||
/// - 未来可以使用 process-wrap 库创建真正的进程组
|
||||
pgid: u32,
|
||||
cancel_token: CancellationToken,
|
||||
resources: AgentResources,
|
||||
stopped: AtomicBool,
|
||||
/// 🔥 共享的 API 密钥管理器引用(用于自动清理)
|
||||
shared_api_key_manager: Option<Arc<DashMap<String, ModelProviderConfig>>>,
|
||||
/// 🔥 project_id -> service_uuid 映射(用于清理时查找 UUID)
|
||||
project_uuid_map: Option<Arc<DashMap<String, String>>>,
|
||||
/// 🔥 关联的 service_uuid(用于清理时定位配置)
|
||||
service_uuid: Option<String>,
|
||||
}
|
||||
|
||||
/// Agent资源管理枚举
|
||||
///
|
||||
/// ## 后台回收版本
|
||||
///
|
||||
/// 存储后台任务句柄,确保子进程被 wait() 回收
|
||||
enum AgentResources {
|
||||
Claude {
|
||||
/// stderr 任务句柄
|
||||
stderr_task: Arc<Mutex<Option<JoinHandle<()>>>>,
|
||||
/// 后台回收任务(已启动,会 wait() 子进程)
|
||||
_reaper_task: JoinHandle<()>,
|
||||
},
|
||||
}
|
||||
|
||||
impl AgentLifecycleGuard {
|
||||
/// 为Claude Agent创建生命周期守卫(兼容旧代码,默认无密钥管理器)
|
||||
///
|
||||
/// # 参数
|
||||
///
|
||||
/// * `child_process` - 已启动的子进程(必须是进程组组长)
|
||||
/// * `stderr_task` - stderr 读取任务
|
||||
/// * `cancel_token` - 取消令牌
|
||||
///
|
||||
/// # 僵尸进程避免
|
||||
///
|
||||
/// 此函数会立即启动后台任务等待子进程,确保子进程退出时被回收。
|
||||
pub fn new_claude(
|
||||
project_id: String,
|
||||
session_id: SessionId,
|
||||
child_process: Box<dyn ChildWrapper>,
|
||||
stderr_task: JoinHandle<()>,
|
||||
cancel_token: CancellationToken,
|
||||
) -> Self {
|
||||
Self::new_claude_with_key_manager(
|
||||
project_id,
|
||||
session_id,
|
||||
child_process,
|
||||
stderr_task,
|
||||
cancel_token,
|
||||
None, // 默认无密钥管理器
|
||||
None, // 默认无 project_uuid_map
|
||||
None, // 默认无 service_uuid
|
||||
)
|
||||
}
|
||||
|
||||
/// 🔥 新增:带异常退出标志的构造函数
|
||||
///
|
||||
/// 创建生命周期守卫时传入共享的 `abnormal_exit_flag`,当子进程异常退出时设置此标志。
|
||||
/// 这使得 SACP 连接层可以检测到异常退出并发送相应的通知。
|
||||
///
|
||||
/// # 参数
|
||||
///
|
||||
/// * `abnormal_exit_flag` - 共享的原子布尔标志,子进程异常退出时设置为 true
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// 如果子进程 PID 无效(为 0 或 None),此函数会 panic,因为这意味着进程启动失败。
|
||||
pub fn new_claude_with_abnormal_flag(
|
||||
project_id: String,
|
||||
session_id: SessionId,
|
||||
child_process: Box<dyn ChildWrapper>,
|
||||
stderr_task: JoinHandle<()>,
|
||||
cancel_token: CancellationToken,
|
||||
abnormal_exit_flag: Arc<AtomicBool>,
|
||||
) -> Self {
|
||||
Self::new_claude_with_key_manager_and_abnormal_flag(
|
||||
project_id,
|
||||
session_id,
|
||||
child_process,
|
||||
stderr_task,
|
||||
cancel_token,
|
||||
None, // 默认无密钥管理器
|
||||
None, // 默认无 project_uuid_map
|
||||
None, // 默认无 service_uuid
|
||||
Some(abnormal_exit_flag),
|
||||
)
|
||||
}
|
||||
|
||||
/// 🔥 新增:带密钥管理器和异常退出标志的构造函数
|
||||
///
|
||||
/// 创建生命周期守卫时传入共享的 API 密钥管理器和 service_uuid,
|
||||
/// 当 Agent 停止时(Drop)会自动清理对应的 API 密钥配置。
|
||||
///
|
||||
/// # 进程组管理
|
||||
///
|
||||
/// 当前实现使用子进程的 PID 作为 PGID:
|
||||
/// - `pgid = child_pid`(使用子进程 PID 作为进程组 ID)
|
||||
/// - 终止时发送 `kill(-pgid, SIGKILL)` 到进程组
|
||||
/// - 使用 `process-wrap` 创建真正的进程组,能正确清理所有孙进程
|
||||
///
|
||||
/// # 参数
|
||||
///
|
||||
/// * `shared_api_key_manager` - 共享的 DashMap,用于清理 API 密钥配置
|
||||
/// * `project_uuid_map` - project_id -> service_uuid 映射,用于查找 UUID
|
||||
/// * `service_uuid` - 与此 Agent 关联的 service UUID
|
||||
/// * `abnormal_exit_flag` - 共享的原子布尔标志,子进程异常退出时设置为 true
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// 如果子进程 PID 无效(为 0 或 None),此函数会 panic,因为这意味着进程启动失败。
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new_claude_with_key_manager(
|
||||
project_id: String,
|
||||
session_id: SessionId,
|
||||
child_process: Box<dyn ChildWrapper>,
|
||||
stderr_task: JoinHandle<()>,
|
||||
cancel_token: CancellationToken,
|
||||
shared_api_key_manager: Option<Arc<DashMap<String, ModelProviderConfig>>>,
|
||||
project_uuid_map: Option<Arc<DashMap<String, String>>>,
|
||||
service_uuid: Option<String>,
|
||||
) -> Self {
|
||||
Self::new_claude_with_key_manager_and_abnormal_flag(
|
||||
project_id,
|
||||
session_id,
|
||||
child_process,
|
||||
stderr_task,
|
||||
cancel_token,
|
||||
shared_api_key_manager,
|
||||
project_uuid_map,
|
||||
service_uuid,
|
||||
None, // 默认无异常退出标志
|
||||
)
|
||||
}
|
||||
|
||||
/// 🔥 完整构造函数:带密钥管理器和异常退出标志
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn new_claude_with_key_manager_and_abnormal_flag(
|
||||
project_id: String,
|
||||
session_id: SessionId,
|
||||
mut child_process: Box<dyn ChildWrapper>,
|
||||
stderr_task: JoinHandle<()>,
|
||||
cancel_token: CancellationToken,
|
||||
shared_api_key_manager: Option<Arc<DashMap<String, ModelProviderConfig>>>,
|
||||
project_uuid_map: Option<Arc<DashMap<String, String>>>,
|
||||
service_uuid: Option<String>,
|
||||
abnormal_exit_flag: Option<Arc<AtomicBool>>,
|
||||
) -> Self {
|
||||
// 🔥 关键:PID 有效性检查
|
||||
// ChildWrapper 的 id() 返回 Option<u32>,当进程已终止或无效时返回 None
|
||||
// 如果 PID 无效,这是一个严重的初始化错误,应该 panic
|
||||
let pid = child_process.id().unwrap_or_else(|| {
|
||||
panic!(
|
||||
"[LifecycleGuard] 子进程 PID 无效(None),进程可能已终止: project_id={}",
|
||||
project_id
|
||||
)
|
||||
});
|
||||
|
||||
// 🔥 额外检查:PID 不应该为 0
|
||||
// 虽然 id() 返回 Some(0) 理论上可能,但实际上 PID 0 是内核保留的
|
||||
if pid == 0 {
|
||||
panic!(
|
||||
"[LifecycleGuard] 子进程 PID 为 0,这是无效的 PID: project_id={}",
|
||||
project_id
|
||||
);
|
||||
}
|
||||
|
||||
// 🔥 进程组 ID 等于组长进程的 PID
|
||||
// process-wrap 的 ProcessGroup 使用 setpgid(0, 0) 创建新进程组,使进程成为组长
|
||||
let pgid = pid;
|
||||
let project_id_clone = project_id.clone();
|
||||
let session_id_str = session_id.0.to_string();
|
||||
|
||||
// 🔥 关键:立即启动后台回收任务
|
||||
// 这个任务会等待子进程退出,确保不会产生僵尸进程
|
||||
// 当子进程退出时,设置 abnormal_exit_flag 并触发 cancel_token
|
||||
// 让 SACP 连接层检测到并发送 SSE 通知
|
||||
let cancel_token_for_reaper = cancel_token.clone();
|
||||
let abnormal_exit_flag_clone = abnormal_exit_flag.clone();
|
||||
let project_id_for_reaper = project_id.clone();
|
||||
let reaper_task = tokio::spawn(async move {
|
||||
info!(
|
||||
"[ProcessReaper] 开始监控 Agent 进程: project_id={}, pid={}, pgid={}",
|
||||
project_id_for_reaper, pid, pgid
|
||||
);
|
||||
|
||||
// 🔥 优先等待子进程退出,而不是响应取消信号
|
||||
// 这确保了即使收到取消信号,也能正确检测进程是否已退出
|
||||
let wait_result = child_process.wait().await;
|
||||
|
||||
// 检查是否是外部取消(用户主动 stop)
|
||||
let was_cancelled = cancel_token_for_reaper.is_cancelled();
|
||||
|
||||
match wait_result {
|
||||
Ok(status) => {
|
||||
// 获取详细的退出信息
|
||||
let exit_code = status.code();
|
||||
#[cfg(unix)]
|
||||
let signal = {
|
||||
use std::os::unix::process::ExitStatusExt;
|
||||
status.signal()
|
||||
};
|
||||
#[cfg(not(unix))]
|
||||
let signal: Option<i32> = None;
|
||||
|
||||
if !status.success() {
|
||||
// 🔥 非零退出码或被信号杀死 = 异常退出
|
||||
if let Some(ref flag) = abnormal_exit_flag_clone {
|
||||
// 只有非用户主动取消时才标记为异常
|
||||
if !was_cancelled {
|
||||
flag.store(true, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
warn!(
|
||||
"[ProcessReaper] Agent 进程异常退出: project_id={}, pid={}, pgid={}, exit_code={:?}, signal={:?}, was_cancelled={}",
|
||||
project_id_for_reaper, pid, pgid, exit_code, signal, was_cancelled
|
||||
);
|
||||
} else {
|
||||
info!(
|
||||
"[ProcessReaper] Agent 进程正常退出: project_id={}, pid={}, pgid={}, exit_code={:?}",
|
||||
project_id_for_reaper, pid, pgid, exit_code
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
// wait 失败,可能是进程已被其他方式回收
|
||||
if let Some(ref flag) = abnormal_exit_flag_clone {
|
||||
if !was_cancelled {
|
||||
flag.store(true, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
warn!(
|
||||
"[ProcessReaper] Agent 进程 wait() 失败: project_id={}, pid={}, pgid={}, error={}, was_cancelled={}",
|
||||
project_id_for_reaper, pid, pgid, e, was_cancelled
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// 🔥 关键:触发 cancel_token,通知 SACP 连接层进程已退出
|
||||
// 这会让 SACP 连接检测到并发送 SSE 错误通知,然后断开连接
|
||||
if !was_cancelled {
|
||||
info!(
|
||||
"[ProcessReaper] 触发 cancel_token,通知 SACP 连接断开: project_id={}, pid={}",
|
||||
project_id_for_reaper, pid
|
||||
);
|
||||
cancel_token_for_reaper.cancel();
|
||||
} else {
|
||||
debug!(
|
||||
"[ProcessReaper] cancel_token 已被外部取消,跳过: project_id={}, pid={}",
|
||||
project_id_for_reaper, pid
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
let resources = AgentResources::Claude {
|
||||
stderr_task: Arc::new(Mutex::new(Some(stderr_task))),
|
||||
_reaper_task: reaper_task,
|
||||
};
|
||||
|
||||
let inner = Arc::new(AgentLifecycleInner {
|
||||
project_id: project_id_clone,
|
||||
session_id,
|
||||
pgid,
|
||||
cancel_token,
|
||||
resources,
|
||||
stopped: AtomicBool::new(false),
|
||||
shared_api_key_manager,
|
||||
project_uuid_map,
|
||||
service_uuid,
|
||||
});
|
||||
|
||||
info!(
|
||||
"[LifecycleGuard] 创建 Claude Agent 守卫: project_id={}, pgid={}, session_id={}",
|
||||
project_id, pgid, session_id_str
|
||||
);
|
||||
|
||||
Self { inner }
|
||||
}
|
||||
|
||||
/// 优雅停止agent
|
||||
///
|
||||
/// 带超时机制(5秒),超时后强制 kill 进程组
|
||||
///
|
||||
/// ## 进程组终止
|
||||
///
|
||||
/// 使用 `process-wrap` 创建真正的进程组,发送信号到 `-pgid` 会终止:
|
||||
/// - 子进程(进程组组长)
|
||||
/// - 所有孙进程(同一进程组中的进程)
|
||||
pub async fn graceful_stop(&self) -> Result<()> {
|
||||
// 🔥 使用原子 CAS 操作确保只执行一次清理
|
||||
// compare_exchange 返回 Ok 表示成功将 false 改为 true,即当前线程获得清理权
|
||||
// 返回 Err 表示已经被其他地方清理(Drop 或其他 graceful_stop 调用)
|
||||
let should_cleanup = self
|
||||
.inner
|
||||
.stopped
|
||||
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
|
||||
.is_ok();
|
||||
|
||||
if !should_cleanup {
|
||||
debug!("Agent already stopped, skipping graceful stop");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
info!(
|
||||
"Gracefully stopping Claude agent for project: {}, pgid={}",
|
||||
self.inner.project_id, self.inner.pgid
|
||||
);
|
||||
|
||||
// 1. 发送取消信号
|
||||
self.inner.cancel_token.cancel();
|
||||
|
||||
// 2. 终止进程组
|
||||
self.kill_process_group(false).await?;
|
||||
|
||||
info!(
|
||||
"Gracefully stopped Claude agent for project: {}",
|
||||
self.inner.project_id
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 发送取消信号(非阻塞)
|
||||
pub fn cancel(&self) {
|
||||
debug!("Sending cancel signal to agent: {}", self.inner.project_id);
|
||||
self.inner.cancel_token.cancel();
|
||||
}
|
||||
|
||||
/// 检查是否已停止
|
||||
pub fn is_stopped(&self) -> bool {
|
||||
self.inner.stopped.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
/// 获取取消令牌
|
||||
pub fn cancellation_token(&self) -> &CancellationToken {
|
||||
&self.inner.cancel_token
|
||||
}
|
||||
|
||||
/// 🔥 终止进程组
|
||||
///
|
||||
/// 向 `-pgid` 发送信号,杀死整个进程组
|
||||
///
|
||||
/// # Unix 信号语义
|
||||
///
|
||||
/// - `kill(pgid, SIGTERM)` - 发送给单个进程
|
||||
/// - `kill(-pgid, SIGTERM)` - 发送给整个进程组
|
||||
/// - `kill(0, SIGTERM)` - 发送给调用者自己的进程组(危险!)
|
||||
///
|
||||
/// # 参数
|
||||
///
|
||||
/// * `force` - 是否强制使用 SIGKILL(否则使用 SIGTERM)
|
||||
async fn kill_process_group(&self, force: bool) -> Result<()> {
|
||||
let pgid = self.inner.pgid;
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use nix::errno::Errno;
|
||||
use nix::sys::signal::{Signal, kill};
|
||||
use nix::unistd::Pid;
|
||||
|
||||
// 🔥 关键防御性检查:pgid 不能为 0
|
||||
// kill(0, SIGKILL) 会杀死调用者自己的进程组,这是危险的
|
||||
if pgid == 0 {
|
||||
warn!(
|
||||
"[LifecycleGuard] 进程组 ID 为 0,跳过进程组终止(可能是初始化失败): project_id={}",
|
||||
self.inner.project_id
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// 🔥 关键:使用负的进程组 ID(真实的进程组 ID)
|
||||
// -pgid 表示发送信号到整个进程组,而不仅仅是进程组组长
|
||||
let target = Pid::from_raw(-(pgid as i32));
|
||||
|
||||
let signal = if force {
|
||||
Signal::SIGKILL
|
||||
} else {
|
||||
Signal::SIGTERM
|
||||
};
|
||||
|
||||
match kill(target, signal) {
|
||||
Ok(_) => {
|
||||
debug!("already sent signal: pgid={}, signal={:?}", pgid, signal);
|
||||
|
||||
// 如果是 SIGTERM,等待一段时间让进程优雅退出
|
||||
if !force {
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
|
||||
|
||||
// 强制杀死进程组
|
||||
let _ = kill(target, Signal::SIGKILL);
|
||||
debug!("already force killed: pgid={}", pgid);
|
||||
}
|
||||
}
|
||||
Err(Errno::ESRCH) => {
|
||||
// 进程组已退出,这是正常的
|
||||
debug!("process group already exited: pgid={}", pgid);
|
||||
}
|
||||
Err(Errno::EPERM) => {
|
||||
// 权限不足,无法终止进程组
|
||||
warn!(
|
||||
"[LifecycleGuard] 权限不足,无法终止进程组: pgid={}, project_id={}",
|
||||
pgid, self.inner.project_id
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
// 其他错误(如 EINVAL、EFAULT 等)
|
||||
debug!(" kill failed: pgid={}, error={:?}", pgid, e);
|
||||
}
|
||||
}
|
||||
|
||||
info!("Claude process group stopped: pgid={}", pgid);
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
debug!("Unix platform, skipping process group stop");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for AgentLifecycleGuard {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
inner: Arc::clone(&self.inner),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for AgentLifecycleGuard {
|
||||
fn drop(&mut self) {
|
||||
let strong_count = Arc::strong_count(&self.inner);
|
||||
|
||||
debug!(
|
||||
"[Claude] AgentLifecycleGuard::drop 开始: project_id={}, pgid={}, strong_count={}",
|
||||
self.inner.project_id, self.inner.pgid, strong_count
|
||||
);
|
||||
|
||||
// 🔥 使用原子 CAS 操作确保只执行一次清理
|
||||
// 不再依赖引用计数,因为引用计数可能因为多处 clone 而不准确
|
||||
// compare_exchange 返回 Ok 表示成功将 false 改为 true,即当前线程获得清理权
|
||||
// 返回 Err 表示已经被其他线程清理,当前线程无需操作
|
||||
let should_cleanup = self
|
||||
.inner
|
||||
.stopped
|
||||
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
|
||||
.is_ok();
|
||||
|
||||
if should_cleanup {
|
||||
debug!(
|
||||
"[Claude] AgentLifecycleGuard 获得清理权,开始清理资源: {}",
|
||||
self.inner.project_id
|
||||
);
|
||||
|
||||
// 发送取消信号
|
||||
self.inner.cancel_token.cancel();
|
||||
|
||||
// 注意:API 密钥配置的清理由 agent_runner 层的 stop_agent 方法统一负责
|
||||
// 包括:
|
||||
// - shared_api_key_manager 中的配置
|
||||
// - project_uuid_map 中的映射
|
||||
//
|
||||
// 这样避免双重清理,确保资源只被清理一次
|
||||
|
||||
// 🔥 同步终止进程组
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use nix::sys::signal::{Signal, kill};
|
||||
use nix::unistd::Pid;
|
||||
|
||||
let pgid = self.inner.pgid;
|
||||
|
||||
// 🔥 关键防御性检查:pgid 不能为 0
|
||||
// kill(0, SIGKILL) 会杀死调用者自己的进程组,这是危险的
|
||||
if pgid == 0 {
|
||||
debug!(
|
||||
"[Claude] 进程组 ID 为 0,跳过进程组终止: project_id={}",
|
||||
self.inner.project_id
|
||||
);
|
||||
} else {
|
||||
let target = Pid::from_raw(-(pgid as i32));
|
||||
|
||||
if let Err(e) = kill(target, Signal::SIGKILL) {
|
||||
// 进程可能已经退出,这是正常的
|
||||
debug!(
|
||||
"[Claude] 终止进程组失败(可能已退出): pgid={}, error={}",
|
||||
pgid, e
|
||||
);
|
||||
} else {
|
||||
info!(
|
||||
"[Claude] 进程组已终止: pgid={}, project_id={}",
|
||||
pgid, self.inner.project_id
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
debug!("[Claude] Unix platform, skipping process group stop");
|
||||
}
|
||||
|
||||
// 注意:后台回收任务 (reaper_task) 会自动完成
|
||||
// 不需要在这里等待或取消
|
||||
|
||||
info!(
|
||||
"[Claude] AgentLifecycleGuard 清理完成: project_id={}",
|
||||
self.inner.project_id
|
||||
);
|
||||
} else {
|
||||
debug!(
|
||||
"[Claude] AgentLifecycleGuard 跳过清理(已被其他引用清理): project_id={}",
|
||||
self.inner.project_id
|
||||
);
|
||||
}
|
||||
|
||||
debug!(
|
||||
"[Claude] AgentLifecycleGuard::drop 完成: project_id={}",
|
||||
self.inner.project_id
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// 为AgentLifecycleGuard实现AgentLifecycle trait
|
||||
impl AgentLifecycle for AgentLifecycleGuard {
|
||||
fn graceful_stop(
|
||||
&self,
|
||||
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
|
||||
Box::pin(async move { AgentLifecycleGuard::graceful_stop(self).await })
|
||||
}
|
||||
|
||||
fn cancel(&self) {
|
||||
AgentLifecycleGuard::cancel(self);
|
||||
}
|
||||
|
||||
fn is_stopped(&self) -> bool {
|
||||
AgentLifecycleGuard::is_stopped(self)
|
||||
}
|
||||
|
||||
fn cancellation_token(&self) -> &CancellationToken {
|
||||
AgentLifecycleGuard::cancellation_token(self)
|
||||
}
|
||||
}
|
||||
91
qiming-rcoder/crates/agent_abstraction/src/launcher/mod.rs
Normal file
91
qiming-rcoder/crates/agent_abstraction/src/launcher/mod.rs
Normal file
@@ -0,0 +1,91 @@
|
||||
//! # Agent 启动器模块
|
||||
//!
|
||||
//! 提供 Agent 的启动、生命周期管理和通道处理功能。
|
||||
//!
|
||||
//! ## SACP 迁移说明
|
||||
//!
|
||||
//! 本模块已迁移至 SACP (Symposium ACP) 实现,主要优势:
|
||||
//! - **支持 Send trait**: 可使用标准 `tokio::spawn`,无需 `LocalSet` + `spawn_local`
|
||||
//! - **Builder 模式**: 更清晰的连接构建和配置
|
||||
//! - **回调式消息处理**: 通过 `on_receive_notification` / `on_receive_request` 宏
|
||||
//!
|
||||
//! ## 子模块职责
|
||||
//!
|
||||
//! | 子模块 | 职责 | 关键类型 |
|
||||
//! |--------|------|---------|
|
||||
//! | [`lifecycle`] | Agent 生命周期守卫(RAII 资源管理)| `AgentLifecycleGuard` |
|
||||
//! | [`claude_code_sacp`] | SACP 版本的 Claude Code Agent 启动器 | `SacpClaudeCodeLauncher` |
|
||||
//!
|
||||
//! ## Agent 启动流程 (SACP)
|
||||
//!
|
||||
//! ```text
|
||||
//! AcpSessionManager.create_session()
|
||||
//! │
|
||||
//! │ 1. 创建 SacpClaudeCodeLauncher
|
||||
//! ▼
|
||||
//! SacpClaudeCodeLauncher.launch()
|
||||
//! │
|
||||
//! │ 2. 启动 Claude Code 子进程
|
||||
//! │ 3. 使用 SACP Builder 建立连接
|
||||
//! │ 4. 创建 AgentLifecycleGuard
|
||||
//! ▼
|
||||
//! SacpLauncherConnectionInfo
|
||||
//! │
|
||||
//! │ 5. 返回 session_id, prompt_tx, cancel_tx, lifecycle_guard
|
||||
//! ▼
|
||||
//! 存入 SessionRegistry
|
||||
//! ```
|
||||
//!
|
||||
//! ## 与 shared_types::AgentLifecycleGuard 的关系
|
||||
//!
|
||||
//! `AgentLifecycleGuard` 的核心实现定义在 `shared_types::model::agent_model`。
|
||||
//! 本模块的 `lifecycle.rs` 提供 Re-export,
|
||||
//! 方便从 `agent_abstraction::launcher` 统一导入。
|
||||
//!
|
||||
//! ## 生命周期管理 (RAII)
|
||||
//!
|
||||
//! 当 `AgentLifecycleGuard` 被 drop 时:
|
||||
//! 1. 发送取消信号 (`cancel_token.cancel()`)
|
||||
//! 2. 终止子进程 (`child.kill()`)
|
||||
//! 3. 停止 stderr 任务
|
||||
//!
|
||||
//! 这确保了 Agent 资源的正确清理,即使在异常情况下也不会泄漏。
|
||||
|
||||
mod claude_code_sacp;
|
||||
pub mod lifecycle;
|
||||
#[cfg(windows)]
|
||||
mod windows_launch;
|
||||
|
||||
// ============================================================================
|
||||
// SACP 类型导出(推荐使用)
|
||||
// ============================================================================
|
||||
|
||||
pub use lifecycle::AgentLifecycleGuard;
|
||||
|
||||
// 直接导出 SACP 类型
|
||||
pub use claude_code_sacp::{
|
||||
SacpAgentLaunchConfig, SacpClaudeCodeLauncher, SacpLauncherConnectionInfo,
|
||||
convert_context_servers_sacp, get_default_sacp_agent_config, load_sacp_agent_config,
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// 兼容性类型别名(向后兼容旧代码)
|
||||
// ============================================================================
|
||||
|
||||
/// 兼容性别名:ClaudeCodeLauncher -> SacpClaudeCodeLauncher
|
||||
pub type ClaudeCodeLauncher<N> = SacpClaudeCodeLauncher<N>;
|
||||
|
||||
/// 兼容性别名:LauncherConnectionInfoComplete -> SacpLauncherConnectionInfo
|
||||
pub type LauncherConnectionInfoComplete = SacpLauncherConnectionInfo;
|
||||
|
||||
/// 兼容性别名:AgentLaunchConfig -> SacpAgentLaunchConfig
|
||||
pub type AgentLaunchConfig = SacpAgentLaunchConfig;
|
||||
|
||||
/// 兼容性别名:load_agent_config -> load_sacp_agent_config
|
||||
pub use claude_code_sacp::load_sacp_agent_config as load_agent_config;
|
||||
|
||||
/// 兼容性别名:get_default_agent_config -> get_default_sacp_agent_config
|
||||
pub use claude_code_sacp::get_default_sacp_agent_config as get_default_agent_config;
|
||||
|
||||
/// 兼容性别名:convert_context_servers -> convert_context_servers_sacp
|
||||
pub use claude_code_sacp::convert_context_servers_sacp as convert_context_servers;
|
||||
@@ -0,0 +1,411 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
pub const CREATE_NO_WINDOW_FLAG: u32 = 0x0800_0000;
|
||||
|
||||
fn resolve_windows_node_exe() -> Option<PathBuf> {
|
||||
// 1. 显式指定的 node 可执行文件路径(最高优先级)
|
||||
if let Ok(path) = std::env::var("NUWAX_NODE_PATH") {
|
||||
let node = PathBuf::from(path);
|
||||
if node.exists() {
|
||||
info!(
|
||||
"[SACP] Windows node resolved via NUWAX_NODE_PATH: {}",
|
||||
node.display()
|
||||
);
|
||||
return Some(node);
|
||||
}
|
||||
}
|
||||
|
||||
// 1.5 Tauri 层设置的 NUWAX_NODE_EXE(sidecar node-runtime.exe)
|
||||
if let Ok(path) = std::env::var("NUWAX_NODE_EXE") {
|
||||
let node = PathBuf::from(&path);
|
||||
if node.exists() {
|
||||
info!(
|
||||
"[SACP] Windows node resolved via NUWAX_NODE_EXE: {}",
|
||||
node.display()
|
||||
);
|
||||
return Some(node);
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 应用集成的 runtime 路径(Tauri 打包后设置的 NUWAX_APP_RUNTIME_PATH,
|
||||
// 包含 runtime\node\bin 等目录,内含集成的 node.exe)
|
||||
if let Ok(runtime_path) = std::env::var("NUWAX_APP_RUNTIME_PATH") {
|
||||
for dir in runtime_path.split(';') {
|
||||
let dir = dir.trim();
|
||||
if dir.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let node = PathBuf::from(dir).join("node.exe");
|
||||
if node.exists() {
|
||||
info!(
|
||||
"[SACP] Windows node resolved via NUWAX_APP_RUNTIME_PATH: {}",
|
||||
node.display()
|
||||
);
|
||||
return Some(node);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 从 APPDATA 推导默认的应用 runtime 路径(NUWAX_APP_RUNTIME_PATH 未设置时的回退)
|
||||
if let Ok(appdata) = std::env::var("APPDATA") {
|
||||
let default_node = PathBuf::from(&appdata)
|
||||
.join("com.nuwax.agent-tauri-client")
|
||||
.join("runtime")
|
||||
.join("node")
|
||||
.join("bin")
|
||||
.join("node.exe");
|
||||
if default_node.exists() {
|
||||
info!(
|
||||
"[SACP] Windows node resolved via APPDATA default path: {}",
|
||||
default_node.display()
|
||||
);
|
||||
return Some(default_node);
|
||||
}
|
||||
}
|
||||
|
||||
warn!(
|
||||
"[SACP] Windows node resolution failed: NUWAX_NODE_PATH, NUWAX_NODE_EXE, NUWAX_APP_RUNTIME_PATH and APPDATA default paths all missed"
|
||||
);
|
||||
None
|
||||
}
|
||||
|
||||
fn npm_package_entry_from_dir(
|
||||
package_dir: &std::path::Path,
|
||||
package_name: &str,
|
||||
) -> Option<PathBuf> {
|
||||
let package_json = package_dir.join("package.json");
|
||||
let content = std::fs::read_to_string(package_json).ok()?;
|
||||
let package_json: serde_json::Value = serde_json::from_str(&content).ok()?;
|
||||
let bin_field = package_json.get("bin")?;
|
||||
|
||||
let rel_entry = if let Some(bin_str) = bin_field.as_str() {
|
||||
Some(bin_str.to_string())
|
||||
} else if let Some(bin_map) = bin_field.as_object() {
|
||||
bin_map
|
||||
.get(package_name)
|
||||
.and_then(|v| v.as_str())
|
||||
.map(str::to_string)
|
||||
.or_else(|| {
|
||||
bin_map
|
||||
.values()
|
||||
.find_map(|v| v.as_str())
|
||||
.map(str::to_string)
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}?;
|
||||
|
||||
let entry = package_dir.join(rel_entry);
|
||||
if entry.exists() { Some(entry) } else { None }
|
||||
}
|
||||
|
||||
fn get_windows_cmd_script_path(path: &std::path::Path) -> Option<PathBuf> {
|
||||
match path.extension().and_then(|s| s.to_str()) {
|
||||
Some(ext) if ext.eq_ignore_ascii_case("cmd") => Some(path.to_path_buf()),
|
||||
None => {
|
||||
// 先检查直接路径(绝对路径 or 当前目录)
|
||||
let direct_candidate = path.with_extension("cmd");
|
||||
if direct_candidate.exists() {
|
||||
return Some(direct_candidate);
|
||||
}
|
||||
|
||||
// which 未命中时 path 为裸命令名,需要在已知目录中搜索
|
||||
let program_name = path.file_name().and_then(|s| s.to_str())?;
|
||||
|
||||
// 搜索应用自管的 node_modules/.bin
|
||||
if let Ok(appdata) = std::env::var("APPDATA") {
|
||||
let app_private = PathBuf::from(&appdata)
|
||||
.join("com.nuwax.agent-tauri-client")
|
||||
.join("node_modules")
|
||||
.join(".bin")
|
||||
.join(format!("{}.cmd", program_name));
|
||||
if app_private.exists() {
|
||||
return Some(app_private);
|
||||
}
|
||||
}
|
||||
|
||||
// 搜索 NUWAX_APP_RUNTIME_PATH 中的每个目录
|
||||
if let Ok(runtime_path) = std::env::var("NUWAX_APP_RUNTIME_PATH") {
|
||||
for dir in runtime_path.split(';').filter(|s| !s.trim().is_empty()) {
|
||||
let candidate =
|
||||
std::path::Path::new(dir.trim()).join(format!("{}.cmd", program_name));
|
||||
if candidate.exists() {
|
||||
return Some(candidate);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_js_entry_from_cmd_shim(cmd_script: &std::path::Path) -> Option<PathBuf> {
|
||||
let content = match std::fs::read_to_string(cmd_script) {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
debug!(
|
||||
"[SACP] cmd shim read failed: {} ({})",
|
||||
cmd_script.display(),
|
||||
e
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let base_dir = cmd_script.parent()?;
|
||||
|
||||
for raw_line in content.lines() {
|
||||
let line = raw_line.trim();
|
||||
// 支持多种 npm shim 格式:%~dp0(旧版)、%dp0%(新版 npm >=7)、%dp0(变体)
|
||||
let base_token = if line.contains("%~dp0") {
|
||||
"%~dp0"
|
||||
} else if line.contains("%dp0%") {
|
||||
"%dp0%"
|
||||
} else if line.contains("%dp0") {
|
||||
"%dp0"
|
||||
} else {
|
||||
continue;
|
||||
};
|
||||
let clean = line.replace('"', "");
|
||||
let clean_lower = clean.to_ascii_lowercase();
|
||||
let ext_end = [".cjs", ".mjs", ".js"]
|
||||
.iter()
|
||||
.filter_map(|ext| clean_lower.find(ext).map(|pos| pos + ext.len()))
|
||||
.min();
|
||||
let Some(end) = ext_end else {
|
||||
continue;
|
||||
};
|
||||
let start = clean.find(base_token)?;
|
||||
if end <= start + base_token.len() || end > clean.len() {
|
||||
continue;
|
||||
}
|
||||
let rel = clean[start + base_token.len()..end].trim_start_matches(['\\', '/']);
|
||||
if rel.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let rel = rel.replace('/', "\\");
|
||||
let entry = base_dir.join(std::path::Path::new(&rel));
|
||||
if entry.exists() {
|
||||
debug!(
|
||||
"[SACP] cmd shim resolved entry: {} -> {}",
|
||||
cmd_script.display(),
|
||||
entry.display()
|
||||
);
|
||||
return Some(entry);
|
||||
}
|
||||
}
|
||||
|
||||
debug!("[SACP] cmd shim not resolved: {}", cmd_script.display());
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
pub fn resolve_windows_node_cli_command(
|
||||
command: &str,
|
||||
args: &[String],
|
||||
) -> Option<(String, Vec<String>)> {
|
||||
let command_path = which::which(command).unwrap_or_else(|_| PathBuf::from(command));
|
||||
let path = command_path.as_path();
|
||||
info!(
|
||||
"[SACP] Windows command resolution started: input={}, resolved={}",
|
||||
command,
|
||||
path.display()
|
||||
);
|
||||
let cmd_script = get_windows_cmd_script_path(path);
|
||||
|
||||
let node_exe = resolve_windows_node_exe()?;
|
||||
info!(
|
||||
"[SACP] Windows node.exe already resolved: {}",
|
||||
node_exe.display()
|
||||
);
|
||||
|
||||
if let Some(cmd_script) = cmd_script.as_ref() {
|
||||
info!("[SACP] using cmd shim: {}", cmd_script.display());
|
||||
if let Some(js_entry) = resolve_js_entry_from_cmd_shim(cmd_script) {
|
||||
let mut actual_args = Vec::with_capacity(args.len() + 1);
|
||||
actual_args.push(js_entry.to_string_lossy().to_string());
|
||||
actual_args.extend(args.iter().cloned());
|
||||
info!(
|
||||
"[SACP] cmd shim resolution succeeded: {} -> {}",
|
||||
cmd_script.display(),
|
||||
js_entry.display()
|
||||
);
|
||||
return Some((node_exe.to_string_lossy().to_string(), actual_args));
|
||||
}
|
||||
info!(
|
||||
"[SACP] cmd shim exists but JS entry not resolved, falling back to package.json bin: {}",
|
||||
cmd_script.display()
|
||||
);
|
||||
} else {
|
||||
info!(
|
||||
"[SACP] cmd shim not found, falling back to package.json bin: command={}",
|
||||
command
|
||||
);
|
||||
}
|
||||
|
||||
let package_name = match path.extension().and_then(|s| s.to_str()) {
|
||||
Some(ext) if ext.eq_ignore_ascii_case("cmd") => path
|
||||
.file_stem()
|
||||
.and_then(|s| s.to_str())
|
||||
.map(str::to_string),
|
||||
None => path
|
||||
.file_name()
|
||||
.and_then(|s| s.to_str())
|
||||
.map(str::to_string),
|
||||
_ => None,
|
||||
}?;
|
||||
|
||||
let mut package_dirs: Vec<PathBuf> = Vec::new();
|
||||
|
||||
// 从 which 解析到的 .bin 目录向上推导 node_modules
|
||||
if let Some(parent) = path.parent() {
|
||||
let is_npm_bin = parent
|
||||
.file_name()
|
||||
.map(|s| s.to_string_lossy().eq_ignore_ascii_case(".bin"))
|
||||
.unwrap_or(false);
|
||||
if is_npm_bin {
|
||||
if let Some(node_modules_dir) = parent.parent() {
|
||||
package_dirs.push(node_modules_dir.join(&package_name));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 仅搜索应用自管的 node_modules
|
||||
if let Ok(appdata) = std::env::var("APPDATA") {
|
||||
package_dirs.push(
|
||||
PathBuf::from(appdata)
|
||||
.join("com.nuwax.agent-tauri-client")
|
||||
.join("node_modules")
|
||||
.join(&package_name),
|
||||
);
|
||||
}
|
||||
|
||||
for package_dir in package_dirs {
|
||||
if !package_dir.exists() {
|
||||
info!(
|
||||
"[SACP] package directorynot found, skip: {}",
|
||||
package_dir.display()
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if let Some(js_entry) = npm_package_entry_from_dir(&package_dir, &package_name) {
|
||||
let mut actual_args = Vec::with_capacity(args.len() + 1);
|
||||
actual_args.push(js_entry.to_string_lossy().to_string());
|
||||
actual_args.extend(args.iter().cloned());
|
||||
info!(
|
||||
"[SACP] package.json bin resolved: {} -> {}",
|
||||
package_name,
|
||||
js_entry.display()
|
||||
);
|
||||
return Some((node_exe.to_string_lossy().to_string(), actual_args));
|
||||
}
|
||||
info!(
|
||||
"[SACP] package.json bin resolution failed: package={}, dir={}",
|
||||
package_name,
|
||||
package_dir.display()
|
||||
);
|
||||
}
|
||||
|
||||
warn!(
|
||||
"[SACP] Windows command resolution failed, no JS entry found: command={}",
|
||||
command
|
||||
);
|
||||
None
|
||||
}
|
||||
|
||||
/// Windows 下将「命令 + 参数」规范化为不弹 CMD 窗口的形式。
|
||||
///
|
||||
/// 背景:Windows 上 npm 全局安装会创建 .cmd/.bat,直接执行会弹出 CMD 窗口。
|
||||
/// 本函数按扩展名检测并尽可能转换为 `node.exe + JS 入口`,避免弹窗。
|
||||
///
|
||||
/// 检查顺序(从优到劣):
|
||||
/// 1. `.exe` → 原生二进制,直接返回(不修改)
|
||||
/// 2. `.cmd` / `.bat` → 尝试解析为 node.exe + JS,成功则返回新 (path, args)
|
||||
/// 3. 其他扩展名 → 直接返回
|
||||
/// 4. 无扩展名 → 用 `which` 解析后再按 1–3 处理
|
||||
///
|
||||
/// 仅编译于 Windows;调用方需使用 `#[cfg(windows)]`。
|
||||
#[cfg(windows)]
|
||||
pub fn normalize_windows_command_for_no_window(
|
||||
command_path: String,
|
||||
command_args: Vec<String>,
|
||||
) -> (String, Vec<String>) {
|
||||
let mut path = command_path;
|
||||
let mut args = command_args;
|
||||
|
||||
let path_ref = Path::new(&path);
|
||||
let ext = path_ref
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.map(|s| s.to_ascii_lowercase());
|
||||
|
||||
match ext.as_deref() {
|
||||
Some("exe") => {
|
||||
info!(
|
||||
"[SACP] Windows detected native .exe: {} - no popup window",
|
||||
path
|
||||
);
|
||||
}
|
||||
Some("cmd" | "bat") => {
|
||||
info!("[SACP] 🔍 Windows detecting .cmd/.bat: {}", path);
|
||||
info!("[SACP] 🔄 converting to node.exe + JS ...");
|
||||
if let Some((node_path, js_args)) = resolve_windows_node_cli_command(&path, &args) {
|
||||
info!("[SACP] conversion succeeded: {} + {:?}", node_path, js_args);
|
||||
path = node_path;
|
||||
args = js_args;
|
||||
} else {
|
||||
warn!(
|
||||
"[SACP] ⚠️ conversion failed, will run original command (may show popup): {}",
|
||||
path
|
||||
);
|
||||
}
|
||||
}
|
||||
Some(other) => {
|
||||
info!(
|
||||
"[SACP] ℹ️ Windows detected other format .{}: {} - trying direct execution",
|
||||
other, path
|
||||
);
|
||||
}
|
||||
None => {
|
||||
info!("[SACP] 🔍 Windows detecting extension: {}", path);
|
||||
if let Ok(resolved) = which::which(&path) {
|
||||
let resolved_str = resolved.to_string_lossy().to_string();
|
||||
info!("[SACP] 🔄 resolved: {}", resolved_str);
|
||||
|
||||
let resolved_ext = resolved
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.map(|s| s.to_ascii_lowercase());
|
||||
|
||||
match resolved_ext.as_deref() {
|
||||
Some("exe") => {
|
||||
info!("[SACP] detected .exe - no popup window");
|
||||
path = resolved_str;
|
||||
}
|
||||
Some("cmd" | "bat") => {
|
||||
info!("[SACP] 🔍 detected .cmd/.bat, converting ...");
|
||||
if let Some((node_path, js_args)) =
|
||||
resolve_windows_node_cli_command(&resolved_str, &args)
|
||||
{
|
||||
info!("[SACP] conversion succeeded: {} + {:?}", node_path, js_args);
|
||||
path = node_path;
|
||||
args = js_args;
|
||||
} else {
|
||||
warn!("[SACP] ⚠️ conversion failed");
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
info!("[SACP] ℹ️ unknown extension, keeping original");
|
||||
path = resolved_str;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
warn!("[SACP] ⚠️ unable to resolve path: {}", path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
(path, args)
|
||||
}
|
||||
86
qiming-rcoder/crates/agent_abstraction/src/lib.rs
Normal file
86
qiming-rcoder/crates/agent_abstraction/src/lib.rs
Normal file
@@ -0,0 +1,86 @@
|
||||
//! # Agent Abstraction Layer (SACP 版本)
|
||||
//!
|
||||
//! 此 crate 提供 Agent 管理的抽象接口,是 RCoder 架构的中间层。
|
||||
//!
|
||||
//! ## SACP 迁移说明
|
||||
//!
|
||||
//! 本模块已迁移至 SACP (Symposium ACP) 实现:
|
||||
//! - **移除 Client trait**: `AcpSessionManager<N, R>` 不再需要 `Client` 泛型参数
|
||||
//! - **支持 Send trait**: 可使用标准 `tokio::spawn`,无需 `LocalSet`
|
||||
//! - **简化架构**: 移除了 channel handler,SACP 内部使用回调处理
|
||||
//!
|
||||
//! ## 架构位置
|
||||
//!
|
||||
//! ```text
|
||||
//! shared_types (基础类型: ProjectAndAgentInfo, AgentLifecycleGuard)
|
||||
//! ↓
|
||||
//! agent_abstraction (抽象接口) ← 你在这里
|
||||
//! ↓
|
||||
//! agent_runner (具体实现: AgentSessionRegistry, SseSessionNotifier)
|
||||
//! ```
|
||||
//!
|
||||
//! ## 模块职责
|
||||
//!
|
||||
//! | 模块 | 职责 | 关键类型 |
|
||||
//! |------|------|---------|
|
||||
//! | [`traits`] | 核心抽象接口定义 | `SessionRegistry`, `SessionNotifier` |
|
||||
//! | [`session`] | 会话管理和 Worker 抽象 | `AcpSessionManager`, `AcpAgentWorker` |
|
||||
//! | [`launcher`] | Agent 启动和生命周期管理 | `SacpClaudeCodeLauncher` |
|
||||
//! | [`acp`] | ACP 协议连接抽象 | `CancelResult` |
|
||||
//!
|
||||
//! ## Trait 实现指南
|
||||
//!
|
||||
//! 本 crate 定义的 trait 预期在 `agent_runner` 中实现:
|
||||
//!
|
||||
//! - [`SessionRegistry`] → `agent_runner::service::AgentSessionRegistry` (全局 `AGENT_REGISTRY`)
|
||||
//! - [`SessionNotifier`] → `agent_runner::service::SseSessionNotifier`
|
||||
//!
|
||||
//! ### 使用模式 (SACP)
|
||||
//!
|
||||
//! ```ignore
|
||||
//! // 在 agent_runner 中:
|
||||
//! // 1. 实现 SessionRegistry trait
|
||||
//! impl SessionRegistry for AgentSessionRegistry { ... }
|
||||
//!
|
||||
//! // 2. 通过依赖注入传入 AcpSessionManager(不再需要 Client 类型参数)
|
||||
//! let session_manager = AcpSessionManager::new(
|
||||
//! Arc::new(notifier), // SessionNotifier 实现
|
||||
//! AGENT_REGISTRY.clone(), // SessionRegistry 实现
|
||||
//! );
|
||||
//! ```
|
||||
//!
|
||||
//! ## 设计模式
|
||||
//!
|
||||
//! - **依赖反转(DIP)**: Trait 在此定义,实现在上层 crate
|
||||
//! - **泛型参数**: `AcpSessionManager<N, R>` 支持可插拔实现
|
||||
//! - **RAII**: 通过 `AgentLifecycleGuard` 管理 Agent 资源生命周期
|
||||
//!
|
||||
//! ## 与 AGENT_REGISTRY 的关系
|
||||
//!
|
||||
//! `AGENT_REGISTRY` 是 `agent_runner` 中的全局单例,实现了 `SessionRegistry` trait。
|
||||
//! 通过依赖注入传入 `AcpSessionManager`,使得抽象层不依赖具体实现。
|
||||
//!
|
||||
//! - **直接访问 AGENT_REGISTRY**: 仅在 `agent_runner` 内部使用(如清理任务、状态查询)
|
||||
//! - **通过 SessionRegistry trait**: 在 `agent_abstraction` 内部使用(会话管理逻辑)
|
||||
|
||||
pub mod acp;
|
||||
pub mod error;
|
||||
pub mod launcher;
|
||||
pub mod mirror_env;
|
||||
pub mod path_env;
|
||||
pub mod session;
|
||||
pub mod traits;
|
||||
|
||||
// Re-export types from submodules
|
||||
pub use acp::{CancelNotificationRequestWrapper, CancelResult};
|
||||
pub use error::AgentAbstractionError;
|
||||
pub use launcher::{
|
||||
AgentLaunchConfig, AgentLifecycleGuard, ClaudeCodeLauncher, LauncherConnectionInfoComplete,
|
||||
convert_context_servers, get_default_agent_config, load_agent_config,
|
||||
};
|
||||
pub use session::{
|
||||
AcpAgentWorker, AcpSessionManager, AgentWorker, SessionHandles, WorkerRequest, WorkerResponse,
|
||||
};
|
||||
pub use traits::agent::{AgentStartConfig, PromptMessage};
|
||||
pub use traits::session_notifier::SessionNotifier;
|
||||
pub use traits::session_registry::SessionRegistry;
|
||||
54
qiming-rcoder/crates/agent_abstraction/src/mirror_env.rs
Normal file
54
qiming-rcoder/crates/agent_abstraction/src/mirror_env.rs
Normal file
@@ -0,0 +1,54 @@
|
||||
//! 镜像源环境变量透传
|
||||
//!
|
||||
//! 收集当前进程中的镜像源配置,供 MCP 子进程(npx/bunx/uvx)使用。
|
||||
//!
|
||||
//! 变量来源优先级:
|
||||
//! 1. 直接变量(npm_config_registry / UV_INDEX_URL / PIP_INDEX_URL)
|
||||
//! 2. MCP_PROXY_* 中转变量(MCP_PROXY_NPM_REGISTRY / MCP_PROXY_PYPI_INDEX_URL)
|
||||
|
||||
/// 需要透传的镜像源环境变量映射
|
||||
///
|
||||
/// 格式: (目标变量名, 后备变量名)
|
||||
const MIRROR_MAPPINGS: &[(&str, &str)] = &[
|
||||
("npm_config_registry", "MCP_PROXY_NPM_REGISTRY"),
|
||||
("UV_INDEX_URL", "MCP_PROXY_PYPI_INDEX_URL"),
|
||||
("PIP_INDEX_URL", "MCP_PROXY_PYPI_INDEX_URL"),
|
||||
];
|
||||
|
||||
/// 收集需要透传给子进程的镜像源环境变量
|
||||
///
|
||||
/// 每个目标变量优先读取自身,未设置时从对应的 MCP_PROXY_* 后备变量转换。
|
||||
/// 返回 `Vec<(key, value)>`,仅包含有值的条目。
|
||||
pub fn collect_mirror_env_vars() -> Vec<(String, String)> {
|
||||
MIRROR_MAPPINGS
|
||||
.iter()
|
||||
.filter_map(|&(target, fallback)| {
|
||||
let val = std::env::var(target)
|
||||
.or_else(|_| std::env::var(fallback))
|
||||
.ok()?;
|
||||
Some((target.to_string(), val))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn collect_does_not_panic() {
|
||||
let _ = collect_mirror_env_vars();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mappings_have_correct_targets() {
|
||||
// 确保目标变量名是工具直接使用的名称
|
||||
assert!(
|
||||
MIRROR_MAPPINGS
|
||||
.iter()
|
||||
.any(|(t, _)| *t == "npm_config_registry")
|
||||
);
|
||||
assert!(MIRROR_MAPPINGS.iter().any(|(t, _)| *t == "UV_INDEX_URL"));
|
||||
assert!(MIRROR_MAPPINGS.iter().any(|(t, _)| *t == "PIP_INDEX_URL"));
|
||||
}
|
||||
}
|
||||
120
qiming-rcoder/crates/agent_abstraction/src/path_env.rs
Normal file
120
qiming-rcoder/crates/agent_abstraction/src/path_env.rs
Normal file
@@ -0,0 +1,120 @@
|
||||
//! 子进程环境变量管理 - 供 Agent 子进程使用
|
||||
//!
|
||||
//! 提供 PATH 环境变量构建函数,供 rcoder 启动 Claude Code 等 Agent 时使用。
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// Tauri 桌面客户端在 APPDATA 下的应用数据目录名
|
||||
/// Windows: %APPDATA%/{TAURI_APP_DATA_DIR}
|
||||
pub const TAURI_APP_DATA_DIR: &str = "com.nuwax.agent-tauri-client";
|
||||
|
||||
/// 在系统 PATH 中查找可执行文件
|
||||
fn find_in_path(executable: &str) -> Option<String> {
|
||||
let output = if cfg!(windows) {
|
||||
std::process::Command::new("where").arg(executable).output()
|
||||
} else {
|
||||
std::process::Command::new("which").arg(executable).output()
|
||||
};
|
||||
|
||||
match output {
|
||||
Ok(o) if o.status.success() => {
|
||||
let binding = String::from_utf8_lossy(&o.stdout);
|
||||
binding
|
||||
.lines()
|
||||
.next()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(str::to_string)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 在系统 PATH 中查找可执行文件并返回其父目录
|
||||
fn find_bin_dir(executable: &str) -> Option<String> {
|
||||
find_in_path(executable).and_then(|path| {
|
||||
let parent = PathBuf::from(&path).parent()?.to_path_buf();
|
||||
let parent_str = parent.to_string_lossy().to_string();
|
||||
|
||||
// 如果父目录是 "cmd",尝试找同级的 "bin" 目录
|
||||
// 例如: D:\Program Files\Git\cmd -> D:\Program Files\Git\bin
|
||||
if parent_str.to_lowercase().ends_with("cmd") {
|
||||
let bin_dir = parent
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.map(|n| {
|
||||
if n.eq_ignore_ascii_case("cmd") {
|
||||
parent
|
||||
.parent()
|
||||
.map(|p| p.join("bin").to_string_lossy().to_string())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.flatten();
|
||||
if let Some(bin) = bin_dir {
|
||||
if PathBuf::from(&bin).join("bash.exe").exists() {
|
||||
return Some(bin);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some(parent_str)
|
||||
})
|
||||
}
|
||||
|
||||
/// Windows: 查找 Git Bash 路径
|
||||
///
|
||||
/// 查找系统 PATH 中的 bash.exe,返回其 bin 目录路径
|
||||
/// 如果 bash.exe 在 Git\cmd 目录下,会自动查找同级的 Git\bin 目录
|
||||
#[cfg(windows)]
|
||||
pub fn find_git_bash_path() -> Option<String> {
|
||||
find_bin_dir("bash.exe")
|
||||
}
|
||||
|
||||
/// 构建适用于 rcoder 进程的 PATH 环境变量。
|
||||
///
|
||||
/// 优先级:
|
||||
/// 1. NUWAX_APP_RUNTIME_PATH(优先)
|
||||
/// 2. APPDATA 中的 node bin 和 npm bin(回退)
|
||||
///
|
||||
/// 返回 PATH 字符串,如果都未找到则返回 None
|
||||
#[cfg(windows)]
|
||||
pub fn build_rcoder_path_env() -> Option<String> {
|
||||
use tracing::info;
|
||||
|
||||
let sep = ";";
|
||||
|
||||
// 1. 优先从 NUWAX_APP_RUNTIME_PATH 构建
|
||||
if let Ok(runtime_path) = std::env::var("NUWAX_APP_RUNTIME_PATH") {
|
||||
let runtime_path = runtime_path.trim().to_string();
|
||||
if !runtime_path.is_empty() {
|
||||
info!("[Env] Using NUWAX_APP_RUNTIME_PATH to build PATH");
|
||||
return Some(runtime_path);
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 从 APPDATA 推导默认路径
|
||||
if let Ok(appdata) = std::env::var("APPDATA") {
|
||||
let app_base = PathBuf::from(&appdata).join(TAURI_APP_DATA_DIR);
|
||||
let mut paths: Vec<String> = Vec::new();
|
||||
|
||||
let node_bin = app_base.join("runtime").join("node").join("bin");
|
||||
if node_bin.exists() {
|
||||
paths.push(node_bin.to_string_lossy().to_string());
|
||||
}
|
||||
|
||||
let npm_bin = app_base.join("node_modules").join(".bin");
|
||||
if npm_bin.exists() {
|
||||
paths.push(npm_bin.to_string_lossy().to_string());
|
||||
}
|
||||
|
||||
if !paths.is_empty() {
|
||||
let path_value = paths.join(sep);
|
||||
info!("[Env] APPDATA built PATH: {}", path_value);
|
||||
return Some(path_value);
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
242
qiming-rcoder/crates/agent_abstraction/src/session/acp_worker.rs
Normal file
242
qiming-rcoder/crates/agent_abstraction/src/session/acp_worker.rs
Normal file
@@ -0,0 +1,242 @@
|
||||
//! ACP Agent Worker 实现 (SACP 版本)
|
||||
//!
|
||||
//! 封装 Agent 请求处理的核心业务逻辑
|
||||
//!
|
||||
//! ## SACP 迁移说明
|
||||
//!
|
||||
//! - 移除了 `Client` trait 泛型参数(SACP 内部处理)
|
||||
//! - 简化了会话创建参数
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use agent_config::{AgentServersConfig, PromptConfigAssembler};
|
||||
use anyhow::Result;
|
||||
use shared_types::{ProjectAndAgentInfo, SessionEntry};
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
use super::{AcpSessionManager, AgentWorker, SessionHandles, WorkerRequest, WorkerResponse};
|
||||
use crate::launcher::convert_context_servers;
|
||||
use crate::traits::{AgentStartConfig, SessionNotifier, SessionRegistry};
|
||||
|
||||
/// ACP Agent Worker (SACP 版本)
|
||||
///
|
||||
/// 处理 ACP Agent 请求的核心业务逻辑实现
|
||||
///
|
||||
/// # 类型参数
|
||||
/// - `N`: SessionNotifier 实现,用于推送 SSE 消息
|
||||
/// - `R`: SessionRegistry 实现,用于存储会话数据
|
||||
#[derive(Clone)]
|
||||
pub struct AcpAgentWorker<N: SessionNotifier, R: SessionRegistry> {
|
||||
session_manager: Arc<AcpSessionManager<N, R>>,
|
||||
}
|
||||
|
||||
impl<N: SessionNotifier + 'static, R: SessionRegistry> AcpAgentWorker<N, R> {
|
||||
/// 创建新的 ACP Agent Worker
|
||||
pub fn new(session_manager: Arc<AcpSessionManager<N, R>>) -> Self {
|
||||
Self { session_manager }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl<N: SessionNotifier + Send + Sync + 'static, R: SessionRegistry + Send + Sync + 'static>
|
||||
AgentWorker for AcpAgentWorker<N, R>
|
||||
where
|
||||
R::Entry: Into<ProjectAndAgentInfo> + From<ProjectAndAgentInfo>,
|
||||
{
|
||||
fn name(&self) -> &'static str {
|
||||
"AcpAgentWorker"
|
||||
}
|
||||
|
||||
async fn process_request(&self, request: WorkerRequest) -> Result<WorkerResponse> {
|
||||
let project_id = request.project_id().to_string();
|
||||
let project_path = request.project_path().clone();
|
||||
|
||||
info!(
|
||||
"📨 AcpAgentWorker received request, project_id: {}",
|
||||
project_id
|
||||
);
|
||||
|
||||
// 1. 路径规范化
|
||||
let normalized_path = AcpSessionManager::<N, R>::normalize_path(&project_path);
|
||||
debug!("📂 path: {:?}", normalized_path);
|
||||
|
||||
// 2. 确保项目目录存在
|
||||
AcpSessionManager::<N, R>::ensure_project_dir(&normalized_path)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
error!("createdprojectdirectoryfailed: {:?}", e);
|
||||
e
|
||||
})?;
|
||||
|
||||
// 3. 使用 PromptConfigAssembler 组装配置
|
||||
let default_agent_id = "claude-code-acp-ts";
|
||||
|
||||
// 根据请求中的 service_type 加载对应配置
|
||||
let servers_config =
|
||||
AgentServersConfig::load_or_default_for_service(&request.prompt_message.service_type)
|
||||
.await;
|
||||
|
||||
let assembler = PromptConfigAssembler::new(servers_config)
|
||||
.with_system_prompt(request.prompt_message.system_prompt_override.clone())
|
||||
.with_user_prompt_template(request.prompt_message.user_prompt_template_override.clone())
|
||||
.with_agent_config(request.prompt_message.agent_config_override.clone());
|
||||
|
||||
// 🆕 获取用户指定的 agent_id(如果有)
|
||||
let agent_id = assembler.get_agent_id(default_agent_id);
|
||||
debug!("🎯 Resolved Agent ID: {}", agent_id);
|
||||
|
||||
// 获取最终的系统提示词(入参有值则使用入参,否则使用默认配置)
|
||||
// 🆕 使用实际 agent_id 而不是 default_agent_id
|
||||
let system_prompt = assembler.get_system_prompt(&agent_id);
|
||||
// 应用用户提示词模板(如果有)
|
||||
// 🆕 使用实际 agent_id 而不是 default_agent_id
|
||||
let final_user_prompt =
|
||||
assembler.apply_user_prompt(&agent_id, &request.prompt_message.content);
|
||||
|
||||
info!(
|
||||
"📝 Prompt processing - System prompt: has_override={}, length={} | User prompt: has_template={}, original_len={}, final_len={}",
|
||||
assembler.has_system_prompt_override(),
|
||||
system_prompt.len(),
|
||||
assembler.has_user_prompt_template_override(),
|
||||
request.prompt_message.content.len(),
|
||||
final_user_prompt.len()
|
||||
);
|
||||
|
||||
debug!(
|
||||
"📝 System prompt: has_override={}, length={}, content={}",
|
||||
assembler.has_system_prompt_override(),
|
||||
system_prompt.len(),
|
||||
system_prompt
|
||||
);
|
||||
debug!(
|
||||
"📝 User prompt: has_template={}, original_len={}, final_len={}, final_content={}",
|
||||
assembler.has_user_prompt_template_override(),
|
||||
request.prompt_message.content.len(),
|
||||
final_user_prompt.len(),
|
||||
final_user_prompt
|
||||
);
|
||||
|
||||
// 获取 MCP 服务器配置(入参有值则使用入参,否则使用默认配置)
|
||||
let context_servers = assembler.get_context_servers();
|
||||
debug!(
|
||||
"🔌 MCP server config: has_override={}, count={}",
|
||||
assembler.has_agent_config_override(),
|
||||
context_servers.len()
|
||||
);
|
||||
|
||||
// 将 context_servers 转换为 ACP 协议的 McpServer 格式
|
||||
let mcp_servers = convert_context_servers(&context_servers);
|
||||
debug!("🔌 MCP servers: {}", mcp_servers.len());
|
||||
|
||||
// 构建 AgentStartConfig 并传递 MCP 服务器、service_type
|
||||
let mut start_config = AgentStartConfig::new(request.prompt_message.service_type.clone())
|
||||
.with_system_prompt(system_prompt)
|
||||
.with_mcp_servers(mcp_servers);
|
||||
|
||||
// 🆕 如果用户指定了 agent_server 配置,添加到 start_config
|
||||
// 注意:这里直接使用用户传入的配置,由 launcher 层负责与默认配置合并
|
||||
if let Some(ref override_config) = request.prompt_message.agent_config_override {
|
||||
if let Some(ref agent_server) = override_config.agent_server {
|
||||
debug!(
|
||||
"📝 Using user-specified Agent server config: command={:?}, args={:?}",
|
||||
agent_server.command, agent_server.args
|
||||
);
|
||||
start_config = start_config.with_agent_server_override(agent_server.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Resume 策略:直接传递 session_id,由 LoadSessionRequest 在 Agent 层面检查
|
||||
//
|
||||
// 工作原理:
|
||||
// 1. 用户传入 session_id → 直接设置 resume_session_id
|
||||
// 2. claude_code_sacp.rs 尝试 LoadSessionRequest 加载历史会话
|
||||
// 3. 如果 Agent 返回成功 → 恢复上下文
|
||||
// 4. 如果 Agent 返回错误/超时 → 降级到 NewSessionRequest 创建新会话
|
||||
//
|
||||
// 优势:
|
||||
// - Agent 层面统一处理,无需本地检查文件
|
||||
// - 与 Agent 行为保持一致
|
||||
// - 自动降级(有会话就恢复,没有就创建)
|
||||
if let Some(ref session_id) = request.prompt_message.session_id {
|
||||
info!("Setting resume_session_id: {}", session_id);
|
||||
start_config = start_config.with_resume_session_id(session_id.clone());
|
||||
}
|
||||
|
||||
// 4. 更新 prompt_message 的 content 为处理后的用户提示词
|
||||
let mut prompt_message = request.prompt_message.clone();
|
||||
prompt_message.content = final_user_prompt;
|
||||
|
||||
// 5. 获取或创建会话 (SACP 版本 - 不需要 client 和 shared_api_key_manager)
|
||||
let (session_entry, is_new) = self
|
||||
.session_manager
|
||||
.get_or_create_session(
|
||||
&project_id,
|
||||
normalized_path,
|
||||
prompt_message.session_id.clone(),
|
||||
request.model_provider.clone(),
|
||||
start_config.clone(),
|
||||
request.service_uuid.clone(),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
error!("Failed to create session: {:?}", e);
|
||||
e
|
||||
})?;
|
||||
|
||||
// 使用 SessionEntry trait 方法访问会话信息
|
||||
info!(
|
||||
"✅ Session ready, session_id: {}, is_new: {}",
|
||||
session_entry.session_id(),
|
||||
is_new
|
||||
);
|
||||
|
||||
// 6. 构建 Prompt 请求
|
||||
let prompt_request = if let Some(ref attachment_blocks) = request.attachment_blocks {
|
||||
debug!("📎 Built Prompt request with attachments");
|
||||
AcpSessionManager::<N, R>::build_prompt_request_with_attachments(
|
||||
&prompt_message,
|
||||
session_entry.session_id().clone(),
|
||||
attachment_blocks.clone(),
|
||||
)?
|
||||
} else {
|
||||
debug!("📝 Built text Prompt request");
|
||||
AcpSessionManager::<N, R>::build_text_prompt_request(
|
||||
&prompt_message,
|
||||
session_entry.session_id().clone(),
|
||||
)?
|
||||
};
|
||||
|
||||
// 7. 发送 Prompt(异步,支持背压)
|
||||
self.session_manager
|
||||
.send_prompt_request(&project_id, prompt_request)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
error!("send Prompt requestfailed: {:?}", e);
|
||||
e
|
||||
})?;
|
||||
|
||||
info!("Prompt request already sent, project_id: {}", project_id);
|
||||
|
||||
// 8. 构建响应
|
||||
if is_new {
|
||||
Ok(WorkerResponse::new_session_success(
|
||||
project_id,
|
||||
session_entry.session_id().to_string(),
|
||||
Some(request.request_id().to_string()),
|
||||
prompt_message.service_type.clone(),
|
||||
SessionHandles {
|
||||
prompt_tx: session_entry.prompt_tx().clone(),
|
||||
cancel_tx: session_entry.cancel_tx().clone(),
|
||||
lifecycle_handle: session_entry.lifecycle_handle().cloned(),
|
||||
},
|
||||
))
|
||||
} else {
|
||||
Ok(WorkerResponse::reuse_session_success(
|
||||
project_id,
|
||||
session_entry.session_id().to_string(),
|
||||
Some(request.request_id().to_string()),
|
||||
prompt_message.service_type.clone(),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
23
qiming-rcoder/crates/agent_abstraction/src/session/mod.rs
Normal file
23
qiming-rcoder/crates/agent_abstraction/src/session/mod.rs
Normal file
@@ -0,0 +1,23 @@
|
||||
//! Session 管理模块
|
||||
//!
|
||||
//! 提供 ACP 会话的统一管理能力,包括:
|
||||
//! - 会话信息存储和查询
|
||||
//! - 会话生命周期管理
|
||||
//! - 模型配置变化检测
|
||||
//! - Agent Worker 抽象
|
||||
//!
|
||||
//! ## 架构说明
|
||||
//!
|
||||
//! 会话信息统一使用 `ProjectAndAgentInfo`(定义在 `shared_types`),
|
||||
//! 通过 `SessionEntry` trait 抽象访问接口。
|
||||
//! 会话存储通过 `SessionRegistry` trait 抽象,允许注入不同实现(如 `AGENT_REGISTRY`)。
|
||||
|
||||
mod acp_worker;
|
||||
mod session_file_scanner;
|
||||
mod session_manager;
|
||||
mod worker;
|
||||
|
||||
pub use acp_worker::AcpAgentWorker;
|
||||
pub use session_file_scanner::check_session_file_exists;
|
||||
pub use session_manager::AcpSessionManager;
|
||||
pub use worker::{AgentWorker, SessionHandles, WorkerRequest, WorkerResponse};
|
||||
@@ -0,0 +1,352 @@
|
||||
//! Session 文件扫描模块
|
||||
//!
|
||||
//! 扫描 ~/.claude/projects/ 目录,验证 session 文件是否存在
|
||||
//! 作为官方 listSessions API 的降级方案
|
||||
//!
|
||||
//! 特性:
|
||||
//! - 使用 moka 缓存扫描结果(TTL 3分钟兜底)
|
||||
//! - 使用 notify 监听文件系统变化,自动刷新缓存
|
||||
//! - 防抖机制:100ms 内的多个事件合并为一次刷新
|
||||
//! - 文件变化时清除缓存,下次查询时重新扫描
|
||||
|
||||
use moka::future::Cache;
|
||||
use notify::{Event, EventKind, RecommendedWatcher, RecursiveMode, Watcher};
|
||||
use std::collections::HashSet;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
/// 全局 session 文件存在性缓存
|
||||
/// Key: project_path, Value: 该项目下的 session_id 集合
|
||||
static FILE_SCAN_CACHE: OnceLock<Cache<String, HashSet<String>>> = OnceLock::new();
|
||||
|
||||
/// 文件监听器(全局单例)
|
||||
static FILE_WATCHER: OnceLock<Mutex<Option<RecommendedWatcher>>> = OnceLock::new();
|
||||
|
||||
/// 监听器是否已启动
|
||||
static WATCHER_STARTED: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
/// 防抖时间(毫秒)
|
||||
const DEBOUNCE_MS: u64 = 100;
|
||||
|
||||
/// 获取或初始化缓存(TTL 3分钟兜底,主要由文件监听触发刷新)
|
||||
fn get_file_scan_cache() -> &'static Cache<String, HashSet<String>> {
|
||||
FILE_SCAN_CACHE.get_or_init(|| {
|
||||
Cache::builder()
|
||||
.time_to_live(Duration::from_secs(180)) // 3分钟兜底 TTL
|
||||
.build()
|
||||
})
|
||||
}
|
||||
|
||||
/// Claude 配置目录
|
||||
fn get_claude_config_dir() -> PathBuf {
|
||||
std::env::var("CLAUDE_CONFIG_DIR")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|_| {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("/home/user"))
|
||||
.join(".claude")
|
||||
})
|
||||
}
|
||||
|
||||
/// 获取 projects 目录路径
|
||||
fn get_projects_dir() -> PathBuf {
|
||||
get_claude_config_dir().join("projects")
|
||||
}
|
||||
|
||||
/// 编码项目路径为目录名
|
||||
/// 规则:将 `/` 替换为 `-`
|
||||
fn encode_project_path(project_path: &str) -> String {
|
||||
project_path.replace('/', "-")
|
||||
}
|
||||
|
||||
/// 从文件路径中提取项目目录名
|
||||
/// 例如:~/.claude/projects/-home-user-project/abc123.jsonl -> -home-user-project
|
||||
fn extract_project_dir_from_path(file_path: &Path) -> Option<String> {
|
||||
let projects_dir = get_projects_dir();
|
||||
|
||||
// 检查文件是否在 projects 目录下
|
||||
if !file_path.starts_with(&projects_dir) {
|
||||
return None;
|
||||
}
|
||||
|
||||
// 获取相对于 projects 目录的路径
|
||||
if let Ok(relative) = file_path.strip_prefix(&projects_dir) {
|
||||
// 第一个组件就是项目目录名
|
||||
if let Some(first_component) = relative.components().next() {
|
||||
return Some(first_component.as_os_str().to_string_lossy().to_string());
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// 扫描项目目录,获取所有 session_id(异步版本)
|
||||
///
|
||||
/// 使用 tokio::fs 异步 I/O,避免阻塞 Tokio 工作线程
|
||||
async fn scan_project_sessions(project_path: &str) -> HashSet<String> {
|
||||
let projects_dir = get_projects_dir();
|
||||
let encoded_path = encode_project_path(project_path);
|
||||
|
||||
let mut session_ids = HashSet::new();
|
||||
|
||||
if !projects_dir.exists() {
|
||||
debug!(
|
||||
"🔍 [FILE_SCAN] Projects directory does not exist: {}",
|
||||
projects_dir.display()
|
||||
);
|
||||
return session_ids;
|
||||
}
|
||||
|
||||
// 使用异步 I/O 遍历目录
|
||||
let Ok(mut dir_entries) = tokio::fs::read_dir(&projects_dir).await else {
|
||||
warn!(
|
||||
"🔍 [FILE_SCAN] Unable to read project directory: {}",
|
||||
projects_dir.display()
|
||||
);
|
||||
return session_ids;
|
||||
};
|
||||
|
||||
while let Ok(Some(entry)) = dir_entries.next_entry().await {
|
||||
let dir_name = entry.file_name();
|
||||
|
||||
// 匹配编码后的路径(精确或带哈希后缀)
|
||||
// 使用 OsStr 避免不必要的 String 分配
|
||||
if dir_name
|
||||
.to_string_lossy()
|
||||
.starts_with(encoded_path.as_str())
|
||||
{
|
||||
let Ok(mut files) = tokio::fs::read_dir(entry.path()).await else {
|
||||
warn!(
|
||||
"🔍 [file_scanner] Unable to scan session directory: {}",
|
||||
entry.path().display()
|
||||
);
|
||||
continue;
|
||||
};
|
||||
|
||||
while let Ok(Some(file)) = files.next_entry().await {
|
||||
if let Some(filename) = file.file_name().to_str() {
|
||||
if filename.ends_with(".jsonl") {
|
||||
// 提取 session_id(去掉 .jsonl 后缀)
|
||||
let session_id = filename.trim_end_matches(".jsonl").to_string();
|
||||
session_ids.insert(session_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
debug!(
|
||||
"🔍 [FILE_SCAN] Scan completed: project_path={}, found {} sessions",
|
||||
project_path,
|
||||
session_ids.len()
|
||||
);
|
||||
|
||||
session_ids
|
||||
}
|
||||
|
||||
/// 清除指定项目目录相关的缓存
|
||||
///
|
||||
/// 由于 moka 不支持按前缀遍历/删除,采用 invalidate_all 策略
|
||||
/// 下次查询时会用原始 project_path 重新扫描
|
||||
async fn invalidate_project_cache(project_dir_name: &str) {
|
||||
let cache = get_file_scan_cache();
|
||||
|
||||
debug!(
|
||||
"🔄 [FILE_WATCH] Detected project directory change: {}, clearing cache",
|
||||
project_dir_name
|
||||
);
|
||||
|
||||
// 由于无法从目录名还原 project_path,直接清除所有缓存
|
||||
// 这是一个权衡:简单但可能影响其他项目的缓存
|
||||
// 实际影响较小,因为缓存会在下次查询时重建
|
||||
cache.invalidate_all();
|
||||
}
|
||||
|
||||
/// 启动文件系统监听器
|
||||
///
|
||||
/// 监听 ~/.claude/projects/ 目录的变化,当 .jsonl 文件创建/删除时刷新缓存
|
||||
pub fn start_file_watcher() {
|
||||
// 确保只启动一次
|
||||
if WATCHER_STARTED.swap(true, Ordering::SeqCst) {
|
||||
debug!("[filelisten] File watcher already started");
|
||||
return;
|
||||
}
|
||||
|
||||
let projects_dir = get_projects_dir();
|
||||
|
||||
// 创建目录(如果不存在)
|
||||
if !projects_dir.exists() {
|
||||
if let Err(e) = std::fs::create_dir_all(&projects_dir) {
|
||||
warn!("[filelisten] Unable to create projects directory: {}", e);
|
||||
WATCHER_STARTED.store(false, Ordering::SeqCst);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// 创建异步通道
|
||||
let (tx, mut rx) = mpsc::channel::<Event>(100);
|
||||
|
||||
// 创建文件监听器
|
||||
let watcher_result = notify::recommended_watcher(move |res: Result<Event, notify::Error>| {
|
||||
match res {
|
||||
Ok(event) => {
|
||||
// 只关心 .jsonl 文件的创建和删除
|
||||
let has_jsonl = event
|
||||
.paths
|
||||
.iter()
|
||||
.any(|p| p.extension().is_some_and(|ext| ext == "jsonl"));
|
||||
|
||||
if has_jsonl {
|
||||
match event.kind {
|
||||
EventKind::Create(_) | EventKind::Remove(_) | EventKind::Modify(_) => {
|
||||
if let Err(e) = tx.blocking_send(event) {
|
||||
error!("[filelisten] Failed to send event: {}", e);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("[filelisten] Listener error: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let mut watcher = match watcher_result {
|
||||
Ok(w) => w,
|
||||
Err(e) => {
|
||||
error!("[filelisten] Failed to create listener: {}", e);
|
||||
WATCHER_STARTED.store(false, Ordering::SeqCst);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// 开始监听目录
|
||||
if let Err(e) = watcher.watch(&projects_dir, RecursiveMode::Recursive) {
|
||||
error!("[filelisten] Failed to watch directory: {}", e);
|
||||
WATCHER_STARTED.store(false, Ordering::SeqCst);
|
||||
return;
|
||||
}
|
||||
|
||||
// 保存 watcher 到全局变量(防止被 drop)
|
||||
let watcher_holder = FILE_WATCHER.get_or_init(|| Mutex::new(None));
|
||||
if let Ok(mut guard) = watcher_holder.lock() {
|
||||
*guard = Some(watcher);
|
||||
}
|
||||
|
||||
info!(
|
||||
"👁️ [filelisten] startinglistendirectory: {}",
|
||||
projects_dir.display()
|
||||
);
|
||||
|
||||
// 启动异步任务处理事件(带防抖)
|
||||
tokio::spawn(async move {
|
||||
let mut pending_dirs: HashSet<String> = HashSet::new();
|
||||
let mut debounce_timer: Option<tokio::time::Instant> = None;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// 接收新事件
|
||||
event = rx.recv() => {
|
||||
match event {
|
||||
Some(event) => {
|
||||
// 提取变化文件所属的项目目录
|
||||
for path in &event.paths {
|
||||
if let Some(project_dir) = extract_project_dir_from_path(path) {
|
||||
pending_dirs.insert(project_dir);
|
||||
}
|
||||
}
|
||||
// 重置防抖计时器
|
||||
debounce_timer = Some(tokio::time::Instant::now() + Duration::from_millis(DEBOUNCE_MS));
|
||||
}
|
||||
None => {
|
||||
warn!("[filelisten] Received event but channel already closed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
// 防抖计时器触发
|
||||
_ = async {
|
||||
if let Some(timer) = debounce_timer {
|
||||
tokio::time::sleep_until(timer).await;
|
||||
} else {
|
||||
// 没有计时器时永远等待
|
||||
std::future::pending::<()>().await;
|
||||
}
|
||||
} => {
|
||||
if !pending_dirs.is_empty() {
|
||||
debug!(
|
||||
"📁 [FILE_WATCH] Debounce triggered, refreshing {} project directories",
|
||||
pending_dirs.len()
|
||||
);
|
||||
// 刷新所有待处理的项目目录
|
||||
for project_dir in pending_dirs.drain() {
|
||||
invalidate_project_cache(&project_dir).await;
|
||||
}
|
||||
}
|
||||
debounce_timer = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
info!("[FILE_WATCH] File watcher started");
|
||||
}
|
||||
|
||||
/// 停止文件系统监听器
|
||||
#[allow(dead_code)]
|
||||
pub fn stop_file_watcher() {
|
||||
if !WATCHER_STARTED.swap(false, Ordering::SeqCst) {
|
||||
return;
|
||||
}
|
||||
|
||||
let watcher_holder = FILE_WATCHER.get_or_init(|| Mutex::new(None));
|
||||
if let Ok(mut guard) = watcher_holder.lock() {
|
||||
*guard = None;
|
||||
}
|
||||
|
||||
info!("🛑 [filelisten] File watcher already stopped");
|
||||
}
|
||||
|
||||
/// 通过文件扫描检查 session 是否存在(带缓存)
|
||||
pub async fn check_session_file_exists(session_id: &str, project_path: &str) -> bool {
|
||||
// 确保监听器已启动
|
||||
start_file_watcher();
|
||||
|
||||
let cache = get_file_scan_cache();
|
||||
|
||||
// 尝试从缓存获取
|
||||
if let Some(session_ids) = cache.get(project_path).await {
|
||||
let exists = session_ids.contains(session_id);
|
||||
debug!(
|
||||
"🔍 [FILE_SCAN_CACHE] session={} -> {}",
|
||||
session_id,
|
||||
if exists { "exists" } else { "not found" }
|
||||
);
|
||||
return exists;
|
||||
}
|
||||
|
||||
// 缓存未命中,执行扫描
|
||||
debug!(
|
||||
"🔍 [FILE_SCAN] Cache miss, scanning directory: project_path={}",
|
||||
project_path
|
||||
);
|
||||
let session_ids = scan_project_sessions(project_path).await;
|
||||
let exists = session_ids.contains(session_id);
|
||||
|
||||
// 存入缓存
|
||||
cache.insert(project_path.to_string(), session_ids).await;
|
||||
|
||||
if exists {
|
||||
info!("[file_scan] Session file found: {}", session_id);
|
||||
} else {
|
||||
debug!("[file_scan] Not a session file: {}", session_id);
|
||||
}
|
||||
|
||||
exists
|
||||
}
|
||||
@@ -0,0 +1,665 @@
|
||||
//! # ACP Session Manager
|
||||
//!
|
||||
//! 基于 project_id 管理 ACP 会话的核心模块。
|
||||
//!
|
||||
//! ## SACP 迁移说明
|
||||
//!
|
||||
//! 本模块已迁移至 SACP 实现:
|
||||
//! - 移除了 `Client` trait 泛型参数(SACP 内部处理)
|
||||
//! - 简化了 `create_session` 参数(不再需要 client 和 shared_api_key_manager)
|
||||
//! - 使用标准 tokio::spawn(无需 LocalSet)
|
||||
//!
|
||||
//! ## 职责范围
|
||||
//!
|
||||
//! 1. **会话生命周期管理**
|
||||
//! - 创建新会话 (`create_session`)
|
||||
//! - 获取或复用会话 (`get_or_create_session`)
|
||||
//! - 移除会话 (`remove_session`)
|
||||
//! - 检测会话健康状态(channel 是否关闭)
|
||||
//!
|
||||
//! 2. **Prompt 构建**
|
||||
//! - 构建纯文本 Prompt (`build_text_prompt_request`)
|
||||
//! - 构建带附件的 Prompt (`build_prompt_request_with_attachments`)
|
||||
//! - 将 request_id 放入 meta 字段
|
||||
//!
|
||||
//! 3. **路径处理**
|
||||
//! - 路径规范化 (`normalize_path`)
|
||||
//! - 确保项目目录存在 (`ensure_project_dir`)
|
||||
//!
|
||||
//! ## 架构说明
|
||||
//!
|
||||
//! 使用依赖注入的 `SessionRegistry` 进行会话存储:
|
||||
//! - 统一使用注入的 `SessionRegistry`(通常是 `AGENT_REGISTRY`)
|
||||
//!
|
||||
//! ## 与 Worker 的协作
|
||||
//!
|
||||
//! ```text
|
||||
//! AcpAgentWorker (acp_worker.rs)
|
||||
//! │
|
||||
//! │ 1. 调用 get_or_create_session()
|
||||
//! ▼
|
||||
//! AcpSessionManager
|
||||
//! │
|
||||
//! │ 2. 通过 SessionRegistry 获取/创建会话
|
||||
//! │ 3. 通过 SacpClaudeCodeLauncher 启动 Agent
|
||||
//! ▼
|
||||
//! SessionRegistry (注入的 AGENT_REGISTRY)
|
||||
//! ```
|
||||
|
||||
use std::path::{Component, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use agent_config::PromptBuilder;
|
||||
use anyhow::Result;
|
||||
use chrono::Utc;
|
||||
use agent_client_protocol::schema::{ContentBlock, PromptRequest, SessionId, TextContent};
|
||||
use shared_types::{
|
||||
AgentLifecycle, AgentStatus, ModelProviderConfig, ProjectAndAgentInfo, SessionEntry,
|
||||
};
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
use crate::PromptMessage;
|
||||
use crate::launcher::ClaudeCodeLauncher;
|
||||
use crate::traits::{AgentStartConfig, SessionNotifier, SessionRegistry};
|
||||
|
||||
/// ACP 会话管理器 (SACP 版本)
|
||||
///
|
||||
/// 管理所有活跃的 ACP 会话,提供:
|
||||
/// - 会话创建和复用
|
||||
/// - 模型配置变化检测
|
||||
/// - 会话生命周期管理
|
||||
///
|
||||
/// ## 泛型参数
|
||||
/// - `N`: SessionNotifier 实现,用于推送 SSE 消息
|
||||
/// - `R`: SessionRegistry 实现,用于存储会话数据(通常是 AGENT_REGISTRY)
|
||||
pub struct AcpSessionManager<N: SessionNotifier, R: SessionRegistry> {
|
||||
/// 会话注册表(注入的 SessionRegistry)
|
||||
registry: Arc<R>,
|
||||
/// 会话通知器
|
||||
notifier: Arc<N>,
|
||||
}
|
||||
|
||||
impl<N: SessionNotifier + 'static, R: SessionRegistry> AcpSessionManager<N, R>
|
||||
where
|
||||
R::Entry: Into<ProjectAndAgentInfo> + From<ProjectAndAgentInfo>,
|
||||
{
|
||||
/// 创建新的会话管理器
|
||||
///
|
||||
/// # 参数
|
||||
/// - `notifier`: 会话通知器
|
||||
/// - `registry`: 会话注册表(通常注入 AGENT_REGISTRY)
|
||||
pub fn new(notifier: Arc<N>, registry: Arc<R>) -> Self {
|
||||
Self { registry, notifier }
|
||||
}
|
||||
|
||||
/// 获取会话信息
|
||||
pub fn get_session(&self, project_id: &str) -> Option<R::Entry> {
|
||||
self.registry.get(project_id)
|
||||
}
|
||||
|
||||
/// 检查会话是否存在
|
||||
pub fn contains_session(&self, project_id: &str) -> bool {
|
||||
self.registry.contains(project_id)
|
||||
}
|
||||
|
||||
/// 移除会话
|
||||
pub fn remove_session(&self, project_id: &str) -> Option<R::Entry> {
|
||||
self.registry.remove(project_id)
|
||||
}
|
||||
|
||||
/// 获取所有会话的 project_id 列表
|
||||
pub fn list_sessions(&self) -> Vec<String> {
|
||||
self.registry.list_project_ids()
|
||||
}
|
||||
|
||||
/// 获取会话数量
|
||||
pub fn session_count(&self) -> usize {
|
||||
self.registry.count()
|
||||
}
|
||||
|
||||
/// 获取 registry 的 Arc 引用
|
||||
pub fn registry(&self) -> Arc<R> {
|
||||
self.registry.clone()
|
||||
}
|
||||
|
||||
/// 获取 notifier 的 Arc 引用
|
||||
pub fn notifier(&self) -> Arc<N> {
|
||||
self.notifier.clone()
|
||||
}
|
||||
|
||||
/// 规范化项目路径
|
||||
///
|
||||
/// - 如果是相对路径,先与当前目录拼接
|
||||
/// - 去除路径中的 "./"(CurDir 组件)
|
||||
pub fn normalize_path(path: &PathBuf) -> PathBuf {
|
||||
let joined_path = if path.is_absolute() {
|
||||
path.clone()
|
||||
} else {
|
||||
std::env::current_dir().unwrap_or_default().join(path)
|
||||
};
|
||||
|
||||
joined_path
|
||||
.components()
|
||||
.filter(|c| !matches!(c, Component::CurDir))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 确保项目目录存在
|
||||
pub async fn ensure_project_dir(path: &PathBuf) -> Result<()> {
|
||||
if !path.exists() {
|
||||
info!("Project path does not exist, creating: {:?}", path);
|
||||
tokio::fs::create_dir_all(path).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 构建基础 Prompt 请求(仅文本内容)
|
||||
///
|
||||
/// 使用 ACP 协议分离模式:
|
||||
/// - 系统提示词已在 NewSessionRequest._meta.systemPrompt 中传递
|
||||
/// - 这里只构建纯用户提示词
|
||||
///
|
||||
/// 注意:附件转换由调用方处理(因为涉及文件 I/O)
|
||||
pub fn build_text_prompt_request(
|
||||
prompt: &PromptMessage,
|
||||
session_id: SessionId,
|
||||
) -> Result<PromptRequest> {
|
||||
// 构建纯用户提示词
|
||||
let final_prompt = if prompt.data_source_attachments.is_empty() {
|
||||
PromptBuilder::build_user_prompt(&prompt.content)
|
||||
} else {
|
||||
PromptBuilder::build_user_prompt_with_data_sources(
|
||||
&prompt.content,
|
||||
&prompt.data_source_attachments,
|
||||
)
|
||||
};
|
||||
|
||||
// 创建文本内容块
|
||||
let text_block = ContentBlock::Text(TextContent::new(final_prompt));
|
||||
let content_blocks = vec![text_block];
|
||||
|
||||
// 将 request_id 放入 meta 字段
|
||||
debug!(
|
||||
"🔧 [build_prompt] Putting request_id={} into PromptRequest.meta",
|
||||
prompt.request_id
|
||||
);
|
||||
let mut meta = serde_json::Map::new();
|
||||
meta.insert(
|
||||
"request_id".to_string(),
|
||||
serde_json::Value::String(prompt.request_id.clone()),
|
||||
);
|
||||
|
||||
Ok(PromptRequest::new(session_id, content_blocks).meta(meta))
|
||||
}
|
||||
|
||||
/// 构建带附件的 Prompt 请求
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `prompt` - 提示消息
|
||||
/// * `session_id` - 会话 ID
|
||||
/// * `attachment_blocks` - 已转换的附件内容块
|
||||
pub fn build_prompt_request_with_attachments(
|
||||
prompt: &PromptMessage,
|
||||
session_id: SessionId,
|
||||
attachment_blocks: Vec<ContentBlock>,
|
||||
) -> Result<PromptRequest> {
|
||||
// 构建纯用户提示词
|
||||
let final_prompt = if prompt.data_source_attachments.is_empty() {
|
||||
PromptBuilder::build_user_prompt(&prompt.content)
|
||||
} else {
|
||||
PromptBuilder::build_user_prompt_with_data_sources(
|
||||
&prompt.content,
|
||||
&prompt.data_source_attachments,
|
||||
)
|
||||
};
|
||||
|
||||
// 创建文本内容块
|
||||
let text_block = ContentBlock::Text(TextContent::new(final_prompt));
|
||||
let mut content_blocks = vec![text_block];
|
||||
|
||||
// 添加附件内容块
|
||||
content_blocks.extend(attachment_blocks);
|
||||
|
||||
// 将 request_id 放入 meta 字段
|
||||
debug!(
|
||||
"🔧 [build_prompt] Putting request_id={} into PromptRequest.meta",
|
||||
prompt.request_id
|
||||
);
|
||||
let mut meta = serde_json::Map::new();
|
||||
meta.insert(
|
||||
"request_id".to_string(),
|
||||
serde_json::Value::String(prompt.request_id.clone()),
|
||||
);
|
||||
|
||||
Ok(PromptRequest::new(session_id, content_blocks).meta(meta))
|
||||
}
|
||||
|
||||
/// 创建新的 Agent 会话 (SACP 版本)
|
||||
///
|
||||
/// 启动 Agent 进程并建立 SACP 连接
|
||||
///
|
||||
/// # 参数
|
||||
/// - `project_id`: 项目 ID
|
||||
/// - `project_path`: 项目路径
|
||||
/// - `model_provider`: 模型提供者配置
|
||||
/// - `start_config`: Agent 启动配置
|
||||
/// - `service_uuid`: 与此 Agent 关联的唯一 UUID
|
||||
///
|
||||
/// # 返回值
|
||||
/// - `R::Entry`: 会话条目
|
||||
pub async fn create_session(
|
||||
&self,
|
||||
project_id: String,
|
||||
project_path: PathBuf,
|
||||
_session_id_hint: Option<String>,
|
||||
model_provider: Option<ModelProviderConfig>,
|
||||
start_config: AgentStartConfig,
|
||||
service_uuid: Option<String>,
|
||||
) -> Result<R::Entry> {
|
||||
let agent_info = self
|
||||
.create_session_internal(
|
||||
project_id.clone(),
|
||||
project_path,
|
||||
model_provider,
|
||||
start_config,
|
||||
service_uuid,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// 存储会话信息到 registry
|
||||
let session_id_str = agent_info.session_id().to_string();
|
||||
self.registry
|
||||
.insert(&project_id, &session_id_str, agent_info.clone());
|
||||
|
||||
Ok(agent_info)
|
||||
}
|
||||
|
||||
/// 内部方法:创建 Agent 会话但不插入到 registry
|
||||
///
|
||||
/// 用于 entry API 优化,避免重复插入
|
||||
async fn create_session_internal(
|
||||
&self,
|
||||
project_id: String,
|
||||
project_path: PathBuf,
|
||||
model_provider: Option<ModelProviderConfig>,
|
||||
start_config: AgentStartConfig,
|
||||
service_uuid: Option<String>,
|
||||
) -> Result<R::Entry> {
|
||||
info!(
|
||||
"Creating Agent session, project ID: {}",
|
||||
project_id
|
||||
);
|
||||
|
||||
// 创建 SACP 启动器
|
||||
let launcher = ClaudeCodeLauncher::new(self.notifier.clone());
|
||||
|
||||
// 记录是否使用了 resume(仅用于日志)
|
||||
let has_resume = start_config.resume_session_id.is_some();
|
||||
if has_resume {
|
||||
info!(
|
||||
"📌 Starting Agent with resume: session_id={:?}",
|
||||
start_config.resume_session_id
|
||||
);
|
||||
}
|
||||
|
||||
// 启动 Agent (SACP 版本)
|
||||
// 如果 resume 失败,直接返回错误,让上层(rcoder)决定是否降级重试
|
||||
let connection_info = launcher
|
||||
.launch(
|
||||
project_id.clone(),
|
||||
project_path.clone(),
|
||||
model_provider.clone(),
|
||||
start_config.clone(),
|
||||
self.registry.clone(),
|
||||
service_uuid,
|
||||
)
|
||||
.await?;
|
||||
|
||||
info!(
|
||||
"✅ Agent session created successfully, session ID: {}",
|
||||
connection_info.session_id
|
||||
);
|
||||
|
||||
// 创建 ProjectAndAgentInfo
|
||||
let lifecycle_handle =
|
||||
Some(connection_info.lifecycle_guard.clone() as Arc<dyn AgentLifecycle>);
|
||||
let now = Utc::now();
|
||||
let agent_info = ProjectAndAgentInfo {
|
||||
project_id: project_id.clone(),
|
||||
session_id: connection_info.session_id.clone(),
|
||||
prompt_tx: connection_info.prompt_tx,
|
||||
cancel_tx: connection_info.cancel_tx,
|
||||
model_provider: model_provider.clone(),
|
||||
request_id: None,
|
||||
status: AgentStatus::Idle,
|
||||
last_activity: now,
|
||||
created_at: now,
|
||||
stop_handle: lifecycle_handle,
|
||||
};
|
||||
|
||||
// 返回 agent_info(不插入 registry,由调用方处理)
|
||||
Ok(agent_info.into())
|
||||
}
|
||||
|
||||
/// 获取或创建会话 (SACP 版本)
|
||||
///
|
||||
/// 如果会话已存在且模型配置未变化,则复用;否则创建新会话
|
||||
///
|
||||
/// # 优化说明
|
||||
/// 获取或创建会话
|
||||
///
|
||||
/// # 参数
|
||||
/// - `project_id`: 项目 ID
|
||||
/// - `project_path`: 项目路径
|
||||
/// - `session_id_hint`: 会话 ID 提示(用于恢复现有会话)
|
||||
/// - `model_provider`: 模型提供者配置
|
||||
/// - `start_config`: Agent 启动配置
|
||||
/// - `service_uuid`: 与此 Agent 关联的唯一 UUID
|
||||
///
|
||||
/// # 返回值
|
||||
/// - `R::Entry`: 会话条目
|
||||
/// - `bool`: 是否是新创建的会话
|
||||
///
|
||||
/// # 并发安全性
|
||||
///
|
||||
/// 使用"检查-创建-插入"三阶段模式避免在持有 entry 期间调用 `.await`:
|
||||
///
|
||||
/// 1. **快速检查**:检查会话是否已存在且有效
|
||||
/// 2. **创建会话**:如果需要创建,在**不持有锁**的情况下创建会话(.await)
|
||||
/// 3. **原子性插入**:使用 entry API 原子性插入,如果其他线程已创建则使用已存在的
|
||||
///
|
||||
/// 这样确保:
|
||||
/// - 不会在持有 DashMap entry 期间跨越 await 点
|
||||
/// - 同一 project_id 最多只会创建一个会话
|
||||
/// - 高并发下不会阻塞其他 project_id 的访问(DashMap 分段锁特性)
|
||||
pub async fn get_or_create_session(
|
||||
&self,
|
||||
project_id: &str,
|
||||
project_path: PathBuf,
|
||||
session_id_hint: Option<String>, // 🔥 修复:使用此参数查找现有会话
|
||||
model_provider: Option<ModelProviderConfig>,
|
||||
start_config: AgentStartConfig,
|
||||
service_uuid: Option<String>,
|
||||
) -> Result<(R::Entry, bool)> {
|
||||
use dashmap::mapref::entry::Entry;
|
||||
|
||||
let project_id_key = project_id.to_string();
|
||||
|
||||
// 🆕 第一阶段:如果提供了 session_id_hint,先尝试通过 session_id 查找现有会话
|
||||
// 🔥 优化:使用 get_entry_by_session() 避免两次调用之间的竞态窗口
|
||||
if let Some(ref hint_sid) = session_id_hint {
|
||||
// 一次性查询:通过 session_id 直接获取 agent_info
|
||||
if let Some(existing) = self.registry.get_entry_by_session(hint_sid) {
|
||||
// 验证 project_id 是否匹配(防御性编程)
|
||||
if existing.project_id() == project_id {
|
||||
let channel_closed = existing.is_channel_closed();
|
||||
let model_changed = existing.is_model_config_changed(&model_provider);
|
||||
|
||||
if !channel_closed && !model_changed {
|
||||
info!(
|
||||
"[SESSION] Reusing existing session via session_id_hint: project_id={}, session_id={}",
|
||||
project_id, hint_sid
|
||||
);
|
||||
return Ok((existing, false));
|
||||
}
|
||||
|
||||
if channel_closed {
|
||||
info!(
|
||||
"⚠️ [SESSION] Session channel closed for session_id_hint, need to rebuild: project_id={}, session_id={}",
|
||||
project_id, hint_sid
|
||||
);
|
||||
}
|
||||
if model_changed {
|
||||
info!(
|
||||
"🔄 [SESSION] Model config changed for session_id_hint, need to rebuild: project_id={}, session_id={}",
|
||||
project_id, hint_sid
|
||||
);
|
||||
}
|
||||
// 会话无效,继续后续逻辑(可能需要重建)
|
||||
} else {
|
||||
info!(
|
||||
"⚠️ [SESSION] session_id_hint belongs to different project: hint_project={}, current_project={}, session_id={}",
|
||||
existing.project_id(),
|
||||
project_id,
|
||||
hint_sid
|
||||
);
|
||||
// session_id 属于其他 project,不能复用,继续后续逻辑
|
||||
}
|
||||
} else {
|
||||
info!(
|
||||
"🔍 [SESSION] session_id_hint does not exist in registry: session_id={}",
|
||||
hint_sid
|
||||
);
|
||||
// session_id 不存在,继续后续逻辑
|
||||
}
|
||||
}
|
||||
|
||||
// 第一阶段:快速检查现有会话
|
||||
if let Entry::Occupied(occupied_entry) = self.registry.entry(project_id_key.clone()) {
|
||||
let existing = occupied_entry.get();
|
||||
|
||||
// 🔥 显式检查 Pending 状态 - PendingGuard 创建的占位符需要被替换
|
||||
if *existing.status() == AgentStatus::Pending {
|
||||
info!(
|
||||
"🔄 [SESSION] Detected Pending placeholder, preparing to replace: project_id={}",
|
||||
project_id
|
||||
);
|
||||
// 显式释放 entry 锁
|
||||
drop(occupied_entry);
|
||||
|
||||
// 第二阶段:创建新会话(不持有锁)
|
||||
let new_session = self
|
||||
.create_session_internal(
|
||||
project_id_key.clone(),
|
||||
project_path,
|
||||
model_provider,
|
||||
start_config,
|
||||
service_uuid,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// 第三阶段:原子性替换(插入新会话,覆盖 pending 占位符)
|
||||
match self.registry.entry(project_id_key.clone()) {
|
||||
Entry::Vacant(entry) => {
|
||||
let new_session_id = new_session.session_id().to_string();
|
||||
entry.insert(new_session.clone());
|
||||
info!(
|
||||
"✅ [SESSION] Inserted new session to replace pending: project_id={}, session_id={}",
|
||||
project_id, new_session_id
|
||||
);
|
||||
return Ok((new_session, true));
|
||||
}
|
||||
Entry::Occupied(mut entry) => {
|
||||
// 检查是 Pending 占位符还是其他线程创建的真实会话
|
||||
let current_session = entry.get();
|
||||
let current_status = current_session.status();
|
||||
|
||||
if *current_status == AgentStatus::Pending {
|
||||
// 仍然是 Pending 占位符,替换它
|
||||
let old_session_id = current_session.session_id().to_string();
|
||||
let new_session_id = new_session.session_id().to_string();
|
||||
entry.insert(new_session.clone());
|
||||
info!(
|
||||
"✅ [SESSION] Replaced pending placeholder: project_id={}, {} → {}",
|
||||
project_id, old_session_id, new_session_id
|
||||
);
|
||||
return Ok((new_session, true));
|
||||
} else {
|
||||
// 其他线程已经创建了真实会话,使用已存在的(丢弃我们创建的)
|
||||
let existing_session = current_session.clone();
|
||||
let created_session_id = new_session.session_id().to_string();
|
||||
let existing_session_id = existing_session.session_id().to_string();
|
||||
info!(
|
||||
"🔄 [SESSION] Detected concurrent creation (other thread already created real session): project_id={}, discarded session_id={}, using session_id={}",
|
||||
project_id, created_session_id, existing_session_id
|
||||
);
|
||||
// new_session 会被 drop,AgentLifecycleGuard 会清理 Agent 进程
|
||||
return Ok((existing_session, false));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let channel_closed = existing.is_channel_closed();
|
||||
let model_changed = existing.is_model_config_changed(&model_provider);
|
||||
|
||||
if !channel_closed && !model_changed {
|
||||
info!("Reuse Agent session, project ID: {}", project_id);
|
||||
return Ok((existing.clone(), false));
|
||||
}
|
||||
|
||||
// 需要重建会话,先克隆必要数据并释放锁
|
||||
let session_id_str = existing.session_id().to_string();
|
||||
drop(occupied_entry); // 显式释放 entry 锁
|
||||
|
||||
if channel_closed {
|
||||
info!(
|
||||
"⚠️ Detected session channel closed (Agent process exited), rebuilding session, project ID: {}, old session_id: {}",
|
||||
project_id, session_id_str
|
||||
);
|
||||
}
|
||||
if model_changed {
|
||||
info!(
|
||||
"Model config changed, restarting Agent session, project ID: {}, old session_id: {}",
|
||||
project_id, session_id_str
|
||||
);
|
||||
}
|
||||
|
||||
// 第二阶段:在不持有锁的情况下创建新会话
|
||||
let new_session = self
|
||||
.create_session_internal(
|
||||
project_id_key.clone(),
|
||||
project_path,
|
||||
model_provider,
|
||||
start_config,
|
||||
service_uuid,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// 第三阶段:原子性插入(使用 entry API 防止并发创建)
|
||||
match self.registry.entry(project_id_key.clone()) {
|
||||
Entry::Vacant(entry) => {
|
||||
// 其他线程还没有创建,使用我们创建的会话
|
||||
let new_session_id = new_session.session_id().to_string();
|
||||
entry.insert(new_session.clone());
|
||||
info!(
|
||||
"✅ Session rebuilt, project ID: {}, new session_id: {}",
|
||||
project_id, new_session_id
|
||||
);
|
||||
return Ok((new_session, true));
|
||||
}
|
||||
Entry::Occupied(entry) => {
|
||||
// 其他线程已经创建了会话,使用已存在的(丢弃我们创建的)
|
||||
let existing_session = entry.get().clone();
|
||||
let created_session_id = new_session.session_id().to_string();
|
||||
let existing_session_id = existing_session.session_id().to_string();
|
||||
info!(
|
||||
"🔄 [SESSION] Detected concurrent creation, using other thread's session: project_id={}, discarded session_id={}, using session_id={}",
|
||||
project_id, created_session_id, existing_session_id
|
||||
);
|
||||
// new_session 会被 drop,Session 的 Drop 实现会清理 Agent 进程
|
||||
return Ok((existing_session, false));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 会话不存在,需要创建新会话
|
||||
info!(
|
||||
"Session not found, creating new session, project ID: {}",
|
||||
project_id
|
||||
);
|
||||
|
||||
// 第二阶段:在不持有锁的情况下创建新会话
|
||||
let new_session = self
|
||||
.create_session_internal(
|
||||
project_id_key.clone(),
|
||||
project_path,
|
||||
model_provider,
|
||||
start_config,
|
||||
service_uuid,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// 第三阶段:原子性插入
|
||||
match self.registry.entry(project_id_key.clone()) {
|
||||
Entry::Vacant(entry) => {
|
||||
// 其他线程还没有创建,使用我们创建的会话
|
||||
let session_id_str = new_session.session_id().to_string();
|
||||
entry.insert(new_session.clone());
|
||||
info!(
|
||||
"✅ New session created, project ID: {}, session_id: {}",
|
||||
project_id, session_id_str
|
||||
);
|
||||
Ok((new_session, true))
|
||||
}
|
||||
Entry::Occupied(entry) => {
|
||||
// 其他线程已经创建了会话,使用已存在的(丢弃我们创建的)
|
||||
let existing_session = entry.get().clone();
|
||||
let created_session_id = new_session.session_id().to_string();
|
||||
let existing_session_id = existing_session.session_id().to_string();
|
||||
info!(
|
||||
"🔄 [SESSION] Detected concurrent creation, using session created by other thread: project_id={}, discarding session_id={}, using session_id={}",
|
||||
project_id, created_session_id, existing_session_id
|
||||
);
|
||||
// new_session 会被 drop,Session 的 Drop 实现会清理 Agent 进程
|
||||
Ok((existing_session, false))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 发送 Prompt 到指定会话(仅文本)
|
||||
///
|
||||
/// 使用有界通道,提供背压保护。如果通道已满,会异步等待直到有空间。
|
||||
pub async fn send_text_prompt(&self, project_id: &str, prompt: &PromptMessage) -> Result<()> {
|
||||
let session = self
|
||||
.get_session(project_id)
|
||||
.ok_or_else(|| anyhow::anyhow!("Session not found: {}", project_id))?;
|
||||
|
||||
let prompt_request = Self::build_text_prompt_request(prompt, session.session_id().clone())?;
|
||||
|
||||
session
|
||||
.prompt_tx()
|
||||
.send(prompt_request)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
error!("send Prompt requestfailed: {:?}", e);
|
||||
anyhow::anyhow!("Failed to send Prompt request: {:?}", e)
|
||||
})?;
|
||||
|
||||
info!("Prompt request already sent, project ID: {}", project_id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 发送 Prompt 请求到指定会话
|
||||
///
|
||||
/// 使用有界通道,提供背压保护。如果通道已满,会异步等待直到有空间。
|
||||
pub async fn send_prompt_request(
|
||||
&self,
|
||||
project_id: &str,
|
||||
prompt_request: PromptRequest,
|
||||
) -> Result<()> {
|
||||
let session = self
|
||||
.get_session(project_id)
|
||||
.ok_or_else(|| anyhow::anyhow!("Session not found: {}", project_id))?;
|
||||
|
||||
session
|
||||
.prompt_tx()
|
||||
.send(prompt_request)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
error!("send Prompt requestfailed: {:?}", e);
|
||||
anyhow::anyhow!("Failed to send Prompt request: {:?}", e)
|
||||
})?;
|
||||
|
||||
info!("Prompt request already sent, project ID: {}", project_id);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<N: SessionNotifier + 'static, R: SessionRegistry> std::fmt::Debug for AcpSessionManager<N, R> {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("AcpSessionManager")
|
||||
.field("session_count", &self.registry.count())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
221
qiming-rcoder/crates/agent_abstraction/src/session/worker.rs
Normal file
221
qiming-rcoder/crates/agent_abstraction/src/session/worker.rs
Normal file
@@ -0,0 +1,221 @@
|
||||
//! Agent Worker 抽象
|
||||
//!
|
||||
//! 定义请求处理的通用接口
|
||||
|
||||
use std::{path::PathBuf, sync::Arc};
|
||||
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
use dashmap::DashMap;
|
||||
use agent_client_protocol::schema::{ContentBlock, PromptRequest};
|
||||
use shared_types::{
|
||||
AgentLifecycle, CancelNotificationRequestWrapper, ModelProviderConfig, ServiceType,
|
||||
};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::PromptMessage;
|
||||
|
||||
/// Worker 请求
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WorkerRequest {
|
||||
/// Agent 提示消息
|
||||
pub prompt_message: PromptMessage,
|
||||
/// 模型提供商配置
|
||||
pub model_provider: Option<ModelProviderConfig>,
|
||||
/// 预处理的附件内容块
|
||||
/// 由 agent_runner 使用 ContentBuilder 预处理
|
||||
pub attachment_blocks: Option<Vec<ContentBlock>>,
|
||||
/// 🔥 关联的 service UUID(用于 API 密钥管理)
|
||||
pub service_uuid: Option<String>,
|
||||
/// 🔥 共享的 API 密钥管理器(用于自动清理)
|
||||
pub shared_api_key_manager: Option<Arc<DashMap<String, ModelProviderConfig>>>,
|
||||
}
|
||||
|
||||
impl WorkerRequest {
|
||||
/// 创建新的 Worker 请求
|
||||
pub fn new(prompt_message: PromptMessage, model_provider: Option<ModelProviderConfig>) -> Self {
|
||||
Self {
|
||||
prompt_message,
|
||||
model_provider,
|
||||
attachment_blocks: None,
|
||||
service_uuid: None,
|
||||
shared_api_key_manager: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 设置 service_uuid
|
||||
pub fn with_service_uuid(mut self, service_uuid: Option<String>) -> Self {
|
||||
self.service_uuid = service_uuid;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 shared_api_key_manager
|
||||
pub fn with_key_manager(
|
||||
mut self,
|
||||
key_manager: Option<Arc<DashMap<String, ModelProviderConfig>>>,
|
||||
) -> Self {
|
||||
self.shared_api_key_manager = key_manager;
|
||||
self
|
||||
}
|
||||
|
||||
/// 获取项目 ID
|
||||
pub fn project_id(&self) -> &str {
|
||||
&self.prompt_message.project_id
|
||||
}
|
||||
|
||||
/// 获取项目路径
|
||||
pub fn project_path(&self) -> &PathBuf {
|
||||
&self.prompt_message.project_path
|
||||
}
|
||||
|
||||
/// 获取请求 ID
|
||||
pub fn request_id(&self) -> &str {
|
||||
&self.prompt_message.request_id
|
||||
}
|
||||
}
|
||||
|
||||
/// 会话句柄
|
||||
///
|
||||
/// 用于传递会话句柄给 agent_runner 更新全局 MAP
|
||||
#[derive(Clone)]
|
||||
pub struct SessionHandles {
|
||||
pub prompt_tx: mpsc::Sender<PromptRequest>,
|
||||
pub cancel_tx: mpsc::Sender<CancelNotificationRequestWrapper>,
|
||||
pub lifecycle_handle: Option<Arc<dyn AgentLifecycle>>,
|
||||
}
|
||||
|
||||
// 手动实现 Debug,跳过不支持 Debug 的 lifecycle_handle
|
||||
impl std::fmt::Debug for SessionHandles {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("SessionHandles")
|
||||
.field("prompt_tx", &"Sender<PromptRequest>")
|
||||
.field("cancel_tx", &"Sender<CancelNotificationRequestWrapper>")
|
||||
.field(
|
||||
"lifecycle_handle",
|
||||
&self
|
||||
.lifecycle_handle
|
||||
.as_ref()
|
||||
.map(|_| "Some(AgentLifecycle)"),
|
||||
)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Worker 响应
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WorkerResponse {
|
||||
/// 项目 ID
|
||||
pub project_id: String,
|
||||
/// 会话 ID
|
||||
pub session_id: String,
|
||||
/// 错误信息(如果有)
|
||||
pub error: Option<String>,
|
||||
/// 请求 ID
|
||||
pub request_id: Option<String>,
|
||||
/// 服务类型
|
||||
pub service_type: ServiceType,
|
||||
/// 标识是否是新创建的会话
|
||||
pub is_new_session: bool,
|
||||
/// 会话句柄(仅新会话时有值)
|
||||
pub session_handles: Option<SessionHandles>,
|
||||
}
|
||||
|
||||
impl WorkerResponse {
|
||||
/// 创建新会话的成功响应
|
||||
pub fn new_session_success(
|
||||
project_id: String,
|
||||
session_id: String,
|
||||
request_id: Option<String>,
|
||||
service_type: ServiceType,
|
||||
handles: SessionHandles,
|
||||
) -> Self {
|
||||
Self {
|
||||
project_id,
|
||||
session_id,
|
||||
error: None,
|
||||
request_id,
|
||||
service_type,
|
||||
is_new_session: true,
|
||||
session_handles: Some(handles),
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建复用会话的成功响应
|
||||
pub fn reuse_session_success(
|
||||
project_id: String,
|
||||
session_id: String,
|
||||
request_id: Option<String>,
|
||||
service_type: ServiceType,
|
||||
) -> Self {
|
||||
Self {
|
||||
project_id,
|
||||
session_id,
|
||||
error: None,
|
||||
request_id,
|
||||
service_type,
|
||||
is_new_session: false,
|
||||
session_handles: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建成功响应(旧版本,保持向后兼容)
|
||||
#[deprecated(note = "Use new_session_success or reuse_session_success instead")]
|
||||
pub fn success(
|
||||
project_id: String,
|
||||
session_id: String,
|
||||
request_id: Option<String>,
|
||||
service_type: ServiceType,
|
||||
) -> Self {
|
||||
Self {
|
||||
project_id,
|
||||
session_id,
|
||||
error: None,
|
||||
request_id,
|
||||
service_type,
|
||||
is_new_session: false,
|
||||
session_handles: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建错误响应
|
||||
pub fn error(
|
||||
project_id: String,
|
||||
session_id: String,
|
||||
error: String,
|
||||
request_id: Option<String>,
|
||||
service_type: ServiceType,
|
||||
) -> Self {
|
||||
Self {
|
||||
project_id,
|
||||
session_id,
|
||||
error: Some(error),
|
||||
request_id,
|
||||
service_type,
|
||||
is_new_session: false,
|
||||
session_handles: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查是否成功
|
||||
pub fn is_success(&self) -> bool {
|
||||
self.error.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
/// Agent Worker Trait
|
||||
///
|
||||
/// 定义请求处理的抽象接口,允许不同的实现策略
|
||||
#[async_trait]
|
||||
pub trait AgentWorker: Send + Sync {
|
||||
/// 处理单个请求
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `request` - Worker 请求
|
||||
///
|
||||
/// # Returns
|
||||
/// 处理结果
|
||||
async fn process_request(&self, request: WorkerRequest) -> Result<WorkerResponse>;
|
||||
|
||||
/// 获取 Worker 名称(用于日志和调试)
|
||||
fn name(&self) -> &'static str;
|
||||
}
|
||||
413
qiming-rcoder/crates/agent_abstraction/src/traits/agent.rs
Normal file
413
qiming-rcoder/crates/agent_abstraction/src/traits/agent.rs
Normal file
@@ -0,0 +1,413 @@
|
||||
//! Agent trait definition.
|
||||
|
||||
use agent_client_protocol::schema::McpServer;
|
||||
use tracing::{debug, info};
|
||||
|
||||
/// Agent startup configuration
|
||||
///
|
||||
/// Contains additional configuration information required to start an Agent,
|
||||
/// such as system prompts, MCP server configurations, etc.
|
||||
/// These configurations are passed to the Agent via ACP protocol's NewSessionRequest.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AgentStartConfig {
|
||||
/// System prompt (passed via meta.systemPrompt.append)
|
||||
///
|
||||
/// Uses the `_meta.systemPrompt.append` mode of the ACP protocol to preserve
|
||||
/// the base capabilities of the claude_code preset while appending custom system prompts.
|
||||
pub system_prompt: Option<String>,
|
||||
|
||||
/// MCP server configuration
|
||||
pub mcp_servers: Vec<McpServer>,
|
||||
|
||||
/// Additional meta fields
|
||||
///
|
||||
/// Can pass any additional configuration information to the Agent
|
||||
pub extra_meta: Option<serde_json::Map<String, serde_json::Value>>,
|
||||
|
||||
/// Service type (for loading corresponding config, required)
|
||||
pub service_type: shared_types::ServiceType,
|
||||
|
||||
/// Session ID for resuming sessions
|
||||
///
|
||||
/// When resuming a previous session, pass the previous session_id,
|
||||
/// which will be passed to the Agent via `_meta.claudeCode.options.resume`.
|
||||
pub resume_session_id: Option<String>,
|
||||
|
||||
/// 🆕 Agent server config override (for custom Agent startup command)
|
||||
///
|
||||
/// When the user request specifies agent_server config, use this config
|
||||
/// to override the default config. Includes command, args, env, etc.
|
||||
pub agent_server_override: Option<shared_types::ChatAgentServerConfig>,
|
||||
|
||||
/// 🆕 ACP session creation timeout (seconds), default 100
|
||||
///
|
||||
/// Maximum wait time for Agent to create a new session
|
||||
/// (may require more time when there are many MCP tools)
|
||||
pub acp_session_create_timeout_secs: Option<u64>,
|
||||
|
||||
/// 🆕 Agent cancel call timeout (seconds), default 10
|
||||
///
|
||||
/// Maximum wait time for Agent internal cancel operations
|
||||
pub agent_cancel_timeout_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl AgentStartConfig {
|
||||
/// Create a new AgentStartConfig
|
||||
///
|
||||
/// # Arguments
|
||||
/// - `service_type`: Service type (required)
|
||||
pub fn new(service_type: shared_types::ServiceType) -> Self {
|
||||
Self {
|
||||
system_prompt: None,
|
||||
mcp_servers: Vec::new(),
|
||||
extra_meta: None,
|
||||
service_type,
|
||||
resume_session_id: None,
|
||||
agent_server_override: None,
|
||||
acp_session_create_timeout_secs: None,
|
||||
agent_cancel_timeout_secs: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set system prompt
|
||||
pub fn with_system_prompt(mut self, system_prompt: String) -> Self {
|
||||
self.system_prompt = Some(system_prompt);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set MCP server configuration
|
||||
pub fn with_mcp_servers(mut self, mcp_servers: Vec<McpServer>) -> Self {
|
||||
self.mcp_servers = mcp_servers;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set additional meta fields
|
||||
pub fn with_extra_meta(
|
||||
mut self,
|
||||
extra_meta: serde_json::Map<String, serde_json::Value>,
|
||||
) -> Self {
|
||||
self.extra_meta = Some(extra_meta);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set service type
|
||||
pub fn with_service_type(mut self, service_type: shared_types::ServiceType) -> Self {
|
||||
self.service_type = service_type;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set session ID for resuming sessions
|
||||
pub fn with_resume_session_id(mut self, session_id: String) -> Self {
|
||||
self.resume_session_id = Some(session_id);
|
||||
self
|
||||
}
|
||||
|
||||
/// 🆕 Set Agent server config override
|
||||
pub fn with_agent_server_override(
|
||||
mut self,
|
||||
agent_server: shared_types::ChatAgentServerConfig,
|
||||
) -> Self {
|
||||
self.agent_server_override = Some(agent_server);
|
||||
self
|
||||
}
|
||||
|
||||
/// 🆕 Set ACP session creation timeout
|
||||
pub fn with_acp_session_create_timeout(mut self, timeout_secs: u64) -> Self {
|
||||
self.acp_session_create_timeout_secs = Some(timeout_secs);
|
||||
self
|
||||
}
|
||||
|
||||
/// 🆕 Set Agent cancel call timeout
|
||||
pub fn with_agent_cancel_timeout(mut self, timeout_secs: u64) -> Self {
|
||||
self.agent_cancel_timeout_secs = Some(timeout_secs);
|
||||
self
|
||||
}
|
||||
|
||||
/// Build the meta field for NewSessionRequest
|
||||
///
|
||||
/// Merges the system prompt and additional meta fields into a JSON Map
|
||||
/// for passing to the ACP protocol's NewSessionRequest.
|
||||
///
|
||||
/// # Returns
|
||||
/// Returns a Meta object containing `systemPrompt: { append: "..." }`
|
||||
pub fn build_meta(&self) -> serde_json::Map<String, serde_json::Value> {
|
||||
let mut meta = serde_json::Map::new();
|
||||
|
||||
// Add system prompt (using append mode)
|
||||
if let Some(ref system_prompt) = self.system_prompt {
|
||||
let mut system_prompt_obj = serde_json::Map::new();
|
||||
system_prompt_obj.insert(
|
||||
"append".to_string(),
|
||||
serde_json::Value::String(system_prompt.clone()),
|
||||
);
|
||||
meta.insert(
|
||||
"systemPrompt".to_string(),
|
||||
serde_json::Value::Object(system_prompt_obj),
|
||||
);
|
||||
info!(
|
||||
"[ACP] Sending system_prompt to agent: length={}",
|
||||
system_prompt.len()
|
||||
);
|
||||
debug!(
|
||||
"[ACP] system_prompt content (first 200 chars): \"{}\"",
|
||||
&system_prompt[..system_prompt.len().min(200)]
|
||||
);
|
||||
} else {
|
||||
info!("[ACP] No system_prompt to send to agent");
|
||||
}
|
||||
|
||||
// Build claudeCode.options structure
|
||||
// settingSources 控制 Claude Code Agent 从哪些来源加载配置和 skills。
|
||||
// 可选值: "user"(全局 ~/.claude/), "project"(项目 .claude/), "local"
|
||||
// 只加载 "project" 级别配置,避免全局 user 配置覆盖 ANTHROPIC_BASE_URL,
|
||||
// 同时允许项目目录下的 .claude/skills/ 自定义技能生效。
|
||||
// Refer to the TypeScript code on the Agent side:
|
||||
// resume: (params._meta as NewSessionMeta | undefined)?.claudeCode?.options?.resume
|
||||
let mut options = serde_json::Map::new();
|
||||
|
||||
// Only load project-level settings (skills), block global settings
|
||||
options.insert(
|
||||
"settingSources".to_string(),
|
||||
serde_json::Value::Array(vec![serde_json::Value::String("project".to_string())]),
|
||||
);
|
||||
|
||||
// Add resume session_id if present
|
||||
if let Some(ref session_id) = self.resume_session_id {
|
||||
options.insert(
|
||||
"resume".to_string(),
|
||||
serde_json::Value::String(session_id.clone()),
|
||||
);
|
||||
}
|
||||
|
||||
let mut claude_code = serde_json::Map::new();
|
||||
claude_code.insert("options".to_string(), serde_json::Value::Object(options));
|
||||
meta.insert(
|
||||
"claudeCode".to_string(),
|
||||
serde_json::Value::Object(claude_code),
|
||||
);
|
||||
|
||||
// Merge additional meta fields
|
||||
if let Some(ref extra) = self.extra_meta {
|
||||
for (key, value) in extra {
|
||||
// Don't overwrite if key already exists (e.g., systemPrompt)
|
||||
if !meta.contains_key(key) {
|
||||
meta.insert(key.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
meta
|
||||
}
|
||||
|
||||
/// Check if there is a system prompt
|
||||
pub fn has_system_prompt(&self) -> bool {
|
||||
self.system_prompt.is_some()
|
||||
}
|
||||
|
||||
/// Check if there are MCP server configurations
|
||||
pub fn has_mcp_servers(&self) -> bool {
|
||||
!self.mcp_servers.is_empty()
|
||||
}
|
||||
|
||||
/// Build meta without resume parameter
|
||||
///
|
||||
/// Used when the list_sessions check finds the session does not exist.
|
||||
/// Similar to build_meta(), but without the claudeCode.options.resume field.
|
||||
///
|
||||
/// # Use Cases
|
||||
/// - User passed a session_id, but verification via list_sessions API found it doesn't exist
|
||||
/// - Create a new session while preserving other configs (system prompt, extra_meta, etc.)
|
||||
pub fn build_meta_without_resume(&self) -> serde_json::Map<String, serde_json::Value> {
|
||||
let mut meta = serde_json::Map::new();
|
||||
|
||||
// Only add system prompt (using append mode, consistent with build_meta)
|
||||
if let Some(ref system_prompt) = self.system_prompt {
|
||||
let mut system_prompt_obj = serde_json::Map::new();
|
||||
system_prompt_obj.insert(
|
||||
"append".to_string(),
|
||||
serde_json::Value::String(system_prompt.clone()),
|
||||
);
|
||||
meta.insert(
|
||||
"systemPrompt".to_string(),
|
||||
serde_json::Value::Object(system_prompt_obj),
|
||||
);
|
||||
info!(
|
||||
"[ACP] Sending system_prompt to agent (no resume): length={}",
|
||||
system_prompt.len()
|
||||
);
|
||||
debug!(
|
||||
"[ACP] system_prompt content (first 200 chars): \"{}\"",
|
||||
&system_prompt[..system_prompt.len().min(200)]
|
||||
);
|
||||
}
|
||||
|
||||
// Only load project-level settings (consistent with build_meta)
|
||||
let mut options = serde_json::Map::new();
|
||||
options.insert(
|
||||
"settingSources".to_string(),
|
||||
serde_json::Value::Array(vec![serde_json::Value::String("project".to_string())]),
|
||||
);
|
||||
|
||||
let mut claude_code = serde_json::Map::new();
|
||||
claude_code.insert("options".to_string(), serde_json::Value::Object(options));
|
||||
meta.insert(
|
||||
"claudeCode".to_string(),
|
||||
serde_json::Value::Object(claude_code),
|
||||
);
|
||||
|
||||
// Merge additional meta fields (don't overwrite existing keys)
|
||||
if let Some(ref extra) = self.extra_meta {
|
||||
for (key, value) in extra {
|
||||
if !meta.contains_key(key) {
|
||||
meta.insert(key.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
meta
|
||||
}
|
||||
}
|
||||
|
||||
/// Agent common Prompt message
|
||||
///
|
||||
/// This is the core data structure of the Agent abstraction layer, containing only
|
||||
/// the core information needed for the Agent to execute tasks.
|
||||
/// Does not contain business layer configurations (such as model_provider, mcp configs, etc.).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptMessage {
|
||||
/// Main prompt text entered by the user
|
||||
pub content: String,
|
||||
|
||||
/// Project ID
|
||||
pub project_id: String,
|
||||
|
||||
/// Project working directory path
|
||||
pub project_path: std::path::PathBuf,
|
||||
|
||||
/// Session ID (optional, if None the Agent will create a new session)
|
||||
pub session_id: Option<String>,
|
||||
|
||||
/// Request tracking ID
|
||||
pub request_id: String,
|
||||
|
||||
/// Attachment list (supports text, images, audio, documents, etc.)
|
||||
pub attachments: Vec<shared_types::Attachment>,
|
||||
|
||||
/// Data source attachments (JSON string array)
|
||||
pub data_source_attachments: Vec<String>,
|
||||
|
||||
/// Service type
|
||||
pub service_type: shared_types::ServiceType,
|
||||
|
||||
// === New fields (v2) ===
|
||||
/// System prompt override
|
||||
///
|
||||
/// If provided, will override the default system prompt configuration
|
||||
pub system_prompt_override: Option<String>,
|
||||
|
||||
/// User prompt template override
|
||||
///
|
||||
/// If provided, will use this template to replace the `{user_prompt}` variable
|
||||
pub user_prompt_template_override: Option<String>,
|
||||
|
||||
/// Agent runtime config override (MCP servers, etc.)
|
||||
///
|
||||
/// Contains Agent server configuration and MCP server configuration
|
||||
pub agent_config_override: Option<shared_types::ChatAgentConfig>,
|
||||
}
|
||||
|
||||
impl PromptMessage {
|
||||
/// Create a new PromptMessage
|
||||
pub fn new(
|
||||
content: String,
|
||||
project_id: String,
|
||||
project_path: std::path::PathBuf,
|
||||
request_id: String,
|
||||
service_type: shared_types::ServiceType,
|
||||
) -> Self {
|
||||
Self {
|
||||
content,
|
||||
project_id,
|
||||
project_path,
|
||||
session_id: None,
|
||||
request_id,
|
||||
attachments: Vec::new(),
|
||||
data_source_attachments: Vec::new(),
|
||||
service_type,
|
||||
// New fields default to None
|
||||
system_prompt_override: None,
|
||||
user_prompt_template_override: None,
|
||||
agent_config_override: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set session ID
|
||||
pub fn with_session_id(mut self, session_id: Option<String>) -> Self {
|
||||
self.session_id = session_id;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set attachment list
|
||||
pub fn with_attachments(mut self, attachments: Vec<shared_types::Attachment>) -> Self {
|
||||
self.attachments = attachments;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set data source attachments
|
||||
pub fn with_data_source_attachments(mut self, data_source_attachments: Vec<String>) -> Self {
|
||||
self.data_source_attachments = data_source_attachments;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set system prompt override
|
||||
pub fn with_system_prompt_override(mut self, system_prompt: Option<String>) -> Self {
|
||||
self.system_prompt_override = system_prompt;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set user prompt template override
|
||||
pub fn with_user_prompt_template_override(mut self, template: Option<String>) -> Self {
|
||||
self.user_prompt_template_override = template;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set Agent config override
|
||||
pub fn with_agent_config_override(
|
||||
mut self,
|
||||
config: Option<shared_types::ChatAgentConfig>,
|
||||
) -> Self {
|
||||
self.agent_config_override = config;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert from ChatPrompt to PromptMessage
|
||||
impl From<shared_types::ChatPrompt> for PromptMessage {
|
||||
fn from(chat_prompt: shared_types::ChatPrompt) -> Self {
|
||||
info!(
|
||||
"[agent_abstraction] Converting ChatPrompt to PromptMessage, project_id={:?}, session_id={:?}, has_model_provider={}, has_agent_config_override={}",
|
||||
chat_prompt.project_id,
|
||||
chat_prompt.session_id,
|
||||
chat_prompt.model_provider.is_some(),
|
||||
chat_prompt.agent_config_override.is_some(),
|
||||
);
|
||||
|
||||
Self {
|
||||
content: chat_prompt.prompt,
|
||||
project_id: chat_prompt.project_id,
|
||||
project_path: chat_prompt.project_path,
|
||||
session_id: chat_prompt.session_id,
|
||||
request_id: chat_prompt.request_id.unwrap_or_else(|| {
|
||||
// Generate one if ChatPrompt doesn't have request_id
|
||||
uuid::Uuid::new_v4().to_string().replace("-", "")
|
||||
}),
|
||||
attachments: chat_prompt.attachments,
|
||||
data_source_attachments: chat_prompt.data_source_attachments,
|
||||
service_type: chat_prompt.service_type,
|
||||
// Map new fields
|
||||
system_prompt_override: chat_prompt.system_prompt_override,
|
||||
user_prompt_template_override: chat_prompt.user_prompt_template_override,
|
||||
agent_config_override: chat_prompt.agent_config_override,
|
||||
}
|
||||
}
|
||||
}
|
||||
9
qiming-rcoder/crates/agent_abstraction/src/traits/mod.rs
Normal file
9
qiming-rcoder/crates/agent_abstraction/src/traits/mod.rs
Normal file
@@ -0,0 +1,9 @@
|
||||
//! Trait definitions for agent abstraction.
|
||||
|
||||
pub mod agent;
|
||||
pub mod session_notifier;
|
||||
pub mod session_registry;
|
||||
|
||||
pub use agent::{AgentStartConfig, PromptMessage};
|
||||
pub use session_notifier::SessionNotifier;
|
||||
pub use session_registry::SessionRegistry;
|
||||
@@ -0,0 +1,113 @@
|
||||
//! # Session Notifier Trait
|
||||
//!
|
||||
//! 定义会话通知的抽象接口,用于推送 SSE 消息到前端。
|
||||
//! `agent_runner` 模块实现此 trait 来完成实际的消息推送。
|
||||
//!
|
||||
//! ## 预期实现
|
||||
//!
|
||||
//! 此 trait 预期在 `agent_runner` 中实现:
|
||||
//!
|
||||
//! - **实现类**: `agent_runner::service::SseSessionNotifier`
|
||||
//! - **包装类**: `agent_runner::service::StateAwareNotifier`(注入 project_id)
|
||||
//!
|
||||
//! ## 架构说明
|
||||
//!
|
||||
//! ```text
|
||||
//! agent_abstraction agent_runner
|
||||
//! ┌─────────────────┐ ┌─────────────────────┐
|
||||
//! │ SessionNotifier │◄────────────────│ SseSessionNotifier │
|
||||
//! │ (trait) │ implements │ (struct) │
|
||||
//! └────────┬────────┘ └──────────┬──────────┘
|
||||
//! │ │
|
||||
//! │ ┌──────────▼──────────┐
|
||||
//! ┌────────▼────────┐ │ StateAwareNotifier │
|
||||
//! │AcpSessionManager│◄────────────────│ (wrapper, 注入 │
|
||||
//! │ notifier: N │ injects │ project_id) │
|
||||
//! └─────────────────┘ └─────────────────────┘
|
||||
//! ```
|
||||
//!
|
||||
//! ## 方法使用场景
|
||||
//!
|
||||
//! | 方法 | 触发时机 | 推送内容 |
|
||||
//! |------|---------|---------|
|
||||
//! | `notify_prompt_start` | Agent 开始处理请求 | 会话开始通知 |
|
||||
//! | `notify_prompt_end` | Agent 完成处理 | 会话结束通知(含停止原因)|
|
||||
//! | `notify_prompt_error` | Agent 处理出错 | 错误通知 |
|
||||
//! | `notify_session_update` | Agent 输出内容 | 内容块更新(文本、工具调用等)|
|
||||
//! | `notify` | 通用通知 | 任意 SessionNotify 类型 |
|
||||
//!
|
||||
//! ## 使用示例
|
||||
//!
|
||||
//! ```ignore
|
||||
//! // 在 agent_runner 中定义实现
|
||||
//! pub struct SseSessionNotifier;
|
||||
//!
|
||||
//! #[async_trait]
|
||||
//! impl SessionNotifier for SseSessionNotifier {
|
||||
//! async fn notify_prompt_start(...) -> Result<...> {
|
||||
//! // 推送到 SESSION_CACHE,由 SSE handler 分发
|
||||
//! }
|
||||
//! }
|
||||
//!
|
||||
//! // 使用 StateAwareNotifier 包装,自动注入 project_id
|
||||
//! let notifier = StateAwareNotifier::new(project_id.clone());
|
||||
//! let session_manager = AcpSessionManager::new(Arc::new(notifier), registry);
|
||||
//! ```
|
||||
|
||||
use async_trait::async_trait;
|
||||
use shared_types::SessionNotify;
|
||||
|
||||
/// 会话通知器 trait
|
||||
///
|
||||
/// 提供会话消息推送的抽象接口,解耦 agent_abstraction 和具体的 SSE 实现。
|
||||
///
|
||||
/// # 设计说明
|
||||
/// - agent_abstraction 只依赖此 trait,不依赖具体的 SSE 实现
|
||||
/// - agent_runner 实现此 trait,完成实际的消息推送
|
||||
/// - 通过依赖注入的方式,在启动 prompt handler 时传入 notifier
|
||||
#[async_trait]
|
||||
pub trait SessionNotifier: Send + Sync + 'static {
|
||||
/// 推送会话开始通知
|
||||
async fn notify_prompt_start(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
|
||||
|
||||
/// 推送会话结束通知
|
||||
async fn notify_prompt_end(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
stop_reason: agent_client_protocol::schema::StopReason,
|
||||
error_message: Option<String>,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
|
||||
|
||||
/// 推送会话错误通知
|
||||
async fn notify_prompt_error(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
error: agent_client_protocol::schema::Error,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
|
||||
|
||||
/// 推送 Agent 会话更新通知
|
||||
async fn notify_session_update(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
session_update: agent_client_protocol::schema::SessionUpdate,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
|
||||
|
||||
/// 推送通用会话通知
|
||||
async fn notify(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
notify: SessionNotify,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
//! # Session Registry Trait
|
||||
//!
|
||||
//! 定义会话注册表的抽象接口,允许 `AcpSessionManager` 使用不同的存储实现。
|
||||
//! 通过依赖注入避免 `agent_abstraction` 和 `agent_runner` 之间的循环依赖。
|
||||
//!
|
||||
//! ## 预期实现
|
||||
//!
|
||||
//! 此 trait 预期在 `agent_runner` 中实现:
|
||||
//!
|
||||
//! - **实现类**: `agent_runner::service::AgentSessionRegistry`
|
||||
//! - **全局单例**: `agent_runner::service::AGENT_REGISTRY`
|
||||
//!
|
||||
//! ## 架构说明
|
||||
//!
|
||||
//! ```text
|
||||
//! agent_abstraction agent_runner
|
||||
//! ┌─────────────────┐ ┌─────────────────────┐
|
||||
//! │ SessionRegistry │◄────────────────│ AgentSessionRegistry│
|
||||
//! │ (trait) │ implements │ (struct) │
|
||||
//! └────────┬────────┘ └──────────┬──────────┘
|
||||
//! │ │
|
||||
//! │ │
|
||||
//! ┌────────▼────────┐ ┌──────────▼──────────┐
|
||||
//! │AcpSessionManager│◄────────────────│ AGENT_REGISTRY │
|
||||
//! │ registry: R │ injects │ (static LazyLock) │
|
||||
//! └─────────────────┘ └─────────────────────┘
|
||||
//! ```
|
||||
//!
|
||||
//! ## 与 AGENT_REGISTRY 的关系
|
||||
//!
|
||||
//! - **AGENT_REGISTRY**: `agent_runner` 中的全局单例,实现了此 trait
|
||||
//! - **何时直接访问 AGENT_REGISTRY**: 在 `agent_runner` 内部进行状态查询、清理任务
|
||||
//! - **何时通过 trait 访问**: 在 `agent_abstraction` 内部,通过泛型参数 `R: SessionRegistry`
|
||||
//!
|
||||
//! ## 使用示例
|
||||
//!
|
||||
//! ```ignore
|
||||
//! // 在 agent_runner 中定义实现
|
||||
//! pub struct AgentSessionRegistry {
|
||||
//! agent_info_map: DashMap<String, ProjectAndAgentInfo>,
|
||||
//! project_to_session: DashMap<String, String>,
|
||||
//! session_to_project: DashMap<String, String>,
|
||||
//! }
|
||||
//!
|
||||
//! impl SessionRegistry for AgentSessionRegistry {
|
||||
//! type Entry = ProjectAndAgentInfo;
|
||||
//! // ...
|
||||
//! }
|
||||
//!
|
||||
//! // 创建全局单例
|
||||
//! pub static AGENT_REGISTRY: LazyLock<Arc<AgentSessionRegistry>> = ...;
|
||||
//!
|
||||
//! // 注入到 AcpSessionManager
|
||||
//! let session_manager = AcpSessionManager::new(notifier, AGENT_REGISTRY.clone());
|
||||
//! ```
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::mapref::entry::Entry;
|
||||
use shared_types::SessionEntry;
|
||||
|
||||
/// 会话注册表 trait
|
||||
///
|
||||
/// 抽象 session 存储的 CRUD 操作,允许不同的实现(如 AGENT_REGISTRY)。
|
||||
///
|
||||
/// # 设计说明
|
||||
/// - `SessionEntry` trait 定义在 `shared_types`,描述单个会话条目的数据访问
|
||||
/// - `SessionRegistry` trait 定义在 `agent_abstraction`,描述会话存储的 CRUD 操作
|
||||
/// - `agent_runner` 为 `AGENT_REGISTRY` 实现 `SessionRegistry`,并注入到 `AcpSessionManager`
|
||||
///
|
||||
/// # 使用示例
|
||||
/// ```ignore
|
||||
/// // 定义在 agent_runner
|
||||
/// impl SessionRegistry for AgentSessionRegistry {
|
||||
/// type Entry = ProjectAndAgentInfo;
|
||||
/// // ...
|
||||
/// }
|
||||
///
|
||||
/// // 注入到 AcpSessionManager
|
||||
/// let session_manager = AcpSessionManager::new(notifier, Arc::new(registry));
|
||||
/// ```
|
||||
pub trait SessionRegistry: Send + Sync + 'static {
|
||||
/// 会话条目类型
|
||||
type Entry: SessionEntry;
|
||||
|
||||
/// 获取会话
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `project_id` - 项目 ID
|
||||
///
|
||||
/// # Returns
|
||||
/// 如果存在则返回会话条目的克隆,否则返回 None
|
||||
fn get(&self, project_id: &str) -> Option<Self::Entry>;
|
||||
|
||||
/// 插入或更新会话
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `project_id` - 项目 ID
|
||||
/// * `session_id` - 会话 ID
|
||||
/// * `entry` - 会话条目
|
||||
fn insert(&self, project_id: &str, session_id: &str, entry: Self::Entry);
|
||||
|
||||
/// 移除会话
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `project_id` - 项目 ID
|
||||
///
|
||||
/// # Returns
|
||||
/// 如果存在则返回被移除的会话条目,否则返回 None
|
||||
fn remove(&self, project_id: &str) -> Option<Self::Entry>;
|
||||
|
||||
/// 检查会话是否存在
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `project_id` - 项目 ID
|
||||
fn contains(&self, project_id: &str) -> bool;
|
||||
|
||||
/// 🆕 通过 session_id 获取 project_id(反向查询)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `session_id` - 会话 ID
|
||||
///
|
||||
/// # Returns
|
||||
/// 如果 session_id 存在,返回对应的 project_id;否则返回 None
|
||||
fn get_project_by_session(&self, session_id: &str) -> Option<String>;
|
||||
|
||||
/// 🆕 通过 session_id 直接获取会话条目(原子性操作)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `session_id` - 会话 ID
|
||||
///
|
||||
/// # Returns
|
||||
/// 如果 session_id 存在,返回对应的会话条目克隆;否则返回 None
|
||||
///
|
||||
/// # 优势
|
||||
/// - 一次性查询,避免两次调用之间的竞态窗口
|
||||
/// - 内部使用 DashMap 的原子性操作
|
||||
fn get_entry_by_session(&self, session_id: &str) -> Option<Self::Entry>;
|
||||
|
||||
/// 获取所有项目 ID 列表
|
||||
fn list_project_ids(&self) -> Vec<String>;
|
||||
|
||||
/// 获取会话数量
|
||||
fn count(&self) -> usize;
|
||||
|
||||
/// 获取 DashMap entry(用于原子性操作)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `project_id` - 项目 ID
|
||||
///
|
||||
/// # Returns
|
||||
/// DashMap Entry,支持原子性的 get/insert/update 操作
|
||||
///
|
||||
/// # 使用示例
|
||||
/// ```ignore
|
||||
/// use dashmap::mapref::entry::Entry;
|
||||
///
|
||||
/// match registry.entry(project_id.to_string()) {
|
||||
/// Entry::Occupied(mut occupied) => {
|
||||
/// // 已存在,可以检查和更新
|
||||
/// let existing = occupied.get();
|
||||
/// if needs_rebuild {
|
||||
/// occupied.insert(new_value);
|
||||
/// }
|
||||
/// }
|
||||
/// Entry::Vacant(vacant) => {
|
||||
/// // 不存在,可以插入
|
||||
/// vacant.insert(new_value);
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
fn entry(&self, project_id: String) -> Entry<'_, String, Self::Entry>;
|
||||
}
|
||||
|
||||
/// SessionRegistry 的 Arc 包装器实现
|
||||
///
|
||||
/// 允许 `Arc<R>` 作为 `SessionRegistry` 使用
|
||||
impl<R: SessionRegistry> SessionRegistry for Arc<R> {
|
||||
type Entry = R::Entry;
|
||||
|
||||
fn get(&self, project_id: &str) -> Option<Self::Entry> {
|
||||
(**self).get(project_id)
|
||||
}
|
||||
|
||||
fn insert(&self, project_id: &str, session_id: &str, entry: Self::Entry) {
|
||||
(**self).insert(project_id, session_id, entry)
|
||||
}
|
||||
|
||||
fn remove(&self, project_id: &str) -> Option<Self::Entry> {
|
||||
(**self).remove(project_id)
|
||||
}
|
||||
|
||||
fn contains(&self, project_id: &str) -> bool {
|
||||
(**self).contains(project_id)
|
||||
}
|
||||
|
||||
fn get_project_by_session(&self, session_id: &str) -> Option<String> {
|
||||
(**self).get_project_by_session(session_id)
|
||||
}
|
||||
|
||||
fn get_entry_by_session(&self, session_id: &str) -> Option<Self::Entry> {
|
||||
(**self).get_entry_by_session(session_id)
|
||||
}
|
||||
|
||||
fn list_project_ids(&self) -> Vec<String> {
|
||||
(**self).list_project_ids()
|
||||
}
|
||||
|
||||
fn count(&self) -> usize {
|
||||
(**self).count()
|
||||
}
|
||||
|
||||
fn entry(&self, project_id: String) -> Entry<'_, String, Self::Entry> {
|
||||
(**self).entry(project_id)
|
||||
}
|
||||
}
|
||||
41
qiming-rcoder/crates/agent_config/Cargo.toml
Normal file
41
qiming-rcoder/crates/agent_config/Cargo.toml
Normal file
@@ -0,0 +1,41 @@
|
||||
[package]
|
||||
name = "agent_config"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
authors = ["Your Name <your.email@example.com>"]
|
||||
license = "MIT OR Apache-2.0"
|
||||
description = "Agent configuration management for RCoder"
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
# 复用 shared_types 中的类型
|
||||
shared_types = { path = "../shared_types" }
|
||||
|
||||
# 异步运行时
|
||||
tokio = { workspace = true, features = ["full", "process"] }
|
||||
|
||||
# 序列化
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
|
||||
# 错误处理
|
||||
thiserror = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
|
||||
# 日志
|
||||
tracing = { workspace = true }
|
||||
|
||||
# 日期时间
|
||||
chrono = { workspace = true, features = ["serde"] }
|
||||
|
||||
# 路径处理
|
||||
path-clean = { workspace = true }
|
||||
|
||||
# 命令执行(跨平台)
|
||||
which = { workspace = true }
|
||||
|
||||
# 异步 trait
|
||||
async-trait = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = { workspace = true }
|
||||
@@ -0,0 +1,90 @@
|
||||
{
|
||||
"agent_servers": {
|
||||
"claude-code-acp-ts": {
|
||||
"agent_id": "claude-code-acp-ts",
|
||||
"agent_type": "claude",
|
||||
"command": "claude-code-acp-ts",
|
||||
"args": [],
|
||||
"env": {
|
||||
"ANTHROPIC_API_KEY": "{MODEL_PROVIDER_API_KEY}",
|
||||
"ANTHROPIC_MODEL": "{MODEL_PROVIDER_DEFAULT_MODEL}",
|
||||
"ANTHROPIC_BASE_URL": "{MODEL_PROVIDER_BASE_URL}",
|
||||
"RUST_LOG": "info",
|
||||
"CLAUDE_CODE_MAX_TOKENS": "32768"
|
||||
},
|
||||
"system_prompt": {
|
||||
"source": "embedded",
|
||||
"template": "",
|
||||
"enabled": true
|
||||
},
|
||||
"user_prompt": {
|
||||
"template": "{user_prompt}",
|
||||
"enabled": false
|
||||
},
|
||||
"installation": {
|
||||
"package_manager": "npm",
|
||||
"package_name": "claude-code-acp-ts@latest",
|
||||
"version": "latest"
|
||||
},
|
||||
"enabled": true,
|
||||
"metadata": {
|
||||
"description": "Claude Code ACP Agent - Computer Agent Runner configuration",
|
||||
"version": "1.0.0",
|
||||
"service_type": "ComputerAgentRunner"
|
||||
}
|
||||
}
|
||||
},
|
||||
"context_servers": {
|
||||
"chrome-devtools": {
|
||||
"source": "custom",
|
||||
"enabled": false,
|
||||
"command": "mcp-proxy",
|
||||
"args": ["convert", "http://127.0.0.1:18099"],
|
||||
"env": {},
|
||||
"metadata": {
|
||||
"description": "Chrome DevTools MCP - 通过共享的 MCP Proxy 服务连接浏览器自动化",
|
||||
"note": "使用 mcp-proxy convert 连接容器启动时自动运行的共享 chrome-devtools-mcp 服务,多个 agent 复用同一个浏览器实例,加快启动速度",
|
||||
"proxy_url": "http://127.0.0.1:18099",
|
||||
"capabilities": [
|
||||
"browser_automation",
|
||||
"page_navigation",
|
||||
"dom_manipulation",
|
||||
"screenshot",
|
||||
"chinese_input_method"
|
||||
]
|
||||
}
|
||||
},
|
||||
"context7": {
|
||||
"source": "custom",
|
||||
"enabled": false,
|
||||
"command": "bunx",
|
||||
"args": ["-y", "@upstash/context7-mcp"],
|
||||
"env": {}
|
||||
},
|
||||
"puppeteer": {
|
||||
"source": "custom",
|
||||
"enabled": false,
|
||||
"command": "npx",
|
||||
"args": ["-y", "@anthropic/mcp-server-puppeteer"],
|
||||
"env": {
|
||||
"PUPPETEER_EXECUTABLE_PATH": "/usr/bin/chromium"
|
||||
},
|
||||
"metadata": {
|
||||
"description": "Puppeteer MCP - Alternative browser automation (disabled by default)",
|
||||
"note": "Enable this if chrome-devtools doesn't meet your needs"
|
||||
}
|
||||
}
|
||||
},
|
||||
"metadata": {
|
||||
"config_name": "computer_agent_default",
|
||||
"description": "Computer Agent Runner 专用配置 - 支持桌面环境和浏览器自动化",
|
||||
"version": "1.0.0",
|
||||
"service_type": "ComputerAgentRunner",
|
||||
"features": {
|
||||
"desktop_environment": "XFCE4 + noVNC",
|
||||
"browser": "Chromium with CDP",
|
||||
"vnc_port": 6080,
|
||||
"cdp_port": 9222
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
{
|
||||
"agent_servers": {
|
||||
"claude-code-acp-ts": {
|
||||
"agent_id": "claude-code-acp-ts",
|
||||
"agent_type": "claude",
|
||||
"command": "claude-code-acp-ts",
|
||||
"args": [],
|
||||
"env": {
|
||||
"ANTHROPIC_API_KEY": "{MODEL_PROVIDER_API_KEY}",
|
||||
"ANTHROPIC_MODEL": "{MODEL_PROVIDER_DEFAULT_MODEL}",
|
||||
"ANTHROPIC_BASE_URL": "{MODEL_PROVIDER_BASE_URL}",
|
||||
"RUST_LOG": "info",
|
||||
"CLAUDE_CODE_MAX_TOKENS": "32768"
|
||||
},
|
||||
"system_prompt": {
|
||||
"source": "embedded",
|
||||
"template": "",
|
||||
"enabled": true
|
||||
},
|
||||
"user_prompt": {
|
||||
"template": "{user_prompt}",
|
||||
"enabled": false
|
||||
},
|
||||
"installation": {
|
||||
"package_manager": "npm",
|
||||
"package_name": "claude-code-acp-ts@latest",
|
||||
"version": "latest"
|
||||
},
|
||||
"enabled": true,
|
||||
"metadata": {
|
||||
"description": "Claude Code ACP Agent - Default configuration",
|
||||
"version": "1.0.0"
|
||||
}
|
||||
}
|
||||
},
|
||||
"context_servers": {
|
||||
"context7": {
|
||||
"source": "custom",
|
||||
"enabled": false,
|
||||
"command": "bunx",
|
||||
"args": ["-y", "@upstash/context7-mcp"],
|
||||
"env": {}
|
||||
},
|
||||
"chrome-devtools": {
|
||||
"source": "custom",
|
||||
"enabled": true,
|
||||
"command": "npx",
|
||||
"args": [
|
||||
"-y",
|
||||
"chrome-devtools-mcp@latest",
|
||||
"--headless",
|
||||
"--isolated",
|
||||
"--executablePath=/usr/bin/chromium",
|
||||
"--chrome-arg=--no-sandbox",
|
||||
"--chrome-arg=--disable-gpu",
|
||||
"--chrome-arg=--disable-dev-shm-usage",
|
||||
"--chrome-arg=--disable-software-rasterizer",
|
||||
"--chrome-arg=--disable-extensions"
|
||||
],
|
||||
"env": {}
|
||||
},
|
||||
"fetch": {
|
||||
"source": "custom",
|
||||
"enabled": true,
|
||||
"command": "uvx",
|
||||
"args": ["mcp-server-fetch"],
|
||||
"env": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
<SYSTEM_INSTRUCTIONS>
|
||||
|
||||
You are a professional frontend project development expert integrated with MCP (Model Context Protocol) tools. You are proficient in modern frontend technology stacks including React, Vue, Vite, TypeScript and other mainstream frameworks and tools. You are designed to identify the framework used by a project and develop based on the existing technology stack, rather than forcibly converting frameworks.
|
||||
|
||||
**Core Capabilities**:
|
||||
• **Framework Identification**: Automatically identify the frontend framework used by the project (React, Vue, etc.)
|
||||
• **Framework Adaptation**: Write code based on the project's current framework, maintaining technology stack consistency
|
||||
• **General Tools**: Vite, TypeScript, Tailwind CSS, ESLint, Prettier
|
||||
• **HTTP Clients**: Axios, Fetch API
|
||||
• **Package Managers**: pnpm, npm, yarn
|
||||
• **Build Tools**: Vite (Hot Module Replacement, Fast Builds)
|
||||
• **Code Standards**: ESLint + Prettier + TypeScript Strict Mode
|
||||
|
||||
**Key Principles**:
|
||||
1. **Prioritize Framework Identification**: Before modifying code, first detect the framework used by the project (through package.json, file structure, etc.)
|
||||
2. **Maintain Technology Stack Consistency**: If the project uses Vue, develop with Vue; if it's React, develop with React
|
||||
3. **No Forced Framework Conversion**: Never convert Vue code to React or React code to Vue
|
||||
4. **Project Development**: Develop new features or fix existing features based on the existing project structure
|
||||
|
||||
<ROLE_DEFINITION>
|
||||
You are a professional frontend development expert proficient in multiple modern frontend frameworks and toolchains. You have access to various MCP tools, including context7 for web search and documentation retrieval.
|
||||
**Technical Capability Scope**:
|
||||
• **Mainstream Frameworks**: React, Vue, Angular, Svelte and other modern frontend frameworks with their ecosystems
|
||||
• **Development Languages**: TypeScript, JavaScript (ES6+), HTML5, CSS3
|
||||
• **Styling Solutions**: Tailwind CSS, CSS Modules, Sass, Less, Styled Components
|
||||
• **Build Tools**: Vite, Webpack, Rollup, esbuild and other modern build tools
|
||||
• **State Management**: State management solutions for each framework (Redux, Pinia, NgRx, Zustand, etc.)
|
||||
• **HTTP Clients**: Axios, Fetch API, HTTP libraries for each framework
|
||||
• **Code Quality Tools**: ESLint, Prettier, TSLint and other code quality tools
|
||||
|
||||
**Core Working Principles**:
|
||||
1. **Identify Framework First**: Must identify the project's framework and technology stack before writing code
|
||||
2. **Respect Existing Technology Stack**: Develop based on the project's existing frameworks and tools without unauthorized changes
|
||||
3. **Maintain Consistency**: Use the project's current framework syntax, conventions, and best practices
|
||||
4. **Use Tools**: Use available MCP tools when they can provide better answers
|
||||
5. **Best Practices**: Follow the latest best practices and design patterns for each framework and tool
|
||||
|
||||
<CODE_FORMAT_RULES>
|
||||
**General Code Standards**:
|
||||
1. Always write code in TypeScript strict mode
|
||||
2. Component files use PascalCase naming, utility functions use camelCase
|
||||
3. Interface types use PascalCase with 'Interface' or 'Type' suffix
|
||||
4. Prefer Tailwind CSS for styling
|
||||
5. API calls use Axios client or Fetch API
|
||||
6. Add JSDoc-style comments for complex logic
|
||||
7. Follow the project's code conventions and file structure conventions
|
||||
8. Ensure code formatting is correct and readable
|
||||
9. Consider error handling and edge cases
|
||||
10. Use appropriate variable and function names
|
||||
11. Leverage Vite's fast builds and hot module replacement
|
||||
12. The 'title' tag in the 'index.html' file in the project root directory should NOT contain any frontend framework names such as: React, Vite, Vue, Antd, Angular, etc.
|
||||
13. **Important: Router Mode Specification**: When developing involving routing, you MUST use hash mode. For example: React Router uses `HashRouter`, Vue Router configures `mode: 'hash'`, Angular Router uses `HashLocationStrategy` from LocationStrategy.
|
||||
14. **Important: Protect Injected Code Blocks**: It is strictly forbidden to delete or modify code blocks surrounded by `DEV-INJECT-START` and `DEV-INJECT-END` markers. These code blocks are automatically injected by development tools and must be completely preserved. When editing code, preserve these markers and all content between them.
|
||||
|
||||
**React Project Specific Standards**:
|
||||
• Follow React functional component best practices, use React.FC types
|
||||
• Use Radix UI component library for building UI
|
||||
• Forms use React Hook Form + Zod for validation
|
||||
• Use React.memo, useCallback, useMemo for performance optimization
|
||||
• Follow React Hooks rules
|
||||
• Routing must use `HashRouter` (from react-router-dom), do not use `BrowserRouter`
|
||||
|
||||
**Vue Project Specific Standards**:
|
||||
• Prefer Composition API (setup syntax sugar)
|
||||
• Use Element Plus or other Vue UI component libraries
|
||||
• Use Pinia for state management
|
||||
• Follow Vue best practices and reactivity system rules
|
||||
• Use computed, watch, ref, reactive and other Composition API features
|
||||
• Vue Router must be configured in hash mode: `createRouter({ history: createWebHashHistory(), ... })`
|
||||
|
||||
<DEVELOPMENT_CONSTRAINTS>
|
||||
**Strictly Prohibited Operations - Absolutely Not Allowed**:
|
||||
|
||||
🚫 **Security Ban** (Highest Priority):
|
||||
- **Absolutely prohibited** from probing, scanning, or accessing private network IP addresses (such as 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8)
|
||||
- **Absolutely prohibited** from attempting to access local services (localhost, 127.0.0.1, 0.0.0.0)
|
||||
- **Absolutely prohibited** from port scanning, network probing, private network service discovery, etc.
|
||||
- **Absolutely prohibited** from hardcoding private network IP addresses or private network addresses in code
|
||||
- **Absolutely prohibited** from using curl, wget, nc, telnet, nmap and other tools to probe private networks
|
||||
- **Absolutely prohibited** from executing any commands or code that may compromise system security
|
||||
- **Absolutely prohibited** from bypassing security restrictions or attempting privilege escalation
|
||||
- **Absolutely prohibited** from executing reverse shells, remote code execution, or other malicious operations
|
||||
- **Core Principle**: All network requests must point to public network services or legitimate API endpoints explicitly provided by the user
|
||||
|
||||
🚫 **Framework Conversion Ban** (Most Important):
|
||||
- **Absolutely prohibited** from rewriting Vue code as React code
|
||||
- **Absolutely prohibited** from rewriting React code as Vue code
|
||||
- **Absolutely prohibited** from arbitrarily switching frameworks in existing projects
|
||||
- **Must Follow**: After identifying the project framework, only use that framework's syntax and APIs
|
||||
- **Core Principle**: Respect the project's existing technology stack and maintain framework consistency
|
||||
|
||||
🚫 **Project Initialization Ban**:
|
||||
- Prohibited from using npm create, npm init
|
||||
- Prohibited from using yarn create, yarn init
|
||||
- Prohibited from using npx create-react-app, npx create-vue
|
||||
- Prohibited from using pnpm create
|
||||
- Prohibited from using any shell commands for project initialization
|
||||
- Prohibited from instructing users on how to use npm dev, npm build and other commands (because the project is a server-deployed service, users do not have permission to execute these)
|
||||
|
||||
🚫 **File/Script Creation Ban**:
|
||||
- **Prohibited** from creating, referencing, or injecting files or scripts named 'dev-monitor.js' in the project
|
||||
|
||||
🚫 **Code Block Protection Ban** (Important):
|
||||
- **Absolutely prohibited** from deleting or modifying code blocks surrounded by `DEV-INJECT-START` and `DEV-INJECT-END` markers
|
||||
- **Absolutely prohibited** from removing these markers or their content when editing code
|
||||
- **Must Follow**: These code blocks are automatically injected by development tools and must be completely preserved
|
||||
- **Core Principle**: When modifying code, if encountering these markers, bypass or preserve all content between these markers
|
||||
|
||||
✅ **Allowed Operation Scope**:
|
||||
- **Primary Task**: Identify the framework used by the project (check package.json, file structure, etc.)
|
||||
- Focus on writing and modifying frontend code files
|
||||
- Create components, pages, style files based on the project framework (.vue for Vue, .tsx/.jsx for React)
|
||||
- Modify existing TypeScript/JavaScript code (maintaining framework syntax)
|
||||
- Write Tailwind CSS or other styles
|
||||
- Use the project's corresponding UI component library (Radix UI for React, Element Plus for Vue)
|
||||
- Code-level modifications to configuration files (such as tsconfig.json, vite.config.ts)
|
||||
- Follow the project's code conventions and file structure
|
||||
- **Only Allowed Access**: Public API endpoints or legitimate external services explicitly provided by the user
|
||||
|
||||
**Core Principles**:
|
||||
- You are a frontend code writing expert, not a project administrator
|
||||
- **Most Important**: Identify and respect the project framework, never arbitrarily convert frameworks
|
||||
- **Security First**: Never execute any operations that may compromise system security
|
||||
- Users are responsible for dependency installation, service startup, and test execution
|
||||
- Always respond in English
|
||||
|
||||
<MCP_TOOL_GUIDANCE>
|
||||
Available MCP tools:
|
||||
- context7: Search the web, retrieve frontend framework documentation (React, Vue, Vite, TypeScript, etc.)
|
||||
|
||||
**Key Tool Usage Rules**:
|
||||
1. **Supported Mainstream Technology Stack**:
|
||||
- Frontend frameworks: React, Vue, Angular, Svelte and their corresponding ecosystems
|
||||
- Build tools: Vite, Webpack, Rollup, esbuild, etc.
|
||||
- Development languages: TypeScript, JavaScript, HTML, CSS
|
||||
- Styling solutions: Tailwind CSS, CSS Modules, Sass, Less, etc.
|
||||
- General tools: Axios, Fetch API, ESLint, Prettier, etc.
|
||||
2. **Existing Project Processing Flow** (Most Important):
|
||||
- **Step 1**: Check package.json to identify the framework and dependencies
|
||||
- **Step 2**: Check file structure to identify project type (.vue = Vue, .tsx/.jsx = React, .component.ts = Angular)
|
||||
- **Step 3**: Write code based on the identified framework, never convert frameworks
|
||||
- **Example**: If "vue" dependency is detected, use Vue syntax; if "react" is detected, use React syntax
|
||||
3. Use context7 to search for corresponding framework documentation, examples, and best practices
|
||||
4. Always verify project structure and framework before writing any code
|
||||
|
||||
**Core Memory**:
|
||||
- Existing project = Identify framework first, then code with corresponding framework syntax
|
||||
- **Never arbitrarily convert frameworks**: Vue projects stay Vue, React projects stay React
|
||||
|
||||
<THINKING_REQUIREMENTS>
|
||||
Before responding, you must follow this exact frontend development workflow:
|
||||
|
||||
**Phase 1: Project Status Detection**
|
||||
1. **Critical First Step**: Check project directory status
|
||||
2. **If It's an Existing Project** (Most Important):
|
||||
- **Step 1**: Immediately read the package.json file
|
||||
- **Step 2**: Check dependencies to identify frontend framework (react, vue, @angular/core, svelte, etc.)
|
||||
- **Step 3**: Check project file structure to identify framework type (.vue, .tsx/.jsx, .component.ts, .svelte, etc.)
|
||||
- **Step 4**: Clearly identify the framework and technology stack used by the project
|
||||
- **Step 5**: Only use that framework's syntax and APIs in all subsequent operations
|
||||
|
||||
**Phase 2: Framework Identification and Confirmation**
|
||||
3. **Framework Identification Indicators**:
|
||||
- Vue projects: Have "vue" dependency in package.json, exist .vue files
|
||||
- React projects: Have "react" dependency in package.json, exist .tsx/.jsx files
|
||||
- Angular projects: Have "@angular/core" dependency in package.json, exist .component.ts files
|
||||
- Svelte projects: Have "svelte" dependency in package.json, exist .svelte files
|
||||
4. **Behavior After Framework Confirmation**:
|
||||
- Vue projects: Use Vue APIs (Composition API or Options API), .vue files, Vue Router, Pinia, etc.
|
||||
- React projects: Use React APIs (Hooks, class components, etc.), .tsx/.jsx files, React Router, Redux/Zustand, etc.
|
||||
- Angular projects: Use Angular APIs, components/services/modules, RxJS, Angular Router, etc.
|
||||
- Svelte projects: Use Svelte syntax, .svelte files, SvelteKit, etc.
|
||||
- **Absolutely Prohibited**: Arbitrarily switching to other framework syntax in any project
|
||||
|
||||
**Phase 3: Development Execution**
|
||||
5. Analyze user's development request in detail
|
||||
6. Determine if context7 needs to be used to search for corresponding framework documentation
|
||||
7. Plan development approach based on the identified framework ecosystem
|
||||
8. Prioritize the framework's best practices and modern development patterns
|
||||
9. Consider framework-specific error handling, state management, component design, etc.
|
||||
10. Follow the project's code conventions and file structure conventions
|
||||
11. **Router Configuration Requirements** (Important):
|
||||
- If routing configuration is involved, hash mode must be used
|
||||
- React projects: Use `HashRouter`
|
||||
- Vue projects: Use `createWebHashHistory()`
|
||||
- Angular projects: Use `HashLocationStrategy`
|
||||
- History mode is strictly prohibited (BrowserRouter, createWebHistory, etc.)
|
||||
12. **MCP Tool Invocation Standards**:
|
||||
- Use context7 to search for corresponding framework documentation and best practices
|
||||
|
||||
**Absolute Rules (Core of the Core)**:
|
||||
⚠️ **Framework Consistency Principle**:
|
||||
- Identify project framework → Only use that framework's syntax and APIs → Never convert to other frameworks
|
||||
- Vue projects stay Vue, React projects stay React, Angular projects stay Angular
|
||||
- **Violating this principle is the most serious error**
|
||||
|
||||
**Checklist**:
|
||||
✓ Have you read package.json?
|
||||
✓ Have you identified the project framework?
|
||||
✓ Have you confirmed using the correct framework syntax?
|
||||
✓ Have you avoided framework conversion?
|
||||
✓ If routing is involved, is hash mode being used?
|
||||
|
||||
</SYSTEM_INSTRUCTIONS>
|
||||
95
qiming-rcoder/crates/agent_config/examples/agents.json
Normal file
95
qiming-rcoder/crates/agent_config/examples/agents.json
Normal file
@@ -0,0 +1,95 @@
|
||||
{
|
||||
"agent_servers": {
|
||||
"claude-code-acp": {
|
||||
"agent_id": "claude-code-acp",
|
||||
"agent_type": "claude",
|
||||
"command": "claude-code-acp",
|
||||
"args": [],
|
||||
"env": {
|
||||
"ANTHROPIC_API_KEY": "{MODEL_PROVIDER_API_KEY}",
|
||||
"ANTHROPIC_BASE_URL": "{MODEL_PROVIDER_BASE_URL}",
|
||||
"ANTHROPIC_MODEL": "{MODEL_PROVIDER_DEFAULT_MODEL}",
|
||||
"RUST_LOG": "info"
|
||||
},
|
||||
"installation": {
|
||||
"package_manager": "npm",
|
||||
"package_name": "@zed-industries/claude-code-acp",
|
||||
"version": "latest"
|
||||
},
|
||||
"enabled": true,
|
||||
"metadata": {
|
||||
"description": "Claude Code ACP Agent",
|
||||
"version": "1.0.0"
|
||||
}
|
||||
},
|
||||
"react-developer": {
|
||||
"agent_id": "react-developer",
|
||||
"agent_type": "claude",
|
||||
"command": "claude-code-acp",
|
||||
"args": [],
|
||||
"env": {
|
||||
"ANTHROPIC_API_KEY": "{MODEL_PROVIDER_API_KEY}"
|
||||
},
|
||||
"system_prompt": {
|
||||
"template": "你是一个专业的 React 开发助手",
|
||||
"enabled": true
|
||||
},
|
||||
"user_prompt": {
|
||||
"template": "作为 React 专家,请帮我解决以下问题:{user_prompt}。请提供现代、可维护的代码示例。",
|
||||
"enabled": true
|
||||
},
|
||||
"enabled": false,
|
||||
"metadata": {
|
||||
"description": "React 专项开发助手"
|
||||
}
|
||||
},
|
||||
"nuwaxcode": {
|
||||
"agent_id": "nuwaxcode",
|
||||
"agent_type": "openai",
|
||||
"command": "nuwaxcode",
|
||||
"args": ["acp"],
|
||||
"env": {
|
||||
"OPENAI_API_KEY": "{MODEL_PROVIDER_API_KEY}",
|
||||
"OPENAI_BASE_URL": "{MODEL_PROVIDER_BASE_URL}",
|
||||
"OPENCODE_MODEL": "{MODEL_PROVIDER_DEFAULT_MODEL}",
|
||||
"RUST_LOG": "info"
|
||||
},
|
||||
"system_prompt": {
|
||||
"source": "embedded",
|
||||
"template": "",
|
||||
"enabled": true
|
||||
},
|
||||
"user_prompt": {
|
||||
"template": "{user_prompt}",
|
||||
"enabled": false
|
||||
},
|
||||
"installation": {
|
||||
"package_manager": "npm",
|
||||
"package_name": "nuwaxcode",
|
||||
"version": "latest"
|
||||
},
|
||||
"enabled": false,
|
||||
"metadata": {
|
||||
"description": "NuwaxCode - OpenAI-compatible coding agent (示例配置)",
|
||||
"version": "1.0.0"
|
||||
}
|
||||
}
|
||||
},
|
||||
"context_servers": {
|
||||
"fetch": {
|
||||
"source": "custom",
|
||||
"enabled": true,
|
||||
"command": "uvx",
|
||||
"args": ["mcp-server-fetch"]
|
||||
},
|
||||
"context7": {
|
||||
"source": "custom",
|
||||
"enabled": true,
|
||||
"command": "npx",
|
||||
"args": ["-y", "@upstash/context7-mcp"],
|
||||
"env": {
|
||||
"NODE_ENV": "production"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,348 @@
|
||||
//! 默认 Agent 配置定义
|
||||
//!
|
||||
//! 此模块从编译时嵌入的 JSON 配置文件加载默认配置。
|
||||
//! 支持两种服务类型的配置:
|
||||
//! - RCoder: `configs/default_agents.json`
|
||||
//! - ComputerAgentRunner: `configs/computer_agent_default.json`
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::LazyLock;
|
||||
|
||||
use crate::types::agent_config::AgentConfig;
|
||||
use crate::types::mcp_config::ContextServerConfig;
|
||||
|
||||
/// 编译时嵌入的 RCoder 默认配置 JSON
|
||||
const DEFAULT_CONFIG_JSON: &str = include_str!("../../configs/default_agents.json");
|
||||
|
||||
/// 编译时嵌入的 ComputerAgentRunner 默认配置 JSON
|
||||
const COMPUTER_AGENT_CONFIG_JSON: &str = include_str!("../../configs/computer_agent_default.json");
|
||||
|
||||
/// Claude Code ACP Agent 的默认 ID
|
||||
pub const CLAUDE_CODE_ACP_AGENT_ID: &str = "claude-code-acp-ts";
|
||||
|
||||
/// 默认配置的内部结构(用于 JSON 反序列化)
|
||||
#[derive(serde::Deserialize)]
|
||||
pub struct EmbeddedConfig {
|
||||
agent_servers: HashMap<String, AgentConfig>,
|
||||
#[serde(default)]
|
||||
context_servers: HashMap<String, ContextServerConfig>,
|
||||
}
|
||||
|
||||
/// 懒加载的默认配置
|
||||
///
|
||||
/// 在首次访问时解析 JSON,之后复用解析结果。
|
||||
/// 如果 JSON 解析失败,程序会 panic(这是编译时嵌入的文件,不应失败)。
|
||||
static DEFAULT_CONFIG: LazyLock<EmbeddedConfig> = LazyLock::new(|| {
|
||||
serde_json::from_str(DEFAULT_CONFIG_JSON).expect(
|
||||
"Failed to parse embedded default_agents.json. This is a bug - please check the JSON syntax.",
|
||||
)
|
||||
});
|
||||
|
||||
/// 懒加载的 ComputerAgentRunner 默认配置
|
||||
///
|
||||
/// 在首次访问时解析 JSON,之后复用解析结果。
|
||||
static COMPUTER_AGENT_CONFIG: LazyLock<EmbeddedConfig> = LazyLock::new(|| {
|
||||
serde_json::from_str(COMPUTER_AGENT_CONFIG_JSON).expect(
|
||||
"Failed to parse embedded computer_agent_default.json. This is a bug - please check the JSON syntax.",
|
||||
)
|
||||
});
|
||||
|
||||
/// 获取默认的 Claude Code ACP Agent 配置
|
||||
///
|
||||
/// 从 `configs/default_agents.json` 加载配置。
|
||||
/// 修改 JSON 文件后重新编译即可生效。
|
||||
///
|
||||
/// # Panics
|
||||
/// 如果 JSON 中不存在 `claude-code-acp-ts` 配置,会 panic。
|
||||
pub fn default_claude_code_agent() -> AgentConfig {
|
||||
DEFAULT_CONFIG
|
||||
.agent_servers
|
||||
.get(CLAUDE_CODE_ACP_AGENT_ID)
|
||||
.cloned()
|
||||
.expect("claude-code-acp-ts not found in default_agents.json")
|
||||
}
|
||||
|
||||
/// 获取默认的 Context Servers 配置
|
||||
///
|
||||
/// 从 `configs/default_agents.json` 加载配置。
|
||||
/// 修改 JSON 文件后重新编译即可生效。
|
||||
pub fn default_context_servers() -> HashMap<String, ContextServerConfig> {
|
||||
DEFAULT_CONFIG.context_servers.clone()
|
||||
}
|
||||
|
||||
/// 获取默认的 Agent Servers 配置
|
||||
///
|
||||
/// 从 `configs/default_agents.json` 加载配置。
|
||||
/// 修改 JSON 文件后重新编译即可生效。
|
||||
pub fn default_agent_servers() -> HashMap<String, AgentConfig> {
|
||||
DEFAULT_CONFIG.agent_servers.clone()
|
||||
}
|
||||
|
||||
/// 获取指定 Agent 的默认配置
|
||||
///
|
||||
/// # 参数
|
||||
/// - `agent_id`: Agent 标识符
|
||||
///
|
||||
/// # 返回
|
||||
/// 如果存在返回配置的克隆,否则返回 None
|
||||
pub fn get_default_agent(agent_id: &str) -> Option<AgentConfig> {
|
||||
DEFAULT_CONFIG.agent_servers.get(agent_id).cloned()
|
||||
}
|
||||
|
||||
/// 获取指定 Context Server 的默认配置
|
||||
///
|
||||
/// # 参数
|
||||
/// - `name`: Context Server 名称
|
||||
///
|
||||
/// # 返回
|
||||
/// 如果存在返回配置的克隆,否则返回 None
|
||||
pub fn get_default_context_server(name: &str) -> Option<ContextServerConfig> {
|
||||
DEFAULT_CONFIG.context_servers.get(name).cloned()
|
||||
}
|
||||
|
||||
/// 根据服务类型获取默认配置
|
||||
///
|
||||
/// # 参数
|
||||
/// - `service_type`: 服务类型(RCoder 或 ComputerAgentRunner)
|
||||
///
|
||||
/// # 返回
|
||||
/// 对应服务类型的默认配置引用
|
||||
pub fn get_default_config_by_service_type(
|
||||
service_type: &shared_types::ServiceType,
|
||||
) -> &'static EmbeddedConfig {
|
||||
match service_type {
|
||||
shared_types::ServiceType::RCoder => &DEFAULT_CONFIG,
|
||||
shared_types::ServiceType::ComputerAgentRunner => &COMPUTER_AGENT_CONFIG,
|
||||
}
|
||||
}
|
||||
|
||||
/// 根据服务类型获取默认 Agent Servers 配置
|
||||
///
|
||||
/// # 参数
|
||||
/// - `service_type`: 服务类型
|
||||
///
|
||||
/// # 返回
|
||||
/// Agent Servers 配置的克隆
|
||||
pub fn default_agent_servers_for_service(
|
||||
service_type: &shared_types::ServiceType,
|
||||
) -> HashMap<String, AgentConfig> {
|
||||
get_default_config_by_service_type(service_type)
|
||||
.agent_servers
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// 根据服务类型获取默认 Context Servers 配置
|
||||
///
|
||||
/// # 参数
|
||||
/// - `service_type`: 服务类型
|
||||
///
|
||||
/// # 返回
|
||||
/// Context Servers 配置的克隆
|
||||
pub fn default_context_servers_for_service(
|
||||
service_type: &shared_types::ServiceType,
|
||||
) -> HashMap<String, ContextServerConfig> {
|
||||
get_default_config_by_service_type(service_type)
|
||||
.context_servers
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// 根据服务类型获取指定 Agent 的默认配置
|
||||
///
|
||||
/// # 参数
|
||||
/// - `agent_id`: Agent 标识符
|
||||
/// - `service_type`: 服务类型
|
||||
///
|
||||
/// # 返回
|
||||
/// 如果存在返回配置的克隆,否则返回 None
|
||||
pub fn get_default_agent_for_service(
|
||||
agent_id: &str,
|
||||
service_type: &shared_types::ServiceType,
|
||||
) -> Option<AgentConfig> {
|
||||
get_default_config_by_service_type(service_type)
|
||||
.agent_servers
|
||||
.get(agent_id)
|
||||
.cloned()
|
||||
}
|
||||
|
||||
/// 根据服务类型获取指定 Context Server 的默认配置
|
||||
///
|
||||
/// # 参数
|
||||
/// - `name`: Context Server 名称
|
||||
/// - `service_type`: 服务类型
|
||||
///
|
||||
/// # 返回
|
||||
/// 如果存在返回配置的克隆,否则返回 None
|
||||
pub fn get_default_context_server_for_service(
|
||||
name: &str,
|
||||
service_type: &shared_types::ServiceType,
|
||||
) -> Option<ContextServerConfig> {
|
||||
get_default_config_by_service_type(service_type)
|
||||
.context_servers
|
||||
.get(name)
|
||||
.cloned()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_default_config_loads_successfully() {
|
||||
// 确保 JSON 能成功解析
|
||||
let _ = &*DEFAULT_CONFIG;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_claude_code_agent() {
|
||||
let agent = default_claude_code_agent();
|
||||
assert_eq!(agent.agent_id, CLAUDE_CODE_ACP_AGENT_ID);
|
||||
assert_eq!(agent.command, "claude-code-acp-ts");
|
||||
assert!(agent.enabled);
|
||||
assert!(agent.env.contains_key("ANTHROPIC_API_KEY"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_context_servers() {
|
||||
let servers = default_context_servers();
|
||||
// 验证 JSON 中定义的 context servers
|
||||
assert!(servers.contains_key("fetch") || servers.contains_key("context7"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_agent_servers() {
|
||||
let agents = default_agent_servers();
|
||||
assert!(agents.contains_key(CLAUDE_CODE_ACP_AGENT_ID));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_default_agent() {
|
||||
let agent = get_default_agent(CLAUDE_CODE_ACP_AGENT_ID);
|
||||
assert!(agent.is_some());
|
||||
assert_eq!(agent.unwrap().agent_id, CLAUDE_CODE_ACP_AGENT_ID);
|
||||
|
||||
let non_existent = get_default_agent("non-existent-agent");
|
||||
assert!(non_existent.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_default_context_server() {
|
||||
// 测试存在的 server
|
||||
let servers = default_context_servers();
|
||||
if let Some(first_server_name) = servers.keys().next() {
|
||||
let server = get_default_context_server(first_server_name);
|
||||
assert!(server.is_some());
|
||||
}
|
||||
|
||||
// 测试不存在的 server
|
||||
let non_existent = get_default_context_server("non-existent-server");
|
||||
assert!(non_existent.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_computer_agent_config_loads_successfully() {
|
||||
// 确保 ComputerAgentRunner 配置能成功解析
|
||||
let _ = &*COMPUTER_AGENT_CONFIG;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_default_config_by_service_type() {
|
||||
// 测试 RCoder 配置
|
||||
let rcoder_config = get_default_config_by_service_type(&shared_types::ServiceType::RCoder);
|
||||
assert!(
|
||||
rcoder_config
|
||||
.agent_servers
|
||||
.contains_key(CLAUDE_CODE_ACP_AGENT_ID)
|
||||
);
|
||||
|
||||
// 测试 ComputerAgentRunner 配置
|
||||
let car_config =
|
||||
get_default_config_by_service_type(&shared_types::ServiceType::ComputerAgentRunner);
|
||||
assert!(
|
||||
car_config
|
||||
.agent_servers
|
||||
.contains_key(CLAUDE_CODE_ACP_AGENT_ID)
|
||||
);
|
||||
|
||||
// ComputerAgentRunner 应该有更多 context servers
|
||||
assert!(
|
||||
car_config.context_servers.contains_key("chrome-devtools")
|
||||
|| car_config.context_servers.contains_key("puppeteer")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_agent_servers_for_service() {
|
||||
// RCoder
|
||||
let rcoder_servers = default_agent_servers_for_service(&shared_types::ServiceType::RCoder);
|
||||
assert!(!rcoder_servers.is_empty());
|
||||
|
||||
// ComputerAgentRunner
|
||||
let car_servers =
|
||||
default_agent_servers_for_service(&shared_types::ServiceType::ComputerAgentRunner);
|
||||
assert!(!car_servers.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_context_servers_for_service() {
|
||||
// RCoder - 基础 MCP servers
|
||||
let rcoder_contexts =
|
||||
default_context_servers_for_service(&shared_types::ServiceType::RCoder);
|
||||
assert!(rcoder_contexts.contains_key("fetch") || rcoder_contexts.contains_key("context7"));
|
||||
|
||||
// ComputerAgentRunner - 增强 MCP servers
|
||||
let car_contexts =
|
||||
default_context_servers_for_service(&shared_types::ServiceType::ComputerAgentRunner);
|
||||
// 至少应该包含浏览器相关的 server
|
||||
assert!(car_contexts.len() >= rcoder_contexts.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_default_agent_for_service() {
|
||||
// 测试 RCoder
|
||||
let agent = get_default_agent_for_service(
|
||||
CLAUDE_CODE_ACP_AGENT_ID,
|
||||
&shared_types::ServiceType::RCoder,
|
||||
);
|
||||
assert!(agent.is_some());
|
||||
|
||||
// 测试 ComputerAgentRunner
|
||||
let agent = get_default_agent_for_service(
|
||||
CLAUDE_CODE_ACP_AGENT_ID,
|
||||
&shared_types::ServiceType::ComputerAgentRunner,
|
||||
);
|
||||
assert!(agent.is_some());
|
||||
|
||||
// 测试不存在的 agent
|
||||
let agent =
|
||||
get_default_agent_for_service("non-existent-agent", &shared_types::ServiceType::RCoder);
|
||||
assert!(agent.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_default_context_server_for_service() {
|
||||
// RCoder context server
|
||||
let rcoder_contexts =
|
||||
default_context_servers_for_service(&shared_types::ServiceType::RCoder);
|
||||
if let Some(first_server_name) = rcoder_contexts.keys().next() {
|
||||
let server = get_default_context_server_for_service(
|
||||
first_server_name,
|
||||
&shared_types::ServiceType::RCoder,
|
||||
);
|
||||
assert!(server.is_some());
|
||||
}
|
||||
|
||||
// ComputerAgentRunner context server
|
||||
let _server = get_default_context_server_for_service(
|
||||
"chrome-devtools",
|
||||
&shared_types::ServiceType::ComputerAgentRunner,
|
||||
);
|
||||
// 可能存在或不存在,取决于配置
|
||||
|
||||
// 测试不存在的 server
|
||||
let server = get_default_context_server_for_service(
|
||||
"non-existent-server",
|
||||
&shared_types::ServiceType::RCoder,
|
||||
);
|
||||
assert!(server.is_none());
|
||||
}
|
||||
}
|
||||
13
qiming-rcoder/crates/agent_config/src/config/mod.rs
Normal file
13
qiming-rcoder/crates/agent_config/src/config/mod.rs
Normal file
@@ -0,0 +1,13 @@
|
||||
//! Configuration management module.
|
||||
//!
|
||||
//! - `default_agent_config`: 默认 Agent 配置定义
|
||||
//! - `servers_config`: 配置文件加载和管理
|
||||
//! - `prompt_assembler`: 提示词配置组装工具
|
||||
|
||||
pub mod default_agent_config;
|
||||
pub mod prompt_assembler;
|
||||
pub mod servers_config;
|
||||
|
||||
pub use default_agent_config::*;
|
||||
pub use prompt_assembler::PromptConfigAssembler;
|
||||
pub use servers_config::*;
|
||||
384
qiming-rcoder/crates/agent_config/src/config/prompt_assembler.rs
Normal file
384
qiming-rcoder/crates/agent_config/src/config/prompt_assembler.rs
Normal file
@@ -0,0 +1,384 @@
|
||||
//! 提示词配置组装工具
|
||||
//!
|
||||
//! 将用户入参组装为内部使用的配置结构。
|
||||
//! 简化设计:直接使用入参,无优先级冲突。
|
||||
|
||||
use shared_types::{ChatAgentConfig, ChatAgentServerConfig};
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::servers_config::AgentServersConfig;
|
||||
use crate::types::agent_config::AgentConfig;
|
||||
use crate::types::mcp_config::ContextServerConfig;
|
||||
|
||||
/// 提示词配置组装器
|
||||
///
|
||||
/// 职责:
|
||||
/// 1. 组装系统提示词(入参 > 默认配置)
|
||||
/// 2. 应用用户提示词模板
|
||||
/// 3. 组装 Agent 服务器配置
|
||||
/// 4. 合并 MCP 服务器配置
|
||||
pub struct PromptConfigAssembler {
|
||||
/// 系统提示词入参
|
||||
system_prompt: Option<String>,
|
||||
/// 用户提示词模板入参
|
||||
user_prompt_template: Option<String>,
|
||||
/// Agent 运行时配置入参
|
||||
agent_config: Option<ChatAgentConfig>,
|
||||
/// 默认配置
|
||||
default_config: AgentServersConfig,
|
||||
}
|
||||
|
||||
impl PromptConfigAssembler {
|
||||
/// 创建新的配置组装器
|
||||
///
|
||||
/// # 参数
|
||||
/// - `default_config`: 默认配置(从文件或内置配置加载)
|
||||
pub fn new(default_config: AgentServersConfig) -> Self {
|
||||
Self {
|
||||
system_prompt: None,
|
||||
user_prompt_template: None,
|
||||
agent_config: None,
|
||||
default_config,
|
||||
}
|
||||
}
|
||||
|
||||
/// 设置系统提示词覆盖
|
||||
pub fn with_system_prompt(mut self, system_prompt: Option<String>) -> Self {
|
||||
self.system_prompt = system_prompt;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置用户提示词模板覆盖
|
||||
pub fn with_user_prompt_template(mut self, template: Option<String>) -> Self {
|
||||
self.user_prompt_template = template;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 Agent 运行时配置覆盖
|
||||
pub fn with_agent_config(mut self, config: Option<ChatAgentConfig>) -> Self {
|
||||
self.agent_config = config;
|
||||
self
|
||||
}
|
||||
|
||||
/// 获取最终的系统提示词
|
||||
///
|
||||
/// 逻辑:入参有值则使用入参,否则使用默认配置
|
||||
pub fn get_system_prompt(&self, agent_id: &str) -> String {
|
||||
// 入参有值且非空,直接使用
|
||||
if let Some(ref sp) = self.system_prompt {
|
||||
if !sp.is_empty() {
|
||||
return sp.clone();
|
||||
}
|
||||
}
|
||||
|
||||
// 使用默认配置
|
||||
self.default_config.get_system_prompt(agent_id)
|
||||
}
|
||||
|
||||
/// 应用用户提示词模板
|
||||
///
|
||||
/// 逻辑:
|
||||
/// 1. 如果有模板入参,使用模板替换 `{user_prompt}`
|
||||
/// 2. 如果没有模板入参,检查默认配置
|
||||
/// 3. 都没有,直接返回原始输入
|
||||
pub fn apply_user_prompt(&self, agent_id: &str, user_input: &str) -> String {
|
||||
// 入参有模板且非空,使用入参模板
|
||||
if let Some(ref template) = self.user_prompt_template {
|
||||
if !template.is_empty() {
|
||||
return template.replace("{user_prompt}", user_input);
|
||||
}
|
||||
}
|
||||
|
||||
// 检查默认配置中的 user_prompt 模板
|
||||
if let Some(agent) = self.default_config.get_agent(agent_id) {
|
||||
if let Some(ref prompt_config) = agent.user_prompt {
|
||||
if prompt_config.enabled {
|
||||
return prompt_config.apply(user_input);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 无模板,直接返回原始输入
|
||||
user_input.to_string()
|
||||
}
|
||||
|
||||
/// 获取最终的 Agent 服务器配置
|
||||
///
|
||||
/// 逻辑:
|
||||
/// 1. 如果入参有 agent_server,与默认配置合并(入参字段覆盖默认值)
|
||||
/// 2. 如果入参没有 agent_server,使用默认配置
|
||||
pub fn get_agent_server_config(&self, default_agent_id: &str) -> AgentConfig {
|
||||
// 获取默认的 Agent 配置
|
||||
let default_agent = self
|
||||
.default_config
|
||||
.get_agent(default_agent_id)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
// 如果入参有 agent_server 配置,合并覆盖
|
||||
if let Some(ref config) = self.agent_config {
|
||||
if let Some(ref agent_server) = config.agent_server {
|
||||
return self.merge_agent_config(&default_agent, agent_server);
|
||||
}
|
||||
}
|
||||
|
||||
// 使用默认配置
|
||||
default_agent
|
||||
}
|
||||
|
||||
/// 合并 Agent 配置(入参覆盖默认值)
|
||||
fn merge_agent_config(
|
||||
&self,
|
||||
default: &AgentConfig,
|
||||
override_config: &ChatAgentServerConfig,
|
||||
) -> AgentConfig {
|
||||
let mut merged = default.clone();
|
||||
|
||||
// agent_id: 入参有值则覆盖
|
||||
if let Some(ref agent_id) = override_config.agent_id {
|
||||
merged.agent_id = agent_id.clone();
|
||||
}
|
||||
|
||||
// command: 入参有值则覆盖
|
||||
if let Some(ref command) = override_config.command {
|
||||
merged.command = command.clone();
|
||||
}
|
||||
|
||||
// args: 入参有值则覆盖(替换而非追加)
|
||||
if let Some(ref args) = override_config.args {
|
||||
merged.args = args.clone();
|
||||
}
|
||||
|
||||
// env: 入参有值则合并(入参优先)
|
||||
if let Some(ref env) = override_config.env {
|
||||
for (key, value) in env {
|
||||
merged.env.insert(key.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// metadata: 入参有值则合并
|
||||
if let Some(ref metadata) = override_config.metadata {
|
||||
for (key, value) in metadata {
|
||||
merged.metadata.insert(key.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
|
||||
merged
|
||||
}
|
||||
|
||||
/// 获取最终的 MCP 服务器配置
|
||||
///
|
||||
/// 逻辑:
|
||||
/// 1. 如果入参明确提供了非空的 context_servers,使用入参
|
||||
/// 2. 否则使用默认配置
|
||||
///
|
||||
/// 注意:即使提供了 agent_config,但 context_servers 为空时,仍使用默认配置
|
||||
pub fn get_context_servers(&self) -> HashMap<String, ContextServerConfig> {
|
||||
// 入参有非空的 MCP 配置,使用入参
|
||||
if let Some(ref config) = self.agent_config {
|
||||
if config.has_context_servers() {
|
||||
return config
|
||||
.context_servers
|
||||
.iter()
|
||||
.map(|(name, chat_config)| {
|
||||
let ctx_config = ContextServerConfig {
|
||||
source: chat_config.source.clone(),
|
||||
enabled: chat_config.enabled,
|
||||
command: chat_config.command.clone(),
|
||||
args: chat_config.args.clone(),
|
||||
env: chat_config.env.clone(),
|
||||
};
|
||||
(name.clone(), ctx_config)
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
}
|
||||
|
||||
// context_servers 为空或未提供,使用默认配置
|
||||
self.default_config.context_servers.clone()
|
||||
}
|
||||
|
||||
/// 获取使用的 Agent ID
|
||||
///
|
||||
/// 逻辑:入参有指定则使用入参,否则使用默认
|
||||
pub fn get_agent_id(&self, default_agent_id: &str) -> String {
|
||||
if let Some(ref config) = self.agent_config {
|
||||
if let Some(ref agent_server) = config.agent_server {
|
||||
return agent_server.get_agent_id().to_string();
|
||||
}
|
||||
}
|
||||
default_agent_id.to_string()
|
||||
}
|
||||
|
||||
/// 检查是否有系统提示词覆盖
|
||||
pub fn has_system_prompt_override(&self) -> bool {
|
||||
self.system_prompt.as_ref().is_some_and(|s| !s.is_empty())
|
||||
}
|
||||
|
||||
/// 检查是否有用户提示词模板覆盖
|
||||
pub fn has_user_prompt_template_override(&self) -> bool {
|
||||
self.user_prompt_template
|
||||
.as_ref()
|
||||
.is_some_and(|s| !s.is_empty())
|
||||
}
|
||||
|
||||
/// 检查是否有 Agent 配置覆盖
|
||||
pub fn has_agent_config_override(&self) -> bool {
|
||||
self.agent_config.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn create_test_default_config() -> AgentServersConfig {
|
||||
AgentServersConfig::default()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_system_prompt_with_override() {
|
||||
let config = create_test_default_config();
|
||||
let assembler = PromptConfigAssembler::new(config)
|
||||
.with_system_prompt(Some("自定义系统提示词".to_string()));
|
||||
|
||||
let result = assembler.get_system_prompt("claude-code-acp-ts");
|
||||
assert_eq!(result, "自定义系统提示词");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_system_prompt_without_override() {
|
||||
let config = create_test_default_config();
|
||||
let assembler = PromptConfigAssembler::new(config);
|
||||
|
||||
let result = assembler.get_system_prompt("claude-code-acp-ts");
|
||||
// 应该返回默认系统提示词
|
||||
assert!(!result.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_system_prompt_empty_override() {
|
||||
let config = create_test_default_config();
|
||||
let assembler = PromptConfigAssembler::new(config).with_system_prompt(Some("".to_string()));
|
||||
|
||||
let result = assembler.get_system_prompt("claude-code-acp-ts");
|
||||
// 空字符串应该回退到默认值
|
||||
assert!(!result.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_prompt_with_template() {
|
||||
let config = create_test_default_config();
|
||||
let assembler = PromptConfigAssembler::new(config)
|
||||
.with_user_prompt_template(Some("请用 Rust 完成:{user_prompt}".to_string()));
|
||||
|
||||
let result = assembler.apply_user_prompt("claude-code-acp-ts", "Hello World");
|
||||
assert_eq!(result, "请用 Rust 完成:Hello World");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_prompt_without_template() {
|
||||
let config = create_test_default_config();
|
||||
let assembler = PromptConfigAssembler::new(config);
|
||||
|
||||
let result = assembler.apply_user_prompt("claude-code-acp-ts", "Hello World");
|
||||
// 没有模板时应该返回原始输入
|
||||
assert_eq!(result, "Hello World");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_servers_with_override() {
|
||||
let config = create_test_default_config();
|
||||
let mut context_servers = HashMap::new();
|
||||
context_servers.insert(
|
||||
"my-mcp".to_string(),
|
||||
shared_types::ChatContextServerConfig {
|
||||
source: "custom".to_string(),
|
||||
enabled: true,
|
||||
command: Some("bunx".to_string()),
|
||||
args: Some(vec!["-y".to_string(), "my-mcp-server".to_string()]),
|
||||
env: None,
|
||||
},
|
||||
);
|
||||
|
||||
let agent_config = ChatAgentConfig {
|
||||
agent_server: None,
|
||||
context_servers,
|
||||
resource_limits: None,
|
||||
};
|
||||
|
||||
let assembler = PromptConfigAssembler::new(config).with_agent_config(Some(agent_config));
|
||||
|
||||
let result = assembler.get_context_servers();
|
||||
assert!(result.contains_key("my-mcp"));
|
||||
assert_eq!(
|
||||
result.get("my-mcp").unwrap().command,
|
||||
Some("bunx".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_servers_without_override() {
|
||||
let config = create_test_default_config();
|
||||
let assembler = PromptConfigAssembler::new(config);
|
||||
|
||||
let result = assembler.get_context_servers();
|
||||
// 应该返回默认的 context servers
|
||||
assert!(!result.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agent_id_with_override() {
|
||||
let config = create_test_default_config();
|
||||
let agent_config = ChatAgentConfig {
|
||||
agent_server: Some(shared_types::ChatAgentServerConfig {
|
||||
agent_id: Some("custom-agent".to_string()),
|
||||
command: None,
|
||||
args: None,
|
||||
env: None,
|
||||
metadata: None,
|
||||
}),
|
||||
context_servers: HashMap::new(),
|
||||
resource_limits: None,
|
||||
};
|
||||
|
||||
let assembler = PromptConfigAssembler::new(config).with_agent_config(Some(agent_config));
|
||||
|
||||
let result = assembler.get_agent_id("claude-code-acp-ts");
|
||||
assert_eq!(result, "custom-agent");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agent_id_without_override() {
|
||||
let config = create_test_default_config();
|
||||
let assembler = PromptConfigAssembler::new(config);
|
||||
|
||||
let result = assembler.get_agent_id("claude-code-acp-ts");
|
||||
assert_eq!(result, "claude-code-acp-ts");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_merge_agent_config_env() {
|
||||
let config = create_test_default_config();
|
||||
let mut override_env = HashMap::new();
|
||||
override_env.insert("NEW_VAR".to_string(), "new_value".to_string());
|
||||
|
||||
let agent_config = ChatAgentConfig {
|
||||
agent_server: Some(shared_types::ChatAgentServerConfig {
|
||||
agent_id: None,
|
||||
command: None,
|
||||
args: None,
|
||||
env: Some(override_env),
|
||||
metadata: None,
|
||||
}),
|
||||
context_servers: HashMap::new(),
|
||||
resource_limits: None,
|
||||
};
|
||||
|
||||
let assembler = PromptConfigAssembler::new(config).with_agent_config(Some(agent_config));
|
||||
|
||||
let result = assembler.get_agent_server_config("claude-code-acp-ts");
|
||||
assert!(result.env.contains_key("NEW_VAR"));
|
||||
assert_eq!(result.env.get("NEW_VAR"), Some(&"new_value".to_string()));
|
||||
}
|
||||
}
|
||||
226
qiming-rcoder/crates/agent_config/src/config/servers_config.rs
Normal file
226
qiming-rcoder/crates/agent_config/src/config/servers_config.rs
Normal file
@@ -0,0 +1,226 @@
|
||||
//! Agent servers configuration structure.
|
||||
//!
|
||||
//! 此模块负责配置的查询和管理。
|
||||
//! 默认配置从 `default_agent_config` 模块加载(来源于 `configs/default_agents.json`)。
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use super::default_agent_config::{
|
||||
default_agent_servers, default_agent_servers_for_service, default_context_servers,
|
||||
default_context_servers_for_service,
|
||||
};
|
||||
use crate::types::agent_config::AgentConfig;
|
||||
use crate::types::error::{ConfigError, Result};
|
||||
use crate::types::mcp_config::ContextServerConfig;
|
||||
use crate::types::system_prompt::DEFAULT_SYSTEM_PROMPT;
|
||||
|
||||
/// Agent servers configuration
|
||||
///
|
||||
/// 包含所有 Agent 和 Context Server 的配置集合。
|
||||
/// 默认配置来自 `configs/default_agents.json`(编译时嵌入)。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentServersConfig {
|
||||
/// Agent servers configuration
|
||||
pub agent_servers: HashMap<String, AgentConfig>,
|
||||
|
||||
/// Context servers configuration (MCP servers)
|
||||
#[serde(default)]
|
||||
pub context_servers: HashMap<String, ContextServerConfig>,
|
||||
}
|
||||
|
||||
impl AgentServersConfig {
|
||||
/// 获取默认配置
|
||||
///
|
||||
/// 返回从 `configs/default_agents.json` 加载的配置。
|
||||
///
|
||||
/// # 已弃用
|
||||
///
|
||||
/// 此方法已弃用,请使用 `load_or_default_for_service` 以支持多服务类型。
|
||||
#[deprecated(
|
||||
since = "0.2.0",
|
||||
note = "请使用 load_or_default_for_service 以支持多服务类型"
|
||||
)]
|
||||
pub async fn load_or_default() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// 根据服务类型加载或使用默认配置
|
||||
///
|
||||
/// # 参数
|
||||
/// - `service_type`: 服务类型(RCoder 或 ComputerAgentRunner)
|
||||
///
|
||||
/// # 返回
|
||||
/// 对应服务类型的默认配置
|
||||
pub async fn load_or_default_for_service(service_type: &shared_types::ServiceType) -> Self {
|
||||
Self::default_for_service(service_type)
|
||||
}
|
||||
|
||||
/// 根据服务类型创建默认配置
|
||||
///
|
||||
/// # 参数
|
||||
/// - `service_type`: 服务类型
|
||||
///
|
||||
/// # 返回
|
||||
/// 对应服务类型的默认配置实例
|
||||
pub fn default_for_service(service_type: &shared_types::ServiceType) -> Self {
|
||||
Self {
|
||||
agent_servers: default_agent_servers_for_service(service_type),
|
||||
context_servers: default_context_servers_for_service(service_type),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从外部文件加载配置
|
||||
///
|
||||
/// 用于加载用户自定义的配置文件,覆盖默认配置。
|
||||
pub async fn from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
|
||||
let path = path.as_ref();
|
||||
if !path.exists() {
|
||||
return Err(ConfigError::file_not_found(path.display()).into());
|
||||
}
|
||||
|
||||
let content = tokio::fs::read_to_string(path)
|
||||
.await
|
||||
.map_err(|e| ConfigError::Read(format!("{}: {}", path.display(), e)))?;
|
||||
|
||||
Self::from_json(&content)
|
||||
}
|
||||
|
||||
/// 从 JSON 字符串加载配置
|
||||
///
|
||||
/// 用于解析用户传入的 JSON 配置。
|
||||
pub fn from_json(json: &str) -> Result<Self> {
|
||||
serde_json::from_str(json).map_err(|e| ConfigError::Deserialization(e.to_string()).into())
|
||||
}
|
||||
|
||||
/// Validate configuration
|
||||
pub fn validate(&self) -> Result<()> {
|
||||
// Validate agent servers
|
||||
for (name, config) in &self.agent_servers {
|
||||
if config.agent_id.is_empty() {
|
||||
return Err(
|
||||
ConfigError::missing_field(format!("agent_servers.{}.agent_id", name)).into(),
|
||||
);
|
||||
}
|
||||
if config.command.is_empty() {
|
||||
return Err(
|
||||
ConfigError::missing_field(format!("agent_servers.{}.command", name)).into(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Validate context servers
|
||||
for (name, config) in &self.context_servers {
|
||||
if config.command.is_none() {
|
||||
return Err(ConfigError::missing_field(format!(
|
||||
"context_servers.{}.command",
|
||||
name
|
||||
))
|
||||
.into());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get all enabled agents
|
||||
pub fn get_enabled_agents(&self) -> Vec<&AgentConfig> {
|
||||
self.agent_servers
|
||||
.values()
|
||||
.filter(|config| config.enabled)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get an agent configuration by ID
|
||||
pub fn get_agent(&self, agent_id: &str) -> Option<&AgentConfig> {
|
||||
self.agent_servers.get(agent_id)
|
||||
}
|
||||
|
||||
/// 获取系统提示词(优先使用配置,否则使用默认值)
|
||||
///
|
||||
/// # 参数
|
||||
/// - `agent_id`: Agent 标识符
|
||||
///
|
||||
/// # 返回
|
||||
/// 根据配置的 source 字段决定返回内容:
|
||||
/// - `source: "embedded"` -> 返回编译时嵌入的默认提示词
|
||||
/// - `source: "custom"` 且 template 非空 -> 返回自定义模板内容
|
||||
/// - 其他情况 -> 回退到编译时嵌入的默认提示词
|
||||
///
|
||||
/// 注意:此方法始终返回有效的系统提示词,不受 enabled 字段影响。
|
||||
/// 如需检查是否启用,请直接调用 `SystemPromptConfig::get_prompt()`。
|
||||
pub fn get_system_prompt(&self, agent_id: &str) -> String {
|
||||
if let Some(agent) = self.get_agent(agent_id) {
|
||||
// 如果有系统提示词配置,使用 get_prompt_or_default() 确保始终返回有效提示词
|
||||
if let Some(ref prompt_config) = agent.system_prompt {
|
||||
return prompt_config.get_prompt_or_default().to_string();
|
||||
}
|
||||
}
|
||||
// 回退到编译时嵌入的默认值
|
||||
DEFAULT_SYSTEM_PROMPT.to_string()
|
||||
}
|
||||
|
||||
/// Get all enabled context servers
|
||||
pub fn get_enabled_context_servers(&self) -> Vec<&ContextServerConfig> {
|
||||
self.context_servers
|
||||
.values()
|
||||
.filter(|config| config.enabled)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get a context server configuration by name
|
||||
pub fn get_context_server(&self, name: &str) -> Option<&ContextServerConfig> {
|
||||
self.context_servers.get(name)
|
||||
}
|
||||
|
||||
/// 创建默认配置
|
||||
///
|
||||
/// 从 `configs/default_agents.json` 加载的配置(编译时嵌入)。
|
||||
/// 修改 JSON 文件后重新编译即可生效。
|
||||
pub fn default() -> Self {
|
||||
Self {
|
||||
agent_servers: default_agent_servers(),
|
||||
context_servers: default_context_servers(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_or_default_for_service() {
|
||||
// 测试 RCoder
|
||||
let rcoder_config =
|
||||
AgentServersConfig::load_or_default_for_service(&shared_types::ServiceType::RCoder)
|
||||
.await;
|
||||
assert!(!rcoder_config.agent_servers.is_empty());
|
||||
|
||||
// 测试 ComputerAgentRunner
|
||||
let car_config = AgentServersConfig::load_or_default_for_service(
|
||||
&shared_types::ServiceType::ComputerAgentRunner,
|
||||
)
|
||||
.await;
|
||||
assert!(!car_config.agent_servers.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_for_service() {
|
||||
// 测试 RCoder
|
||||
let rcoder_config =
|
||||
AgentServersConfig::default_for_service(&shared_types::ServiceType::RCoder);
|
||||
assert!(!rcoder_config.agent_servers.is_empty());
|
||||
|
||||
// 测试 ComputerAgentRunner
|
||||
let car_config = AgentServersConfig::default_for_service(
|
||||
&shared_types::ServiceType::ComputerAgentRunner,
|
||||
);
|
||||
assert!(!car_config.agent_servers.is_empty());
|
||||
|
||||
// 验证配置不同
|
||||
// ComputerAgentRunner 应该有更多或不同的 context servers
|
||||
assert!(car_config.context_servers.len() >= rcoder_config.context_servers.len());
|
||||
}
|
||||
}
|
||||
305
qiming-rcoder/crates/agent_config/src/installer/manager.rs
Normal file
305
qiming-rcoder/crates/agent_config/src/installer/manager.rs
Normal file
@@ -0,0 +1,305 @@
|
||||
//! Agent Installation Manager
|
||||
//!
|
||||
//! Manages agent installation, validation, and updates.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use thiserror::Error;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use super::npm_installer::NpmInstaller;
|
||||
use super::traits::AgentInstaller;
|
||||
use crate::types::installation::{InstallationConfig, PackageManager};
|
||||
|
||||
/// Installation error types
|
||||
#[derive(Error, Debug)]
|
||||
pub enum InstallationError {
|
||||
#[error("Package manager '{0}' is not available")]
|
||||
PackageManagerNotAvailable(String),
|
||||
|
||||
#[error("Package '{package}' not found in {manager}")]
|
||||
PackageNotFound { package: String, manager: String },
|
||||
|
||||
#[error("Installation failed for '{package}': {reason}")]
|
||||
InstallFailed { package: String, reason: String },
|
||||
|
||||
#[error("Validation failed: {0}")]
|
||||
ValidationFailed(String),
|
||||
|
||||
#[error("Command failed: {command} - {reason}")]
|
||||
CommandFailed { command: String, reason: String },
|
||||
|
||||
#[error("IO error: {0}")]
|
||||
IoError(#[from] std::io::Error),
|
||||
|
||||
#[error("Configuration error: {0}")]
|
||||
ConfigError(String),
|
||||
|
||||
#[error("Unsupported package manager: {0}")]
|
||||
UnsupportedPackageManager(String),
|
||||
}
|
||||
|
||||
/// Installation result
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct InstallResult {
|
||||
/// Whether installation was successful
|
||||
pub success: bool,
|
||||
/// Installed version
|
||||
pub version: Option<String>,
|
||||
/// Installation path
|
||||
pub install_path: Option<String>,
|
||||
/// Additional message
|
||||
pub message: String,
|
||||
/// Whether the agent was already installed
|
||||
pub already_installed: bool,
|
||||
}
|
||||
|
||||
impl InstallResult {
|
||||
/// Create a successful installation result
|
||||
pub fn success(version: Option<String>, message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
success: true,
|
||||
version,
|
||||
install_path: None,
|
||||
message: message.into(),
|
||||
already_installed: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a result indicating already installed
|
||||
pub fn already_installed(version: Option<String>) -> Self {
|
||||
Self {
|
||||
success: true,
|
||||
version,
|
||||
install_path: None,
|
||||
message: "Already installed".to_string(),
|
||||
already_installed: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a failed installation result
|
||||
pub fn failed(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
success: false,
|
||||
version: None,
|
||||
install_path: None,
|
||||
message: message.into(),
|
||||
already_installed: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set installation path
|
||||
pub fn with_path(mut self, path: impl Into<String>) -> Self {
|
||||
self.install_path = Some(path.into());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Agent Installation Manager
|
||||
///
|
||||
/// Manages multiple package manager installers and handles
|
||||
/// agent installation lifecycle.
|
||||
pub struct AgentInstallationManager {
|
||||
installers: HashMap<String, Arc<dyn AgentInstaller>>,
|
||||
}
|
||||
|
||||
impl AgentInstallationManager {
|
||||
/// Create a new installation manager with default installers
|
||||
pub fn new() -> Self {
|
||||
let mut installers: HashMap<String, Arc<dyn AgentInstaller>> = HashMap::new();
|
||||
|
||||
// Register npm installer
|
||||
installers.insert("npm".to_string(), Arc::new(NpmInstaller::new()));
|
||||
|
||||
Self { installers }
|
||||
}
|
||||
|
||||
/// Register a custom installer
|
||||
pub fn register_installer(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
installer: Arc<dyn AgentInstaller>,
|
||||
) {
|
||||
self.installers.insert(name.into(), installer);
|
||||
}
|
||||
|
||||
/// Get installer for package manager
|
||||
fn get_installer(&self, pm: &PackageManager) -> Option<Arc<dyn AgentInstaller>> {
|
||||
let name = pm.as_str();
|
||||
self.installers.get(name).cloned()
|
||||
}
|
||||
|
||||
/// Ensure agent is installed
|
||||
///
|
||||
/// This method checks if the agent is already installed and working.
|
||||
/// If not, it attempts to install it.
|
||||
pub async fn ensure_installed(
|
||||
&self,
|
||||
config: &InstallationConfig,
|
||||
command: &str,
|
||||
) -> Result<InstallResult, InstallationError> {
|
||||
info!("Checking Agent installation status: {}", command);
|
||||
|
||||
// First, check if the command already exists
|
||||
if self.is_command_available(command).await {
|
||||
// Validate if it's working correctly
|
||||
if self.validate_installation(config, command).await? {
|
||||
info!("Agent is installed and verified: {}", command);
|
||||
return Ok(InstallResult::already_installed(None));
|
||||
}
|
||||
warn!(
|
||||
"Agent command exists but verification failed, trying reinstall: {}",
|
||||
command
|
||||
);
|
||||
}
|
||||
|
||||
// Not installed or validation failed, try to install
|
||||
info!("Trying to install Agent: {}", command);
|
||||
let result = self.install(config, None).await;
|
||||
|
||||
match &result {
|
||||
Ok(install_result) => {
|
||||
if install_result.success {
|
||||
info!(
|
||||
"Agent installed successfully: {} - {}",
|
||||
command, install_result.message
|
||||
);
|
||||
} else {
|
||||
warn!(
|
||||
"⚠️ Agent Installation failed: {} - {}",
|
||||
command, install_result.message
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("Agent installation error: {} - {}", command, e);
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Install agent using appropriate package manager
|
||||
pub async fn install(
|
||||
&self,
|
||||
config: &InstallationConfig,
|
||||
install_dir: Option<&Path>,
|
||||
) -> Result<InstallResult, InstallationError> {
|
||||
let installer = self.get_installer(&config.package_manager).ok_or_else(|| {
|
||||
InstallationError::UnsupportedPackageManager(
|
||||
config.package_manager.as_str().to_string(),
|
||||
)
|
||||
})?;
|
||||
|
||||
let package_name = config.package_name.as_ref().ok_or_else(|| {
|
||||
InstallationError::ConfigError("Package name is required".to_string())
|
||||
})?;
|
||||
|
||||
info!(
|
||||
"Starting Agent installation: {} via {}",
|
||||
package_name,
|
||||
installer.package_manager_name()
|
||||
);
|
||||
|
||||
installer.install(config, install_dir).await
|
||||
}
|
||||
|
||||
/// Validate agent installation
|
||||
pub async fn validate_installation(
|
||||
&self,
|
||||
config: &InstallationConfig,
|
||||
command: &str,
|
||||
) -> Result<bool, InstallationError> {
|
||||
debug!("Verifying Agent installation: {}", command);
|
||||
|
||||
// If validate_command is specified, use it
|
||||
if let Some(validate_cmd) = &config.validate_command {
|
||||
if !validate_cmd.is_empty() {
|
||||
return self.run_validation_command(validate_cmd).await;
|
||||
}
|
||||
}
|
||||
|
||||
// 使用 which 检查命令是否可用(适用于 Debian/Docker 环境)
|
||||
// which::which 是 Rust crate,不依赖系统 which 命令
|
||||
// 注意:只检查命令是否在 PATH 中,不运行命令避免副作用(如 claude-code-acp-ts --version 会阻塞)
|
||||
if self.is_command_available(command).await {
|
||||
debug!("Command exists in PATH: {}", command);
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
debug!("Command not in PATH: {}", command);
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
/// Check if a command is available in PATH
|
||||
async fn is_command_available(&self, command: &str) -> bool {
|
||||
which::which(command).is_ok()
|
||||
}
|
||||
|
||||
/// Run custom validation command
|
||||
async fn run_validation_command(&self, cmd: &[String]) -> Result<bool, InstallationError> {
|
||||
if cmd.is_empty() {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let program = &cmd[0];
|
||||
let args: Vec<&str> = cmd[1..].iter().map(|s| s.as_str()).collect();
|
||||
|
||||
debug!("Running verification command: {} {:?}", program, args);
|
||||
|
||||
let output = tokio::process::Command::new(program)
|
||||
.args(&args)
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| InstallationError::CommandFailed {
|
||||
command: format!("{} {}", program, args.join(" ")),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
Ok(output.status.success())
|
||||
}
|
||||
|
||||
/// Update agent to latest version
|
||||
pub async fn update(
|
||||
&self,
|
||||
config: &InstallationConfig,
|
||||
install_dir: Option<&Path>,
|
||||
) -> Result<InstallResult, InstallationError> {
|
||||
let installer = self.get_installer(&config.package_manager).ok_or_else(|| {
|
||||
InstallationError::UnsupportedPackageManager(
|
||||
config.package_manager.as_str().to_string(),
|
||||
)
|
||||
})?;
|
||||
|
||||
info!("Updating Agent via {}", installer.package_manager_name());
|
||||
|
||||
installer.update(config, install_dir).await
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AgentInstallationManager {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_command_availability() {
|
||||
let manager = AgentInstallationManager::new();
|
||||
// Test with a command that should exist on most systems
|
||||
let ls_available = manager.is_command_available("ls").await;
|
||||
assert!(ls_available);
|
||||
|
||||
// Test with a command that shouldn't exist
|
||||
let fake_available = manager
|
||||
.is_command_available("this_command_does_not_exist_12345")
|
||||
.await;
|
||||
assert!(!fake_available);
|
||||
}
|
||||
}
|
||||
11
qiming-rcoder/crates/agent_config/src/installer/mod.rs
Normal file
11
qiming-rcoder/crates/agent_config/src/installer/mod.rs
Normal file
@@ -0,0 +1,11 @@
|
||||
//! Agent Installation Management
|
||||
//!
|
||||
//! This module provides automatic installation and validation for agents.
|
||||
|
||||
mod manager;
|
||||
mod npm_installer;
|
||||
mod traits;
|
||||
|
||||
pub use manager::{AgentInstallationManager, InstallResult, InstallationError};
|
||||
pub use npm_installer::NpmInstaller;
|
||||
pub use traits::AgentInstaller;
|
||||
236
qiming-rcoder/crates/agent_config/src/installer/npm_installer.rs
Normal file
236
qiming-rcoder/crates/agent_config/src/installer/npm_installer.rs
Normal file
@@ -0,0 +1,236 @@
|
||||
//! NPM Package Installer
|
||||
//!
|
||||
//! Implements agent installation via npm.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use super::manager::{InstallResult, InstallationError};
|
||||
use super::traits::AgentInstaller;
|
||||
use crate::types::installation::InstallationConfig;
|
||||
|
||||
/// NPM Installer
|
||||
///
|
||||
/// Installs agents via npm global installation.
|
||||
pub struct NpmInstaller {
|
||||
/// Use pnpm instead of npm if available
|
||||
prefer_pnpm: bool,
|
||||
}
|
||||
|
||||
impl NpmInstaller {
|
||||
/// Create a new NPM installer
|
||||
pub fn new() -> Self {
|
||||
Self { prefer_pnpm: false }
|
||||
}
|
||||
|
||||
/// Create NPM installer that prefers pnpm
|
||||
pub fn with_pnpm_preference() -> Self {
|
||||
Self { prefer_pnpm: true }
|
||||
}
|
||||
|
||||
/// Get the npm command to use (npm or pnpm)
|
||||
async fn get_npm_command(&self) -> &'static str {
|
||||
if self.prefer_pnpm && which::which("pnpm").is_ok() {
|
||||
"pnpm"
|
||||
} else if which::which("npm").is_ok() {
|
||||
"npm"
|
||||
} else {
|
||||
"npm" // Default, will fail if not available
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse version from npm output
|
||||
fn parse_installed_version(&self, output: &str) -> Option<String> {
|
||||
// npm list output format: "package@version"
|
||||
for line in output.lines() {
|
||||
if line.contains('@') {
|
||||
if let Some(version) = line.split('@').last() {
|
||||
return Some(version.trim().to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for NpmInstaller {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl AgentInstaller for NpmInstaller {
|
||||
async fn install(
|
||||
&self,
|
||||
config: &InstallationConfig,
|
||||
_install_dir: Option<&Path>,
|
||||
) -> Result<InstallResult, InstallationError> {
|
||||
let npm_cmd = self.get_npm_command().await;
|
||||
|
||||
// Check if npm is available
|
||||
if !self.command_exists(npm_cmd).await {
|
||||
return Err(InstallationError::PackageManagerNotAvailable(
|
||||
npm_cmd.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let package_name = config.package_name.as_ref().ok_or_else(|| {
|
||||
InstallationError::ConfigError("Package name is required".to_string())
|
||||
})?;
|
||||
|
||||
// Build package spec with version
|
||||
let package_spec = if let Some(version) = &config.version {
|
||||
if version == "latest" {
|
||||
package_name.clone()
|
||||
} else {
|
||||
format!("{}@{}", package_name, version)
|
||||
}
|
||||
} else {
|
||||
package_name.clone()
|
||||
};
|
||||
|
||||
info!(" running {} : {}", npm_cmd, package_spec);
|
||||
|
||||
// Run npm install -g
|
||||
// 仅 nuwaxcode 使用官方源(--registry 参数优先级高于 .npmrc)
|
||||
let output = if package_name == "nuwaxcode" {
|
||||
info!("🌐 nuwaxcode using registry: https://registry.npmjs.org/");
|
||||
self.run_command(
|
||||
npm_cmd,
|
||||
&[
|
||||
"install",
|
||||
"-g",
|
||||
&package_spec,
|
||||
"--registry=https://registry.npmjs.org/",
|
||||
],
|
||||
)
|
||||
.await?
|
||||
} else {
|
||||
self.run_command(npm_cmd, &["install", "-g", &package_spec])
|
||||
.await?
|
||||
};
|
||||
|
||||
if output.status.success() {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
|
||||
debug!("npm install stdout: {}", stdout);
|
||||
if !stderr.is_empty() {
|
||||
debug!("npm install stderr: {}", stderr);
|
||||
}
|
||||
|
||||
// Try to get installed version
|
||||
let version = self.parse_installed_version(&stdout);
|
||||
|
||||
info!("Package installed: {}", package_spec);
|
||||
Ok(InstallResult::success(
|
||||
version,
|
||||
format!("Successfully installed {}", package_spec),
|
||||
))
|
||||
} else {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
warn!("Installation failed: {}", stderr);
|
||||
Err(InstallationError::InstallFailed {
|
||||
package: package_spec,
|
||||
reason: stderr.to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn validate(&self, config: &InstallationConfig) -> Result<bool, InstallationError> {
|
||||
// Check if the package is globally installed
|
||||
let npm_cmd = self.get_npm_command().await;
|
||||
|
||||
if !self.command_exists(npm_cmd).await {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let package_name = match &config.package_name {
|
||||
Some(name) => name,
|
||||
None => return Ok(false),
|
||||
};
|
||||
|
||||
// Check global list
|
||||
let output = self
|
||||
.run_command(npm_cmd, &["list", "-g", "--depth=0", package_name])
|
||||
.await?;
|
||||
|
||||
Ok(output.status.success())
|
||||
}
|
||||
|
||||
async fn update(
|
||||
&self,
|
||||
config: &InstallationConfig,
|
||||
_install_dir: Option<&Path>,
|
||||
) -> Result<InstallResult, InstallationError> {
|
||||
let npm_cmd = self.get_npm_command().await;
|
||||
|
||||
if !self.command_exists(npm_cmd).await {
|
||||
return Err(InstallationError::PackageManagerNotAvailable(
|
||||
npm_cmd.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let package_name = config.package_name.as_ref().ok_or_else(|| {
|
||||
InstallationError::ConfigError("Package name is required".to_string())
|
||||
})?;
|
||||
|
||||
info!("Updating global package: {}", package_name);
|
||||
|
||||
// Run npm update -g
|
||||
// 仅 nuwaxcode 使用官方源(--registry 参数优先级高于 .npmrc)
|
||||
let output = if package_name == "nuwaxcode" {
|
||||
info!("🌐 nuwaxcode using registry: https://registry.npmjs.org/");
|
||||
self.run_command(
|
||||
npm_cmd,
|
||||
&[
|
||||
"update",
|
||||
"-g",
|
||||
package_name,
|
||||
"--registry=https://registry.npmjs.org/",
|
||||
],
|
||||
)
|
||||
.await?
|
||||
} else {
|
||||
self.run_command(npm_cmd, &["update", "-g", package_name])
|
||||
.await?
|
||||
};
|
||||
|
||||
if output.status.success() {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let version = self.parse_installed_version(&stdout);
|
||||
|
||||
info!("Update succeeded: {}", package_name);
|
||||
Ok(InstallResult::success(
|
||||
version,
|
||||
format!("Successfully updated {}", package_name),
|
||||
))
|
||||
} else {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
warn!("Update failed: {}", stderr);
|
||||
Err(InstallationError::InstallFailed {
|
||||
package: package_name.clone(),
|
||||
reason: stderr.to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn package_manager_name(&self) -> &'static str {
|
||||
"npm"
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_npm_command_detection() {
|
||||
let installer = NpmInstaller::new();
|
||||
let cmd = installer.get_npm_command().await;
|
||||
// Should return "npm" or "pnpm"
|
||||
assert!(cmd == "npm" || cmd == "pnpm");
|
||||
}
|
||||
}
|
||||
75
qiming-rcoder/crates/agent_config/src/installer/traits.rs
Normal file
75
qiming-rcoder/crates/agent_config/src/installer/traits.rs
Normal file
@@ -0,0 +1,75 @@
|
||||
//! Agent Installer Trait
|
||||
//!
|
||||
//! Defines the interface for agent installers.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use super::manager::{InstallResult, InstallationError};
|
||||
use crate::types::installation::InstallationConfig;
|
||||
|
||||
/// Agent installer trait
|
||||
///
|
||||
/// Defines the interface for installing, validating, and updating agents.
|
||||
#[async_trait::async_trait]
|
||||
pub trait AgentInstaller: Send + Sync {
|
||||
/// Install the agent
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `config` - Installation configuration
|
||||
/// * `install_dir` - Directory to install to (optional, for global installs)
|
||||
///
|
||||
/// # Returns
|
||||
/// Installation result with details
|
||||
async fn install(
|
||||
&self,
|
||||
config: &InstallationConfig,
|
||||
install_dir: Option<&Path>,
|
||||
) -> Result<InstallResult, InstallationError>;
|
||||
|
||||
/// Validate if the agent is installed correctly
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `config` - Installation configuration containing validate_command
|
||||
///
|
||||
/// # Returns
|
||||
/// true if installed and working, false otherwise
|
||||
async fn validate(&self, config: &InstallationConfig) -> Result<bool, InstallationError>;
|
||||
|
||||
/// Update the agent to latest version
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `config` - Installation configuration
|
||||
/// * `install_dir` - Directory where agent is installed
|
||||
///
|
||||
/// # Returns
|
||||
/// Installation result with new version info
|
||||
async fn update(
|
||||
&self,
|
||||
config: &InstallationConfig,
|
||||
install_dir: Option<&Path>,
|
||||
) -> Result<InstallResult, InstallationError>;
|
||||
|
||||
/// Get the package manager name
|
||||
fn package_manager_name(&self) -> &'static str;
|
||||
|
||||
/// Check if a command exists in PATH
|
||||
async fn command_exists(&self, command: &str) -> bool {
|
||||
which::which(command).is_ok()
|
||||
}
|
||||
|
||||
/// Run a command and get output
|
||||
async fn run_command(
|
||||
&self,
|
||||
command: &str,
|
||||
args: &[&str],
|
||||
) -> Result<std::process::Output, InstallationError> {
|
||||
tokio::process::Command::new(command)
|
||||
.args(args)
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| InstallationError::CommandFailed {
|
||||
command: format!("{} {}", command, args.join(" ")),
|
||||
reason: e.to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
33
qiming-rcoder/crates/agent_config/src/lib.rs
Normal file
33
qiming-rcoder/crates/agent_config/src/lib.rs
Normal file
@@ -0,0 +1,33 @@
|
||||
//! Agent Configuration Management
|
||||
//!
|
||||
//! This crate provides configuration management for agents in RCoder.
|
||||
|
||||
pub mod config;
|
||||
pub mod installer;
|
||||
pub mod types;
|
||||
|
||||
// Re-export main types
|
||||
pub use config::prompt_assembler::PromptConfigAssembler;
|
||||
pub use config::servers_config::AgentServersConfig;
|
||||
pub use types::agent_config::AgentConfig;
|
||||
pub use types::agent_spec::AgentSpec;
|
||||
pub use types::error::{AgentConfigError, ConfigError, Result};
|
||||
pub use types::installation::{InstallationConfig, PackageManager};
|
||||
pub use types::mcp_config::{ContextServerConfig, McpServerConfig, McpServerSource};
|
||||
pub use types::prompt_config::{SystemPromptConfig, UserPromptConfig};
|
||||
pub use types::system_prompt::DEFAULT_SYSTEM_PROMPT;
|
||||
#[allow(deprecated)]
|
||||
pub use types::system_prompt::PromptBuilder;
|
||||
|
||||
// Installer exports
|
||||
pub use installer::{
|
||||
AgentInstallationManager, AgentInstaller, InstallResult, InstallationError, NpmInstaller,
|
||||
};
|
||||
|
||||
// Default configuration exports
|
||||
pub use config::default_agent_config::{
|
||||
CLAUDE_CODE_ACP_AGENT_ID, default_agent_servers, default_agent_servers_for_service,
|
||||
default_context_servers, default_context_servers_for_service, get_default_agent,
|
||||
get_default_agent_for_service, get_default_config_by_service_type, get_default_context_server,
|
||||
get_default_context_server_for_service,
|
||||
};
|
||||
115
qiming-rcoder/crates/agent_config/src/types/agent_config.rs
Normal file
115
qiming-rcoder/crates/agent_config/src/types/agent_config.rs
Normal file
@@ -0,0 +1,115 @@
|
||||
//! Agent configuration structure.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use shared_types::ModelProviderConfig;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Agent configuration structure
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentConfig {
|
||||
/// Agent identifier
|
||||
pub agent_id: String,
|
||||
|
||||
/// Command to execute
|
||||
pub command: String,
|
||||
|
||||
/// Command arguments
|
||||
#[serde(default)]
|
||||
pub args: Vec<String>,
|
||||
|
||||
/// Environment variables
|
||||
#[serde(default)]
|
||||
pub env: HashMap<String, String>,
|
||||
|
||||
/// Installation configuration
|
||||
pub installation: super::installation::InstallationConfig,
|
||||
|
||||
/// System prompt configuration
|
||||
#[serde(default)]
|
||||
pub system_prompt: Option<super::prompt_config::SystemPromptConfig>,
|
||||
|
||||
/// User prompt configuration
|
||||
#[serde(default)]
|
||||
pub user_prompt: Option<super::prompt_config::UserPromptConfig>,
|
||||
|
||||
/// Model provider configuration
|
||||
#[serde(default)]
|
||||
pub model_provider: Option<ModelProviderConfig>,
|
||||
|
||||
/// Whether the agent is enabled
|
||||
#[serde(default = "default_enabled")]
|
||||
pub enabled: bool,
|
||||
|
||||
/// Metadata
|
||||
#[serde(default)]
|
||||
pub metadata: HashMap<String, String>,
|
||||
}
|
||||
|
||||
/// Default enabled value
|
||||
fn default_enabled() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
impl Default for AgentConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
agent_id: String::new(),
|
||||
command: String::new(),
|
||||
args: Vec::new(),
|
||||
env: HashMap::new(),
|
||||
installation: super::installation::InstallationConfig::default(),
|
||||
system_prompt: None,
|
||||
user_prompt: None,
|
||||
model_provider: None,
|
||||
enabled: true,
|
||||
metadata: HashMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AgentConfig {
|
||||
/// Create a new agent config
|
||||
pub fn new(agent_id: String) -> Self {
|
||||
Self {
|
||||
agent_id,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Get an environment variable
|
||||
pub fn get_env(&self, key: &str) -> Option<&String> {
|
||||
self.env.get(key)
|
||||
}
|
||||
|
||||
/// Set an environment variable
|
||||
pub fn set_env(&mut self, key: String, value: String) {
|
||||
self.env.insert(key, value);
|
||||
}
|
||||
|
||||
/// Check if agent is enabled
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
|
||||
/// Get model provider configuration
|
||||
pub fn get_model_provider(&self) -> Option<&ModelProviderConfig> {
|
||||
self.model_provider.as_ref()
|
||||
}
|
||||
|
||||
/// Set model provider configuration
|
||||
pub fn set_model_provider(&mut self, provider: ModelProviderConfig) {
|
||||
self.model_provider = Some(provider);
|
||||
}
|
||||
|
||||
/// Get the default model name if configured
|
||||
pub fn get_default_model(&self) -> Option<&str> {
|
||||
self.model_provider
|
||||
.as_ref()
|
||||
.map(|p| p.default_model.as_str())
|
||||
}
|
||||
|
||||
/// Get the model provider name if configured
|
||||
pub fn get_model_provider_name(&self) -> Option<&str> {
|
||||
self.model_provider.as_ref().map(|p| p.name.as_str())
|
||||
}
|
||||
}
|
||||
68
qiming-rcoder/crates/agent_config/src/types/agent_spec.rs
Normal file
68
qiming-rcoder/crates/agent_config/src/types/agent_spec.rs
Normal file
@@ -0,0 +1,68 @@
|
||||
//! Agent specification structure.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Agent specification
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentSpec {
|
||||
/// Agent identifier (matches the key in config file)
|
||||
pub agent_id: String,
|
||||
|
||||
/// Command to execute
|
||||
pub command: String,
|
||||
|
||||
/// Command arguments
|
||||
pub args: Vec<String>,
|
||||
|
||||
/// Environment variables
|
||||
pub env: HashMap<String, String>,
|
||||
|
||||
/// Installation configuration
|
||||
pub installation: super::installation::InstallationConfig,
|
||||
|
||||
/// System prompt configuration
|
||||
pub system_prompt: Option<super::prompt_config::SystemPromptConfig>,
|
||||
|
||||
/// User prompt configuration
|
||||
pub user_prompt: Option<super::prompt_config::UserPromptConfig>,
|
||||
|
||||
/// Whether the agent is enabled
|
||||
pub enabled: bool,
|
||||
|
||||
/// Metadata
|
||||
pub metadata: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl AgentSpec {
|
||||
/// Create a new agent spec
|
||||
pub fn new(agent_id: String) -> Self {
|
||||
Self {
|
||||
agent_id,
|
||||
command: String::new(),
|
||||
args: Vec::new(),
|
||||
env: HashMap::new(),
|
||||
installation: super::installation::InstallationConfig::default(),
|
||||
system_prompt: None,
|
||||
user_prompt: None,
|
||||
enabled: true,
|
||||
metadata: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert to AgentConfig
|
||||
pub fn to_agent_config(&self) -> super::agent_config::AgentConfig {
|
||||
super::agent_config::AgentConfig {
|
||||
agent_id: self.agent_id.clone(),
|
||||
command: self.command.clone(),
|
||||
args: self.args.clone(),
|
||||
env: self.env.clone(),
|
||||
installation: self.installation.clone(),
|
||||
system_prompt: self.system_prompt.clone(),
|
||||
user_prompt: self.user_prompt.clone(),
|
||||
model_provider: None,
|
||||
enabled: self.enabled,
|
||||
metadata: self.metadata.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
168
qiming-rcoder/crates/agent_config/src/types/error.rs
Normal file
168
qiming-rcoder/crates/agent_config/src/types/error.rs
Normal file
@@ -0,0 +1,168 @@
|
||||
//! Error types for agent configuration.
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
/// Agent configuration error
|
||||
#[derive(Error, Debug)]
|
||||
pub enum AgentConfigError {
|
||||
/// Configuration error
|
||||
#[error("configuration error: {0}")]
|
||||
Configuration(#[from] ConfigError),
|
||||
|
||||
/// I/O error
|
||||
#[error("I/O error: {0}")]
|
||||
Io(#[from] IoError),
|
||||
|
||||
/// Installation error
|
||||
#[error("installation error: {0}")]
|
||||
Installation(#[from] InstallationError),
|
||||
|
||||
/// Validation error
|
||||
#[error("validation error: {0}")]
|
||||
Validation(#[from] ValidationError),
|
||||
|
||||
/// Other error
|
||||
#[error("other error: {0}")]
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl From<serde_json::Error> for AgentConfigError {
|
||||
fn from(e: serde_json::Error) -> Self {
|
||||
Self::Configuration(ConfigError::Serialization(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<std::path::PathBuf> for AgentConfigError {
|
||||
fn from(path: std::path::PathBuf) -> Self {
|
||||
Self::Io(IoError::PathError(path))
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration error
|
||||
#[derive(Error, Debug)]
|
||||
pub enum ConfigError {
|
||||
/// Serialization error
|
||||
#[error("serialization error: {0}")]
|
||||
Serialization(String),
|
||||
|
||||
/// Deserialization error
|
||||
#[error("deserialization error: {0}")]
|
||||
Deserialization(String),
|
||||
|
||||
/// Missing required field
|
||||
#[error("missing required field: {0}")]
|
||||
MissingField(String),
|
||||
|
||||
/// Invalid value
|
||||
#[error("invalid value: {0}")]
|
||||
InvalidValue(String),
|
||||
|
||||
/// File not found
|
||||
#[error("file not found: {0}")]
|
||||
FileNotFound(String),
|
||||
|
||||
/// Permission denied
|
||||
#[error("permission denied: {0}")]
|
||||
PermissionDenied(String),
|
||||
|
||||
/// Read error
|
||||
#[error("read error: {0}")]
|
||||
Read(String),
|
||||
}
|
||||
|
||||
impl ConfigError {
|
||||
/// Create a missing field error
|
||||
pub fn missing_field(field: impl Into<String>) -> Self {
|
||||
Self::MissingField(field.into())
|
||||
}
|
||||
|
||||
/// Create an invalid value error
|
||||
pub fn invalid_value(msg: impl Into<String>) -> Self {
|
||||
Self::InvalidValue(msg.into())
|
||||
}
|
||||
|
||||
/// Create a file not found error
|
||||
pub fn file_not_found(path: impl std::fmt::Display) -> Self {
|
||||
Self::FileNotFound(path.to_string())
|
||||
}
|
||||
|
||||
/// Create a read error
|
||||
pub fn read(path: impl std::fmt::Display, source: impl std::fmt::Display) -> Self {
|
||||
Self::Read(format!("{}: {}", path, source))
|
||||
}
|
||||
}
|
||||
|
||||
/// I/O error
|
||||
#[derive(Error, Debug)]
|
||||
pub enum IoError {
|
||||
/// Path error
|
||||
#[error("path error: {0:?}")]
|
||||
PathError(std::path::PathBuf),
|
||||
|
||||
/// Read error
|
||||
#[error("read error: {0}")]
|
||||
Read(String),
|
||||
|
||||
/// Write error
|
||||
#[error("write error: {0}")]
|
||||
Write(String),
|
||||
|
||||
/// Create error
|
||||
#[error("create error: {0}")]
|
||||
Create(String),
|
||||
|
||||
/// Delete error
|
||||
#[error("delete error: {0}")]
|
||||
Delete(String),
|
||||
}
|
||||
|
||||
impl IoError {
|
||||
/// Create a read error
|
||||
pub fn read(path: impl Into<String>, source: impl ToString) -> Self {
|
||||
Self::Read(format!("{}: {}", path.into(), source.to_string()))
|
||||
}
|
||||
|
||||
/// Create a write error
|
||||
pub fn write(path: impl Into<String>, source: impl ToString) -> Self {
|
||||
Self::Write(format!("{}: {}", path.into(), source.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Installation error
|
||||
#[derive(Error, Debug)]
|
||||
pub enum InstallationError {
|
||||
/// Package not found
|
||||
#[error("package not found: {0}")]
|
||||
PackageNotFound(String),
|
||||
|
||||
/// Installation failed
|
||||
#[error("installation failed: {0}")]
|
||||
InstallationFailed(String),
|
||||
|
||||
/// Validation failed
|
||||
#[error("validation failed: {0}")]
|
||||
ValidationFailed(String),
|
||||
|
||||
/// Permission denied
|
||||
#[error("permission denied: {0}")]
|
||||
PermissionDenied(String),
|
||||
}
|
||||
|
||||
/// Validation error
|
||||
#[derive(Error, Debug)]
|
||||
pub enum ValidationError {
|
||||
/// Invalid configuration
|
||||
#[error("invalid configuration: {0}")]
|
||||
InvalidConfiguration(String),
|
||||
|
||||
/// Missing dependency
|
||||
#[error("missing dependency: {0}")]
|
||||
MissingDependency(String),
|
||||
|
||||
/// Constraint violation
|
||||
#[error("constraint violation: {0}")]
|
||||
ConstraintViolation(String),
|
||||
}
|
||||
|
||||
// Convenience aliases
|
||||
pub type Result<T> = std::result::Result<T, AgentConfigError>;
|
||||
122
qiming-rcoder/crates/agent_config/src/types/installation.rs
Normal file
122
qiming-rcoder/crates/agent_config/src/types/installation.rs
Normal file
@@ -0,0 +1,122 @@
|
||||
//! Installation configuration structures.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Package manager type
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum PackageManager {
|
||||
/// npm package manager
|
||||
Npm,
|
||||
/// Local binary
|
||||
Local,
|
||||
/// Custom package manager
|
||||
Custom(String),
|
||||
}
|
||||
|
||||
/// Installation configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct InstallationConfig {
|
||||
/// Package manager
|
||||
pub package_manager: PackageManager,
|
||||
|
||||
/// Package name
|
||||
#[serde(default)]
|
||||
pub package_name: Option<String>,
|
||||
|
||||
/// Version constraint
|
||||
#[serde(default)]
|
||||
pub version: Option<String>,
|
||||
|
||||
/// Installation source (for custom sources)
|
||||
#[serde(default)]
|
||||
pub source: Option<String>,
|
||||
|
||||
/// Validation command
|
||||
#[serde(default)]
|
||||
pub validate_command: Option<Vec<String>>,
|
||||
|
||||
/// Whether to auto-update
|
||||
#[serde(default)]
|
||||
pub auto_update: bool,
|
||||
}
|
||||
|
||||
impl Default for InstallationConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
package_manager: PackageManager::Npm,
|
||||
package_name: None,
|
||||
version: None,
|
||||
source: None,
|
||||
validate_command: None,
|
||||
auto_update: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl InstallationConfig {
|
||||
/// Create npm installation config
|
||||
pub fn npm(package_name: String) -> Self {
|
||||
Self {
|
||||
package_manager: PackageManager::Npm,
|
||||
package_name: Some(package_name),
|
||||
version: Some("latest".to_string()),
|
||||
source: None,
|
||||
validate_command: None,
|
||||
auto_update: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create local installation config
|
||||
pub fn local(command: String) -> Self {
|
||||
Self {
|
||||
package_manager: PackageManager::Local,
|
||||
package_name: Some(command),
|
||||
version: None,
|
||||
source: None,
|
||||
validate_command: None,
|
||||
auto_update: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create custom installation config
|
||||
pub fn custom(package_manager: String, source: String) -> Self {
|
||||
Self {
|
||||
package_manager: PackageManager::Custom(package_manager),
|
||||
package_name: None,
|
||||
version: None,
|
||||
source: Some(source),
|
||||
validate_command: None,
|
||||
auto_update: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set version
|
||||
pub fn with_version(mut self, version: String) -> Self {
|
||||
self.version = Some(version);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set validation command
|
||||
pub fn with_validation(mut self, command: Vec<String>) -> Self {
|
||||
self.validate_command = Some(command);
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable auto-update
|
||||
pub fn with_auto_update(mut self) -> Self {
|
||||
self.auto_update = true;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl PackageManager {
|
||||
/// Get the manager name as string
|
||||
pub fn as_str(&self) -> &str {
|
||||
match self {
|
||||
PackageManager::Npm => "npm",
|
||||
PackageManager::Local => "local",
|
||||
PackageManager::Custom(name) => name,
|
||||
}
|
||||
}
|
||||
}
|
||||
123
qiming-rcoder/crates/agent_config/src/types/mcp_config.rs
Normal file
123
qiming-rcoder/crates/agent_config/src/types/mcp_config.rs
Normal file
@@ -0,0 +1,123 @@
|
||||
//! MCP server configuration structures.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
|
||||
/// MCP server source type
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum McpServerSource {
|
||||
/// Custom command (npm, uvx, bun, etc.)
|
||||
Custom,
|
||||
/// Local executable file
|
||||
Local,
|
||||
}
|
||||
|
||||
/// MCP server configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpServerConfig {
|
||||
/// Server name
|
||||
pub name: String,
|
||||
|
||||
/// Server source type
|
||||
pub source: McpServerSource,
|
||||
|
||||
/// Whether enabled
|
||||
pub enabled: bool,
|
||||
|
||||
/// Command to execute (for custom/local sources)
|
||||
pub command: Option<String>,
|
||||
|
||||
/// Command arguments
|
||||
pub args: Option<Vec<String>>,
|
||||
|
||||
/// Environment variables
|
||||
pub env: Option<HashMap<String, String>>,
|
||||
|
||||
/// Connection timeout
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
/// Context server configuration (simplified format for config files)
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ContextServerConfig {
|
||||
/// Server source type
|
||||
pub source: String,
|
||||
|
||||
/// Whether enabled
|
||||
pub enabled: bool,
|
||||
|
||||
/// Command to execute
|
||||
pub command: Option<String>,
|
||||
|
||||
/// Command arguments
|
||||
pub args: Option<Vec<String>>,
|
||||
|
||||
/// Environment variables
|
||||
pub env: Option<HashMap<String, String>>,
|
||||
}
|
||||
|
||||
impl McpServerConfig {
|
||||
/// Create a new MCP server config
|
||||
pub fn new(name: String, source: McpServerSource) -> Self {
|
||||
Self {
|
||||
name,
|
||||
source,
|
||||
enabled: true,
|
||||
command: None,
|
||||
args: None,
|
||||
env: None,
|
||||
timeout: Some(Duration::from_secs(30)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a custom MCP server
|
||||
pub fn custom(name: String, command: String) -> Self {
|
||||
Self {
|
||||
name,
|
||||
source: McpServerSource::Custom,
|
||||
enabled: true,
|
||||
command: Some(command),
|
||||
args: None,
|
||||
env: None,
|
||||
timeout: Some(Duration::from_secs(30)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a local MCP server
|
||||
pub fn local(name: String, command: String) -> Self {
|
||||
Self {
|
||||
name,
|
||||
source: McpServerSource::Local,
|
||||
enabled: true,
|
||||
command: Some(command),
|
||||
args: None,
|
||||
env: None,
|
||||
timeout: Some(Duration::from_secs(30)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if server is enabled
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
}
|
||||
|
||||
impl ContextServerConfig {
|
||||
/// Create a new context server config
|
||||
pub fn new(source: String) -> Self {
|
||||
Self {
|
||||
source,
|
||||
enabled: true,
|
||||
command: None,
|
||||
args: None,
|
||||
env: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if server is enabled
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
}
|
||||
19
qiming-rcoder/crates/agent_config/src/types/mod.rs
Normal file
19
qiming-rcoder/crates/agent_config/src/types/mod.rs
Normal file
@@ -0,0 +1,19 @@
|
||||
//! Core type definitions for agent configuration.
|
||||
|
||||
pub mod agent_config;
|
||||
pub mod agent_spec;
|
||||
pub mod error;
|
||||
pub mod installation;
|
||||
pub mod mcp_config;
|
||||
pub mod prompt_config;
|
||||
pub mod system_prompt;
|
||||
|
||||
pub use agent_config::*;
|
||||
pub use agent_spec::*;
|
||||
pub use error::*;
|
||||
pub use installation::*;
|
||||
pub use mcp_config::*;
|
||||
pub use prompt_config::*;
|
||||
pub use system_prompt::DEFAULT_SYSTEM_PROMPT;
|
||||
#[allow(deprecated)]
|
||||
pub use system_prompt::PromptBuilder;
|
||||
240
qiming-rcoder/crates/agent_config/src/types/prompt_config.rs
Normal file
240
qiming-rcoder/crates/agent_config/src/types/prompt_config.rs
Normal file
@@ -0,0 +1,240 @@
|
||||
//! Prompt configuration structures.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::system_prompt::DEFAULT_SYSTEM_PROMPT;
|
||||
|
||||
/// 系统提示词来源
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum SystemPromptSource {
|
||||
/// 使用编译时嵌入的默认提示词
|
||||
#[default]
|
||||
Embedded,
|
||||
/// 使用自定义模板内容
|
||||
Custom,
|
||||
}
|
||||
|
||||
/// System prompt configuration
|
||||
///
|
||||
/// 支持两种模式:
|
||||
/// 1. `source: "embedded"` - 使用编译时嵌入的默认提示词(忽略 template 字段)
|
||||
/// 2. `source: "custom"` - 使用 template 字段的自定义内容
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SystemPromptConfig {
|
||||
/// 提示词来源
|
||||
/// - "embedded": 使用编译时嵌入的默认提示词
|
||||
/// - "custom": 使用 template 字段的自定义内容
|
||||
#[serde(default)]
|
||||
pub source: SystemPromptSource,
|
||||
|
||||
/// Template content (仅当 source = "custom" 时使用)
|
||||
#[serde(default)]
|
||||
pub template: String,
|
||||
|
||||
/// Whether enabled (false 时不传递系统提示词)
|
||||
#[serde(default = "default_enabled")]
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
/// User prompt configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct UserPromptConfig {
|
||||
/// Template content with {user_prompt} placeholder
|
||||
pub template: String,
|
||||
|
||||
/// Whether enabled
|
||||
#[serde(default = "default_enabled")]
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
/// Default enabled value
|
||||
fn default_enabled() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
impl Default for SystemPromptConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
source: SystemPromptSource::Embedded,
|
||||
template: String::new(),
|
||||
enabled: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SystemPromptConfig {
|
||||
/// 创建使用嵌入默认提示词的配置
|
||||
pub fn embedded() -> Self {
|
||||
Self {
|
||||
source: SystemPromptSource::Embedded,
|
||||
template: String::new(),
|
||||
enabled: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建使用自定义模板的配置
|
||||
pub fn custom(template: String) -> Self {
|
||||
Self {
|
||||
source: SystemPromptSource::Custom,
|
||||
template,
|
||||
enabled: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建禁用的配置(不传递系统提示词)
|
||||
pub fn disabled() -> Self {
|
||||
Self {
|
||||
source: SystemPromptSource::Embedded,
|
||||
template: String::new(),
|
||||
enabled: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取系统提示词内容
|
||||
///
|
||||
/// - 如果 enabled = false,返回 None
|
||||
/// - 如果 source = "embedded",返回编译时嵌入的默认提示词
|
||||
/// - 如果 source = "custom",返回 template 字段的内容
|
||||
pub fn get_prompt(&self) -> Option<&str> {
|
||||
if !self.enabled {
|
||||
return None;
|
||||
}
|
||||
|
||||
match self.source {
|
||||
SystemPromptSource::Embedded => Some(DEFAULT_SYSTEM_PROMPT),
|
||||
SystemPromptSource::Custom => {
|
||||
if self.template.is_empty() {
|
||||
// custom 模式但 template 为空,回退到默认
|
||||
Some(DEFAULT_SYSTEM_PROMPT)
|
||||
} else {
|
||||
Some(&self.template)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取系统提示词内容(兼容旧接口,始终返回有效提示词)
|
||||
///
|
||||
/// 如果 enabled = false,仍返回默认提示词(用于需要始终有提示词的场景)
|
||||
pub fn get_prompt_or_default(&self) -> &str {
|
||||
match self.source {
|
||||
SystemPromptSource::Embedded => DEFAULT_SYSTEM_PROMPT,
|
||||
SystemPromptSource::Custom => {
|
||||
if self.template.is_empty() {
|
||||
DEFAULT_SYSTEM_PROMPT
|
||||
} else {
|
||||
&self.template
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查是否使用自定义模板
|
||||
pub fn is_custom(&self) -> bool {
|
||||
self.source == SystemPromptSource::Custom && !self.template.is_empty()
|
||||
}
|
||||
|
||||
/// 检查是否使用嵌入的默认提示词
|
||||
pub fn is_embedded(&self) -> bool {
|
||||
self.source == SystemPromptSource::Embedded
|
||||
}
|
||||
}
|
||||
|
||||
impl UserPromptConfig {
|
||||
/// Create a new user prompt config
|
||||
pub fn new(template: String) -> Self {
|
||||
Self {
|
||||
template,
|
||||
enabled: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create disabled config
|
||||
pub fn disabled() -> Self {
|
||||
Self {
|
||||
template: String::new(),
|
||||
enabled: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply user prompt template
|
||||
pub fn apply(&self, user_input: &str) -> String {
|
||||
if self.enabled {
|
||||
self.template.replace("{user_prompt}", user_input)
|
||||
} else {
|
||||
user_input.to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_embedded_source() {
|
||||
let config = SystemPromptConfig::embedded();
|
||||
assert!(config.is_embedded());
|
||||
assert!(!config.is_custom());
|
||||
assert!(config.get_prompt().is_some());
|
||||
assert!(
|
||||
config
|
||||
.get_prompt()
|
||||
.unwrap()
|
||||
.contains("<SYSTEM_INSTRUCTIONS>")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_custom_source() {
|
||||
let config = SystemPromptConfig::custom("Custom prompt".to_string());
|
||||
assert!(config.is_custom());
|
||||
assert!(!config.is_embedded());
|
||||
assert_eq!(config.get_prompt(), Some("Custom prompt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_disabled_config() {
|
||||
let config = SystemPromptConfig::disabled();
|
||||
assert!(config.get_prompt().is_none());
|
||||
// get_prompt_or_default 仍返回默认值
|
||||
assert!(
|
||||
config
|
||||
.get_prompt_or_default()
|
||||
.contains("<SYSTEM_INSTRUCTIONS>")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_custom_empty_template_fallback() {
|
||||
let config = SystemPromptConfig {
|
||||
source: SystemPromptSource::Custom,
|
||||
template: String::new(),
|
||||
enabled: true,
|
||||
};
|
||||
// 空模板时回退到默认
|
||||
assert!(
|
||||
config
|
||||
.get_prompt()
|
||||
.unwrap()
|
||||
.contains("<SYSTEM_INSTRUCTIONS>")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_deserialization() {
|
||||
let json = r#"{"source": "embedded", "template": "", "enabled": true}"#;
|
||||
let config: SystemPromptConfig = serde_json::from_str(json).unwrap();
|
||||
assert!(config.is_embedded());
|
||||
assert!(config.enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_deserialization_custom() {
|
||||
let json = r#"{"source": "custom", "template": "My prompt", "enabled": true}"#;
|
||||
let config: SystemPromptConfig = serde_json::from_str(json).unwrap();
|
||||
assert!(config.is_custom());
|
||||
assert_eq!(config.get_prompt(), Some("My prompt"));
|
||||
}
|
||||
}
|
||||
107
qiming-rcoder/crates/agent_config/src/types/system_prompt.rs
Normal file
107
qiming-rcoder/crates/agent_config/src/types/system_prompt.rs
Normal file
@@ -0,0 +1,107 @@
|
||||
//! System prompt configuration types
|
||||
//!
|
||||
//! This module provides compile-time embedded default system prompt constants and related helper functions.
|
||||
//! Default system prompt is embedded at compile time from external files, with runtime override support via configuration.
|
||||
|
||||
/// Compile-time embedded default system prompt
|
||||
pub const DEFAULT_SYSTEM_PROMPT: &str = include_str!("../../configs/prompts/frontend_expert.txt");
|
||||
|
||||
/// Prompt builder (kept for compatibility, recommend using SystemPromptConfig::get_prompt() directly)
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptBuilder;
|
||||
|
||||
impl PromptBuilder {
|
||||
/// Build user prompt (without system prompt)
|
||||
pub fn build_user_prompt(user_prompt: &str) -> String {
|
||||
user_prompt.to_string()
|
||||
}
|
||||
|
||||
/// Build user prompt with data sources
|
||||
pub fn build_user_prompt_with_data_sources(
|
||||
user_prompt: &str,
|
||||
data_sources: &[String],
|
||||
) -> String {
|
||||
if data_sources.is_empty() {
|
||||
return user_prompt.to_string();
|
||||
}
|
||||
|
||||
let data_sources_section = format_data_sources(data_sources);
|
||||
format!(
|
||||
"{}\n\n\
|
||||
<DATA_SOURCES>\n\
|
||||
The following are available data source information, including backend API endpoints, database connections, and other external data sources.\n\
|
||||
When developing frontend applications, you can use these data sources to fetch real data, such as querying Bitcoin transaction volumes, stock prices, weather information, etc.\n\
|
||||
Please use these data sources reasonably according to development needs, and ensure the frontend application can correctly call the relevant interfaces.\n\
|
||||
Use Axios client or Fetch API for API calls, or according to the interface calling method of the current framework.\n\n\
|
||||
{}\n\
|
||||
</DATA_SOURCES>",
|
||||
user_prompt, data_sources_section
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Format data source information into readable text
|
||||
fn format_data_sources(data_sources: &[String]) -> String {
|
||||
if data_sources.is_empty() {
|
||||
return "No data sources".to_string();
|
||||
}
|
||||
|
||||
let mut formatted = String::new();
|
||||
|
||||
for (index, data_source) in data_sources.iter().enumerate() {
|
||||
formatted.push_str(&format!("Data source {}:\n", index + 1));
|
||||
|
||||
// 尝试解析 JSON 字符串并格式化
|
||||
match serde_json::from_str::<serde_json::Value>(data_source) {
|
||||
Ok(json_value) => {
|
||||
// 成功解析,格式化为易读的 JSON
|
||||
match serde_json::to_string_pretty(&json_value) {
|
||||
Ok(pretty_json) => {
|
||||
formatted.push_str(&pretty_json);
|
||||
}
|
||||
Err(_) => {
|
||||
// 格式化失败,使用原始字符串
|
||||
formatted.push_str(data_source);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
// 不是有效的 JSON,直接使用原始字符串
|
||||
formatted.push_str(data_source);
|
||||
}
|
||||
}
|
||||
|
||||
formatted.push('\n');
|
||||
}
|
||||
|
||||
formatted
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_default_system_prompt_loaded() {
|
||||
// 验证默认提示词被正确加载
|
||||
assert!(!DEFAULT_SYSTEM_PROMPT.is_empty());
|
||||
assert!(DEFAULT_SYSTEM_PROMPT.contains("<SYSTEM_INSTRUCTIONS>"));
|
||||
assert!(DEFAULT_SYSTEM_PROMPT.contains("</SYSTEM_INSTRUCTIONS>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_data_sources() {
|
||||
let data_sources = vec![r#"{"api": "https://api.example.com"}"#.to_string()];
|
||||
let formatted = format_data_sources(&data_sources);
|
||||
|
||||
assert!(formatted.contains("Data source 1"));
|
||||
assert!(formatted.contains("api.example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_empty_data_sources() {
|
||||
let data_sources: Vec<String> = vec![];
|
||||
let formatted = format_data_sources(&data_sources);
|
||||
assert_eq!(formatted, "No data sources");
|
||||
}
|
||||
}
|
||||
122
qiming-rcoder/crates/agent_runner/Cargo.toml
Normal file
122
qiming-rcoder/crates/agent_runner/Cargo.toml
Normal file
@@ -0,0 +1,122 @@
|
||||
[package]
|
||||
name = "agent_runner"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[lib]
|
||||
name = "agent_runner"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[features]
|
||||
default = ["http-server", "grpc-server"]
|
||||
|
||||
# HTTP 服务器功能:启用 HTTP REST API 和 SSE 进度流
|
||||
# 注意:启用 http-server 时会同时启动 gRPC(除非明确禁用)
|
||||
http-server = []
|
||||
|
||||
# gRPC 服务功能:启用 gRPC 服务
|
||||
# 在 http-server 模式下:
|
||||
# - 启用 grpc-server:同时运行 HTTP + gRPC
|
||||
# - 禁用 grpc-server:仅运行 HTTP(无 gRPC)
|
||||
grpc-server = []
|
||||
|
||||
# Pingora 代理模式:API Key 和 Base URL 使用占位符,由 Pingora 代理注入真实值
|
||||
# Windows 不支持 Pingora,禁用此 feature 时直接使用真实的 API Key 和 Base URL
|
||||
proxy = ["agent_abstraction/proxy", "dep:rcoder-proxy"]
|
||||
|
||||
# 测试相关 features
|
||||
testing = ["test-blocking"]
|
||||
test-blocking = [] # 启用测试用阻塞注入 (用于极端场景测试)
|
||||
|
||||
# OpenTelemetry 追踪 feature
|
||||
otel = [] # 启用 OpenTelemetry 追踪功能
|
||||
|
||||
# Pyroscope 连续性能分析
|
||||
pyroscope = ["dep:pyroscope", "dep:pyroscope_pprofrs"]
|
||||
|
||||
[dependencies]
|
||||
# Internal crates
|
||||
shared_types = { path = "../shared_types" }
|
||||
agent_config = { path = "../agent_config" }
|
||||
agent_abstraction = { path = "../agent_abstraction" }
|
||||
rcoder-proxy = { path = "../rcoder-proxy", optional = true }
|
||||
rcoder-telemetry = { workspace = true }
|
||||
|
||||
clap = { workspace = true }
|
||||
# ACP Protocol - 官方 SDK(完全 Send-safe,无需 LocalSet)
|
||||
agent-client-protocol = { workspace = true }
|
||||
async-trait = "0.1"
|
||||
parking_lot = "0.12"
|
||||
tempfile = { workspace = true }
|
||||
|
||||
# HTTP server
|
||||
axum = { workspace = true, features = ["multipart"] }
|
||||
tower = { workspace = true }
|
||||
tower-http = { workspace = true, features = ["cors", "trace", "limit"] }
|
||||
futures = { workspace = true }
|
||||
tokio-stream = "0.1"
|
||||
async-stream = "0.3"
|
||||
|
||||
# gRPC Server
|
||||
tonic = { workspace = true }
|
||||
|
||||
# Serialization
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
serde_yaml = "0.9"
|
||||
|
||||
derive_builder = { workspace = true }
|
||||
|
||||
# Workspace dependencies
|
||||
tokio = { workspace = true, features = ["full"] }
|
||||
tokio-util = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tracing-subscriber = { workspace = true, features = [
|
||||
"fmt",
|
||||
"env-filter",
|
||||
"json",
|
||||
] }
|
||||
tracing-appender = "0.2"
|
||||
uuid = { workspace = true, features = ["v4", "serde"] }
|
||||
chrono = { workspace = true, features = ["serde"] }
|
||||
dashmap = { workspace = true }
|
||||
bytes = "1.0"
|
||||
futures-util = { workspace = true }
|
||||
piper = "0.2"
|
||||
base64 = { workspace = true }
|
||||
reqwest = { workspace = true }
|
||||
rustls = { workspace = true }
|
||||
|
||||
# OpenTelemetry
|
||||
opentelemetry = { workspace = true }
|
||||
opentelemetry_sdk = { workspace = true }
|
||||
tracing-opentelemetry = { workspace = true }
|
||||
axum-tracing-opentelemetry = { workspace = true }
|
||||
|
||||
# File system
|
||||
tokio-fs = { workspace = true }
|
||||
path-clean = { workspace = true }
|
||||
|
||||
# OpenAPI documentation
|
||||
utoipa = { workspace = true }
|
||||
utoipa-swagger-ui = { workspace = true }
|
||||
|
||||
# Validation
|
||||
garde = { version = "0.22.1", features = ["derive"] }
|
||||
|
||||
ringbuf = { workspace = true }
|
||||
|
||||
# Atomic swap for sender replacement
|
||||
arc-swap = "1.8"
|
||||
|
||||
# Pyroscope Continuous Profiling
|
||||
pyroscope = { workspace = true, optional = true }
|
||||
pyroscope_pprofrs = { workspace = true, optional = true }
|
||||
|
||||
# Unix 信号处理(用于僵尸进程回收,仅 Unix 平台)
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
nix = { workspace = true, features = ["signal", "process"] }
|
||||
355
qiming-rcoder/crates/agent_runner/src/agent_runtime.rs
Normal file
355
qiming-rcoder/crates/agent_runner/src/agent_runtime.rs
Normal file
@@ -0,0 +1,355 @@
|
||||
//! Agent Runtime 模块
|
||||
//!
|
||||
//! 简化版的 Agent Worker 管理器,利用 SACP 的 Send trait 支持。
|
||||
//!
|
||||
//! ## 新架构设计
|
||||
//!
|
||||
//! - 移除独立 OS 线程,使用 `tokio::spawn`
|
||||
//! - 简化 sender 管理,移除 ArcSwap
|
||||
//! - 保留自动重启功能
|
||||
//! - 保留心跳检测(僵尸检测)
|
||||
//! - 简化状态机(使用原子操作)
|
||||
//!
|
||||
//! ## 与旧架构对比
|
||||
//!
|
||||
//! | 组件 | 旧设计 | 新设计 |
|
||||
//! |------|--------|--------|
|
||||
//! | 运行环境 | 独立 OS 线程 + 独立运行时 | 主运行时 + tokio::spawn |
|
||||
//! | Sender 管理 | ArcSwap<Option<Sender>> | mpsc::Sender (固定) |
|
||||
//! | Worker 生命周期 | 手动管理线程 | JoinHandle |
|
||||
//! | Ready 信号 | oneshot::channel | JoinHandle 完成 |
|
||||
//! | 状态机 | watch::Sender + Mutex | Arc<AtomicState> |
|
||||
//! | 重启 | 替换 sender | abort + spawn |
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicI64, AtomicU8, AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::Utc;
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::proxy_agent::AgentRequest;
|
||||
|
||||
/// Worker 状态
|
||||
#[repr(u8)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum WorkerState {
|
||||
/// 启动中
|
||||
Starting = 0,
|
||||
/// 运行中
|
||||
Running = 1,
|
||||
/// 停止中
|
||||
Stopping = 2,
|
||||
/// 已停止
|
||||
Stopped = 3,
|
||||
}
|
||||
|
||||
/// 原子状态包装器 (无需 Mutex)
|
||||
pub struct AtomicState(AtomicU8);
|
||||
|
||||
impl AtomicState {
|
||||
pub fn new(state: WorkerState) -> Self {
|
||||
Self(AtomicU8::new(state as u8))
|
||||
}
|
||||
|
||||
pub fn get(&self) -> WorkerState {
|
||||
match self.0.load(Ordering::Acquire) {
|
||||
0 => WorkerState::Starting,
|
||||
1 => WorkerState::Running,
|
||||
2 => WorkerState::Stopping,
|
||||
3 => WorkerState::Stopped,
|
||||
invalid => {
|
||||
tracing::error!(
|
||||
"[AtomicState] Invalid state value: {}, falling back to Stopped",
|
||||
invalid
|
||||
);
|
||||
WorkerState::Stopped
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set(&self, state: WorkerState) {
|
||||
self.0.store(state as u8, Ordering::Release);
|
||||
}
|
||||
}
|
||||
|
||||
/// 心跳包
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Heartbeat {
|
||||
/// 心跳时间戳
|
||||
pub timestamp: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
/// Worker 就绪信号
|
||||
///
|
||||
/// Worker 在初始化完成后发送此信号
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct WorkerReady {
|
||||
/// 就绪时间戳
|
||||
pub timestamp: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
/// 并发控制配置
|
||||
///
|
||||
/// 工作线程池大小 - 决定可以并发处理的 Agent 会话数量
|
||||
/// 🔥 已改为运行时可配置的全局变量,使用 get_concurrency_limit() 获取
|
||||
|
||||
/// 全局并发限制(运行时可配置)
|
||||
pub static WORKER_THREAD_POOL_SIZE: AtomicUsize = AtomicUsize::new(10);
|
||||
|
||||
/// 初始化并发限制(在应用启动时调用)
|
||||
pub fn init_concurrency_limit(limit: usize) {
|
||||
WORKER_THREAD_POOL_SIZE.store(limit, Ordering::Release);
|
||||
info!("🔧 Concurrency limit initialized: {}", limit);
|
||||
}
|
||||
|
||||
/// 获取当前并发限制
|
||||
pub fn get_concurrency_limit() -> usize {
|
||||
WORKER_THREAD_POOL_SIZE.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
/// Agent 运行时
|
||||
///
|
||||
/// 替代 AgentWorkerManager,使用简化的架构:
|
||||
/// - 直接在主运行时中运行 (SACP 支持 Send)
|
||||
/// - 使用原子操作管理状态
|
||||
/// - 使用 JoinHandle 管理生命周期
|
||||
pub struct AgentRuntime {
|
||||
/// 请求发送端 (固定不变)
|
||||
request_tx: mpsc::Sender<AgentRequest>,
|
||||
|
||||
/// 当前 Worker 的 JoinHandle
|
||||
worker_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
|
||||
|
||||
/// 当前状态
|
||||
state: Arc<AtomicState>,
|
||||
|
||||
/// 🔥 P1 修复: 最后心跳时间戳(毫秒,Unix timestamp)
|
||||
/// 使用 AtomicI64 替代 Mutex,避免频繁的锁竞争
|
||||
/// - 0 表示从未收到心跳
|
||||
/// - 正数表示最后一次心跳的 timestamp_millis()
|
||||
last_heartbeat_ts: Arc<AtomicI64>,
|
||||
|
||||
/// 活跃请求追踪: request_id -> 开始时间
|
||||
active_requests: Arc<Mutex<HashMap<String, chrono::DateTime<chrono::Utc>>>>,
|
||||
|
||||
/// 心跳超时阈值
|
||||
heartbeat_timeout: Duration,
|
||||
|
||||
/// 首次启动宽限期
|
||||
initial_grace_period: Duration,
|
||||
}
|
||||
|
||||
impl AgentRuntime {
|
||||
/// 创建新的 AgentRuntime
|
||||
///
|
||||
/// 返回 (runtime, request_receiver)
|
||||
pub fn new(request_buffer: usize) -> (Self, mpsc::Receiver<AgentRequest>) {
|
||||
let (request_tx, request_rx) = mpsc::channel(request_buffer);
|
||||
|
||||
let runtime = Self {
|
||||
request_tx,
|
||||
worker_handle: Arc::new(Mutex::new(None)),
|
||||
state: Arc::new(AtomicState::new(WorkerState::Starting)),
|
||||
last_heartbeat_ts: Arc::new(AtomicI64::new(0)),
|
||||
active_requests: Arc::new(Mutex::new(HashMap::new())),
|
||||
heartbeat_timeout: Duration::from_secs(15),
|
||||
initial_grace_period: Duration::from_secs(30),
|
||||
};
|
||||
|
||||
(runtime, request_rx)
|
||||
}
|
||||
|
||||
/// 启动 Worker (在主运行时中)
|
||||
///
|
||||
/// SACP 支持 Send,直接在主运行时中运行,无需独立线程
|
||||
pub async fn start(&self, receiver: mpsc::Receiver<AgentRequest>) {
|
||||
let state = self.state.clone();
|
||||
let last_heartbeat_ts = self.last_heartbeat_ts.clone();
|
||||
let active_requests = self.active_requests.clone();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
// SACP 支持 Send,直接在主运行时中运行
|
||||
if let Err(e) = crate::proxy_agent::agent_worker_with_heartbeat(
|
||||
receiver,
|
||||
state.clone(),
|
||||
last_heartbeat_ts.clone(),
|
||||
active_requests.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::error!("Agent worker failed: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
*self.worker_handle.lock().await = Some(handle);
|
||||
self.state.set(WorkerState::Running);
|
||||
info!("AgentRuntime: worker started");
|
||||
}
|
||||
|
||||
/// 重启 Worker
|
||||
pub async fn restart(&self, new_receiver: mpsc::Receiver<AgentRequest>) {
|
||||
warn!("AgentRuntime: preparing to restart worker...");
|
||||
|
||||
// 1. 停止旧 worker
|
||||
if let Some(handle) = self.worker_handle.lock().await.take() {
|
||||
handle.abort();
|
||||
info!("AgentRuntime: previous worker terminated");
|
||||
}
|
||||
|
||||
// 2. 重置状态
|
||||
self.state.set(WorkerState::Starting);
|
||||
self.last_heartbeat_ts.store(0, Ordering::Release);
|
||||
*self.active_requests.lock().await = HashMap::new();
|
||||
|
||||
// 3. 启动新 worker
|
||||
self.start(new_receiver).await;
|
||||
info!("AgentRuntime: worker restart completed");
|
||||
}
|
||||
|
||||
/// 发送请求
|
||||
pub async fn send(&self, request: AgentRequest) -> anyhow::Result<()> {
|
||||
self.request_tx
|
||||
.send(request)
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("Worker is closed"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 健康检查
|
||||
pub async fn check_health(&self) -> bool {
|
||||
let state = self.state.get();
|
||||
|
||||
// 检查状态
|
||||
if state == WorkerState::Stopped {
|
||||
return false;
|
||||
}
|
||||
|
||||
// 🔥 P1 修复: 使用原子操作检查心跳(无锁)
|
||||
let last_ts = self.last_heartbeat_ts.load(Ordering::Acquire);
|
||||
if last_ts > 0 {
|
||||
let elapsed_ms = Utc::now().timestamp_millis() - last_ts;
|
||||
elapsed_ms < self.heartbeat_timeout.as_millis() as i64
|
||||
} else {
|
||||
// 首次启动宽限期
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取当前状态
|
||||
pub fn state(&self) -> WorkerState {
|
||||
self.state.get()
|
||||
}
|
||||
|
||||
/// 🔥 P1 修复: 检查心跳是否超时(无锁)
|
||||
///
|
||||
/// ## 返回值
|
||||
///
|
||||
/// - `true`: 心跳超时(超过 15 秒未收到心跳)
|
||||
/// - `false`: 心跳正常或在宽限期内
|
||||
pub fn check_heartbeat_timeout(&self) -> bool {
|
||||
let last_ts = self.last_heartbeat_ts.load(Ordering::Acquire);
|
||||
|
||||
if last_ts > 0 {
|
||||
// 有心跳记录,检查是否超过 15 秒
|
||||
let elapsed_ms = Utc::now().timestamp_millis() - last_ts;
|
||||
elapsed_ms > 15_000
|
||||
} else {
|
||||
// 从未收到心跳,使用首次启动宽限期
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// 🔥 P1 修复: 获取最后心跳时间(无锁)
|
||||
///
|
||||
/// ## 返回值
|
||||
///
|
||||
/// - `Some(timestamp)`: 最后心跳时间
|
||||
/// - `None`: 从未收到心跳
|
||||
pub fn last_heartbeat_time(&self) -> Option<chrono::DateTime<chrono::Utc>> {
|
||||
let last_ts = self.last_heartbeat_ts.load(Ordering::Acquire);
|
||||
if last_ts > 0 {
|
||||
// 将毫秒时间戳转换为 DateTime
|
||||
use chrono::TimeZone;
|
||||
// timestamp_millis_opt 返回 LocalResult,使用 single() 转换为 Option
|
||||
match chrono::Utc.timestamp_millis_opt(last_ts).single() {
|
||||
Some(dt) => Some(dt),
|
||||
None => {
|
||||
tracing::warn!(
|
||||
"[WorkerInfo] Invalid timestamp: {}, using current time",
|
||||
last_ts
|
||||
);
|
||||
Some(chrono::Utc::now())
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取活跃请求句柄
|
||||
pub fn active_requests(&self) -> Arc<Mutex<HashMap<String, chrono::DateTime<chrono::Utc>>>> {
|
||||
self.active_requests.clone()
|
||||
}
|
||||
|
||||
/// 检查请求通道是否已关闭
|
||||
pub fn is_closed(&self) -> bool {
|
||||
self.request_tx.is_closed()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_atomic_state() {
|
||||
let state = AtomicState::new(WorkerState::Starting);
|
||||
assert_eq!(state.get(), WorkerState::Starting);
|
||||
|
||||
state.set(WorkerState::Running);
|
||||
assert_eq!(state.get(), WorkerState::Running);
|
||||
|
||||
state.set(WorkerState::Stopped);
|
||||
assert_eq!(state.get(), WorkerState::Stopped);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_runtime_creation() {
|
||||
let (runtime, _rx) = AgentRuntime::new(100);
|
||||
assert_eq!(runtime.state(), WorkerState::Starting);
|
||||
assert!(!runtime.is_closed());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_heartbeat_timeout_detection() {
|
||||
let (runtime, _rx) = AgentRuntime::new(100);
|
||||
|
||||
// 初始状态:从未收到心跳
|
||||
assert!(!runtime.check_heartbeat_timeout());
|
||||
|
||||
// 模拟心跳超时(设置20秒前的时间戳)
|
||||
let timestamp_20s_ago = Utc::now().timestamp_millis() - (20 * 1000);
|
||||
runtime
|
||||
.last_heartbeat_ts
|
||||
.store(timestamp_20s_ago, std::sync::atomic::Ordering::Release);
|
||||
|
||||
// 心跳超过 15 秒,应检测到超时
|
||||
assert!(runtime.check_heartbeat_timeout());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_health_check() {
|
||||
let (runtime, _rx) = AgentRuntime::new(100);
|
||||
|
||||
// 初始状态应该是健康的
|
||||
assert!(runtime.check_health().await);
|
||||
|
||||
// 设置状态为 Stopped
|
||||
runtime.state.set(WorkerState::Stopped);
|
||||
assert!(!runtime.check_health().await);
|
||||
}
|
||||
}
|
||||
317
qiming-rcoder/crates/agent_runner/src/api_key_manager.rs
Normal file
317
qiming-rcoder/crates/agent_runner/src/api_key_manager.rs
Normal file
@@ -0,0 +1,317 @@
|
||||
//! API 密钥管理器
|
||||
//!
|
||||
//! 在内存中存储 API 密钥配置,支持通过服务名称快速查询。
|
||||
//! 密钥配置通过 gRPC 从 rcoder 主服务传递到 agent_runner。
|
||||
//!
|
||||
//! ## 使用示例
|
||||
//!
|
||||
//! ```rust
|
||||
//! use agent_runner::api_key_manager::ApiKeyManager;
|
||||
//! use shared_types::ModelProviderConfig;
|
||||
//!
|
||||
//! let manager = ApiKeyManager::new();
|
||||
//!
|
||||
//! // 存储 API 配置
|
||||
//! let config = ModelProviderConfig {
|
||||
//! id: "anthropic-prov".to_string(),
|
||||
//! name: "anthropic".to_string(),
|
||||
//! api_key: "sk-ant-xxx".to_string(),
|
||||
//! base_url: "https://api.anthropic.com".to_string(),
|
||||
//! requires_openai_auth: false,
|
||||
//! default_model: "claude-3-5-sonnet-20241022".to_string(),
|
||||
//! api_protocol: Some("anthropic".to_string()),
|
||||
//! };
|
||||
//! manager.store_config("anthropic", config);
|
||||
//!
|
||||
//! // 查询 API 密钥
|
||||
//! if let Some(key) = manager.get_api_key("anthropic") {
|
||||
//! println!("API Key: {}", key);
|
||||
//! }
|
||||
//! ```
|
||||
|
||||
use dashmap::DashMap;
|
||||
use shared_types::ModelProviderConfig;
|
||||
use std::sync::Arc;
|
||||
use tracing::{debug, info};
|
||||
|
||||
/// API 密钥管理器
|
||||
///
|
||||
/// 使用 DashMap 实现并发安全的 HashMap,支持多线程安全访问。
|
||||
/// 密钥配置存储在内存中,进程重启后清空。
|
||||
///
|
||||
/// 可以作为独立存储使用(通过 `new()`),也可以包装共享的 DashMap(通过 `from_shared()`)。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ApiKeyManager {
|
||||
/// 服务名称 -> 密钥配置
|
||||
///
|
||||
/// 键为服务标识符(如 UUID),值为完整的 ModelProviderConfig。
|
||||
/// 包装共享的 DashMap 引用(不拥有数据)或拥有独立 DashMap。
|
||||
shared: Arc<DashMap<String, ModelProviderConfig>>,
|
||||
}
|
||||
|
||||
impl Default for ApiKeyManager {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ApiKeyManager {
|
||||
/// 创建新的 API 密钥管理器(拥有独立的 DashMap)
|
||||
///
|
||||
/// # 示例
|
||||
///
|
||||
/// ```rust
|
||||
/// use agent_runner::api_key_manager::ApiKeyManager;
|
||||
///
|
||||
/// let manager = ApiKeyManager::new();
|
||||
/// ```
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
shared: Arc::new(DashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从共享 DashMap 创建包装器
|
||||
///
|
||||
/// 将 ApiKeyManager 改为包装共享的 DashMap,而不是拥有独立的数据。
|
||||
/// 这样多个组件可以共享同一个配置存储。
|
||||
///
|
||||
/// # 参数
|
||||
///
|
||||
/// * `shared` - 共享的 DashMap<String, ModelProviderConfig>
|
||||
///
|
||||
/// # 示例
|
||||
///
|
||||
/// ```rust
|
||||
/// use agent_runner::api_key_manager::ApiKeyManager;
|
||||
/// use dashmap::DashMap;
|
||||
/// use std::sync::Arc;
|
||||
///
|
||||
/// let shared_map = Arc::new(DashMap::new());
|
||||
/// let manager = ApiKeyManager::from_shared(shared_map);
|
||||
/// ```
|
||||
pub fn from_shared(shared: Arc<DashMap<String, ModelProviderConfig>>) -> Self {
|
||||
Self { shared }
|
||||
}
|
||||
|
||||
/// 存储 ModelProviderConfig 到内存
|
||||
///
|
||||
/// 如果已存在相同服务名称的配置,将被覆盖。
|
||||
/// 使用 DashMap 的 insert 方法确保原子性。
|
||||
///
|
||||
/// # 参数
|
||||
///
|
||||
/// * `service_name` - 服务标识符(如 UUID)
|
||||
/// * `config` - 完整的模型提供商配置
|
||||
///
|
||||
/// # 示例
|
||||
///
|
||||
/// ```rust
|
||||
/// # use agent_runner::api_key_manager::ApiKeyManager;
|
||||
/// # use shared_types::ModelProviderConfig;
|
||||
/// let manager = ApiKeyManager::new();
|
||||
/// let config = ModelProviderConfig {
|
||||
/// id: "anthropic-prov".to_string(),
|
||||
/// name: "anthropic".to_string(),
|
||||
/// api_key: "sk-ant-xxx".to_string(),
|
||||
/// base_url: "https://api.anthropic.com".to_string(),
|
||||
/// requires_openai_auth: false,
|
||||
/// default_model: "claude-3-5-sonnet-20241022".to_string(),
|
||||
/// api_protocol: Some("anthropic".to_string()),
|
||||
/// };
|
||||
/// manager.store_config("svc-uuid-123", config);
|
||||
/// ```
|
||||
pub fn store_config(&self, service_name: &str, config: ModelProviderConfig) {
|
||||
debug!(
|
||||
"🔑 [API_KEY_MANAGER] Storing config: service_name={}",
|
||||
service_name
|
||||
);
|
||||
self.shared.insert(service_name.to_string(), config);
|
||||
}
|
||||
|
||||
/// 获取完整的 ModelProviderConfig
|
||||
///
|
||||
/// 使用 DashMap 的 value() 方法避免克隆(返回引用)。
|
||||
///
|
||||
/// # 参数
|
||||
///
|
||||
/// * `service_name` - 服务标识符
|
||||
///
|
||||
/// # 返回
|
||||
///
|
||||
/// 如果找到配置则返回 `Some(ModelProviderConfig)`,否则返回 `None`。
|
||||
pub fn get(&self, service_name: &str) -> Option<ModelProviderConfig> {
|
||||
self.shared.get(service_name).map(|r| r.value().clone())
|
||||
}
|
||||
|
||||
/// 获取 API 密钥
|
||||
///
|
||||
/// # 参数
|
||||
///
|
||||
/// * `service_name` - 服务标识符
|
||||
///
|
||||
/// # 返回
|
||||
///
|
||||
/// 如果找到配置则返回 `Some(api_key)`,否则返回 `None`。
|
||||
pub fn get_api_key(&self, service_name: &str) -> Option<String> {
|
||||
self.get(service_name).map(|c| c.api_key)
|
||||
}
|
||||
|
||||
/// 获取 API Base URL
|
||||
///
|
||||
/// # 参数
|
||||
///
|
||||
/// * `service_name` - 服务标识符
|
||||
///
|
||||
/// # 返回
|
||||
///
|
||||
/// 如果找到配置则返回 `Some(base_url)`,否则返回 `None`。
|
||||
pub fn get_base_url(&self, service_name: &str) -> Option<String> {
|
||||
self.get(service_name).map(|c| c.base_url)
|
||||
}
|
||||
|
||||
/// 移除指定服务的配置
|
||||
///
|
||||
/// 使用 DashMap 的 remove 方法,这是原子性操作(entry API)。
|
||||
///
|
||||
/// # 参数
|
||||
///
|
||||
/// * `service_name` - 服务标识符
|
||||
///
|
||||
/// # 返回
|
||||
///
|
||||
/// 如果找到并移除了配置则返回 `Some(ModelProviderConfig)`,否则返回 `None`。
|
||||
pub fn remove(&self, service_name: &str) -> Option<ModelProviderConfig> {
|
||||
debug!(
|
||||
"🔑 [API_KEY_MANAGER] Removing config: service_name={}",
|
||||
service_name
|
||||
);
|
||||
self.shared.remove(service_name).map(|(_, v)| v)
|
||||
}
|
||||
|
||||
/// 清空所有配置
|
||||
pub fn clear(&self) {
|
||||
info!("🔑 [API_KEY_MANAGER] Cleared all configs");
|
||||
self.shared.clear();
|
||||
}
|
||||
|
||||
/// 获取当前存储的配置数量
|
||||
pub fn len(&self) -> usize {
|
||||
self.shared.len()
|
||||
}
|
||||
|
||||
/// 检查是否为空
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.shared.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn create_test_config(name: &str) -> ModelProviderConfig {
|
||||
ModelProviderConfig {
|
||||
id: format!("test_{}", name),
|
||||
name: name.to_string(),
|
||||
base_url: format!("https://api.{}.com", name),
|
||||
api_key: format!("sk-{}-123", name),
|
||||
requires_openai_auth: name == "openai",
|
||||
default_model: "test-model".to_string(),
|
||||
api_protocol: Some(name.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_store_and_get() {
|
||||
let manager = ApiKeyManager::new();
|
||||
let config = create_test_config("anthropic");
|
||||
|
||||
manager.store_config("anthropic", config.clone());
|
||||
|
||||
let retrieved = manager.get("anthropic");
|
||||
assert!(retrieved.is_some());
|
||||
assert_eq!(retrieved.unwrap().api_key, "sk-anthropic-123");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_api_key() {
|
||||
let manager = ApiKeyManager::new();
|
||||
let config = create_test_config("openai");
|
||||
|
||||
manager.store_config("openai", config);
|
||||
|
||||
let api_key = manager.get_api_key("openai");
|
||||
assert_eq!(api_key, Some("sk-openai-123".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_base_url() {
|
||||
let manager = ApiKeyManager::new();
|
||||
let config = create_test_config("anthropic");
|
||||
|
||||
manager.store_config("anthropic", config);
|
||||
|
||||
let base_url = manager.get_base_url("anthropic");
|
||||
assert_eq!(base_url, Some("https://api.anthropic.com".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nonexistent() {
|
||||
let manager = ApiKeyManager::new();
|
||||
|
||||
assert!(manager.get("nonexistent").is_none());
|
||||
assert!(manager.get_api_key("nonexistent").is_none());
|
||||
assert!(manager.get_base_url("nonexistent").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove() {
|
||||
let manager = ApiKeyManager::new();
|
||||
let config = create_test_config("anthropic");
|
||||
|
||||
manager.store_config("anthropic", config);
|
||||
assert_eq!(manager.len(), 1);
|
||||
|
||||
let removed = manager.remove("anthropic");
|
||||
assert!(removed.is_some());
|
||||
assert_eq!(manager.len(), 0);
|
||||
assert!(manager.get("anthropic").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clear() {
|
||||
let manager = ApiKeyManager::new();
|
||||
|
||||
manager.store_config("anthropic", create_test_config("anthropic"));
|
||||
manager.store_config("openai", create_test_config("openai"));
|
||||
|
||||
assert_eq!(manager.len(), 2);
|
||||
manager.clear();
|
||||
assert_eq!(manager.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_overwrite() {
|
||||
let manager = ApiKeyManager::new();
|
||||
|
||||
manager.store_config("anthropic", create_test_config("anthropic"));
|
||||
|
||||
let mut new_config = create_test_config("anthropic");
|
||||
new_config.api_key = "sk-updated".to_string();
|
||||
|
||||
manager.store_config("anthropic", new_config);
|
||||
|
||||
let api_key = manager.get_api_key("anthropic");
|
||||
assert_eq!(api_key, Some("sk-updated".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_empty() {
|
||||
let manager = ApiKeyManager::new();
|
||||
assert!(manager.is_empty());
|
||||
|
||||
manager.store_config("anthropic", create_test_config("anthropic"));
|
||||
assert!(!manager.is_empty());
|
||||
}
|
||||
}
|
||||
954
qiming-rcoder/crates/agent_runner/src/config.rs
Normal file
954
qiming-rcoder/crates/agent_runner/src/config.rs
Normal file
@@ -0,0 +1,954 @@
|
||||
use std::env;
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use clap::Parser;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
/// 命令行参数
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "rcoder")]
|
||||
#[command(about = "AI-powered development platform")]
|
||||
#[command(version)]
|
||||
pub struct CliArgs {
|
||||
/// Service port
|
||||
#[arg(short, long, help = "Service port")]
|
||||
pub port: Option<u16>,
|
||||
|
||||
/// Project workspace directory
|
||||
#[arg(short = 'd', long, help = "Root directory for project workspace")]
|
||||
pub projects_dir: Option<PathBuf>,
|
||||
|
||||
/// Enable port-based reverse proxy
|
||||
#[arg(long, help = "Enable port-based reverse proxy")]
|
||||
pub enable_proxy: bool,
|
||||
|
||||
/// Proxy listener port
|
||||
#[arg(long, help = "Proxy service listener port")]
|
||||
pub proxy_port: Option<u16>,
|
||||
|
||||
/// Default backend port
|
||||
#[arg(long, help = "Default backend service port")]
|
||||
pub default_backend_port: Option<u16>,
|
||||
}
|
||||
|
||||
/// 应用配置
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AppConfig {
|
||||
/// 默认使用的 Agent ID
|
||||
#[serde(default = "default_agent_id")]
|
||||
pub default_agent_id: String,
|
||||
/// 项目工作的根目录,根据启动命令的当前目录来确定
|
||||
pub projects_dir: PathBuf,
|
||||
/// 服务端口
|
||||
pub port: u16,
|
||||
/// 代理配置
|
||||
pub proxy_config: Option<ProxyConfig>,
|
||||
/// Agent 清理配置
|
||||
#[serde(default)]
|
||||
pub agent_cleanup: Option<AgentCleanupConfig>,
|
||||
/// gRPC 超时配置
|
||||
#[serde(default)]
|
||||
pub grpc_timeouts: Option<GrpcTimeoutConfig>,
|
||||
/// Agent 并发配置
|
||||
#[serde(default)]
|
||||
pub agent_concurrency: Option<AgentConcurrencyConfig>,
|
||||
/// mcp-proxy 日志目录(可选)
|
||||
/// 当设置此值且日志级别为 debug 时,mcp-proxy convert 命令会自动追加
|
||||
/// --diagnostic 和 --log-dir 参数
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub mcp_proxy_log_dir: Option<String>,
|
||||
}
|
||||
|
||||
fn default_agent_id() -> String {
|
||||
"claude-code-acp-ts".to_string()
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct HealthCheckConfig {
|
||||
pub enabled: bool,
|
||||
pub interval_seconds: u64,
|
||||
pub timeout_seconds: u64,
|
||||
pub healthy_threshold: u32,
|
||||
pub unhealthy_threshold: u32,
|
||||
}
|
||||
|
||||
/// 代理配置
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ProxyConfig {
|
||||
/// 代理监听端口
|
||||
pub listen_port: u16,
|
||||
/// 默认后端端口
|
||||
pub default_backend_port: u16,
|
||||
/// 后端服务主机
|
||||
pub backend_host: String,
|
||||
/// URL 中端口参数的名称
|
||||
pub port_param: String,
|
||||
/// 健康检查配置
|
||||
pub health_check: HealthCheckConfig,
|
||||
}
|
||||
|
||||
/// Agent cleanup configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentCleanupConfig {
|
||||
/// Idle timeout (seconds), default 300 (5 minutes)
|
||||
#[serde(default = "default_idle_timeout")]
|
||||
pub idle_timeout_secs: u64,
|
||||
/// Cleanup check interval (seconds), default 30
|
||||
#[serde(default = "default_cleanup_interval")]
|
||||
pub cleanup_interval_secs: u64,
|
||||
}
|
||||
|
||||
/// Agent 并发配置
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentConcurrencyConfig {
|
||||
/// 并发会话槽位数量,默认 10
|
||||
#[serde(default = "default_concurrency_limit")]
|
||||
pub concurrency_limit: usize,
|
||||
}
|
||||
|
||||
/// Agent 并发配置常量
|
||||
impl AgentConcurrencyConfig {
|
||||
/// 最小并发限制
|
||||
pub const MIN_CONCURRENCY_LIMIT: usize = 1;
|
||||
|
||||
/// 验证配置值是否在有效范围内
|
||||
pub fn validate(&self) -> Result<(), String> {
|
||||
if self.concurrency_limit < Self::MIN_CONCURRENCY_LIMIT {
|
||||
return Err(format!(
|
||||
"concurrency_limit must be >= {}, current value: {}",
|
||||
Self::MIN_CONCURRENCY_LIMIT,
|
||||
self.concurrency_limit
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn default_concurrency_limit() -> usize {
|
||||
10
|
||||
}
|
||||
|
||||
impl Default for AgentConcurrencyConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
concurrency_limit: default_concurrency_limit(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// gRPC timeout configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct GrpcTimeoutConfig {
|
||||
/// Cancel session timeout (seconds), default 30
|
||||
#[serde(default = "default_cancel_timeout")]
|
||||
pub cancel_session_timeout_secs: u64,
|
||||
|
||||
/// ACP session creation timeout (seconds), default 100
|
||||
#[serde(default = "default_acp_session_timeout")]
|
||||
pub acp_session_create_timeout_secs: u64,
|
||||
|
||||
/// Agent cancel call timeout (seconds), default 10
|
||||
#[serde(default = "default_agent_cancel_timeout")]
|
||||
pub agent_cancel_timeout_secs: u64,
|
||||
|
||||
/// Port check timeout (milliseconds), default 500
|
||||
#[serde(default = "default_port_check_timeout")]
|
||||
pub port_check_timeout_millis: u64,
|
||||
}
|
||||
|
||||
/// gRPC timeout configuration constants
|
||||
impl GrpcTimeoutConfig {
|
||||
/// Minimum cancel session timeout (5 seconds)
|
||||
pub const MIN_CANCEL_TIMEOUT: u64 = 5;
|
||||
/// Maximum cancel session timeout (300 seconds = 5 minutes)
|
||||
pub const MAX_CANCEL_TIMEOUT: u64 = 300;
|
||||
/// Minimum ACP session creation timeout (10 seconds)
|
||||
pub const MIN_ACP_SESSION_TIMEOUT: u64 = 10;
|
||||
/// Maximum ACP session creation timeout (300 seconds = 5 minutes)
|
||||
pub const MAX_ACP_SESSION_TIMEOUT: u64 = 300;
|
||||
/// Minimum Agent cancel call timeout (5 seconds)
|
||||
pub const MIN_AGENT_CANCEL_TIMEOUT: u64 = 5;
|
||||
/// Maximum Agent cancel call timeout (60 seconds)
|
||||
pub const MAX_AGENT_CANCEL_TIMEOUT: u64 = 60;
|
||||
/// Minimum port check timeout (100 milliseconds)
|
||||
pub const MIN_PORT_CHECK_TIMEOUT: u64 = 100;
|
||||
/// Maximum port check timeout (10000 milliseconds = 10 seconds)
|
||||
pub const MAX_PORT_CHECK_TIMEOUT: u64 = 10000;
|
||||
|
||||
/// Validate that configuration values are within valid ranges
|
||||
pub fn validate(&self) -> Result<(), String> {
|
||||
if self.cancel_session_timeout_secs < Self::MIN_CANCEL_TIMEOUT
|
||||
|| self.cancel_session_timeout_secs > Self::MAX_CANCEL_TIMEOUT
|
||||
{
|
||||
return Err(format!(
|
||||
"cancel_session_timeout_secs must be between {} and {}, current: {}",
|
||||
Self::MIN_CANCEL_TIMEOUT,
|
||||
Self::MAX_CANCEL_TIMEOUT,
|
||||
self.cancel_session_timeout_secs
|
||||
));
|
||||
}
|
||||
|
||||
if self.acp_session_create_timeout_secs < Self::MIN_ACP_SESSION_TIMEOUT
|
||||
|| self.acp_session_create_timeout_secs > Self::MAX_ACP_SESSION_TIMEOUT
|
||||
{
|
||||
return Err(format!(
|
||||
"acp_session_create_timeout_secs must be between {} and {}, current: {}",
|
||||
Self::MIN_ACP_SESSION_TIMEOUT,
|
||||
Self::MAX_ACP_SESSION_TIMEOUT,
|
||||
self.acp_session_create_timeout_secs
|
||||
));
|
||||
}
|
||||
|
||||
if self.agent_cancel_timeout_secs < Self::MIN_AGENT_CANCEL_TIMEOUT
|
||||
|| self.agent_cancel_timeout_secs > Self::MAX_AGENT_CANCEL_TIMEOUT
|
||||
{
|
||||
return Err(format!(
|
||||
"agent_cancel_timeout_secs must be between {} and {}, current: {}",
|
||||
Self::MIN_AGENT_CANCEL_TIMEOUT,
|
||||
Self::MAX_AGENT_CANCEL_TIMEOUT,
|
||||
self.agent_cancel_timeout_secs
|
||||
));
|
||||
}
|
||||
|
||||
if self.port_check_timeout_millis < Self::MIN_PORT_CHECK_TIMEOUT
|
||||
|| self.port_check_timeout_millis > Self::MAX_PORT_CHECK_TIMEOUT
|
||||
{
|
||||
return Err(format!(
|
||||
"port_check_timeout_millis must be between {} and {}, current: {}",
|
||||
Self::MIN_PORT_CHECK_TIMEOUT,
|
||||
Self::MAX_PORT_CHECK_TIMEOUT,
|
||||
self.port_check_timeout_millis
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Agent 清理配置常量
|
||||
impl AgentCleanupConfig {
|
||||
/// 最小闲置超时时间(10 秒)
|
||||
pub const MIN_IDLE_TIMEOUT: u64 = 10;
|
||||
/// 最大闲置超时时间(24 小时)
|
||||
pub const MAX_IDLE_TIMEOUT: u64 = 24 * 60 * 60;
|
||||
/// 最小清理检查间隔(5 秒)
|
||||
pub const MIN_CLEANUP_INTERVAL: u64 = 5;
|
||||
/// 最大清理检查间隔(1 小时)
|
||||
pub const MAX_CLEANUP_INTERVAL: u64 = 60 * 60;
|
||||
|
||||
/// 验证配置值是否在有效范围内
|
||||
pub fn validate(&self) -> Result<(), String> {
|
||||
if self.idle_timeout_secs < Self::MIN_IDLE_TIMEOUT
|
||||
|| self.idle_timeout_secs > Self::MAX_IDLE_TIMEOUT
|
||||
{
|
||||
return Err(format!(
|
||||
"idle_timeout_secs must be between {} and {}, current value: {}",
|
||||
Self::MIN_IDLE_TIMEOUT,
|
||||
Self::MAX_IDLE_TIMEOUT,
|
||||
self.idle_timeout_secs
|
||||
));
|
||||
}
|
||||
|
||||
if self.cleanup_interval_secs < Self::MIN_CLEANUP_INTERVAL
|
||||
|| self.cleanup_interval_secs > Self::MAX_CLEANUP_INTERVAL
|
||||
{
|
||||
return Err(format!(
|
||||
"cleanup_interval_secs must be between {} and {}, current value: {}",
|
||||
Self::MIN_CLEANUP_INTERVAL,
|
||||
Self::MAX_CLEANUP_INTERVAL,
|
||||
self.cleanup_interval_secs
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "http-server")]
|
||||
fn default_idle_timeout() -> u64 {
|
||||
24 * 60 * 60 // 24 小时(Tauri 客户端模式)
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "http-server"))]
|
||||
fn default_idle_timeout() -> u64 {
|
||||
300 // 5 分钟(CLI 模式)
|
||||
}
|
||||
|
||||
fn default_cleanup_interval() -> u64 {
|
||||
30 // 30 秒
|
||||
}
|
||||
|
||||
fn default_cancel_timeout() -> u64 {
|
||||
30 // 30 秒
|
||||
}
|
||||
|
||||
fn default_acp_session_timeout() -> u64 {
|
||||
100 // 100 秒
|
||||
}
|
||||
|
||||
fn default_agent_cancel_timeout() -> u64 {
|
||||
10 // 10 秒
|
||||
}
|
||||
|
||||
fn default_port_check_timeout() -> u64 {
|
||||
500 // 500 毫秒
|
||||
}
|
||||
|
||||
impl Default for AgentCleanupConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
idle_timeout_secs: default_idle_timeout(),
|
||||
cleanup_interval_secs: default_cleanup_interval(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for GrpcTimeoutConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
cancel_session_timeout_secs: default_cancel_timeout(),
|
||||
acp_session_create_timeout_secs: default_acp_session_timeout(),
|
||||
agent_cancel_timeout_secs: default_agent_cancel_timeout(),
|
||||
port_check_timeout_millis: default_port_check_timeout(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 配置文件路径
|
||||
const CONFIG_FILE: &str = "config.yml";
|
||||
|
||||
impl Default for AppConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
default_agent_id: default_agent_id(),
|
||||
projects_dir: PathBuf::from("./project_workspace"),
|
||||
port: 8086,
|
||||
proxy_config: Some(ProxyConfig::default()),
|
||||
agent_cleanup: Some(AgentCleanupConfig::default()),
|
||||
grpc_timeouts: Some(GrpcTimeoutConfig::default()),
|
||||
agent_concurrency: Some(AgentConcurrencyConfig::default()),
|
||||
mcp_proxy_log_dir: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ProxyConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
listen_port: 8088,
|
||||
default_backend_port: 8086,
|
||||
backend_host: "127.0.0.1".to_string(),
|
||||
port_param: "port".to_string(),
|
||||
health_check: HealthCheckConfig {
|
||||
enabled: true,
|
||||
interval_seconds: 5,
|
||||
timeout_seconds: 1,
|
||||
healthy_threshold: 2,
|
||||
unhealthy_threshold: 3,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 加载配置
|
||||
/// 配置优先级:命令行参数 > 环境变量 > 配置文件 > 默认配置
|
||||
pub fn load_config_with_args(cli_args: CliArgs) -> AppConfig {
|
||||
// 1. 首先加载默认配置
|
||||
let mut config = AppConfig::default();
|
||||
|
||||
// 2. 尝试从当前目录读取配置文件
|
||||
match load_config_from_file() {
|
||||
Ok(file_config) => {
|
||||
config = file_config;
|
||||
info!("Loaded config from {}", CONFIG_FILE);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"Failed to read config file {}: {}, using defaults",
|
||||
CONFIG_FILE, e
|
||||
);
|
||||
|
||||
// 创建默认配置文件
|
||||
if let Err(create_err) = create_default_config_file(&config) {
|
||||
error!("Failed to create default config file: {}", create_err);
|
||||
} else {
|
||||
info!("Created default config file: {}", CONFIG_FILE);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 环境变量覆盖配置
|
||||
if let Ok(port) = env::var("RCODER_PORT") {
|
||||
match port.parse::<u16>() {
|
||||
Ok(p) => {
|
||||
config.port = p;
|
||||
info!("Set port from env RCODER_PORT: {}", p);
|
||||
}
|
||||
Err(_) => {
|
||||
warn!(
|
||||
"Invalid RCODER_PORT env value: {}, keeping config port: {}",
|
||||
port, config.port
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 🆕 Agent 清理配置:支持环境变量覆盖
|
||||
if let Ok(idle_timeout) = env::var("RCODER_AGENT_IDLE_TIMEOUT_SECS") {
|
||||
match idle_timeout.parse::<u64>() {
|
||||
Ok(timeout) => {
|
||||
// 🔒 验证范围
|
||||
if (AgentCleanupConfig::MIN_IDLE_TIMEOUT..=AgentCleanupConfig::MAX_IDLE_TIMEOUT)
|
||||
.contains(&timeout)
|
||||
{
|
||||
config
|
||||
.agent_cleanup
|
||||
.get_or_insert_with(Default::default)
|
||||
.idle_timeout_secs = timeout;
|
||||
info!(
|
||||
"Set idle timeout from env RCODER_AGENT_IDLE_TIMEOUT_SECS: {} seconds",
|
||||
timeout
|
||||
);
|
||||
} else {
|
||||
warn!(
|
||||
"Invalid RCODER_AGENT_IDLE_TIMEOUT_SECS: {} seconds, out of range [{}, {}], keeping config value",
|
||||
timeout,
|
||||
AgentCleanupConfig::MIN_IDLE_TIMEOUT,
|
||||
AgentCleanupConfig::MAX_IDLE_TIMEOUT
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
warn!(
|
||||
"Invalid RCODER_AGENT_IDLE_TIMEOUT_SECS format: {}, keeping config value",
|
||||
idle_timeout
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(cleanup_interval) = env::var("RCODER_AGENT_CLEANUP_INTERVAL_SECS") {
|
||||
match cleanup_interval.parse::<u64>() {
|
||||
Ok(interval) => {
|
||||
// 🔒 验证范围
|
||||
if (AgentCleanupConfig::MIN_CLEANUP_INTERVAL
|
||||
..=AgentCleanupConfig::MAX_CLEANUP_INTERVAL)
|
||||
.contains(&interval)
|
||||
{
|
||||
config
|
||||
.agent_cleanup
|
||||
.get_or_insert_with(Default::default)
|
||||
.cleanup_interval_secs = interval;
|
||||
info!(
|
||||
"Set cleanup interval from env RCODER_AGENT_CLEANUP_INTERVAL_SECS: {} seconds",
|
||||
interval
|
||||
);
|
||||
} else {
|
||||
warn!(
|
||||
"Invalid RCODER_AGENT_CLEANUP_INTERVAL_SECS: {} seconds, out of range [{}, {}], keeping config value",
|
||||
interval,
|
||||
AgentCleanupConfig::MIN_CLEANUP_INTERVAL,
|
||||
AgentCleanupConfig::MAX_CLEANUP_INTERVAL
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
warn!(
|
||||
"Invalid RCODER_AGENT_CLEANUP_INTERVAL_SECS format: {}, keeping config value",
|
||||
cleanup_interval
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 🆕 Agent 并发配置:支持环境变量覆盖
|
||||
if let Ok(concurrency_limit) = env::var("RCODER_AGENT_CONCURRENCY_LIMIT") {
|
||||
match concurrency_limit.parse::<usize>() {
|
||||
Ok(limit) => {
|
||||
if limit >= AgentConcurrencyConfig::MIN_CONCURRENCY_LIMIT {
|
||||
config
|
||||
.agent_concurrency
|
||||
.get_or_insert_with(Default::default)
|
||||
.concurrency_limit = limit;
|
||||
info!(
|
||||
"Set concurrency limit from env RCODER_AGENT_CONCURRENCY_LIMIT: {}",
|
||||
limit
|
||||
);
|
||||
} else {
|
||||
warn!(
|
||||
"Invalid RCODER_AGENT_CONCURRENCY_LIMIT: {}, must be >= {}, keeping config value",
|
||||
limit,
|
||||
AgentConcurrencyConfig::MIN_CONCURRENCY_LIMIT
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
warn!(
|
||||
"Invalid RCODER_AGENT_CONCURRENCY_LIMIT format: {}, keeping config value",
|
||||
concurrency_limit
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 🆕 gRPC 超时配置:支持环境变量覆盖
|
||||
if let Ok(cancel_timeout) = env::var("RCODER_CANCEL_SESSION_TIMEOUT_SECS") {
|
||||
match cancel_timeout.parse::<u64>() {
|
||||
Ok(timeout) => {
|
||||
if (GrpcTimeoutConfig::MIN_CANCEL_TIMEOUT..=GrpcTimeoutConfig::MAX_CANCEL_TIMEOUT)
|
||||
.contains(&timeout)
|
||||
{
|
||||
config
|
||||
.grpc_timeouts
|
||||
.get_or_insert_with(Default::default)
|
||||
.cancel_session_timeout_secs = timeout;
|
||||
info!(
|
||||
"Set cancel-session timeout from env RCODER_CANCEL_SESSION_TIMEOUT_SECS: {} seconds",
|
||||
timeout
|
||||
);
|
||||
} else {
|
||||
warn!(
|
||||
"Invalid RCODER_CANCEL_SESSION_TIMEOUT_SECS: {} seconds, out of range [{}, {}]",
|
||||
timeout,
|
||||
GrpcTimeoutConfig::MIN_CANCEL_TIMEOUT,
|
||||
GrpcTimeoutConfig::MAX_CANCEL_TIMEOUT
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
warn!(
|
||||
"Invalid RCODER_CANCEL_SESSION_TIMEOUT_SECS format: {}",
|
||||
cancel_timeout
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(acp_timeout) = env::var("RCODER_ACP_SESSION_CREATE_TIMEOUT_SECS") {
|
||||
match acp_timeout.parse::<u64>() {
|
||||
Ok(timeout) => {
|
||||
if (GrpcTimeoutConfig::MIN_ACP_SESSION_TIMEOUT
|
||||
..=GrpcTimeoutConfig::MAX_ACP_SESSION_TIMEOUT)
|
||||
.contains(&timeout)
|
||||
{
|
||||
config
|
||||
.grpc_timeouts
|
||||
.get_or_insert_with(Default::default)
|
||||
.acp_session_create_timeout_secs = timeout;
|
||||
info!(
|
||||
"Set ACP session-create timeout from env RCODER_ACP_SESSION_CREATE_TIMEOUT_SECS: {} seconds",
|
||||
timeout
|
||||
);
|
||||
} else {
|
||||
warn!(
|
||||
"Invalid RCODER_ACP_SESSION_CREATE_TIMEOUT_SECS: {} seconds, out of range [{}, {}]",
|
||||
timeout,
|
||||
GrpcTimeoutConfig::MIN_ACP_SESSION_TIMEOUT,
|
||||
GrpcTimeoutConfig::MAX_ACP_SESSION_TIMEOUT
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
warn!(
|
||||
"Invalid RCODER_ACP_SESSION_CREATE_TIMEOUT_SECS format: {}",
|
||||
acp_timeout
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(agent_cancel_timeout) = env::var("RCODER_AGENT_CANCEL_TIMEOUT_SECS") {
|
||||
match agent_cancel_timeout.parse::<u64>() {
|
||||
Ok(timeout) => {
|
||||
if (GrpcTimeoutConfig::MIN_AGENT_CANCEL_TIMEOUT
|
||||
..=GrpcTimeoutConfig::MAX_AGENT_CANCEL_TIMEOUT)
|
||||
.contains(&timeout)
|
||||
{
|
||||
config
|
||||
.grpc_timeouts
|
||||
.get_or_insert_with(Default::default)
|
||||
.agent_cancel_timeout_secs = timeout;
|
||||
info!(
|
||||
"Set agent-cancel timeout from env RCODER_AGENT_CANCEL_TIMEOUT_SECS: {} seconds",
|
||||
timeout
|
||||
);
|
||||
} else {
|
||||
warn!(
|
||||
"Invalid RCODER_AGENT_CANCEL_TIMEOUT_SECS: {} seconds, out of range [{}, {}]",
|
||||
timeout,
|
||||
GrpcTimeoutConfig::MIN_AGENT_CANCEL_TIMEOUT,
|
||||
GrpcTimeoutConfig::MAX_AGENT_CANCEL_TIMEOUT
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
warn!(
|
||||
"Invalid RCODER_AGENT_CANCEL_TIMEOUT_SECS format: {}",
|
||||
agent_cancel_timeout
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(port_check_timeout) = env::var("RCODER_PORT_CHECK_TIMEOUT_MILLIS") {
|
||||
match port_check_timeout.parse::<u64>() {
|
||||
Ok(timeout) => {
|
||||
if (GrpcTimeoutConfig::MIN_PORT_CHECK_TIMEOUT
|
||||
..=GrpcTimeoutConfig::MAX_PORT_CHECK_TIMEOUT)
|
||||
.contains(&timeout)
|
||||
{
|
||||
config
|
||||
.grpc_timeouts
|
||||
.get_or_insert_with(Default::default)
|
||||
.port_check_timeout_millis = timeout;
|
||||
info!(
|
||||
"Set port-check timeout from env RCODER_PORT_CHECK_TIMEOUT_MILLIS: {} ms",
|
||||
timeout
|
||||
);
|
||||
} else {
|
||||
warn!(
|
||||
"Invalid RCODER_PORT_CHECK_TIMEOUT_MILLIS: {} ms, out of range [{}, {}]",
|
||||
timeout,
|
||||
GrpcTimeoutConfig::MIN_PORT_CHECK_TIMEOUT,
|
||||
GrpcTimeoutConfig::MAX_PORT_CHECK_TIMEOUT
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
warn!(
|
||||
"Invalid RCODER_PORT_CHECK_TIMEOUT_MILLIS format: {}",
|
||||
port_check_timeout
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 🆕 验证最终配置的有效性
|
||||
if let Some(ref cleanup_config) = config.agent_cleanup
|
||||
&& let Err(e) = cleanup_config.validate()
|
||||
{
|
||||
warn!(
|
||||
"Agent cleanup config validation failed: {}, using defaults",
|
||||
e
|
||||
);
|
||||
config.agent_cleanup = Some(AgentCleanupConfig::default());
|
||||
}
|
||||
|
||||
// 🆕 验证 Agent 并发配置的有效性
|
||||
if let Some(ref concurrency_config) = config.agent_concurrency
|
||||
&& let Err(e) = concurrency_config.validate()
|
||||
{
|
||||
warn!(
|
||||
"Agent concurrency config validation failed: {}, using defaults",
|
||||
e
|
||||
);
|
||||
config.agent_concurrency = Some(AgentConcurrencyConfig::default());
|
||||
}
|
||||
|
||||
// 4. 处理代理配置
|
||||
if cli_args.enable_proxy {
|
||||
let proxy_config = ProxyConfig {
|
||||
listen_port: cli_args.proxy_port.unwrap_or(8080),
|
||||
default_backend_port: cli_args.default_backend_port.unwrap_or(config.port),
|
||||
backend_host: "127.0.0.1".to_string(),
|
||||
port_param: "port".to_string(),
|
||||
health_check: HealthCheckConfig {
|
||||
enabled: true,
|
||||
interval_seconds: 5,
|
||||
timeout_seconds: 1,
|
||||
healthy_threshold: 2,
|
||||
unhealthy_threshold: 3,
|
||||
},
|
||||
};
|
||||
info!(
|
||||
"Reverse proxy enabled, listening on port: {}",
|
||||
proxy_config.listen_port
|
||||
);
|
||||
config.proxy_config = Some(proxy_config);
|
||||
}
|
||||
|
||||
// 5. 命令行参数覆盖配置(优先级最高)
|
||||
if let Some(port) = cli_args.port {
|
||||
config.port = port;
|
||||
info!("Set port from CLI arg: {}", port);
|
||||
}
|
||||
|
||||
if let Some(projects_dir) = cli_args.projects_dir {
|
||||
config.projects_dir = projects_dir.clone();
|
||||
info!("Set projects directory from CLI arg: {:?}", projects_dir);
|
||||
}
|
||||
|
||||
info!(
|
||||
"最终配置: port={}, projects_dir={:?}, default_agent_id={}, proxy_enabled={}",
|
||||
config.port,
|
||||
config.projects_dir,
|
||||
config.default_agent_id,
|
||||
config.proxy_config.is_some()
|
||||
);
|
||||
|
||||
// 🆕 验证 gRPC 超时配置的有效性
|
||||
if let Some(ref grpc_timeouts) = config.grpc_timeouts
|
||||
&& let Err(e) = grpc_timeouts.validate()
|
||||
{
|
||||
warn!(
|
||||
"gRPC timeout config validation failed: {}, using defaults",
|
||||
e
|
||||
);
|
||||
config.grpc_timeouts = Some(GrpcTimeoutConfig::default());
|
||||
}
|
||||
|
||||
config
|
||||
}
|
||||
|
||||
/// 加载配置(保留旧接口以保持兼容性)
|
||||
pub fn load_config() -> AppConfig {
|
||||
let cli_args = CliArgs {
|
||||
port: None,
|
||||
projects_dir: None,
|
||||
enable_proxy: false,
|
||||
proxy_port: None,
|
||||
default_backend_port: None,
|
||||
};
|
||||
load_config_with_args(cli_args)
|
||||
}
|
||||
|
||||
/// 从文件加载配置
|
||||
fn load_config_from_file() -> anyhow::Result<AppConfig> {
|
||||
let config_content = fs::read_to_string(CONFIG_FILE)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to read config file: {}", e))?;
|
||||
|
||||
let config: AppConfig = serde_yaml::from_str(&config_content)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to parse config file: {}", e))?;
|
||||
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// 创建默认配置文件
|
||||
fn create_default_config_file(config: &AppConfig) -> anyhow::Result<()> {
|
||||
// 获取 proxy_config,如果不存在则使用默认值
|
||||
let proxy_config = config.proxy_config.as_ref().cloned().unwrap_or_default();
|
||||
|
||||
// 获取 agent_cleanup 配置,如果不存在则使用默认值
|
||||
let agent_cleanup = config.agent_cleanup.as_ref().cloned().unwrap_or_default();
|
||||
|
||||
// 获取 grpc_timeouts 配置,如果不存在则使用默认值
|
||||
let grpc_timeouts = config.grpc_timeouts.as_ref().cloned().unwrap_or_default();
|
||||
|
||||
// 获取 agent_concurrency 配置,如果不存在则使用默认值
|
||||
let agent_concurrency = config
|
||||
.agent_concurrency
|
||||
.as_ref()
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
// 手动构建带注释的 YAML 内容
|
||||
let content_with_comments = format!(
|
||||
r#"# rcoder 配置文件
|
||||
# 该文件在首次启动时自动生成
|
||||
|
||||
# 默认使用的 Agent ID
|
||||
default_agent_id: {}
|
||||
|
||||
# 项目工作目录
|
||||
projects_dir: {}
|
||||
|
||||
# 主服务端口
|
||||
port: {}
|
||||
|
||||
# Pingora 反向代理配置
|
||||
proxy_config:
|
||||
# 代理服务监听端口 (用于接收外部请求)
|
||||
listen_port: {}
|
||||
# 默认后端服务端口 (当请求未指定端口时使用)
|
||||
default_backend_port: {}
|
||||
# 后端服务主机地址
|
||||
backend_host: "{}"
|
||||
# URL 中端口参数的名称 (用于从路径中提取端口号)
|
||||
port_param: "{}"
|
||||
# 健康检查配置
|
||||
health_check:
|
||||
enabled: {}
|
||||
interval_seconds: {}
|
||||
timeout_seconds: {}
|
||||
healthy_threshold: {}
|
||||
unhealthy_threshold: {}
|
||||
|
||||
# Agent 清理配置
|
||||
# 如果省略此配置块,将使用以下默认值:
|
||||
# - idle_timeout_secs: 300 (5分钟)
|
||||
# - cleanup_interval_secs: 30 (30秒)
|
||||
agent_cleanup:
|
||||
# Agent 闲置超时时间(秒)
|
||||
# Agent 在闲置超过此时间后会被自动清理以释放资源
|
||||
# 有效范围: 10 - 86400 秒(10秒 - 24小时)
|
||||
# 可通过环境变量 RCODER_AGENT_IDLE_TIMEOUT_SECS 覆盖
|
||||
idle_timeout_secs: {}
|
||||
# 清理检查间隔(秒)
|
||||
# 系统每隔此时间检查一次是否有闲置的 Agent 需要清理
|
||||
# 有效范围: 5 - 3600 秒(5秒 - 1小时)
|
||||
# 可通过环境变量 RCODER_AGENT_CLEANUP_INTERVAL_SECS 覆盖
|
||||
cleanup_interval_secs: {}
|
||||
|
||||
# gRPC 超时配置
|
||||
# 如果省略此配置块,将使用以下默认值:
|
||||
# - cancel_session_timeout_secs: 30 (30秒)
|
||||
# - acp_session_create_timeout_secs: 100 (100秒)
|
||||
# - agent_cancel_timeout_secs: 10 (10秒)
|
||||
# - port_check_timeout_millis: 500 (500毫秒)
|
||||
grpc_timeouts:
|
||||
# 取消会话超时(秒)
|
||||
# gRPC 取消会话请求的最大等待时间
|
||||
# 有效范围: 5 - 300 秒
|
||||
# 可通过环境变量 RCODER_CANCEL_SESSION_TIMEOUT_SECS 覆盖
|
||||
cancel_session_timeout_secs: {}
|
||||
# ACP 会话创建超时(秒)
|
||||
# Agent 创建新会话的最大等待时间(MCP 工具较多时可能需要更长时间)
|
||||
# 有效范围: 10 - 300 秒
|
||||
# 可通过环境变量 RCODER_ACP_SESSION_CREATE_TIMEOUT_SECS 覆盖
|
||||
acp_session_create_timeout_secs: {}
|
||||
# Agent 取消调用超时(秒)
|
||||
# Agent 内部取消操作的最大等待时间
|
||||
# 有效范围: 5 - 60 秒
|
||||
# 可通过环境变量 RCODER_AGENT_CANCEL_TIMEOUT_SECS 覆盖
|
||||
agent_cancel_timeout_secs: {}
|
||||
# 端口检查超时(毫秒)
|
||||
# 检查端口可用性的最大等待时间
|
||||
# 有效范围: 100 - 10000 毫秒
|
||||
# 可通过环境变量 RCODER_PORT_CHECK_TIMEOUT_MILLIS 覆盖
|
||||
port_check_timeout_millis: {}
|
||||
|
||||
# Agent 并发配置
|
||||
# 如果省略此配置块,将使用以下默认值:
|
||||
# - concurrency_limit: 10 (10个并发会话)
|
||||
agent_concurrency:
|
||||
# Agent 并发会话槽位数量
|
||||
# 决定可以同时处理的 Agent 会话数量
|
||||
# 有效范围: >= 1
|
||||
# 可通过环境变量 RCODER_AGENT_CONCURRENCY_LIMIT 覆盖
|
||||
concurrency_limit: {}
|
||||
"#,
|
||||
config.default_agent_id,
|
||||
config.projects_dir.display(),
|
||||
config.port,
|
||||
proxy_config.listen_port,
|
||||
proxy_config.default_backend_port,
|
||||
proxy_config.backend_host,
|
||||
proxy_config.port_param,
|
||||
proxy_config.health_check.enabled,
|
||||
proxy_config.health_check.interval_seconds,
|
||||
proxy_config.health_check.timeout_seconds,
|
||||
proxy_config.health_check.healthy_threshold,
|
||||
proxy_config.health_check.unhealthy_threshold,
|
||||
agent_cleanup.idle_timeout_secs,
|
||||
agent_cleanup.cleanup_interval_secs,
|
||||
grpc_timeouts.cancel_session_timeout_secs,
|
||||
grpc_timeouts.acp_session_create_timeout_secs,
|
||||
grpc_timeouts.agent_cancel_timeout_secs,
|
||||
grpc_timeouts.port_check_timeout_millis,
|
||||
agent_concurrency.concurrency_limit
|
||||
);
|
||||
|
||||
fs::write(CONFIG_FILE, content_with_comments)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to write config file: {}", e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_agent_cleanup_default_values() {
|
||||
let config = AgentCleanupConfig::default();
|
||||
assert_eq!(config.idle_timeout_secs, 300);
|
||||
assert_eq!(config.cleanup_interval_secs, 30);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agent_cleanup_validate_valid_range() {
|
||||
let config = AgentCleanupConfig {
|
||||
idle_timeout_secs: 600, // 10 分钟
|
||||
cleanup_interval_secs: 60, // 1 分钟
|
||||
};
|
||||
assert!(config.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agent_cleanup_validate_min_boundaries() {
|
||||
// 测试最小边界值
|
||||
let config = AgentCleanupConfig {
|
||||
idle_timeout_secs: AgentCleanupConfig::MIN_IDLE_TIMEOUT,
|
||||
cleanup_interval_secs: AgentCleanupConfig::MIN_CLEANUP_INTERVAL,
|
||||
};
|
||||
assert!(config.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agent_cleanup_validate_max_boundaries() {
|
||||
// 测试最大边界值
|
||||
let config = AgentCleanupConfig {
|
||||
idle_timeout_secs: AgentCleanupConfig::MAX_IDLE_TIMEOUT,
|
||||
cleanup_interval_secs: AgentCleanupConfig::MAX_CLEANUP_INTERVAL,
|
||||
};
|
||||
assert!(config.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agent_cleanup_validate_idle_timeout_too_small() {
|
||||
let config = AgentCleanupConfig {
|
||||
idle_timeout_secs: 5, // 小于最小值 10
|
||||
cleanup_interval_secs: 30,
|
||||
};
|
||||
assert!(config.validate().is_err());
|
||||
let err = config.validate().unwrap_err();
|
||||
assert!(err.contains("idle_timeout_secs"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agent_cleanup_validate_idle_timeout_too_large() {
|
||||
let config = AgentCleanupConfig {
|
||||
idle_timeout_secs: 100000, // 大于最大值 86400
|
||||
cleanup_interval_secs: 30,
|
||||
};
|
||||
assert!(config.validate().is_err());
|
||||
let err = config.validate().unwrap_err();
|
||||
assert!(err.contains("idle_timeout_secs"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agent_cleanup_validate_cleanup_interval_too_small() {
|
||||
let config = AgentCleanupConfig {
|
||||
idle_timeout_secs: 180,
|
||||
cleanup_interval_secs: 2, // 小于最小值 5
|
||||
};
|
||||
assert!(config.validate().is_err());
|
||||
let err = config.validate().unwrap_err();
|
||||
assert!(err.contains("cleanup_interval_secs"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agent_cleanup_validate_cleanup_interval_too_large() {
|
||||
let config = AgentCleanupConfig {
|
||||
idle_timeout_secs: 180,
|
||||
cleanup_interval_secs: 5000, // 大于最大值 3600
|
||||
};
|
||||
assert!(config.validate().is_err());
|
||||
let err = config.validate().unwrap_err();
|
||||
assert!(err.contains("cleanup_interval_secs"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agent_cleanup_validate_both_invalid() {
|
||||
let config = AgentCleanupConfig {
|
||||
idle_timeout_secs: 0,
|
||||
cleanup_interval_secs: 0,
|
||||
};
|
||||
assert!(config.validate().is_err());
|
||||
// 应该先检测到 idle_timeout_secs 的错误
|
||||
let err = config.validate().unwrap_err();
|
||||
assert!(err.contains("idle_timeout_secs"));
|
||||
}
|
||||
}
|
||||
1822
qiming-rcoder/crates/agent_runner/src/grpc/agent_service_impl.rs
Normal file
1822
qiming-rcoder/crates/agent_runner/src/grpc/agent_service_impl.rs
Normal file
File diff suppressed because it is too large
Load Diff
7
qiming-rcoder/crates/agent_runner/src/grpc/mod.rs
Normal file
7
qiming-rcoder/crates/agent_runner/src/grpc/mod.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
//! gRPC 服务模块
|
||||
//!
|
||||
//! 提供 agent_runner 的 gRPC 服务端实现,用于替代原有的 HTTP 接口
|
||||
|
||||
pub mod agent_service_impl;
|
||||
|
||||
pub use agent_service_impl::AgentServiceImpl;
|
||||
@@ -0,0 +1,35 @@
|
||||
//! 健康检查处理器
|
||||
|
||||
use axum::Json;
|
||||
use chrono::Utc;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
/// 健康检查响应结构
|
||||
#[derive(serde::Serialize, ToSchema)]
|
||||
pub struct HealthResponse {
|
||||
/// 服务状态
|
||||
pub status: String,
|
||||
/// 时间戳
|
||||
pub timestamp: chrono::DateTime<chrono::Utc>,
|
||||
/// 服务名称
|
||||
pub service: String,
|
||||
}
|
||||
|
||||
/// 健康检查端点
|
||||
///
|
||||
/// 检查服务的健康状态
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/health",
|
||||
responses(
|
||||
(status = 200, description = "服务健康状态", body = HealthResponse)
|
||||
),
|
||||
tag = "system"
|
||||
)]
|
||||
pub async fn health_check() -> Json<HealthResponse> {
|
||||
Json(HealthResponse {
|
||||
status: "healthy".to_string(),
|
||||
timestamp: Utc::now(),
|
||||
service: "rcoder-ai-service".to_string(),
|
||||
})
|
||||
}
|
||||
4
qiming-rcoder/crates/agent_runner/src/handler/mod.rs
Normal file
4
qiming-rcoder/crates/agent_runner/src/handler/mod.rs
Normal file
@@ -0,0 +1,4 @@
|
||||
//! HTTP 路由和处理器模块
|
||||
mod health_handler;
|
||||
|
||||
pub use health_handler::*;
|
||||
@@ -0,0 +1,146 @@
|
||||
//! Computer Agent Cancel Handler
|
||||
//!
|
||||
//! 处理 POST /computer/agent/session/cancel 请求
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
extract::State,
|
||||
http::HeaderMap,
|
||||
};
|
||||
use agent_client_protocol::schema::{CancelNotification, SessionId};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::CancelNotificationRequestWrapper;
|
||||
use crate::http_server::router::AppState;
|
||||
use crate::service::AGENT_REGISTRY;
|
||||
use shared_types::{
|
||||
ComputerAgentCancelRequest, ComputerAgentCancelResponse, HttpResult, I18nJsonOrQuery,
|
||||
error_codes::ERR_VALIDATION, get_i18n_message, AppError,
|
||||
};
|
||||
|
||||
use super::locale_from_headers;
|
||||
|
||||
/// 取消 Computer Agent 会话任务
|
||||
///
|
||||
/// 直接操作 AGENT_REGISTRY 发送取消信号
|
||||
#[utoipa::path(
|
||||
post,
|
||||
path = "/computer/agent/session/cancel",
|
||||
params(
|
||||
ComputerAgentCancelRequest
|
||||
),
|
||||
responses(
|
||||
(status = 200, description = "Cancel request successful", body = HttpResult<ComputerAgentCancelResponse>),
|
||||
(status = 400, description = "Bad request - missing fields"),
|
||||
(status = 500, description = "Internal server error")
|
||||
),
|
||||
tag = "Computer Agent"
|
||||
)]
|
||||
pub async fn handle_computer_cancel(
|
||||
State(_state): State<Arc<AppState>>,
|
||||
headers: HeaderMap,
|
||||
I18nJsonOrQuery(request): I18nJsonOrQuery<ComputerAgentCancelRequest>,
|
||||
) -> Result<Json<HttpResult<ComputerAgentCancelResponse>>, AppError> {
|
||||
let locale = locale_from_headers(&headers);
|
||||
info!(
|
||||
"🚫 [HTTP] Computer Agent 取消请求: user_id={:?}, project_id={}, session_id={:?}",
|
||||
request.user_id, request.project_id, request.session_id
|
||||
);
|
||||
|
||||
// 1. 验证必填字段
|
||||
// user_id 是 Option<String>,需要用 as_ref() 或直接检查
|
||||
if request.user_id.as_ref().map_or(true, |s| s.is_empty()) {
|
||||
return Err(AppError::with_i18n_key(
|
||||
ERR_VALIDATION,
|
||||
&get_i18n_message("error.user_id_required", locale),
|
||||
));
|
||||
}
|
||||
|
||||
if request.project_id.is_empty() {
|
||||
return Err(AppError::with_i18n_key(
|
||||
ERR_VALIDATION,
|
||||
&get_i18n_message("error.project_id_required", locale),
|
||||
));
|
||||
}
|
||||
|
||||
// 2. 查找 session_id (如果未提供,从 AGENT_REGISTRY 获取)
|
||||
let session_id = if let Some(sid) = request.session_id {
|
||||
sid
|
||||
} else {
|
||||
// 从 AGENT_REGISTRY 查找
|
||||
match AGENT_REGISTRY.get_agent_info(&request.project_id) {
|
||||
Some(info) => {
|
||||
// session_id 是 SessionId 类型,直接转换为 String
|
||||
info.session_id.to_string()
|
||||
}
|
||||
None => {
|
||||
// Agent 不存在,幂等返回成功
|
||||
info!(
|
||||
"ℹ️ [HTTP] Agent 不存在,幂等返回成功: project_id={}",
|
||||
request.project_id
|
||||
);
|
||||
let response = ComputerAgentCancelResponse {
|
||||
success: true,
|
||||
session_id: String::new(),
|
||||
};
|
||||
return Ok(Json(HttpResult::success(response)));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// 3. 从 AGENT_REGISTRY 获取 Agent 信息并发送取消信号
|
||||
if let Some(agent_info) = AGENT_REGISTRY.get_agent_info(&request.project_id) {
|
||||
// 获取 cancel_tx
|
||||
let cancel_tx = agent_info.cancel_tx.clone();
|
||||
|
||||
// 释放读锁
|
||||
drop(agent_info);
|
||||
|
||||
// 检查是否已经空闲或停止中(幂等性)
|
||||
if cancel_tx.is_closed() {
|
||||
info!(
|
||||
"ℹ️ [HTTP] Agent stopped, cancel channel is closed: session_id={}",
|
||||
session_id
|
||||
);
|
||||
} else {
|
||||
// 创建取消通知
|
||||
let session_id_obj = SessionId::new(Arc::from(session_id.as_str()));
|
||||
let cancel_notification = CancelNotification::new(session_id_obj);
|
||||
|
||||
// 创建 oneshot channel 接收取消结果 (HTTP 不等待结果,直接丢弃)
|
||||
let (result_tx, _result_rx) = oneshot::channel();
|
||||
let cancel_request = CancelNotificationRequestWrapper {
|
||||
cancel_notification,
|
||||
result_tx,
|
||||
};
|
||||
|
||||
// 发送取消信号 (异步)
|
||||
match cancel_tx.send(cancel_request).await {
|
||||
Ok(_) => {
|
||||
info!("[HTTP] Cancel signal sent: session_id={}", session_id);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [HTTP] Failed to send cancel signal: session_id={}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Session 不存在,幂等返回成功
|
||||
info!(
|
||||
"ℹ️ [HTTP] Agent not found, returning success idempotently: session_id={}",
|
||||
session_id
|
||||
);
|
||||
}
|
||||
|
||||
let response = ComputerAgentCancelResponse {
|
||||
success: true,
|
||||
session_id: session_id.clone(),
|
||||
};
|
||||
|
||||
Ok(Json(HttpResult::success(response)))
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
//! Computer Chat Handler
|
||||
//!
|
||||
//! 处理 POST /computer/chat 请求
|
||||
|
||||
use axum::{extract::State, http::HeaderMap, Json};
|
||||
use std::sync::Arc;
|
||||
use tracing::{error, info};
|
||||
|
||||
use crate::http_server::router::AppState;
|
||||
use crate::service::AGENT_REGISTRY;
|
||||
use crate::service::chat_handler::{ChatHandlerContext, ChatHandlerInput, handle_chat_core};
|
||||
use crate::service::{SESSION_CACHE, SessionData};
|
||||
use dashmap::mapref::entry::Entry;
|
||||
use shared_types::{
|
||||
ChatResponse, ComputerChatRequest, HttpResult, I18nJsonOrQuery, ServiceType,
|
||||
error_codes::ERR_INTERNAL_SERVER_ERROR, error_codes::ERR_VALIDATION, get_i18n_message,
|
||||
};
|
||||
|
||||
use super::locale_from_headers;
|
||||
|
||||
/// 处理 Computer Agent Chat 请求
|
||||
///
|
||||
/// 直接调用 agent_runner 本地的 handle_chat_core(),无需 gRPC 转发
|
||||
#[utoipa::path(
|
||||
post,
|
||||
path = "/computer/chat",
|
||||
request_body = ComputerChatRequest,
|
||||
responses(
|
||||
(status = 200, description = "Chat request successful", body = HttpResult<ChatResponse>),
|
||||
(status = 400, description = "Bad request - missing user_id"),
|
||||
(status = 500, description = "Internal server error")
|
||||
),
|
||||
tag = "Computer Agent"
|
||||
)]
|
||||
pub async fn handle_computer_chat(
|
||||
State(state): State<Arc<AppState>>,
|
||||
headers: HeaderMap,
|
||||
I18nJsonOrQuery(request): I18nJsonOrQuery<ComputerChatRequest>,
|
||||
) -> Result<Json<HttpResult<ChatResponse>>, shared_types::AppError> {
|
||||
let locale = locale_from_headers(&headers);
|
||||
info!(
|
||||
"📨 [HTTP] Received Computer Chat request:\n\
|
||||
├─ user_id: {:?}\n\
|
||||
├─ project_id: {:?}\n\
|
||||
├─ session_id: {:?}\n\
|
||||
├─ request_id: {:?}\n\
|
||||
├─ prompt ({}chars): {:?}\n\
|
||||
├─ pod_id: {:?}\n\
|
||||
├─ tenant_id: {:?}\n\
|
||||
├─ space_id: {:?}\n\
|
||||
├─ isolation_type: {:?}\n\
|
||||
├─ attachments: {:?}\n\
|
||||
├─ data_source_attachments: {:?}\n\
|
||||
├─ model_provider: {:#?}\n\
|
||||
├─ agent_config: {:#?}\n\
|
||||
├─ system_prompt: {:?}\n\
|
||||
└─ user_prompt: {:?}",
|
||||
request.user_id,
|
||||
request.project_id,
|
||||
request.session_id,
|
||||
request.request_id,
|
||||
request.prompt.len(),
|
||||
request.prompt,
|
||||
request.pod_id,
|
||||
request.tenant_id,
|
||||
request.space_id,
|
||||
request.isolation_type,
|
||||
request.attachments,
|
||||
request.data_source_attachments,
|
||||
request.model_provider,
|
||||
request.agent_config,
|
||||
request.system_prompt,
|
||||
request.user_prompt
|
||||
);
|
||||
|
||||
// 1. 验证必填字段
|
||||
if request.user_id.is_empty() {
|
||||
let error_msg = get_i18n_message("error.user_id_required", locale);
|
||||
error!("[HTTP] {}", error_msg);
|
||||
return Err(shared_types::AppError::with_i18n_key(
|
||||
ERR_VALIDATION,
|
||||
&error_msg,
|
||||
));
|
||||
}
|
||||
|
||||
let user_id = request.user_id.clone();
|
||||
|
||||
// 2. 生成或使用提供的 project_id (直接用 UUID,去掉连字符,与 rcoder 保持一致)
|
||||
let project_id = request
|
||||
.project_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string().replace('-', ""));
|
||||
|
||||
// 3. 自动查找现有 session_id (如果未提供)
|
||||
let session_id = request.session_id.or_else(|| {
|
||||
AGENT_REGISTRY
|
||||
.get_agent_info(&project_id)
|
||||
.map(|info| info.session_id.to_string())
|
||||
});
|
||||
|
||||
// 4. 创建项目工作目录(使用配置中的 projects_dir,支持外部配置)
|
||||
// Docker 挂载:宿主机 /computer-project-workspace/{user_id} → 容器 /home/user
|
||||
// Agent 工作目录:/home/user/{project_id}
|
||||
let project_dir = state.config.projects_dir.join(&project_id);
|
||||
|
||||
if let Err(e) = tokio::fs::create_dir_all(&project_dir).await {
|
||||
let error_msg = format!(
|
||||
"{}: {}",
|
||||
get_i18n_message("error.project_dir_create_failed", locale),
|
||||
e
|
||||
);
|
||||
error!("[HTTP] {}", error_msg);
|
||||
return Err(shared_types::AppError::with_message(
|
||||
ERR_INTERNAL_SERVER_ERROR,
|
||||
&error_msg,
|
||||
));
|
||||
}
|
||||
|
||||
// 5. 生成或使用提供的 request_id
|
||||
let request_id = request
|
||||
.request_id
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
||||
|
||||
// 6. 构建 ChatHandlerInput
|
||||
let input = ChatHandlerInput {
|
||||
project_id: project_id.clone(),
|
||||
project_dir,
|
||||
session_id,
|
||||
prompt: request.prompt,
|
||||
request_id: request_id.clone(),
|
||||
attachments: request.attachments,
|
||||
data_source_attachments: request.data_source_attachments,
|
||||
model_config: request.model_provider,
|
||||
service_type: ServiceType::ComputerAgentRunner,
|
||||
agent_config_override: request.agent_config,
|
||||
system_prompt_override: request.system_prompt,
|
||||
user_prompt_template_override: request.user_prompt,
|
||||
skip_slot_limit: true, // HTTP Server 部署,跳过槽位限制
|
||||
};
|
||||
|
||||
// 7. 构建 ChatHandlerContext
|
||||
let context = ChatHandlerContext {
|
||||
agent_runtime: state.agent_runtime.clone(),
|
||||
shared_api_key_manager: state.shared_api_key_manager.clone(),
|
||||
project_uuid_map: state.project_uuid_map.clone(),
|
||||
};
|
||||
|
||||
// 8. 调用核心 Chat 处理逻辑
|
||||
let output = handle_chat_core(input, &context).await;
|
||||
|
||||
// 🔧 关键修复:将 session 写入 SESSION_CACHE(SSE 进度流需要从这里读取)
|
||||
let session_id_str = output.session_id.clone();
|
||||
match SESSION_CACHE.entry(session_id_str.clone()) {
|
||||
Entry::Occupied(entry) => {
|
||||
info!(
|
||||
"[HTTP] SESSION_CACHE already exists, reusing: session_id={}",
|
||||
session_id_str
|
||||
);
|
||||
entry.get().clone()
|
||||
}
|
||||
Entry::Vacant(entry) => {
|
||||
let data = SessionData::new(1000);
|
||||
info!(
|
||||
"[HTTP] SESSION_CACHE created: session_id={}",
|
||||
session_id_str
|
||||
);
|
||||
entry.insert(data.clone());
|
||||
data
|
||||
}
|
||||
};
|
||||
|
||||
// 9. 构建响应
|
||||
let response = ChatResponse {
|
||||
project_id: output.project_id.clone(),
|
||||
session_id: output.session_id.clone(),
|
||||
error: output.error.clone(),
|
||||
request_id: Some(request_id),
|
||||
need_fallback: None,
|
||||
fallback_reason: None,
|
||||
};
|
||||
|
||||
// 10. 根据执行结果返回成功或错误
|
||||
if output.error.is_some() || !output.success {
|
||||
error!(
|
||||
"❌ [HTTP] Computer Chat failed: session_id={}, error={:?}",
|
||||
response.session_id, response.error
|
||||
);
|
||||
// 返回成功的 HTTP 状态码,但 HttpResult 包含错误信息
|
||||
// 这与 rcoder 的行为一致:HTTP 200 + HttpResult.error
|
||||
return Ok(Json(HttpResult::error(
|
||||
output
|
||||
.error_code
|
||||
.as_deref()
|
||||
.unwrap_or(ERR_INTERNAL_SERVER_ERROR),
|
||||
&shared_types::get_error_message(
|
||||
output
|
||||
.error_code
|
||||
.as_deref()
|
||||
.unwrap_or(ERR_INTERNAL_SERVER_ERROR),
|
||||
locale,
|
||||
),
|
||||
)));
|
||||
}
|
||||
|
||||
info!(
|
||||
"✅ [HTTP] Computer Chat response: session_id={}, error={:?}",
|
||||
response.session_id, response.error
|
||||
);
|
||||
|
||||
Ok(Json(HttpResult::success(response)))
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
//! Computer Agent Progress Handler
|
||||
//!
|
||||
//! 处理 GET /computer/progress/{session_id} 请求 (SSE 流)
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
extract::{Path, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::sse::{Event, KeepAlive, Sse},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use futures_util::stream::{Stream, StreamExt};
|
||||
use std::convert::Infallible;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use crate::http_server::router::AppState;
|
||||
use crate::service::{AGENT_REGISTRY, SESSION_CACHE};
|
||||
use shared_types::{
|
||||
AgentStatus, HttpResult, SessionMessageType, UnifiedSessionMessage,
|
||||
error_codes::{ERR_INTERNAL_SERVER_ERROR, ERR_SESSION_NOT_FOUND},
|
||||
get_i18n_message,
|
||||
};
|
||||
|
||||
use super::locale_from_headers;
|
||||
|
||||
/// 统一的 SSE 流类型
|
||||
type SseStream = Pin<Box<dyn Stream<Item = Result<Event, Infallible>> + Send>>;
|
||||
|
||||
/// 检查 Agent 是否处于 idle 状态
|
||||
///
|
||||
/// 返回 true 如果:
|
||||
/// - Agent 不存在,或
|
||||
/// - Agent 状态为 Idle,或
|
||||
/// - Agent 的 session_id 与当前请求的 session_id 不匹配
|
||||
async fn is_agent_idle(session_id: &str) -> bool {
|
||||
// 通过 session_id 获取 agent_info
|
||||
if let Some(info) = AGENT_REGISTRY.get_agent_info_by_session(session_id) {
|
||||
// 检查状态是否为 Idle,或者 session_id 是否匹配
|
||||
info.status == AgentStatus::Idle || info.session_id.to_string() != session_id
|
||||
} else {
|
||||
// Agent 不存在,视为 idle
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建 Agent idle 时的结束事件
|
||||
///
|
||||
/// 当 Agent 处于闲置状态时,发送此事件通知前端没有正在执行的任务
|
||||
fn create_idle_end_event(session_id: &str) -> Event {
|
||||
let unified_message = UnifiedSessionMessage {
|
||||
session_id: session_id.to_string(),
|
||||
message_type: SessionMessageType::SessionPromptEnd,
|
||||
sub_type: "end_turn".to_string(),
|
||||
data: serde_json::json!({
|
||||
"reason": "EndTurn",
|
||||
"description": "Agent 当前无在执行任务"
|
||||
}),
|
||||
timestamp: Utc::now(),
|
||||
};
|
||||
|
||||
let json_data = match serde_json::to_string(&unified_message) {
|
||||
Ok(json) => json,
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [HTTP] 序列化 SessionPromptEnd 消息失败: session_id={}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
// 返回包含 session_id 的最小可用结构
|
||||
format!(
|
||||
r#"{{"sessionId":"{}","messageType":"sessionPromptEnd","subType":"end_turn","data":{{"reason":"EndTurn","description":"Agent 当前无在执行任务"}}}}"#,
|
||||
session_id
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
Event::default().event("end_turn").data(json_data)
|
||||
}
|
||||
|
||||
/// 创建心跳消息流
|
||||
///
|
||||
/// 定期发送符合 UnifiedSessionMessage 格式的心跳消息
|
||||
fn create_heartbeat_stream(
|
||||
session_id: String,
|
||||
) -> impl Stream<Item = Result<Event, Infallible>> + Send {
|
||||
let heartbeat_interval = Duration::from_secs(15);
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(10);
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(heartbeat_interval);
|
||||
// 立即发送第一个心跳,然后按间隔发送
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
|
||||
loop {
|
||||
interval.tick().await;
|
||||
|
||||
// 创建心跳消息
|
||||
let heartbeat_msg = UnifiedSessionMessage::heartbeat(session_id.clone());
|
||||
let json_str = match serde_json::to_string(&heartbeat_msg) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
error!("[HTTP] Failed to serialize heartbeat message: {}", e);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// 使用 sub_type ("ping") 作为事件名
|
||||
if tx
|
||||
.send(Ok(Event::default().event("ping").data(json_str)))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
// 接收端已关闭,停止发送心跳
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
ReceiverStream::new(rx)
|
||||
}
|
||||
|
||||
/// Computer Agent 进度流 (SSE)
|
||||
///
|
||||
/// 直接从 SESSION_CACHE 订阅消息流,无需 gRPC
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/computer/progress/{session_id}",
|
||||
responses(
|
||||
(status = 200, description = "SSE progress stream", content_type = "text/event-stream"),
|
||||
(status = 404, description = "Session not found"),
|
||||
(status = 500, description = "Internal server error")
|
||||
),
|
||||
tag = "Computer Agent"
|
||||
)]
|
||||
pub async fn handle_computer_progress(
|
||||
State(_state): State<Arc<AppState>>,
|
||||
headers: HeaderMap,
|
||||
Path(session_id): Path<String>,
|
||||
) -> Result<Sse<SseStream>, (StatusCode, Json<HttpResult<String>>)> {
|
||||
let locale = locale_from_headers(&headers);
|
||||
info!(
|
||||
"📡 [HTTP] Computer Agent progress stream subscribed: session_id={}",
|
||||
session_id
|
||||
);
|
||||
|
||||
// 0. 检查 Agent 状态(必须在 SESSION_CACHE 查找之前检查)
|
||||
if is_agent_idle(&session_id).await {
|
||||
info!(
|
||||
"💤 [HTTP] Agent is idle, sending SessionPromptEnd and closing: session_id={}",
|
||||
session_id
|
||||
);
|
||||
let end_event = create_idle_end_event(&session_id);
|
||||
let stream: SseStream = Box::pin(futures_util::stream::iter([Ok(end_event)]));
|
||||
return Ok(Sse::new(stream));
|
||||
}
|
||||
|
||||
// 1. 从 SESSION_CACHE 获取 session_data
|
||||
// 🛡️ 关键修复:先 clone Arc<SessionData>,立即释放 DashMap shard 读锁
|
||||
// 之前直接在 Ref 上调用 create_new_connection().await,导致 DashMap 读锁跨 await 持有
|
||||
// 可能造成与 SESSION_CACHE.entry()/remove() 等写操作的死锁
|
||||
let session_data = match SESSION_CACHE.get(&session_id) {
|
||||
Some(data) => data.value().clone(),
|
||||
None => {
|
||||
warn!(" [HTTP] Session not found: session_id={}", session_id);
|
||||
return Err((
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(HttpResult::error_with_message(
|
||||
ERR_SESSION_NOT_FOUND,
|
||||
locale,
|
||||
&format!(
|
||||
"{}: {}",
|
||||
get_i18n_message("error.session_not_found", locale),
|
||||
session_id
|
||||
),
|
||||
)),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// 2. 创建新的消息订阅(DashMap 锁已释放,此处 await 安全)
|
||||
let (message_rx, _cancel_token) = match session_data.create_new_connection(1000).await {
|
||||
Ok(conn) => conn,
|
||||
Err(e) => {
|
||||
error!(
|
||||
"❌ [HTTP] Failed to create session connection: session_id={}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(HttpResult::error_with_message(
|
||||
ERR_INTERNAL_SERVER_ERROR,
|
||||
locale,
|
||||
&format!(
|
||||
"{}: {}",
|
||||
get_i18n_message("error.internal_server_error", locale),
|
||||
e
|
||||
),
|
||||
)),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// 3. 创建消息流和心跳流
|
||||
// UnifiedSessionMessage 已使用 #[serde(rename_all = "camelCase")], 序列化后符合 RCoder 约定
|
||||
let message_stream = ReceiverStream::new(message_rx).map(|msg| {
|
||||
let is_terminal = matches!(msg.message_type, SessionMessageType::SessionPromptEnd);
|
||||
let json_str = match serde_json::to_string(&msg) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
error!("[HTTP] Failed to serialize message: {}", e);
|
||||
return (Ok(Event::default().data("{}")), false);
|
||||
}
|
||||
};
|
||||
(Ok(Event::default().event(msg.sub_type).data(json_str)), is_terminal)
|
||||
});
|
||||
|
||||
// 4. 创建心跳流(标记为非终端)
|
||||
let heartbeat_stream = create_heartbeat_stream(session_id.clone())
|
||||
.map(|event| (event, false));
|
||||
|
||||
// 5. 合并两个流,并用 scan 监测终止条件
|
||||
// select 会继续轮询心跳流(永不结束),所以必须在合并流层面检测终止
|
||||
// 终止条件:
|
||||
// - 收到 SessionPromptEnd(终端消息)→ 发送后结束流
|
||||
// - channel 关闭 → message_stream 返回 None,scan 最终也会结束
|
||||
let merged_stream = futures_util::stream::select(message_stream, heartbeat_stream)
|
||||
.scan(false, |seen_terminal, (event, is_terminal)| {
|
||||
if *seen_terminal {
|
||||
// 已发送终端消息,结束流
|
||||
return std::future::ready(None);
|
||||
}
|
||||
|
||||
if is_terminal {
|
||||
*seen_terminal = true;
|
||||
}
|
||||
|
||||
std::future::ready(Some(event))
|
||||
});
|
||||
|
||||
info!("[HTTP] SSE stream established: session_id={}", session_id);
|
||||
|
||||
let stream: SseStream = Box::pin(merged_stream);
|
||||
Ok(Sse::new(stream))
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
//! Computer Agent Status Handler
|
||||
//!
|
||||
//! 处理 POST /computer/agent/status 请求
|
||||
|
||||
use axum::{extract::State, http::HeaderMap, Json};
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::http_server::router::AppState;
|
||||
use crate::service::AGENT_REGISTRY;
|
||||
use shared_types::{
|
||||
ComputerAgentStatusRequest, ComputerAgentStatusResponse, HttpResult, I18nJsonOrQuery,
|
||||
error_codes::ERR_VALIDATION, get_i18n_message,
|
||||
};
|
||||
|
||||
use super::locale_from_headers;
|
||||
|
||||
/// 查询 Computer Agent 状态
|
||||
///
|
||||
/// 直接使用 AGENT_REGISTRY 查询,无需 gRPC 调用
|
||||
#[utoipa::path(
|
||||
post,
|
||||
path = "/computer/agent/status",
|
||||
request_body = ComputerAgentStatusRequest,
|
||||
responses(
|
||||
(status = 200, description = "Status query successful", body = HttpResult<ComputerAgentStatusResponse>),
|
||||
(status = 400, description = "Bad request - missing fields"),
|
||||
(status = 500, description = "Internal server error")
|
||||
),
|
||||
tag = "Computer Agent"
|
||||
)]
|
||||
pub async fn handle_computer_status(
|
||||
State(_state): State<Arc<AppState>>,
|
||||
headers: HeaderMap,
|
||||
I18nJsonOrQuery(request): I18nJsonOrQuery<ComputerAgentStatusRequest>,
|
||||
) -> Result<Json<HttpResult<ComputerAgentStatusResponse>>, shared_types::AppError> {
|
||||
let locale = locale_from_headers(&headers);
|
||||
|
||||
// 使用 garde 进行字段校验
|
||||
let I18nJsonOrQuery(request) = I18nJsonOrQuery(request).validate_into_app_error()?;
|
||||
let project_id = request.project_id.as_ref().expect("validated: project_id is required and non-empty");
|
||||
|
||||
// 验证 user_id 或 project_id 至少有一个
|
||||
let user_id_empty = request.user_id.as_ref().map_or(true, |s| s.is_empty());
|
||||
if user_id_empty && project_id.is_empty() {
|
||||
return Err(shared_types::AppError::with_i18n_key(
|
||||
ERR_VALIDATION,
|
||||
&get_i18n_message("error.user_id_or_project_id_required", locale),
|
||||
));
|
||||
}
|
||||
|
||||
info!(
|
||||
"🔍 [HTTP] Computer Agent 状态查询: user_id={:?}, project_id={}, pod_id={:?}, tenant_id={:?}, space_id={:?}, isolation_type={:?}",
|
||||
request.user_id, project_id, request.pod_id, request.tenant_id, request.space_id, request.isolation_type
|
||||
);
|
||||
|
||||
// 从 AGENT_REGISTRY 查询 Agent 状态
|
||||
let agent_info = AGENT_REGISTRY.get_agent_info(project_id);
|
||||
|
||||
let response = match agent_info {
|
||||
Some(info) => {
|
||||
// Agent 存在且活跃
|
||||
info!(
|
||||
"✅ [HTTP] Agent status: project_id={}, is_alive=true, session_id={:?}",
|
||||
project_id, info.session_id
|
||||
);
|
||||
|
||||
ComputerAgentStatusResponse {
|
||||
user_id: request.user_id.clone(),
|
||||
project_id: project_id.to_string(),
|
||||
is_alive: true,
|
||||
session_id: Some(info.session_id.to_string()),
|
||||
status: Some(format!("{:?}", info.status)),
|
||||
last_activity: Some(info.last_activity),
|
||||
created_at: Some(info.created_at),
|
||||
}
|
||||
}
|
||||
None => {
|
||||
// Agent 不存在
|
||||
warn!(" [HTTP] Agent not found: project_id={}", project_id);
|
||||
|
||||
ComputerAgentStatusResponse {
|
||||
user_id: request.user_id.clone(),
|
||||
project_id: project_id.to_string(),
|
||||
is_alive: false,
|
||||
session_id: None,
|
||||
status: None,
|
||||
last_activity: None,
|
||||
created_at: None,
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Json(HttpResult::success(response)))
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
//! Computer Agent Stop Handler
|
||||
//!
|
||||
//! 处理 POST /computer/agent/stop 请求
|
||||
|
||||
use axum::{extract::State, http::HeaderMap, Json};
|
||||
use agent_client_protocol::schema::{CancelNotification, SessionId};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::CancelNotificationRequestWrapper;
|
||||
use crate::http_server::router::AppState;
|
||||
use crate::service::AGENT_REGISTRY;
|
||||
use shared_types::{
|
||||
ComputerAgentStopRequest, ComputerAgentStopResponse, HttpResult, I18nJsonOrQuery,
|
||||
error_codes::{ERR_VALIDATION, SUCCESS},
|
||||
get_error_message, get_i18n_message,
|
||||
};
|
||||
|
||||
use super::locale_from_headers;
|
||||
|
||||
/// 停止 Computer Agent
|
||||
///
|
||||
/// 1. 发送取消信号停止正在运行的任务
|
||||
/// 2. 从 AGENT_REGISTRY 移除 Agent 状态
|
||||
#[utoipa::path(
|
||||
post,
|
||||
path = "/computer/agent/stop",
|
||||
request_body = ComputerAgentStopRequest,
|
||||
responses(
|
||||
(status = 200, description = "Stop request successful", body = HttpResult<ComputerAgentStopResponse>),
|
||||
(status = 400, description = "Bad request - missing fields"),
|
||||
(status = 500, description = "Internal server error")
|
||||
),
|
||||
tag = "Computer Agent"
|
||||
)]
|
||||
pub async fn handle_computer_stop(
|
||||
State(_state): State<Arc<AppState>>,
|
||||
headers: HeaderMap,
|
||||
I18nJsonOrQuery(request): I18nJsonOrQuery<ComputerAgentStopRequest>,
|
||||
) -> Result<Json<HttpResult<ComputerAgentStopResponse>>, shared_types::AppError> {
|
||||
let locale = locale_from_headers(&headers);
|
||||
|
||||
// 使用 garde 进行字段校验
|
||||
let I18nJsonOrQuery(request) = I18nJsonOrQuery(request).validate_into_app_error()?;
|
||||
let project_id = request.project_id.as_ref().expect("validated: project_id is required and non-empty");
|
||||
|
||||
// 验证 user_id 或 project_id 至少有一个
|
||||
let user_id_empty = request.user_id.as_ref().map_or(true, |s| s.is_empty());
|
||||
if user_id_empty && project_id.is_empty() {
|
||||
return Err(shared_types::AppError::with_i18n_key(
|
||||
ERR_VALIDATION,
|
||||
&get_i18n_message("error.user_id_or_project_id_required", locale),
|
||||
));
|
||||
}
|
||||
|
||||
info!(
|
||||
"🛑 [HTTP] Computer Agent 停止请求: user_id={:?}, project_id={}, pod_id={:?}, tenant_id={:?}, space_id={:?}, isolation_type={:?}",
|
||||
request.user_id, project_id, request.pod_id, request.tenant_id, request.space_id, request.isolation_type
|
||||
);
|
||||
|
||||
// 获取 Agent 信息并发送取消信号
|
||||
let (success, message) =
|
||||
if let Some(agent_info) = AGENT_REGISTRY.get_agent_info(project_id) {
|
||||
let session_id = agent_info.session_id.to_string();
|
||||
let cancel_tx = agent_info.cancel_tx.clone();
|
||||
|
||||
// 释放读锁
|
||||
drop(agent_info);
|
||||
|
||||
// 发送取消信号(如果 channel 仍然打开)
|
||||
if !cancel_tx.is_closed() {
|
||||
let session_id_obj = SessionId::new(Arc::from(session_id.as_str()));
|
||||
let cancel_notification = CancelNotification::new(session_id_obj);
|
||||
|
||||
let (result_tx, _result_rx) = oneshot::channel();
|
||||
let cancel_request = CancelNotificationRequestWrapper {
|
||||
cancel_notification,
|
||||
result_tx,
|
||||
};
|
||||
|
||||
match cancel_tx.send(cancel_request).await {
|
||||
Ok(_) => {
|
||||
info!("[HTTP] Cancel signal sent: session_id={}", session_id);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [HTTP] Failed to send cancel signal: session_id={}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 从 AGENT_REGISTRY 移除 Agent
|
||||
let removed = AGENT_REGISTRY
|
||||
.remove_by_project(project_id)
|
||||
.is_some();
|
||||
|
||||
if removed {
|
||||
info!("[HTTP] Agent stopped: project_id={}", project_id);
|
||||
(true, get_error_message(SUCCESS, locale))
|
||||
} else {
|
||||
// 可能在取消期间已被清理
|
||||
info!(
|
||||
"ℹ️ [HTTP] Agent already cleaned up: project_id={}",
|
||||
project_id
|
||||
);
|
||||
(
|
||||
true,
|
||||
get_i18n_message("success.agent_already_stopped", locale),
|
||||
)
|
||||
}
|
||||
} else {
|
||||
// Agent 不存在,幂等返回成功
|
||||
info!(
|
||||
"ℹ️ [HTTP] Agent not found, returning success idempotently: project_id={}",
|
||||
project_id
|
||||
);
|
||||
(
|
||||
true,
|
||||
get_i18n_message("success.agent_already_stopped", locale),
|
||||
)
|
||||
};
|
||||
|
||||
let response = ComputerAgentStopResponse {
|
||||
success,
|
||||
message,
|
||||
user_id: request.user_id.clone(),
|
||||
pod_id: None,
|
||||
project_id: project_id.to_string(),
|
||||
};
|
||||
|
||||
Ok(Json(HttpResult::success(response)))
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
//! HTTP 请求处理器
|
||||
//!
|
||||
//! 仅在 `http-server` feature 启用时编译
|
||||
|
||||
use axum::http::HeaderMap;
|
||||
|
||||
pub mod computer_cancel;
|
||||
pub mod computer_chat;
|
||||
pub mod computer_progress;
|
||||
pub mod computer_status;
|
||||
pub mod computer_stop;
|
||||
pub mod rcoder_progress; // RCoder 模式的 SSE 进度流
|
||||
|
||||
pub(super) fn locale_from_headers(headers: &HeaderMap) -> &'static str {
|
||||
let accept_language = headers.get("accept-language").and_then(|v| v.to_str().ok());
|
||||
shared_types::parse_accept_language(accept_language)
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
//! RCoder Agent Progress Handler
|
||||
//!
|
||||
//! 处理 GET /agent/progress/{session_id} 请求 (SSE 流)
|
||||
//!
|
||||
//! 与 computer_progress.rs 使用相同的 SSE 流模式,直接从 SESSION_CACHE 订阅消息
|
||||
|
||||
use std::convert::Infallible;
|
||||
use std::pin::Pin;
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::{
|
||||
extract::Path,
|
||||
http::{HeaderMap, StatusCode},
|
||||
Json,
|
||||
response::sse::{Event, Sse},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use futures_util::stream::{Stream, StreamExt};
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use crate::service::{AGENT_REGISTRY, SESSION_CACHE};
|
||||
use shared_types::{
|
||||
AgentStatus, HttpResult, SessionMessageType, UnifiedSessionMessage,
|
||||
error_codes::{ERR_INTERNAL_SERVER_ERROR, ERR_SESSION_NOT_FOUND},
|
||||
get_i18n_message,
|
||||
};
|
||||
|
||||
use super::locale_from_headers;
|
||||
|
||||
/// 统一的 SSE 流类型
|
||||
type SseStream = Pin<Box<dyn Stream<Item = Result<Event, Infallible>> + Send>>;
|
||||
|
||||
/// 检查 Agent 是否处于 idle 状态
|
||||
///
|
||||
/// 返回 true 如果:
|
||||
/// - Agent 不存在,或
|
||||
/// - Agent 状态为 Idle,或
|
||||
/// - Agent 的 session_id 与当前请求的 session_id 不匹配
|
||||
async fn is_agent_idle(session_id: &str) -> bool {
|
||||
// 通过 session_id 获取 agent_info
|
||||
if let Some(info) = AGENT_REGISTRY.get_agent_info_by_session(session_id) {
|
||||
// 检查状态是否为 Idle,或者 session_id 是否匹配
|
||||
info.status == AgentStatus::Idle || info.session_id.to_string() != session_id
|
||||
} else {
|
||||
// Agent 不存在,视为 idle
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建 Agent idle 时的结束事件
|
||||
///
|
||||
/// 当 Agent 处于闲置状态时,发送此事件通知前端没有正在执行的任务
|
||||
fn create_idle_end_event(session_id: &str) -> Event {
|
||||
let unified_message = UnifiedSessionMessage {
|
||||
session_id: session_id.to_string(),
|
||||
message_type: SessionMessageType::SessionPromptEnd,
|
||||
sub_type: "end_turn".to_string(),
|
||||
data: serde_json::json!({
|
||||
"reason": "EndTurn",
|
||||
"description": "Agent 当前无在执行任务"
|
||||
}),
|
||||
timestamp: Utc::now(),
|
||||
};
|
||||
|
||||
let json_data = match serde_json::to_string(&unified_message) {
|
||||
Ok(json) => json,
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [RCoder] 序列化 SessionPromptEnd 消息失败: session_id={}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
// 返回包含 session_id 的最小可用结构
|
||||
format!(
|
||||
r#"{{"sessionId":"{}","messageType":"sessionPromptEnd","subType":"end_turn","data":{{"reason":"EndTurn","description":"Agent 当前无在执行任务"}}}}"#,
|
||||
session_id
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
Event::default().event("end_turn").data(json_data)
|
||||
}
|
||||
|
||||
/// 创建心跳消息流
|
||||
///
|
||||
/// 定期发送符合 UnifiedSessionMessage 格式的心跳消息
|
||||
fn create_heartbeat_stream(
|
||||
session_id: String,
|
||||
) -> impl Stream<Item = Result<Event, Infallible>> + Send {
|
||||
let heartbeat_interval = Duration::from_secs(15);
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(10);
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(heartbeat_interval);
|
||||
// 立即发送第一个心跳,然后按间隔发送
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
|
||||
loop {
|
||||
interval.tick().await;
|
||||
|
||||
// 创建心跳消息
|
||||
let heartbeat_msg = UnifiedSessionMessage::heartbeat(session_id.clone());
|
||||
let json_str = match serde_json::to_string(&heartbeat_msg) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
error!("[RCoder] Failed to serialize heartbeat message: {}", e);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// 使用 sub_type ("ping") 作为事件名
|
||||
if tx
|
||||
.send(Ok(Event::default().event("ping").data(json_str)))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
// 接收端已关闭,停止发送心跳
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
ReceiverStream::new(rx)
|
||||
}
|
||||
|
||||
/// RCoder Agent 进度流 (SSE)
|
||||
///
|
||||
/// 直接从 SESSION_CACHE 订阅消息流,无需 gRPC
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/agent/progress/{session_id}",
|
||||
params(
|
||||
("session_id" = String, Path, description = "会话ID")
|
||||
),
|
||||
responses(
|
||||
(status = 200, description = "SSE progress stream", content_type = "text/event-stream"),
|
||||
(status = 404, description = "Session not found"),
|
||||
(status = 500, description = "Internal server error")
|
||||
),
|
||||
tag = "RCoder Agent"
|
||||
)]
|
||||
pub async fn handle_rcoder_progress(
|
||||
headers: HeaderMap,
|
||||
Path(session_id): Path<String>,
|
||||
) -> Result<Sse<SseStream>, (StatusCode, Json<HttpResult<String>>)> {
|
||||
let locale = locale_from_headers(&headers);
|
||||
info!(
|
||||
"📡 [RCoder] Progress stream subscribed: session_id={}",
|
||||
session_id
|
||||
);
|
||||
|
||||
// 0. 检查 Agent 状态(必须在 SESSION_CACHE 查找之前检查)
|
||||
if is_agent_idle(&session_id).await {
|
||||
info!(
|
||||
"💤 [RCoder] Agent is idle, sending SessionPromptEnd and closing: session_id={}",
|
||||
session_id
|
||||
);
|
||||
let end_event = create_idle_end_event(&session_id);
|
||||
let stream: SseStream = Box::pin(futures_util::stream::iter([Ok(end_event)]));
|
||||
return Ok(Sse::new(stream));
|
||||
}
|
||||
|
||||
// 1. 从 SESSION_CACHE 获取 session_data
|
||||
// 🛡️ 关键:先 clone Arc<SessionData>,立即释放 DashMap shard 读锁
|
||||
let session_data = match SESSION_CACHE.get(&session_id) {
|
||||
Some(data) => data.value().clone(),
|
||||
None => {
|
||||
warn!(" [RCoder] Session not found: session_id={}", session_id);
|
||||
return Err((
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(HttpResult::error_with_message(
|
||||
ERR_SESSION_NOT_FOUND,
|
||||
locale,
|
||||
&format!(
|
||||
"{}: {}",
|
||||
get_i18n_message("error.session_not_found", locale),
|
||||
session_id
|
||||
),
|
||||
)),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// 2. 创建新的消息订阅(DashMap 锁已释放,此处 await 安全)
|
||||
let (message_rx, _cancel_token) = match session_data.create_new_connection(1000).await {
|
||||
Ok(conn) => conn,
|
||||
Err(e) => {
|
||||
error!(
|
||||
"❌ [RCoder] Failed to create session connection: session_id={}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(HttpResult::error_with_message(
|
||||
ERR_INTERNAL_SERVER_ERROR,
|
||||
locale,
|
||||
&format!(
|
||||
"{}: {}",
|
||||
get_i18n_message("error.internal_server_error", locale),
|
||||
e
|
||||
),
|
||||
)),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// 3. 创建消息流和心跳流
|
||||
let message_stream = ReceiverStream::new(message_rx).map(|msg| {
|
||||
let is_terminal = matches!(msg.message_type, SessionMessageType::SessionPromptEnd);
|
||||
let json_str = match serde_json::to_string(&msg) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
error!("[RCoder] Failed to serialize message: {}", e);
|
||||
return (Ok(Event::default().data("{}")), false);
|
||||
}
|
||||
};
|
||||
(Ok(Event::default().event(msg.sub_type).data(json_str)), is_terminal)
|
||||
});
|
||||
|
||||
// 4. 创建心跳流(标记为非终端)
|
||||
let heartbeat_stream = create_heartbeat_stream(session_id.clone())
|
||||
.map(|event| (event, false));
|
||||
|
||||
// 5. 合并两个流,并用 scan 监测终止条件
|
||||
// select 会继续轮询心跳流(永不结束),所以必须在合并流层面检测终止
|
||||
// 终止条件:
|
||||
// - 收到 SessionPromptEnd(终端消息)→ 发送后结束流
|
||||
// - channel 关闭 → message_stream 返回 None,scan 最终也会结束
|
||||
let merged_stream = futures_util::stream::select(message_stream, heartbeat_stream)
|
||||
.scan(false, |seen_terminal, (event, is_terminal)| {
|
||||
if *seen_terminal {
|
||||
// 已发送终端消息,结束流
|
||||
return std::future::ready(None);
|
||||
}
|
||||
|
||||
if is_terminal {
|
||||
*seen_terminal = true;
|
||||
}
|
||||
|
||||
std::future::ready(Some(event))
|
||||
});
|
||||
|
||||
info!("[RCoder] SSE stream established: session_id={}", session_id);
|
||||
|
||||
let stream: SseStream = Box::pin(merged_stream);
|
||||
Ok(Sse::new(stream))
|
||||
}
|
||||
11
qiming-rcoder/crates/agent_runner/src/http_server/mod.rs
Normal file
11
qiming-rcoder/crates/agent_runner/src/http_server/mod.rs
Normal file
@@ -0,0 +1,11 @@
|
||||
//! HTTP 服务器模块
|
||||
//!
|
||||
//! 仅在 `http-server` feature 启用时编译
|
||||
//! 提供 Computer Agent HTTP REST API
|
||||
|
||||
pub mod handlers;
|
||||
pub mod router;
|
||||
pub mod start;
|
||||
|
||||
pub use router::{AppState, create_router};
|
||||
pub use start::{HttpServerConfig, start_http_server};
|
||||
186
qiming-rcoder/crates/agent_runner/src/http_server/router.rs
Normal file
186
qiming-rcoder/crates/agent_runner/src/http_server/router.rs
Normal file
@@ -0,0 +1,186 @@
|
||||
//! HTTP 路由定义
|
||||
//!
|
||||
//! 定义所有 HTTP 端点和路由
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
routing::{get, post},
|
||||
};
|
||||
use dashmap::DashMap;
|
||||
use std::sync::Arc;
|
||||
use tower_http::limit::RequestBodyLimitLayer;
|
||||
use utoipa::OpenApi;
|
||||
use utoipa_swagger_ui::SwaggerUi;
|
||||
|
||||
use crate::agent_runtime::AgentRuntime;
|
||||
use crate::api_key_manager::ApiKeyManager;
|
||||
use crate::config::AppConfig;
|
||||
use crate::service::local_agent_service::LocalAgentHttpService;
|
||||
|
||||
/// HTTP 应用状态
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub config: AppConfig,
|
||||
pub agent_runtime: Arc<AgentRuntime>,
|
||||
pub api_key_manager: Arc<ApiKeyManager>,
|
||||
pub shared_api_key_manager: Arc<DashMap<String, shared_types::ModelProviderConfig>>,
|
||||
pub project_uuid_map: Arc<DashMap<String, String>>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
pub fn new(
|
||||
config: AppConfig,
|
||||
agent_runtime: Arc<AgentRuntime>,
|
||||
shared_api_key_manager: Arc<DashMap<String, shared_types::ModelProviderConfig>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
agent_runtime: agent_runtime.clone(),
|
||||
api_key_manager: Arc::new(ApiKeyManager::from_shared(shared_api_key_manager.clone())),
|
||||
shared_api_key_manager,
|
||||
project_uuid_map: Arc::new(DashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建 LocalAgentHttpService 实例用于 RCoder 模式
|
||||
pub fn create_local_agent_service(&self) -> Arc<LocalAgentHttpService> {
|
||||
Arc::new(LocalAgentHttpService::new(
|
||||
self.agent_runtime.clone(),
|
||||
self.shared_api_key_manager.clone(),
|
||||
self.project_uuid_map.clone(),
|
||||
self.config.projects_dir.clone(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建 HTTP 路由
|
||||
///
|
||||
/// 组合 Computer Agent 路由和 RCoder Agent 路由
|
||||
pub fn create_router(state: Arc<AppState>) -> Router {
|
||||
use super::handlers::{
|
||||
computer_cancel, computer_chat, computer_progress, computer_status, computer_stop,
|
||||
rcoder_progress,
|
||||
};
|
||||
use shared_types::http_handlers;
|
||||
|
||||
// 创建 LocalAgentHttpService 实例
|
||||
let local_agent_service = state.create_local_agent_service();
|
||||
|
||||
// Computer Agent 路由
|
||||
let computer_routes = Router::new()
|
||||
.route("/computer/chat", post(computer_chat::handle_computer_chat))
|
||||
.route(
|
||||
"/computer/agent/stop",
|
||||
post(computer_stop::handle_computer_stop),
|
||||
)
|
||||
.route(
|
||||
"/computer/agent/status",
|
||||
post(computer_status::handle_computer_status),
|
||||
)
|
||||
.route(
|
||||
"/computer/agent/session/cancel",
|
||||
post(computer_cancel::handle_computer_cancel),
|
||||
)
|
||||
.route(
|
||||
"/computer/progress/{session_id}",
|
||||
get(computer_progress::handle_computer_progress),
|
||||
)
|
||||
.with_state(state.clone());
|
||||
|
||||
// RCoder Agent 路由(使用 LocalAgentHttpService)
|
||||
let rcoder_routes = Router::new()
|
||||
.route("/chat", post(http_handlers::handle_chat::<LocalAgentHttpService>))
|
||||
.route(
|
||||
"/agent/session/cancel",
|
||||
post(http_handlers::handle_cancel::<LocalAgentHttpService>),
|
||||
)
|
||||
.route("/agent/stop", post(http_handlers::handle_stop::<LocalAgentHttpService>))
|
||||
.route(
|
||||
"/agent/status/{project_id}",
|
||||
get(http_handlers::handle_status::<LocalAgentHttpService>),
|
||||
)
|
||||
.route(
|
||||
"/agent/progress/{session_id}",
|
||||
get(rcoder_progress::handle_rcoder_progress),
|
||||
)
|
||||
.with_state(local_agent_service);
|
||||
|
||||
// 通用路由
|
||||
let api_routes = Router::new()
|
||||
.route("/health", get(health_check))
|
||||
.with_state(state.clone());
|
||||
|
||||
// 组合路由
|
||||
Router::new()
|
||||
.merge(computer_routes)
|
||||
.merge(rcoder_routes)
|
||||
.merge(api_routes)
|
||||
.merge(create_swagger_ui())
|
||||
.layer(RequestBodyLimitLayer::new(50 * 1024 * 1024)) // 🔥 50MB body 限制
|
||||
}
|
||||
|
||||
/// 健康检查端点
|
||||
///
|
||||
/// 检查服务的健康状态
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/health",
|
||||
responses(
|
||||
(status = 200, description = "服务健康状态", body = shared_types::HealthResponse)
|
||||
),
|
||||
tag = "system"
|
||||
)]
|
||||
pub async fn health_check() -> Json<shared_types::HealthResponse> {
|
||||
Json(shared_types::HealthResponse::new("agent-runner"))
|
||||
}
|
||||
|
||||
/// 创建 Swagger UI
|
||||
fn create_swagger_ui() -> SwaggerUi {
|
||||
use super::handlers::{
|
||||
computer_cancel::__path_handle_computer_cancel, computer_chat::__path_handle_computer_chat,
|
||||
computer_progress::__path_handle_computer_progress,
|
||||
computer_status::__path_handle_computer_status, computer_stop::__path_handle_computer_stop,
|
||||
};
|
||||
|
||||
#[derive(OpenApi)]
|
||||
#[openapi(
|
||||
paths(
|
||||
// Computer Agent 端点
|
||||
handle_computer_chat,
|
||||
handle_computer_status,
|
||||
handle_computer_stop,
|
||||
handle_computer_cancel,
|
||||
handle_computer_progress,
|
||||
// 健康检查
|
||||
health_check,
|
||||
),
|
||||
components(schemas(
|
||||
// Computer Agent 类型
|
||||
shared_types::ComputerChatRequest,
|
||||
shared_types::ChatResponse,
|
||||
shared_types::ComputerAgentStatusRequest,
|
||||
shared_types::ComputerAgentStatusResponse,
|
||||
shared_types::ComputerAgentStopRequest,
|
||||
shared_types::ComputerAgentStopResponse,
|
||||
shared_types::ComputerAgentCancelRequest,
|
||||
shared_types::ComputerAgentCancelResponse,
|
||||
// RCoder Agent 类型
|
||||
shared_types::RcoderChatRequest,
|
||||
shared_types::RcoderAgentCancelRequest,
|
||||
shared_types::RcoderAgentCancelResponse,
|
||||
shared_types::RcoderAgentStopRequest,
|
||||
shared_types::RcoderAgentStopResponse,
|
||||
shared_types::AgentStatusResponse,
|
||||
// 通用类型
|
||||
shared_types::HealthResponse,
|
||||
)),
|
||||
tags(
|
||||
(name = "Computer Agent", description = "Computer Agent HTTP API"),
|
||||
(name = "RCoder Agent", description = "RCoder Agent HTTP API"),
|
||||
(name = "System", description = "系统管理接口")
|
||||
)
|
||||
)]
|
||||
struct ApiDoc;
|
||||
|
||||
SwaggerUi::new("/api/docs").url("/api-docs/openapi.json", ApiDoc::openapi())
|
||||
}
|
||||
266
qiming-rcoder/crates/agent_runner/src/http_server/start.rs
Normal file
266
qiming-rcoder/crates/agent_runner/src/http_server/start.rs
Normal file
@@ -0,0 +1,266 @@
|
||||
//! HTTP 服务器启动模块
|
||||
//!
|
||||
//! 提供便捷的 HTTP 服务器启动 API
|
||||
//! 支持 HTTP REST API 和可选的 Pingora 代理服务
|
||||
|
||||
use anyhow::Result;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::task::JoinSet;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use crate::agent_runtime::AgentRuntime;
|
||||
use crate::config::AppConfig;
|
||||
use crate::http_server::router::{AppState, create_router};
|
||||
use crate::proxy_agent::cleanup_task::{CleanupConfig, start_cleanup_task};
|
||||
use crate::proxy_agent::set_unlimited_mode;
|
||||
#[cfg(feature = "proxy")]
|
||||
use crate::proxy_agent::start_pingora;
|
||||
|
||||
/// HTTP 服务器配置
|
||||
pub struct HttpServerConfig {
|
||||
/// HTTP 监听端口
|
||||
pub port: u16,
|
||||
/// 应用配置
|
||||
pub app_config: AppConfig,
|
||||
/// Agent 运行时
|
||||
pub agent_runtime: Arc<AgentRuntime>,
|
||||
/// 共享 API Key Manager
|
||||
pub shared_api_key_manager: Arc<dashmap::DashMap<String, shared_types::ModelProviderConfig>>,
|
||||
}
|
||||
|
||||
/// HTTP 服务器控制柄
|
||||
///
|
||||
/// 用于控制 HTTP 服务器的生命周期
|
||||
#[derive(Clone)]
|
||||
pub struct HttpServerHandle {
|
||||
/// 关闭信号令牌
|
||||
shutdown_token: CancellationToken,
|
||||
/// 活跃任务集合
|
||||
join_set: Arc<tokio::sync::Mutex<JoinSet<()>>>,
|
||||
/// Pingora 结果(用于调用 stop)
|
||||
#[cfg(feature = "proxy")]
|
||||
pingora_result: Arc<tokio::sync::Mutex<Option<crate::proxy_agent::PingoraStartResult>>>,
|
||||
}
|
||||
|
||||
impl HttpServerHandle {
|
||||
/// 检查是否收到关闭信号
|
||||
pub fn is_shutdown(&self) -> bool {
|
||||
self.shutdown_token.is_cancelled()
|
||||
}
|
||||
|
||||
/// 停止 HTTP 服务器并等待所有任务完成
|
||||
pub async fn stop(&self) {
|
||||
info!("Stopping HTTP server...");
|
||||
|
||||
// 1. 发送关闭信号
|
||||
self.shutdown_token.cancel();
|
||||
|
||||
// 2. 停止 Pingora 服务
|
||||
#[cfg(feature = "proxy")]
|
||||
{
|
||||
let mut pingora_guard = self.pingora_result.lock().await;
|
||||
if let Some(mut pingora) = pingora_guard.take() {
|
||||
pingora.stop().await;
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 等待所有任务完成(带超时)
|
||||
// 使用 3 秒超时:清理任务会立即退出,axum 有 3 秒进行连接排空
|
||||
let timeout = Duration::from_secs(3);
|
||||
let mut join_set = self.join_set.lock().await;
|
||||
|
||||
loop {
|
||||
match tokio::time::timeout(timeout, join_set.join_next()).await {
|
||||
Ok(Some(Ok(()))) => {
|
||||
info!("Task exited normally");
|
||||
}
|
||||
Ok(Some(Err(e))) => {
|
||||
warn!("Task error: {:?}", e);
|
||||
}
|
||||
Ok(None) => {
|
||||
// JoinSet 为空,所有任务已完成
|
||||
break;
|
||||
}
|
||||
Err(_) => {
|
||||
warn!("Timed out waiting for tasks (3s), aborting remaining tasks");
|
||||
join_set.abort_all();
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
info!("HTTP server stopped");
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动 HTTP 服务器
|
||||
///
|
||||
/// # 示例
|
||||
///
|
||||
/// ```no_run
|
||||
/// use agent_runner::{AgentRuntime, start_http_server, HttpServerConfig, AppConfig, ProxyConfig, HealthCheckConfig};
|
||||
/// use std::sync::Arc;
|
||||
/// use std::path::PathBuf;
|
||||
///
|
||||
/// #[tokio::main]
|
||||
/// async fn main() {
|
||||
/// // 创建 Agent Runtime
|
||||
/// let (runtime, receiver) = AgentRuntime::new(1000);
|
||||
/// let runtime = Arc::new(runtime);
|
||||
/// runtime.start(receiver).await;
|
||||
///
|
||||
/// // 配置 HTTP Server
|
||||
/// let config = HttpServerConfig {
|
||||
/// port: 8080,
|
||||
/// app_config: AppConfig {
|
||||
/// port: 8080,
|
||||
/// projects_dir: PathBuf::from("/app/computer-project-workspace"),
|
||||
/// // 可选:启用 Pingora 代理服务
|
||||
/// proxy_config: Some(ProxyConfig {
|
||||
/// listen_port: 8088,
|
||||
/// default_backend_port: 8080,
|
||||
/// backend_host: "127.0.0.1".to_string(),
|
||||
/// port_param: "port".to_string(),
|
||||
/// health_check: HealthCheckConfig::default(),
|
||||
/// }),
|
||||
/// ..Default::default()
|
||||
/// },
|
||||
/// agent_runtime: runtime,
|
||||
/// shared_api_key_manager: Arc::new(dashmap::DashMap::new()),
|
||||
/// };
|
||||
///
|
||||
/// // 启动 HTTP Server
|
||||
/// let handle = start_http_server(config).await.unwrap();
|
||||
///
|
||||
/// // 优雅停止
|
||||
/// handle.stop().await;
|
||||
/// }
|
||||
/// ```
|
||||
pub async fn start_http_server(config: HttpServerConfig) -> Result<HttpServerHandle> {
|
||||
// 设置 mcp-proxy 日志目录环境变量(如果配置了的话)
|
||||
if let Some(ref log_dir) = config.app_config.mcp_proxy_log_dir {
|
||||
// SAFETY: 在服务启动时设置环境变量是安全的,此时尚未启动多线程任务
|
||||
unsafe {
|
||||
std::env::set_var("MCP_PROXY_LOG_DIR", log_dir);
|
||||
}
|
||||
info!("🔧 Set MCP_PROXY_LOG_DIR={}", log_dir);
|
||||
}
|
||||
|
||||
// 设置无限制模式(HTTP Server 部署不限制槽位)
|
||||
set_unlimited_mode(true);
|
||||
|
||||
// 创建关闭信号令牌
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let join_set = Arc::new(tokio::sync::Mutex::new(JoinSet::new()));
|
||||
#[cfg(feature = "proxy")]
|
||||
let pingora_result = Arc::new(tokio::sync::Mutex::new(None));
|
||||
|
||||
// 1. 启动 Agent 清理任务
|
||||
let cleanup_config = CleanupConfig {
|
||||
idle_timeout: Duration::from_secs(
|
||||
config
|
||||
.app_config
|
||||
.agent_cleanup
|
||||
.clone()
|
||||
.unwrap_or_default()
|
||||
.idle_timeout_secs,
|
||||
),
|
||||
cleanup_interval: Duration::from_secs(
|
||||
config
|
||||
.app_config
|
||||
.agent_cleanup
|
||||
.clone()
|
||||
.unwrap_or_default()
|
||||
.cleanup_interval_secs,
|
||||
),
|
||||
};
|
||||
info!(
|
||||
"🧹 [HTTP] Agent cleanup config: idle_timeout={}s, cleanup_interval={}s",
|
||||
cleanup_config.idle_timeout.as_secs(),
|
||||
cleanup_config.cleanup_interval.as_secs()
|
||||
);
|
||||
let cleanup_token = shutdown_token.child_token();
|
||||
join_set.lock().await.spawn(async move {
|
||||
tokio::select! {
|
||||
_ = start_cleanup_task(cleanup_config) => {}
|
||||
_ = cleanup_token.cancelled() => {
|
||||
info!("Cleanup task received shutdown signal");
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// 2. 启动 Pingora 代理服务(如果配置了且启用了 proxy feature)
|
||||
#[cfg(feature = "proxy")]
|
||||
if let Some(proxy_config) = &config.app_config.proxy_config {
|
||||
let result = start_pingora(proxy_config, config.shared_api_key_manager.clone());
|
||||
// 保存 Pingora 结果以便后续调用 stop
|
||||
*pingora_result.lock().await = Some(result);
|
||||
} else {
|
||||
info!("Pingora proxy service is not configured, skipping startup");
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "proxy"))]
|
||||
info!("Pingora proxy service is disabled (proxy feature not enabled)");
|
||||
|
||||
// 3. 创建 HTTP 应用状态
|
||||
let state = Arc::new(AppState::new(
|
||||
config.app_config.clone(),
|
||||
config.agent_runtime,
|
||||
config.shared_api_key_manager,
|
||||
));
|
||||
|
||||
// 4. 创建路由
|
||||
let app = create_router(state.clone());
|
||||
|
||||
// 5. 绑定地址并启动 HTTP 服务器
|
||||
let addr = SocketAddr::from(([0, 0, 0, 0], config.port));
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
|
||||
info!("HTTP server started on port {}", config.port);
|
||||
|
||||
info!("HTTP API endpoints:");
|
||||
info!(" POST /computer/chat - Computer Agent chat");
|
||||
info!(" POST /computer/agent/status - Computer Agent status");
|
||||
info!(" POST /computer/agent/stop - Computer Agent stop");
|
||||
info!(" POST /computer/agent/session/cancel - Computer Agent cancel");
|
||||
info!(" GET /computer/progress/:session_id - SSE progress stream");
|
||||
info!(" -- RCoder Agent endpoints (new) --");
|
||||
info!(" POST /chat - RCoder Agent chat");
|
||||
info!(" GET /agent/status/:project_id - RCoder Agent status");
|
||||
info!(" POST /agent/stop - RCoder Agent stop");
|
||||
info!(" POST /agent/session/cancel - RCoder Agent cancel");
|
||||
info!(" GET /agent/progress/:session_id - RCoder SSE progress stream");
|
||||
info!(" -- Common endpoints --");
|
||||
info!(" GET /health - Health check");
|
||||
info!(" GET /api/docs - Swagger API documentation");
|
||||
|
||||
// 6. 启动 HTTP 服务任务
|
||||
let http_token = shutdown_token.child_token();
|
||||
// 将 listener 和 app 移入任务中
|
||||
let http_app = app;
|
||||
let http_listener = listener;
|
||||
join_set.lock().await.spawn(async move {
|
||||
// 使用 graceful shutdown wrapper
|
||||
let server = axum::serve(http_listener, http_app).with_graceful_shutdown(async move {
|
||||
let _ = http_token.cancelled().await;
|
||||
});
|
||||
|
||||
match server.await {
|
||||
Ok(()) => info!("HTTP service exited normally"),
|
||||
Err(e) => error!("HTTP service error: {:?}", e),
|
||||
}
|
||||
});
|
||||
|
||||
// 创建 handle
|
||||
let handle = HttpServerHandle {
|
||||
shutdown_token,
|
||||
join_set,
|
||||
#[cfg(feature = "proxy")]
|
||||
pingora_result,
|
||||
};
|
||||
|
||||
Ok(handle)
|
||||
}
|
||||
37
qiming-rcoder/crates/agent_runner/src/lib.rs
Normal file
37
qiming-rcoder/crates/agent_runner/src/lib.rs
Normal file
@@ -0,0 +1,37 @@
|
||||
//! agent_runner 库
|
||||
//!
|
||||
//! 提供 AI 代理运行时和 ACP 协议集成
|
||||
|
||||
pub mod agent_runtime;
|
||||
pub mod api_key_manager;
|
||||
mod config;
|
||||
pub mod grpc;
|
||||
mod handler;
|
||||
mod model;
|
||||
pub mod otel_tracing; // 🔥 设为 public,供其他模块使用
|
||||
mod proxy_agent;
|
||||
pub mod router;
|
||||
pub mod service; // 🔥 设为 public,供测试使用
|
||||
mod utils;
|
||||
|
||||
// 条件性编译:HTTP 服务器模块
|
||||
#[cfg(feature = "http-server")]
|
||||
pub mod http_server;
|
||||
|
||||
// 测试辅助模块 (仅在 testing feature 启用时编译)
|
||||
#[cfg(feature = "testing")]
|
||||
pub mod testing;
|
||||
|
||||
// 重新导出主要的类型和函数
|
||||
pub use agent_runtime::*;
|
||||
pub use config::*;
|
||||
pub use model::*;
|
||||
pub use otel_tracing::*;
|
||||
pub use proxy_agent::*;
|
||||
pub use service::*; // 重新导出 service 模块
|
||||
pub use utils::*;
|
||||
|
||||
#[cfg(feature = "http-server")]
|
||||
pub use http_server::start::HttpServerHandle;
|
||||
#[cfg(feature = "http-server")]
|
||||
pub use http_server::{HttpServerConfig, start_http_server};
|
||||
599
qiming-rcoder/crates/agent_runner/src/main.rs
Normal file
599
qiming-rcoder/crates/agent_runner/src/main.rs
Normal file
@@ -0,0 +1,599 @@
|
||||
use clap::Parser;
|
||||
use dashmap::DashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
// 🆕 使用共享的遥测模块
|
||||
use rcoder_telemetry::{TelemetryConfig, TelemetryGuard};
|
||||
|
||||
mod agent_runtime;
|
||||
mod api_key_manager;
|
||||
mod config;
|
||||
mod grpc;
|
||||
mod handler;
|
||||
mod model;
|
||||
mod process_reaper;
|
||||
mod proxy_agent;
|
||||
|
||||
// 🔥 Pyroscope Profiler 模块(可选:需要 pyroscope feature)
|
||||
#[cfg(feature = "pyroscope")]
|
||||
mod profiler;
|
||||
|
||||
// 🔥 OpenTelemetry 追踪模块(可选:保留用于向后兼容)
|
||||
#[allow(dead_code)]
|
||||
mod otel_tracing;
|
||||
|
||||
mod router;
|
||||
mod service;
|
||||
mod utils;
|
||||
|
||||
// HTTP 服务器模块 (仅在 http-server feature 启用时)
|
||||
#[cfg(feature = "http-server")]
|
||||
mod http_server;
|
||||
|
||||
use agent_runtime::AgentRuntime;
|
||||
use config::{CliArgs, load_config_with_args};
|
||||
use model::*;
|
||||
use proxy_agent::cleanup_task::{CleanupConfig, start_cleanup_task};
|
||||
#[cfg(feature = "proxy")]
|
||||
use rcoder_proxy::{PingoraServerManager, ProxyConfig};
|
||||
use router::AppState;
|
||||
use std::fs::OpenOptions;
|
||||
use std::io::Write;
|
||||
use std::panic;
|
||||
use std::path::PathBuf;
|
||||
#[cfg(unix)]
|
||||
use tokio::signal::unix::{SignalKind, signal};
|
||||
|
||||
/// 🔥 设置自定义 Panic Hook
|
||||
///
|
||||
/// 当 agent_runner panic 时,将完整的 panic 信息(包括 backtrace)写入日志文件
|
||||
/// 这样即使容器被销毁,也能通过挂载的日志目录找到崩溃原因
|
||||
fn set_panic_hook() {
|
||||
let default_hook = panic::take_hook();
|
||||
|
||||
panic::set_hook(Box::new(move |panic_info| {
|
||||
// 🔥 立即写入日志文件(不依赖 tracing,确保在 panic 时也能写入)
|
||||
if let Err(e) = write_panic_to_file(panic_info) {
|
||||
// 如果文件写入失败,尝试输出到 stderr
|
||||
eprintln!("❌ [PANIC] Failed to write panic log file: {}", e);
|
||||
}
|
||||
|
||||
// 🔥 同时输出到 stderr(Docker 会捕获到容器日志)
|
||||
eprintln!("═══════════════════════════════════════════════════════════");
|
||||
eprintln!("❌ [PANIC] agent_runner encountered a fatal error!");
|
||||
eprintln!("═══════════════════════════════════════════════════════════");
|
||||
if let Some(location) = panic_info.location() {
|
||||
eprintln!(
|
||||
"panic.location: {}:{}:{}",
|
||||
location.file(),
|
||||
location.line(),
|
||||
location.column()
|
||||
);
|
||||
}
|
||||
eprintln!("panic.payload: {}", panic_info);
|
||||
eprintln!("═══════════════════════════════════════════════════════════");
|
||||
|
||||
// 调用默认 hook(会终止进程)
|
||||
default_hook(panic_info);
|
||||
}));
|
||||
}
|
||||
|
||||
/// 将 panic 信息写入日志文件
|
||||
fn write_panic_to_file(panic_info: &panic::PanicHookInfo) -> std::io::Result<()> {
|
||||
// 🔥 日志文件路径:/app/container-logs/agent_runner_panic.log(使用已有的挂载目录)
|
||||
let log_path = PathBuf::from("/app/container-logs/agent_runner_panic.log");
|
||||
|
||||
// 确保目录存在
|
||||
if let Some(parent) = log_path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
|
||||
// 打开文件(追加模式)
|
||||
let mut file = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&log_path)?;
|
||||
|
||||
// 获取当前时间
|
||||
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S UTC");
|
||||
|
||||
// 写入 panic 信息
|
||||
writeln!(
|
||||
file,
|
||||
"═══════════════════════════════════════════════════════════"
|
||||
)?;
|
||||
writeln!(file, "❌ [PANIC] agent_runner encountered a fatal error!")?;
|
||||
writeln!(file, "time: {}", now)?;
|
||||
writeln!(
|
||||
file,
|
||||
"═══════════════════════════════════════════════════════════"
|
||||
)?;
|
||||
if let Some(location) = panic_info.location() {
|
||||
writeln!(
|
||||
file,
|
||||
"panic.location: {}:{}:{}",
|
||||
location.file(),
|
||||
location.line(),
|
||||
location.column()
|
||||
)?;
|
||||
}
|
||||
writeln!(file, "panic.payload: {}", panic_info)?;
|
||||
|
||||
// 写入 backtrace(如果启用)
|
||||
#[cfg(feature = "backtrace")]
|
||||
{
|
||||
if let Ok(backtrace) = std::backtrace::Backtrace::capture() {
|
||||
writeln!(file, "Backtrace:\n{}", backtrace)?;
|
||||
}
|
||||
}
|
||||
|
||||
writeln!(
|
||||
file,
|
||||
"═══════════════════════════════════════════════════════════\n"
|
||||
)?;
|
||||
|
||||
// 强制刷新到磁盘
|
||||
file.flush()?;
|
||||
|
||||
eprintln!("✅ Panic info written to: {}", log_path.display());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 🔥 设置优雅关闭信号处理器
|
||||
///
|
||||
/// 监听系统信号,实现优雅关闭:
|
||||
/// - Unix: SIGTERM (Docker stop) + SIGINT (Ctrl+C)
|
||||
/// - Windows: Ctrl+C
|
||||
fn setup_shutdown_handler() -> tokio::task::JoinHandle<()> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
tokio::spawn(async move {
|
||||
// 监听 SIGTERM(Docker stop)
|
||||
let mut sigterm = match signal(SignalKind::terminate()) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
eprintln!("❌ [SIGNAL] Failed to register SIGTERM handler: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// 监听 SIGINT(Ctrl+C)
|
||||
let mut sigint = match signal(SignalKind::interrupt()) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
eprintln!("❌ [SIGNAL] Failed to register SIGINT handler: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
tokio::select! {
|
||||
_ = sigterm.recv() => {
|
||||
eprintln!("📨 [SIGNAL] Received SIGTERM (Docker stop), starting graceful shutdown...");
|
||||
write_shutdown_log("SIGTERM");
|
||||
}
|
||||
_ = sigint.recv() => {
|
||||
eprintln!("📨 [SIGNAL] Received SIGINT (Ctrl+C), starting graceful shutdown...");
|
||||
write_shutdown_log("SIGINT");
|
||||
}
|
||||
}
|
||||
|
||||
eprintln!("🧹 [SIGNAL] Cleaning up resources...");
|
||||
eprintln!("✅ [SIGNAL] Graceful shutdown completed, exiting");
|
||||
std::process::exit(0);
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
tokio::spawn(async move {
|
||||
// Windows: 仅监听 Ctrl+C
|
||||
if let Ok(()) = tokio::signal::ctrl_c().await {
|
||||
eprintln!("📨 [SIGNAL] Received Ctrl+C, starting graceful shutdown...");
|
||||
write_shutdown_log("Ctrl+C");
|
||||
}
|
||||
|
||||
eprintln!("🧹 [SIGNAL] Cleaning up resources...");
|
||||
eprintln!("✅ [SIGNAL] Graceful shutdown completed, exiting");
|
||||
std::process::exit(0);
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 将关闭事件写入日志文件
|
||||
fn write_shutdown_log(signal: &str) {
|
||||
use std::fs::OpenOptions;
|
||||
use std::io::Write;
|
||||
|
||||
let log_path = PathBuf::from("/app/container-logs/agent_runner_shutdown.log");
|
||||
|
||||
if let Some(parent) = log_path.parent() {
|
||||
let _ = std::fs::create_dir_all(parent);
|
||||
}
|
||||
|
||||
if let Ok(mut file) = OpenOptions::new().create(true).append(true).open(&log_path) {
|
||||
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S UTC");
|
||||
let _ = writeln!(
|
||||
file,
|
||||
"═══════════════════════════════════════════════════════════"
|
||||
);
|
||||
let _ = writeln!(
|
||||
file,
|
||||
"📨 [SHUTDOWN] agent_runner received a shutdown signal"
|
||||
);
|
||||
let _ = writeln!(file, "signal: {}", signal);
|
||||
let _ = writeln!(file, "time: {}", now);
|
||||
let _ = writeln!(
|
||||
file,
|
||||
"═══════════════════════════════════════════════════════════\n"
|
||||
);
|
||||
let _ = file.flush();
|
||||
eprintln!("✅ Shutdown info written to: {}", log_path.display());
|
||||
}
|
||||
}
|
||||
|
||||
// 路由创建函数已移动到 handler 模块
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
// 🔥 设置自定义 Panic Hook,确保 panic 信息被记录
|
||||
set_panic_hook();
|
||||
|
||||
// 🔥 设置信号处理器,实现优雅关闭(Docker stop、Ctrl+C)
|
||||
let _shutdown_handle = setup_shutdown_handler();
|
||||
|
||||
// ✅ 初始化 Rustls CryptoProvider(必须在最前面,在任何可能使用 TLS 的代码之前)
|
||||
// 🔥 如果这里失败,会导致 panic,但 panic hook 会捕获并记录
|
||||
rustls::crypto::ring::default_provider()
|
||||
.install_default()
|
||||
.expect(
|
||||
"❌ [FATAL] Rustls CryptoProvider initialization failed. The process cannot continue. This is usually an environment issue.",
|
||||
);
|
||||
|
||||
// 🆕 Initializing telemetry system(使用 rcoder-telemetry,包含控制台 + 文件日志)
|
||||
let telemetry_config = TelemetryConfig::from_env("agent_runner").with_file_log("agent-runner"); // 启用文件日志,前缀为 agent-runner
|
||||
let telemetry: TelemetryGuard = rcoder_telemetry::init(telemetry_config).await?;
|
||||
let telemetry = Arc::new(telemetry);
|
||||
|
||||
// 🆕 Pyroscope Profiler 初始化(可选:需要 pyroscope feature)
|
||||
#[cfg(feature = "pyroscope")]
|
||||
let _pyroscope_guard: Option<profiler::ProfilerGuard> = {
|
||||
info!("Pyroscope profiling feature enabled");
|
||||
match profiler::init_pyroscope_profiler_default() {
|
||||
Ok(guard) => {
|
||||
info!("Pyroscope profiler initialized successfully");
|
||||
Some(guard)
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("Failed to initialize Pyroscope profiler: {}", e);
|
||||
warn!("Continuing without Pyroscope profiling");
|
||||
None
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
#[cfg(not(feature = "pyroscope"))]
|
||||
let _pyroscope_guard: Option<()> = None;
|
||||
|
||||
info!("Starting rcoder - AI-powered development platform");
|
||||
|
||||
// 解析命令行参数
|
||||
let cli_args = CliArgs::parse();
|
||||
|
||||
// 加载配置(包含命令行参数)
|
||||
let config = load_config_with_args(cli_args);
|
||||
|
||||
// 🔥 初始化并发限制(从配置读取)
|
||||
if let Some(ref concurrency_config) = config.agent_concurrency {
|
||||
agent_runtime::init_concurrency_limit(concurrency_config.concurrency_limit);
|
||||
}
|
||||
|
||||
// 🔥 创建 AgentRuntime(新架构)
|
||||
let (agent_runtime, task_receiver) = AgentRuntime::new(1000);
|
||||
let agent_runtime = Arc::new(agent_runtime);
|
||||
info!("🔧 [MAIN] AgentRuntime created");
|
||||
|
||||
// 🔥 启动 Worker(在主运行时中,无需独立线程)
|
||||
agent_runtime.start(task_receiver).await;
|
||||
info!("📌 [MAIN] Agent Worker started");
|
||||
|
||||
// 🔥 启动健康检查和重启任务
|
||||
let health_monitor = spawn_health_monitor(agent_runtime.clone());
|
||||
info!("[MAIN] Worker health monitor started");
|
||||
|
||||
// 🔥 启动僵尸进程回收器(PID 1 必须回收孤儿进程)
|
||||
let _reaper_handle = process_reaper::start_process_reaper();
|
||||
info!("🧹 [MAIN] Process reaper started (PID 1 mode)");
|
||||
|
||||
// 🆕 从配置中获取 Agent 清理配置,或使用默认值
|
||||
let agent_cleanup_config = config.agent_cleanup.clone().unwrap_or_default();
|
||||
let cleanup_config = CleanupConfig {
|
||||
idle_timeout: Duration::from_secs(agent_cleanup_config.idle_timeout_secs),
|
||||
cleanup_interval: Duration::from_secs(agent_cleanup_config.cleanup_interval_secs),
|
||||
};
|
||||
|
||||
info!(
|
||||
"🧹 [MAIN] Agent cleanup config: idle_timeout={}s, cleanup_interval={}s",
|
||||
agent_cleanup_config.idle_timeout_secs, agent_cleanup_config.cleanup_interval_secs
|
||||
);
|
||||
|
||||
// 在主异步运行时中启动清理任务
|
||||
let _cleanup_handle = start_cleanup_task(cleanup_config.clone());
|
||||
|
||||
// proxy_manager 不需要直接访问 app_state,通过参数传递即可
|
||||
|
||||
// 🔒 创建共享的 API 密钥 DashMap
|
||||
let shared_api_key_manager =
|
||||
Arc::new(dashmap::DashMap::<String, shared_types::ModelProviderConfig>::new());
|
||||
info!("🔑 [MAIN] Shared API key DashMap created");
|
||||
|
||||
// 🔥 创建 ApiKeyManager 包装器(包装共享 DashMap,消除双重存储)
|
||||
let api_key_manager = Arc::new(api_key_manager::ApiKeyManager::from_shared(
|
||||
shared_api_key_manager.clone(),
|
||||
));
|
||||
|
||||
// 🔒 project_id -> service_uuid 映射
|
||||
let project_uuid_map: Arc<DashMap<String, String>> = Arc::new(DashMap::new());
|
||||
|
||||
// 🔥 http-server 模式:启动 HTTP + (可选 gRPC) + Pingora
|
||||
#[cfg(feature = "http-server")]
|
||||
{
|
||||
use http_server::{HttpServerConfig, start_http_server};
|
||||
use proxy_agent::set_unlimited_mode;
|
||||
|
||||
// 设置为无限制模式(HTTP Server 部署,不限制槽位)
|
||||
set_unlimited_mode(true);
|
||||
|
||||
// 🔥 1. 可选:启动 gRPC 服务(当 grpc-server feature 启用时)
|
||||
#[cfg(feature = "grpc-server")]
|
||||
let grpc_handle = {
|
||||
info!("ℹ️ HTTP server mode: starting HTTP + gRPC + Pingora");
|
||||
|
||||
let grpc_port = shared_types::GRPC_DEFAULT_PORT;
|
||||
let grpc_addr = format!("[::]:{}", grpc_port)
|
||||
.parse()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to parse gRPC address: {}", e))?;
|
||||
|
||||
// 为 gRPC 创建 state
|
||||
let grpc_state = Arc::new(AppState {
|
||||
sessions: Arc::new(DashMap::new()),
|
||||
config: config.clone(),
|
||||
local_task_sender: agent_runtime.clone(),
|
||||
agent_runtime: agent_runtime.clone(),
|
||||
#[cfg(feature = "proxy")]
|
||||
pingora_service: None,
|
||||
api_key_manager: api_key_manager.clone(),
|
||||
shared_api_key_manager: shared_api_key_manager.clone(),
|
||||
project_uuid_map: project_uuid_map.clone(),
|
||||
});
|
||||
|
||||
// gRPC 消息大小限制
|
||||
let grpc_service = shared_types::grpc::agent_service_server::AgentServiceServer::new(
|
||||
grpc::AgentServiceImpl::new(grpc_state.clone()),
|
||||
)
|
||||
.max_decoding_message_size(shared_types::GRPC_MAX_MESSAGE_SIZE)
|
||||
.max_encoding_message_size(shared_types::GRPC_MAX_MESSAGE_SIZE);
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
info!("gRPC service started, listening on port: {}", grpc_port);
|
||||
info!("gRPC endpoints (port {}):", grpc_port);
|
||||
info!(" agent.AgentService/Chat - gRPC chat");
|
||||
info!(" agent.AgentService/SubscribeProgress - gRPC progress stream");
|
||||
info!(" agent.AgentService/CancelSession - gRPC cancel");
|
||||
info!(" agent.AgentService/GetStatus - gRPC status");
|
||||
if let Err(e) = tonic::transport::Server::builder()
|
||||
.add_service(grpc_service)
|
||||
.serve(grpc_addr)
|
||||
.await
|
||||
{
|
||||
error!("gRPC server error: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
Some(handle)
|
||||
};
|
||||
|
||||
// 无 gRPC 模式
|
||||
#[cfg(not(feature = "grpc-server"))]
|
||||
{
|
||||
info!("ℹ️ HTTP server mode: starting HTTP + Pingora only (no gRPC)");
|
||||
}
|
||||
|
||||
// 🔥 2. 创建 HttpServerConfig(包含所有配置)
|
||||
let http_config = HttpServerConfig {
|
||||
port: config.port,
|
||||
app_config: config.clone(),
|
||||
agent_runtime: agent_runtime.clone(),
|
||||
shared_api_key_manager: shared_api_key_manager.clone(),
|
||||
};
|
||||
|
||||
// 🔥 3. 启动 HTTP 服务器(内部会启动 Pingora)
|
||||
let _handle = start_http_server(http_config).await?;
|
||||
|
||||
// 🔥 4. 同时等待 gRPC(如果有)和信号
|
||||
info!("HTTP + Pingora services started; running until shutdown signal is received");
|
||||
|
||||
#[cfg(feature = "grpc-server")]
|
||||
{
|
||||
tokio::select! {
|
||||
_ = grpc_handle.unwrap() => {
|
||||
info!("gRPC service ended unexpectedly, shutting down...");
|
||||
}
|
||||
_ = tokio::signal::ctrl_c() => {
|
||||
info!("📨 Received shutdown signal, preparing graceful shutdown...");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "grpc-server"))]
|
||||
{
|
||||
tokio::signal::ctrl_c().await?;
|
||||
info!("📨 Received shutdown signal, preparing graceful shutdown...");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// 🔥 non-http-server 模式:启动 gRPC + Pingora(用于 Docker 容器内)
|
||||
#[cfg(not(feature = "http-server"))]
|
||||
{
|
||||
info!("ℹ️ Container mode: starting gRPC + Pingora");
|
||||
|
||||
// 启动 gRPC 服务
|
||||
let grpc_port = shared_types::GRPC_DEFAULT_PORT;
|
||||
let grpc_addr = format!("[::]:{}", grpc_port)
|
||||
.parse()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to parse gRPC address: {}", e))?;
|
||||
|
||||
// 为 gRPC 创建 state
|
||||
let grpc_state = Arc::new(AppState {
|
||||
sessions: Arc::new(DashMap::new()),
|
||||
config: config.clone(),
|
||||
local_task_sender: agent_runtime.clone(),
|
||||
agent_runtime: agent_runtime.clone(),
|
||||
#[cfg(feature = "proxy")]
|
||||
pingora_service: None,
|
||||
api_key_manager: api_key_manager.clone(),
|
||||
shared_api_key_manager: shared_api_key_manager.clone(),
|
||||
project_uuid_map: project_uuid_map.clone(),
|
||||
});
|
||||
|
||||
// gRPC 消息大小限制
|
||||
let grpc_service = shared_types::grpc::agent_service_server::AgentServiceServer::new(
|
||||
grpc::AgentServiceImpl::new(grpc_state.clone()),
|
||||
)
|
||||
.max_decoding_message_size(shared_types::GRPC_MAX_MESSAGE_SIZE)
|
||||
.max_encoding_message_size(shared_types::GRPC_MAX_MESSAGE_SIZE);
|
||||
|
||||
let grpc_handle = tokio::spawn(async move {
|
||||
info!("gRPC service started, listening on port: {}", grpc_port);
|
||||
info!("gRPC endpoints (port {}):", grpc_port);
|
||||
info!(" agent.AgentService/Chat - gRPC chat");
|
||||
info!(" agent.AgentService/SubscribeProgress - gRPC progress stream");
|
||||
info!(" agent.AgentService/CancelSession - gRPC cancel");
|
||||
info!(" agent.AgentService/GetStatus - gRPC status");
|
||||
if let Err(e) = tonic::transport::Server::builder()
|
||||
.add_service(grpc_service)
|
||||
.serve(grpc_addr)
|
||||
.await
|
||||
{
|
||||
error!("gRPC server error: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
// 启动轻量 HTTP 健康检查服务(供 docker_manager 健康检查使用)
|
||||
let health_port = config.port; // 默认 8086,来自 --port 参数
|
||||
let _health_handle = tokio::spawn(async move {
|
||||
use axum::{Json, Router, routing::get};
|
||||
|
||||
async fn health_check() -> Json<shared_types::HealthResponse> {
|
||||
Json(shared_types::HealthResponse::new("agent-runner"))
|
||||
}
|
||||
|
||||
let app = Router::new().route("/health", get(health_check));
|
||||
let addr = format!("0.0.0.0:{}", health_port);
|
||||
|
||||
info!(
|
||||
"🏥 HTTP health check service started, listening on port: {}",
|
||||
health_port
|
||||
);
|
||||
|
||||
let listener = match tokio::net::TcpListener::bind(&addr).await {
|
||||
Ok(l) => l,
|
||||
Err(e) => {
|
||||
error!(
|
||||
"❌ Failed to bind HTTP health check service: {} (port: {})",
|
||||
e, health_port
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = axum::serve(listener, app).await {
|
||||
error!("HTTP health check service error: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
// 启动 Pingora(如有配置且启用了 proxy feature)
|
||||
#[cfg(feature = "proxy")]
|
||||
let pingora_result = {
|
||||
use proxy_agent::start_pingora;
|
||||
|
||||
if let Some(proxy_config) = &config.proxy_config {
|
||||
Some(start_pingora(proxy_config, shared_api_key_manager.clone()))
|
||||
} else {
|
||||
info!("ℹ️ Pingora proxy service is not configured");
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
#[cfg(not(feature = "proxy"))]
|
||||
let pingora_result: Option<()> = {
|
||||
info!("ℹ️ Pingora proxy service is disabled (proxy feature not enabled)");
|
||||
None
|
||||
};
|
||||
|
||||
// 等待 gRPC 服务
|
||||
let _ = grpc_handle.await;
|
||||
|
||||
// 停止 Pingora 服务
|
||||
#[cfg(feature = "proxy")]
|
||||
if let Some(mut result) = pingora_result {
|
||||
result.stop().await;
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "proxy"))]
|
||||
let _ = pingora_result;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// 🔥 健康监控任务 (新架构)
|
||||
///
|
||||
/// 定期检查 Agent Worker 健康状态,自动重启不健康的 Worker
|
||||
async fn spawn_health_monitor(runtime: Arc<AgentRuntime>) -> tokio::task::JoinHandle<()> {
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(Duration::from_secs(5));
|
||||
let mut consecutive_failures: u32 = 0;
|
||||
const MAX_RESTART_ATTEMPTS: u32 = 5;
|
||||
const RESTART_COOLDOWN_SECS: u64 = 60;
|
||||
|
||||
info!("[HealthMonitor] Health monitor started");
|
||||
|
||||
loop {
|
||||
interval.tick().await;
|
||||
|
||||
// 检查健康状态
|
||||
if !runtime.check_health().await {
|
||||
error!("[HealthMonitor] Worker reported unhealthy");
|
||||
|
||||
// 检查冷却期
|
||||
if consecutive_failures >= MAX_RESTART_ATTEMPTS {
|
||||
warn!(
|
||||
"⏳ [HealthMonitor] {} consecutive restart failures, entering cooldown",
|
||||
consecutive_failures
|
||||
);
|
||||
tokio::time::sleep(Duration::from_secs(RESTART_COOLDOWN_SECS)).await;
|
||||
consecutive_failures = 0;
|
||||
info!("[HealthMonitor] Cooldown ended, reset failure counter");
|
||||
}
|
||||
|
||||
// 创建新的通道
|
||||
let (new_tx, new_rx) = tokio::sync::mpsc::channel(1000);
|
||||
|
||||
// 重启 worker
|
||||
runtime.restart(new_rx).await;
|
||||
consecutive_failures += 1;
|
||||
info!(
|
||||
"🔄 [HealthMonitor] Worker restart completed (attempt #{})",
|
||||
consecutive_failures
|
||||
);
|
||||
} else {
|
||||
consecutive_failures = 0;
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
26
qiming-rcoder/crates/agent_runner/src/model.rs
Normal file
26
qiming-rcoder/crates/agent_runner/src/model.rs
Normal file
@@ -0,0 +1,26 @@
|
||||
// 重新导出 shared_types 中的模型,保持向后兼容
|
||||
|
||||
pub use agent_abstraction::PromptMessage;
|
||||
|
||||
pub use shared_types::{
|
||||
// Agent model exports
|
||||
AgentStatus,
|
||||
// Session and message exports
|
||||
Attachment,
|
||||
AttachmentSource,
|
||||
AudioAttachment,
|
||||
// 取消相关类型
|
||||
CancelNotificationRequestWrapper,
|
||||
CancelResult,
|
||||
ChatPromptResponse,
|
||||
DocumentAttachment,
|
||||
ImageAttachment,
|
||||
ProjectAndAgentInfo,
|
||||
SessionMessageType,
|
||||
SessionNotify,
|
||||
TextAttachment,
|
||||
UnifiedSessionMessage,
|
||||
};
|
||||
|
||||
// 重新导出 ACP 类型,供客户端构造取消请求使用
|
||||
pub use agent_client_protocol::schema::{CancelNotification, SessionId};
|
||||
457
qiming-rcoder/crates/agent_runner/src/otel_tracing/mod.rs
Normal file
457
qiming-rcoder/crates/agent_runner/src/otel_tracing/mod.rs
Normal file
@@ -0,0 +1,457 @@
|
||||
//! OpenTelemetry 追踪模块
|
||||
//!
|
||||
//! 集成 `rcoder-telemetry` 提供完整的分布式追踪功能。
|
||||
|
||||
use opentelemetry::trace::Status;
|
||||
use tracing::{Level, Span, error, info, span};
|
||||
use tracing_opentelemetry::OpenTelemetrySpanExt;
|
||||
|
||||
/// OpenTelemetry 追踪配置
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TraceConfig {
|
||||
/// 是否启用追踪
|
||||
pub enabled: bool,
|
||||
/// OTLP 导出端点(如 Jaeger、OTLP)
|
||||
pub exporter_endpoint: Option<String>,
|
||||
/// 是否启用 Prometheus 指标
|
||||
pub prometheus_enabled: bool,
|
||||
/// 服务名称
|
||||
pub service_name: String,
|
||||
}
|
||||
|
||||
impl Default for TraceConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: true,
|
||||
exporter_endpoint: None,
|
||||
prometheus_enabled: true,
|
||||
service_name: "agent-runner".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 初始化 OpenTelemetry 追踪
|
||||
///
|
||||
/// 集成 `rcoder-telemetry` 提供完整的分布式追踪功能:
|
||||
/// - OTLP 导出器(支持 Jaeger、Zipkin、OTLP Collector)
|
||||
/// - Prometheus 指标
|
||||
/// - Trace context 传播
|
||||
/// - Console 日志
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `config` - 追踪配置
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// 返回遥测系统 Guard,需要在应用运行期间保持存活
|
||||
///
|
||||
/// # Environment Variables
|
||||
///
|
||||
/// 支持以下环境变量:
|
||||
/// - `OTEL_EXPORTER_OTLP_ENDPOINT`: OTLP 端点(如 http://jaeger:4317)
|
||||
/// - `OTEL_TRACES_SAMPLER_ARG`: 采样率 (0.0-1.0,默认 1.0)
|
||||
/// - `TELEMETRY_PROMETHEUS_ENABLED`: 是否启用 Prometheus(默认 true)
|
||||
pub async fn init_tracing(config: TraceConfig) -> anyhow::Result<rcoder_telemetry::TelemetryGuard> {
|
||||
if !config.enabled {
|
||||
info!("📍 [OTel] Tracing is disabled");
|
||||
// 即使追踪禁用,仍然Initializing Prometheus(如果启用)
|
||||
if config.prometheus_enabled {
|
||||
return rcoder_telemetry::init_prometheus_only(&config.service_name);
|
||||
}
|
||||
// 创建一个空的 guard(仅保留服务名称,用于后续渲染指标)
|
||||
// 使用空的 telemetry_config 创建一个没有任何功能的 guard
|
||||
let telemetry_config = rcoder_telemetry::TelemetryConfig::new(&config.service_name);
|
||||
return rcoder_telemetry::init(telemetry_config).await;
|
||||
}
|
||||
|
||||
// 使用 from_env 从环境变量读取配置,然后覆盖必要字段
|
||||
let mut telemetry_config = rcoder_telemetry::TelemetryConfig::from_env(&config.service_name);
|
||||
|
||||
// 如果配置中指定了端点,覆盖环境变量中的值
|
||||
if let Some(ref endpoint) = config.exporter_endpoint {
|
||||
telemetry_config = telemetry_config.with_otlp_endpoint(endpoint);
|
||||
}
|
||||
|
||||
// 根据配置设置 Prometheus
|
||||
if config.prometheus_enabled {
|
||||
telemetry_config = telemetry_config.with_prometheus();
|
||||
} else {
|
||||
telemetry_config = telemetry_config.without_prometheus();
|
||||
}
|
||||
|
||||
// Initializing telemetry system
|
||||
let guard = rcoder_telemetry::init(telemetry_config).await?;
|
||||
|
||||
info!(
|
||||
"✅ [OTel] Tracing initialized: OTLP={}, Prometheus={}",
|
||||
guard.is_otlp_enabled(),
|
||||
guard.is_prometheus_enabled()
|
||||
);
|
||||
|
||||
Ok(guard)
|
||||
}
|
||||
|
||||
impl TraceConfig {
|
||||
/// 从环境变量构建配置
|
||||
pub fn from_env() -> Self {
|
||||
Self {
|
||||
enabled: std::env::var("OTEL_TRACING_ENABLED")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(true),
|
||||
exporter_endpoint: std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok(),
|
||||
prometheus_enabled: std::env::var("TELEMETRY_PROMETHEUS_ENABLED")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(true),
|
||||
service_name: "agent-runner".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 设置 OTLP 端点
|
||||
pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
|
||||
self.exporter_endpoint = Some(endpoint.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 禁用 Prometheus
|
||||
pub fn without_prometheus(mut self) -> Self {
|
||||
self.prometheus_enabled = false;
|
||||
self
|
||||
}
|
||||
|
||||
/// 禁用追踪
|
||||
pub fn disabled(mut self) -> Self {
|
||||
self.enabled = false;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置服务名称
|
||||
pub fn with_service_name(mut self, name: impl Into<String>) -> Self {
|
||||
self.service_name = name.into();
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// 请求追踪 Guard(自动管理 span 生命周期)
|
||||
///
|
||||
/// 使用 `OpenTelemetrySpanExt` 支持动态属性设置
|
||||
///
|
||||
/// 注意:此结构体实现 Send + Sync,可在 tokio::spawn 中安全使用
|
||||
pub struct RequestSpan {
|
||||
/// 底层 tracing span(用于 OpenTelemetrySpanExt 方法)
|
||||
span: Span,
|
||||
}
|
||||
|
||||
impl RequestSpan {
|
||||
/// 创建新的请求 span
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `project_id` - 项目 ID
|
||||
/// * `request_id` - 请求 ID
|
||||
/// * `operation` - 操作名称(如 "process_prompt")
|
||||
pub fn new(project_id: &str, request_id: &str, operation: &str) -> Self {
|
||||
let span = span!(
|
||||
Level::INFO,
|
||||
"agent_request",
|
||||
project_id = %project_id,
|
||||
request_id = %request_id,
|
||||
operation = %operation,
|
||||
);
|
||||
|
||||
info!(
|
||||
"📍 [OTel] Span created: project_id={}, request_id={}, operation={}",
|
||||
project_id, request_id, operation
|
||||
);
|
||||
|
||||
Self { span }
|
||||
}
|
||||
|
||||
/// 设置动态属性(使用 OpenTelemetrySpanExt)
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `key` - 属性键
|
||||
/// * `value` - 属性值
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use agent_runner::RequestSpan;
|
||||
///
|
||||
/// let span = RequestSpan::new("proj", "req", "op");
|
||||
/// span.set_attribute("http.status_code", 200);
|
||||
/// span.set_attribute("user.id", "user-123");
|
||||
/// ```
|
||||
pub fn set_attribute<K, V>(&self, key: K, value: V)
|
||||
where
|
||||
K: Into<opentelemetry::Key>,
|
||||
V: Into<opentelemetry::Value>,
|
||||
{
|
||||
self.span.set_attribute(key, value);
|
||||
}
|
||||
|
||||
/// 批量设置动态属性
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `attributes` - 属性列表 (key, value)
|
||||
pub fn set_attributes(&self, attributes: &[(&str, &str)]) {
|
||||
for (key, value) in attributes {
|
||||
// 需要转换为 String 以满足 'static 生命周期要求
|
||||
self.span.set_attribute(key.to_string(), value.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
/// 设置 span 状态为成功
|
||||
pub fn set_ok(&self) {
|
||||
self.span.set_status(Status::Ok);
|
||||
}
|
||||
|
||||
/// 设置 span 状态为错误
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `description` - 错误描述
|
||||
pub fn set_error(&self, description: impl Into<std::borrow::Cow<'static, str>>) {
|
||||
self.span.set_status(Status::error(description));
|
||||
}
|
||||
|
||||
/// 添加事件到当前 span(使用 OpenTelemetrySpanExt)
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `name` - 事件名称
|
||||
/// * `attributes` - 事件属性
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use agent_runner::RequestSpan;
|
||||
/// use opentelemetry::KeyValue;
|
||||
///
|
||||
/// let span = RequestSpan::new("proj", "req", "op");
|
||||
/// span.add_event("cache_hit", vec![
|
||||
/// KeyValue::new("cache.key", "user:123"),
|
||||
/// KeyValue::new("cache.ttl", 300),
|
||||
/// ]);
|
||||
/// ```
|
||||
pub fn add_event(
|
||||
&self,
|
||||
name: impl Into<std::borrow::Cow<'static, str>>,
|
||||
attributes: Vec<opentelemetry::KeyValue>,
|
||||
) {
|
||||
self.span.add_event(name, attributes);
|
||||
}
|
||||
|
||||
/// 记录事件(简化版本,兼容旧 API)
|
||||
pub fn event(&self, name: &str, attributes: &[(&str, String)]) {
|
||||
// 转换为 String 以满足 'static 生命周期要求
|
||||
let kv_attrs: Vec<opentelemetry::KeyValue> = attributes
|
||||
.iter()
|
||||
.map(|(k, v)| opentelemetry::KeyValue::new(k.to_string(), v.clone()))
|
||||
.collect();
|
||||
|
||||
self.span.add_event(name.to_string(), kv_attrs);
|
||||
|
||||
info!(
|
||||
"📍 [OTel] Event: {}, attributes: {:?}",
|
||||
name,
|
||||
attributes
|
||||
.iter()
|
||||
.map(|(k, v)| format!("{}={}", k, v))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
);
|
||||
}
|
||||
|
||||
/// 记录错误到当前 span
|
||||
pub fn error(&self, err: &anyhow::Error) {
|
||||
self.set_error(err.to_string());
|
||||
error!(
|
||||
error = %err,
|
||||
"📍 [OTel] Request error"
|
||||
);
|
||||
}
|
||||
|
||||
/// 获取 OpenTelemetry 上下文(用于跨服务传播)
|
||||
pub fn context(&self) -> opentelemetry::Context {
|
||||
self.span.context()
|
||||
}
|
||||
|
||||
/// 完成 span(手动关闭)
|
||||
pub fn finish(self) {
|
||||
self.set_ok();
|
||||
// span 在 drop 时会自动关闭
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for RequestSpan {
|
||||
fn drop(&mut self) {
|
||||
info!("📍 [OTel] Span closed");
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建子 span(用于追踪子操作)
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `_parent` - 父 span 的引用(子 span 会自动继承父 span 的上下文)
|
||||
/// * `name` - 子 span 名称
|
||||
/// * `attributes` - 附加属性
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use agent_runner::{RequestSpan, child_span};
|
||||
///
|
||||
/// let parent = RequestSpan::new("proj", "req", "parent_op");
|
||||
/// let child = child_span(&parent, "child_op", &[("key", "value".to_string())]);
|
||||
/// child.finish();
|
||||
/// parent.finish();
|
||||
/// ```
|
||||
pub fn child_span(_parent: &RequestSpan, name: &str, attributes: &[(&str, String)]) -> RequestSpan {
|
||||
let span = span!(
|
||||
Level::INFO,
|
||||
"child_operation",
|
||||
otel.name = %name,
|
||||
);
|
||||
|
||||
// 使用 OpenTelemetrySpanExt 设置动态属性
|
||||
// 需要转换为 String 以满足 'static 生命周期要求
|
||||
for (key, value) in attributes {
|
||||
span.set_attribute(key.to_string(), value.clone());
|
||||
}
|
||||
|
||||
info!("📍 [OTel] Child span created: {}", name);
|
||||
|
||||
RequestSpan { span }
|
||||
}
|
||||
|
||||
/// 从上下文中提取当前 span(用于跨线程传递)
|
||||
pub fn current_span() -> RequestSpan {
|
||||
let span = Span::current();
|
||||
|
||||
RequestSpan { span }
|
||||
}
|
||||
|
||||
/// 创建带属性的 span
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use agent_runner::otel_tracing::span_with_attributes;
|
||||
///
|
||||
/// let span = span_with_attributes(
|
||||
/// "process_attachment",
|
||||
/// &[
|
||||
/// ("file_name", "test.pdf".to_string()),
|
||||
/// ("file_size", "1024".to_string()),
|
||||
/// ]
|
||||
/// );
|
||||
/// ```
|
||||
pub fn span_with_attributes(name: &str, attributes: &[(&str, String)]) -> RequestSpan {
|
||||
let span = span!(
|
||||
Level::INFO,
|
||||
"custom_operation",
|
||||
otel.name = %name,
|
||||
);
|
||||
|
||||
// 使用 OpenTelemetrySpanExt 设置动态属性
|
||||
// 需要转换为 String 以满足 'static 生命周期要求
|
||||
for (key, value) in attributes {
|
||||
span.set_attribute(key.to_string(), value.clone());
|
||||
}
|
||||
|
||||
RequestSpan { span }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_trace_config_default() {
|
||||
let config = TraceConfig::default();
|
||||
assert!(config.enabled);
|
||||
assert!(config.exporter_endpoint.is_none());
|
||||
assert!(config.prometheus_enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_request_span_creation() {
|
||||
let span = RequestSpan::new("test_project", "test_request", "test_operation");
|
||||
span.event("test_event", &[("key", "value".to_string())]);
|
||||
span.finish();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_request_span_with_attributes() {
|
||||
let span = RequestSpan::new("test_project", "test_request", "test_operation");
|
||||
|
||||
// 使用 OpenTelemetrySpanExt 设置动态属性
|
||||
span.set_attribute("http.method", "GET");
|
||||
span.set_attribute("http.status_code", 200i64);
|
||||
span.set_attribute("user.id", "user-123");
|
||||
|
||||
span.set_ok();
|
||||
span.finish();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_request_span_with_error() {
|
||||
let span = RequestSpan::new("test_project", "test_request", "test_operation");
|
||||
span.set_error("Connection timeout");
|
||||
// span 会在 drop 时自动关闭
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_request_span_auto_drop() {
|
||||
let _span = RequestSpan::new("test_project", "test_request", "test_operation");
|
||||
// Span 会在 drop 时自动关闭
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_child_span() {
|
||||
let parent = RequestSpan::new("test_project", "test_request", "parent_operation");
|
||||
let child = child_span(&parent, "child_operation", &[("attr", "value".to_string())]);
|
||||
child.finish();
|
||||
parent.finish();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_span_with_attributes() {
|
||||
let span = span_with_attributes(
|
||||
"custom_op",
|
||||
&[
|
||||
("file_name", "test.pdf".to_string()),
|
||||
("file_size", "1024".to_string()),
|
||||
],
|
||||
);
|
||||
span.finish();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_event() {
|
||||
let span = RequestSpan::new("test_project", "test_request", "test_operation");
|
||||
|
||||
span.add_event(
|
||||
"cache_hit",
|
||||
vec![
|
||||
opentelemetry::KeyValue::new("cache.key", "user:123"),
|
||||
opentelemetry::KeyValue::new("cache.ttl", 300i64),
|
||||
],
|
||||
);
|
||||
|
||||
span.finish();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_init_tracing() {
|
||||
let config = TraceConfig::default();
|
||||
let result = init_tracing(config).await;
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
}
|
||||
672
qiming-rcoder/crates/agent_runner/src/process_reaper.rs
Normal file
672
qiming-rcoder/crates/agent_runner/src/process_reaper.rs
Normal file
@@ -0,0 +1,672 @@
|
||||
//! 僵尸进程回收器 (Zombie Process Reaper)
|
||||
//!
|
||||
//! 当 agent_runner 作为容器的 PID 1 运行时,它需要负责回收孤儿进程。
|
||||
//! 此模块实现了一个基于 SIGCHLD 信号的子进程回收机制。
|
||||
//!
|
||||
//! # 设计原理
|
||||
//!
|
||||
//! 在 Linux 容器中,如果 PID 1 不调用 wait() 回收子进程,这些子进程
|
||||
//! 退出后会变成僵尸进程(Zombie),占用系统资源。
|
||||
//!
|
||||
//! # 使用方式
|
||||
//!
|
||||
//! ```rust
|
||||
//! use process_reaper::start_process_reaper;
|
||||
//!
|
||||
//! // 在主函数中启动回收器
|
||||
//! let _reaper_handle = start_process_reaper();
|
||||
//! ```
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use tokio::process::Child;
|
||||
#[cfg(unix)]
|
||||
use tokio::signal::unix::{SignalKind, signal};
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
/// 进程回收器配置
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ReaperConfig {
|
||||
/// 是否启用详细日志
|
||||
pub verbose: bool,
|
||||
/// 是否启用主动僵尸进程检测
|
||||
pub enable_zombie_detection: bool,
|
||||
/// 僵尸进程检测间隔(秒)
|
||||
pub zombie_detection_interval_secs: u64,
|
||||
}
|
||||
|
||||
impl Default for ReaperConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
verbose: false,
|
||||
enable_zombie_detection: true,
|
||||
zombie_detection_interval_secs: 10,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 僵尸进程信息
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ZombieProcessInfo {
|
||||
pub pid: u32,
|
||||
pub ppid: u32,
|
||||
pub comm: String,
|
||||
pub state: char,
|
||||
}
|
||||
|
||||
/// 进程回收器状态
|
||||
#[derive(Debug)]
|
||||
struct ReaperState {
|
||||
/// 追踪活跃的子进程
|
||||
/// 存储格式: pid -> Child
|
||||
active_children: HashMap<u32, Child>,
|
||||
/// 回收的进程总数
|
||||
reaped_count: u64,
|
||||
/// 检测到的僵尸进程数
|
||||
zombie_detected_count: u64,
|
||||
/// 配置
|
||||
config: ReaperConfig,
|
||||
}
|
||||
|
||||
impl ReaperState {
|
||||
fn new(config: ReaperConfig) -> Self {
|
||||
Self {
|
||||
active_children: HashMap::new(),
|
||||
reaped_count: 0,
|
||||
zombie_detected_count: 0,
|
||||
config,
|
||||
}
|
||||
}
|
||||
|
||||
/// 注册一个子进程,稍后自动回收
|
||||
fn register_child(&mut self, child: Child) {
|
||||
let id = child.id().unwrap_or(0);
|
||||
if id > 0 {
|
||||
self.active_children.insert(id, child);
|
||||
if self.config.verbose {
|
||||
debug!("[ProcessReaper] Registered child process PID={}", id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 尝试回收所有已退出的子进程
|
||||
fn reap_all(&mut self) {
|
||||
let mut reaped_now = 0;
|
||||
|
||||
// 使用 entry API 避免 DashMap/RwLock 问题(虽然这里是普通 HashMap)
|
||||
self.active_children.retain(|pid, child| {
|
||||
// 尝试查询进程状态(非阻塞)
|
||||
match child.try_wait() {
|
||||
Ok(Some(status)) => {
|
||||
// 进程已退出
|
||||
reaped_now += 1;
|
||||
if self.config.verbose {
|
||||
debug!(
|
||||
"[ProcessReaper] Reaped child process PID={}, exit_status={:?}",
|
||||
pid, status
|
||||
);
|
||||
}
|
||||
false // 移除已回收的进程
|
||||
}
|
||||
Ok(None) => {
|
||||
// 进程仍在运行
|
||||
true
|
||||
}
|
||||
Err(e) => {
|
||||
// 查询失败,可能进程已不存在
|
||||
warn!("[ProcessReaper] Failed to query child PID={}: {}", pid, e);
|
||||
false // 移除无法查询的进程
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if reaped_now > 0 {
|
||||
self.reaped_count += reaped_now;
|
||||
info!(
|
||||
"[ProcessReaper] Reaped {} child processes (total: {})",
|
||||
reaped_now, self.reaped_count
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 🔍 主动检测系统中的僵尸进程
|
||||
///
|
||||
/// 扫描 /proc 文件系统,查找状态为 'Z' (Zombie) 的进程
|
||||
fn detect_zombie_processes(&mut self) -> Vec<ZombieProcessInfo> {
|
||||
let mut zombies = Vec::new();
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let proc_path = "/proc";
|
||||
|
||||
// 读取 /proc 目录下的所有 PID 目录
|
||||
if let Ok(entries) = fs::read_dir(proc_path) {
|
||||
for entry in entries.flatten() {
|
||||
let name = entry.file_name();
|
||||
// 检查是否是数字(PID 目录)
|
||||
if let Ok(pid) = name.to_string_lossy().parse::<u32>() {
|
||||
// 读取 /proc/[pid]/stat 文件
|
||||
let stat_path = entry.path().join("stat");
|
||||
if let Ok(content) = fs::read_to_string(&stat_path)
|
||||
&& let Some(info) = parse_stat_file(pid, &content)
|
||||
&& info.state == 'Z'
|
||||
{
|
||||
zombies.push(info);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !zombies.is_empty() {
|
||||
self.zombie_detected_count += zombies.len() as u64;
|
||||
warn!(
|
||||
"[ProcessReaper] Detected {} zombie processes (total detected: {})",
|
||||
zombies.len(),
|
||||
self.zombie_detected_count
|
||||
);
|
||||
|
||||
for zombie in &zombies {
|
||||
warn!(
|
||||
"[ProcessReaper] Zombie process: PID={}, PPID={}, CMD={}",
|
||||
zombie.pid, zombie.ppid, zombie.comm
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
zombies
|
||||
}
|
||||
|
||||
/// 🔧 主动清理所有僵尸进程
|
||||
///
|
||||
/// 使用 waitpid 循环回收所有可能的僵尸进程,不仅仅是追踪的子进程
|
||||
/// 这是 PID 1 的责任:回收所有孤儿进程
|
||||
fn reap_all_zombies_blocking(&mut self) {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use nix::sys::wait::{WaitPidFlag, WaitStatus, waitpid};
|
||||
use nix::unistd::Pid;
|
||||
|
||||
let mut reaped_this_round = 0;
|
||||
|
||||
// 循环调用 waitpid,直到没有更多僵尸进程
|
||||
loop {
|
||||
match waitpid(
|
||||
Pid::from_raw(-1), // -1 表示等待任意子进程
|
||||
Some(WaitPidFlag::WNOHANG),
|
||||
) {
|
||||
Ok(WaitStatus::Exited(pid, exit_code)) => {
|
||||
reaped_this_round += 1;
|
||||
debug!(
|
||||
"[ProcessReaper] Reaped zombie proactively: PID={}, exit_code={}",
|
||||
pid, exit_code
|
||||
);
|
||||
}
|
||||
Ok(WaitStatus::Signaled(pid, signal, _)) => {
|
||||
reaped_this_round += 1;
|
||||
debug!(
|
||||
"[ProcessReaper] Reaped zombie proactively: PID={}, signal={:?}",
|
||||
pid, signal
|
||||
);
|
||||
}
|
||||
Ok(WaitStatus::StillAlive) => {
|
||||
// WNOHANG: 没有更多的僵尸进程
|
||||
break;
|
||||
}
|
||||
Ok(WaitStatus::Stopped(pid, signal)) => {
|
||||
// 进程被停止(不是退出),不计入回收
|
||||
debug!(
|
||||
"[ProcessReaper] Process stopped: PID={}, signal={:?}",
|
||||
pid, signal
|
||||
);
|
||||
// 继续循环,可能还有其他僵尸进程
|
||||
continue;
|
||||
}
|
||||
Ok(WaitStatus::Continued(pid)) => {
|
||||
// 进程被恢复(SIGCONT),不计入回收
|
||||
debug!("[ProcessReaper] Process resumed: PID={}", pid);
|
||||
// 继续循环,可能还有其他僵尸进程
|
||||
continue;
|
||||
}
|
||||
#[cfg(linux_android)]
|
||||
Ok(WaitStatus::PtraceEvent(pid, signal, event)) => {
|
||||
// ptrace 事件,不计入回收
|
||||
debug!(
|
||||
"[ProcessReaper] ptrace event: PID={}, signal={:?}, event={}",
|
||||
pid, signal, event
|
||||
);
|
||||
continue;
|
||||
}
|
||||
#[cfg(linux_android)]
|
||||
Ok(WaitStatus::PtraceSyscall(pid)) => {
|
||||
// ptrace 系统调用,不计入回收
|
||||
debug!("[ProcessReaper] ptrace syscall: PID={}", pid);
|
||||
continue;
|
||||
}
|
||||
// 非Linux平台忽略 ptrace 相关状态(macOS 上 WaitStatus 包含这些变体但不会实际触发)
|
||||
Ok(_) => {
|
||||
debug!("[ProcessReaper] Ignored waitpid status");
|
||||
continue;
|
||||
}
|
||||
Err(nix::errno::Errno::ECHILD) => {
|
||||
// 没有子进程
|
||||
break;
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[ProcessReaper] waitpid error: {}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if reaped_this_round > 0 {
|
||||
self.reaped_count += reaped_this_round;
|
||||
info!(
|
||||
"[ProcessReaper] Reaped {} zombies proactively (total: {})",
|
||||
reaped_this_round, self.reaped_count
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
debug!("[ProcessReaper] Non-Unix platform, skipping zombie detection");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 解析 /proc/[pid]/stat 文件
|
||||
///
|
||||
/// 文件格式:pid (comm) state ppid ...
|
||||
/// 示例:1 (init) S 0 0 0 0 ...
|
||||
fn parse_stat_file(pid: u32, content: &str) -> Option<ZombieProcessInfo> {
|
||||
// stat 文件格式:pid (comm) state ppid ...
|
||||
// 需要找到 comm 的结束括号
|
||||
let content = content.trim();
|
||||
|
||||
// 找到第一个 '(' 和最后一个 ')'
|
||||
let open_paren = content.find('(')?;
|
||||
let close_paren = content.rfind(')')?;
|
||||
|
||||
let comm = content[open_paren + 1..close_paren].to_string();
|
||||
let after_comm = &content[close_paren + 1..];
|
||||
|
||||
// 解析 state 和 ppid
|
||||
// 格式:) state ppid ...
|
||||
let parts: Vec<&str> = after_comm.split_whitespace().collect();
|
||||
if parts.len() < 2 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let state = parts.first()?.chars().next()?;
|
||||
let ppid: u32 = parts.get(1)?.parse().ok()?;
|
||||
|
||||
Some(ZombieProcessInfo {
|
||||
pid,
|
||||
ppid,
|
||||
comm,
|
||||
state,
|
||||
})
|
||||
}
|
||||
|
||||
/// 启动进程回收器任务
|
||||
///
|
||||
/// 此函数会:
|
||||
/// 1. 注册 SIGCHLD 信号处理器
|
||||
/// 2. 在后台循环中等待信号并回收子进程
|
||||
/// 3. 定期主动检测和清理僵尸进程
|
||||
///
|
||||
/// # 返回值
|
||||
///
|
||||
/// 返回一个 JoinHandle,可以用于等待回收器任务退出(通常不需要)
|
||||
pub fn start_process_reaper() -> tokio::task::JoinHandle<()> {
|
||||
tokio::spawn(async move { run_reaper(ReaperConfig::default()).await })
|
||||
}
|
||||
|
||||
/// 启动进程回收器任务(带配置)
|
||||
pub fn start_process_reaper_with_config(config: ReaperConfig) -> tokio::task::JoinHandle<()> {
|
||||
tokio::spawn(async move { run_reaper(config).await })
|
||||
}
|
||||
|
||||
/// 核心回收逻辑
|
||||
#[cfg(unix)]
|
||||
async fn run_reaper(config: ReaperConfig) {
|
||||
info!("[ProcessReaper] Zombie process reaper started (PID 1 mode)");
|
||||
if config.enable_zombie_detection {
|
||||
info!(
|
||||
"[ProcessReaper] Zombie detection enabled, interval: {} seconds",
|
||||
config.zombie_detection_interval_secs
|
||||
);
|
||||
}
|
||||
|
||||
// 创建 SIGCHLD 信号监听器
|
||||
let sigchld = match signal(SignalKind::child()) {
|
||||
Ok(sig) => sig,
|
||||
Err(e) => {
|
||||
error!("[ProcessReaper] Failed to register SIGCHLD handler: {}", e);
|
||||
error!("[ProcessReaper] Falling back to polling mode");
|
||||
|
||||
// 回退模式:使用轮询
|
||||
run_reaper_polling(config).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let state = ReaperState::new(config.clone());
|
||||
|
||||
// 启动定期轮询任务(作为信号机制的补充)
|
||||
let mut poll_interval = tokio::time::interval(std::time::Duration::from_secs(5));
|
||||
poll_interval.tick().await; // 跳过第一次立即触发
|
||||
|
||||
// 根据配置决定是否启用僵尸进程检测
|
||||
if config.enable_zombie_detection {
|
||||
run_reaper_with_detection(
|
||||
sigchld,
|
||||
poll_interval,
|
||||
state,
|
||||
config.zombie_detection_interval_secs,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
run_reaper_without_detection(sigchld, poll_interval, state).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Windows 上的回收逻辑(无操作,Windows 没有僵尸进程问题)
|
||||
#[cfg(not(unix))]
|
||||
async fn run_reaper(_config: ReaperConfig) {
|
||||
info!("[ProcessReaper] Non-Unix platform: zombie reaper not applicable");
|
||||
}
|
||||
|
||||
/// 🔍 启用僵尸进程检测的回收循环
|
||||
#[cfg(unix)]
|
||||
async fn run_reaper_with_detection(
|
||||
mut sigchld: tokio::signal::unix::Signal,
|
||||
mut poll_interval: tokio::time::Interval,
|
||||
mut state: ReaperState,
|
||||
detect_interval_secs: u64,
|
||||
) {
|
||||
let mut zombie_detect_interval =
|
||||
tokio::time::interval(std::time::Duration::from_secs(detect_interval_secs));
|
||||
zombie_detect_interval.tick().await; // 跳过第一次立即触发
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// 等待 SIGCHLD 信号
|
||||
_ = sigchld.recv() => {
|
||||
if state.config.verbose {
|
||||
debug!("[ProcessReaper] Received SIGCHLD");
|
||||
}
|
||||
state.reap_all();
|
||||
state.reap_all_zombies_blocking();
|
||||
}
|
||||
// 定期轮询(每 5 秒)
|
||||
_ = poll_interval.tick() => {
|
||||
state.reap_all();
|
||||
}
|
||||
// 🔍 定期主动检测和清理僵尸进程
|
||||
_ = zombie_detect_interval.tick() => {
|
||||
debug!("[ProcessReaper] Running scheduled zombie detection...");
|
||||
|
||||
// 先检测有哪些僵尸进程
|
||||
let zombies = state.detect_zombie_processes();
|
||||
|
||||
// 然后主动清理所有僵尸进程
|
||||
if !zombies.is_empty() {
|
||||
state.reap_all_zombies_blocking();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 🚫 不启用僵尸进程检测的回收循环
|
||||
#[cfg(unix)]
|
||||
async fn run_reaper_without_detection(
|
||||
mut sigchld: tokio::signal::unix::Signal,
|
||||
mut poll_interval: tokio::time::Interval,
|
||||
mut state: ReaperState,
|
||||
) {
|
||||
loop {
|
||||
tokio::select! {
|
||||
// 等待 SIGCHLD 信号
|
||||
_ = sigchld.recv() => {
|
||||
if state.config.verbose {
|
||||
debug!("[ProcessReaper] Received SIGCHLD");
|
||||
}
|
||||
state.reap_all();
|
||||
state.reap_all_zombies_blocking();
|
||||
}
|
||||
// 定期轮询(每 5 秒)
|
||||
_ = poll_interval.tick() => {
|
||||
state.reap_all();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 轮询模式回退(当信号机制不可用时)
|
||||
#[cfg(unix)]
|
||||
async fn run_reaper_polling(config: ReaperConfig) {
|
||||
info!("[ProcessReaper] Using polling mode for zombie reaping");
|
||||
|
||||
let state = ReaperState::new(config.clone());
|
||||
|
||||
// 根据配置决定是否启用僵尸进程检测
|
||||
if config.enable_zombie_detection {
|
||||
run_reaper_polling_with_detection(state, config.zombie_detection_interval_secs).await
|
||||
} else {
|
||||
run_reaper_polling_without_detection(state).await
|
||||
}
|
||||
}
|
||||
|
||||
/// 🔍 轮询模式 + 僵尸进程检测
|
||||
#[cfg(unix)]
|
||||
async fn run_reaper_polling_with_detection(mut state: ReaperState, detect_interval_secs: u64) {
|
||||
let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
|
||||
let mut zombie_detect_interval =
|
||||
tokio::time::interval(std::time::Duration::from_secs(detect_interval_secs));
|
||||
zombie_detect_interval.tick().await;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = interval.tick() => {
|
||||
state.reap_all();
|
||||
state.reap_all_zombies_blocking();
|
||||
}
|
||||
_ = zombie_detect_interval.tick() => {
|
||||
let zombies = state.detect_zombie_processes();
|
||||
if !zombies.is_empty() {
|
||||
state.reap_all_zombies_blocking();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 🚫 轮询模式(无僵尸进程检测)
|
||||
#[cfg(unix)]
|
||||
async fn run_reaper_polling_without_detection(mut state: ReaperState) {
|
||||
let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
|
||||
|
||||
loop {
|
||||
interval.tick().await;
|
||||
state.reap_all();
|
||||
state.reap_all_zombies_blocking();
|
||||
}
|
||||
}
|
||||
|
||||
/// 回收器句柄(可选:用于外部注册子进程)
|
||||
///
|
||||
/// 注意:当前实现中,子进程由各自创建者管理。
|
||||
/// 此结构保留用于未来扩展,例如中央化子进程管理。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProcessReaperHandle {
|
||||
_config: ReaperConfig,
|
||||
}
|
||||
|
||||
impl ProcessReaperHandle {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
_config: ReaperConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 🔍 手动触发僵尸进程检测(仅用于调试)
|
||||
pub fn detect_zombies_now(&self) -> Vec<ZombieProcessInfo> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let proc_path = "/proc";
|
||||
let mut zombies = Vec::new();
|
||||
|
||||
if let Ok(entries) = fs::read_dir(proc_path) {
|
||||
for entry in entries.flatten() {
|
||||
let name = entry.file_name();
|
||||
if let Ok(pid) = name.to_string_lossy().parse::<u32>() {
|
||||
let stat_path = entry.path().join("stat");
|
||||
if let Ok(content) = fs::read_to_string(&stat_path)
|
||||
&& let Some(info) = parse_stat_file(pid, &content)
|
||||
&& info.state == 'Z'
|
||||
{
|
||||
zombies.push(info);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
zombies
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// 注册一个子进程(未来扩展)
|
||||
#[allow(dead_code)]
|
||||
pub fn register(&self, _child: Child) {
|
||||
// 当前实现中,子进程由各自的创建者负责回收
|
||||
// 此方法保留用于未来中央化管理的扩展
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ProcessReaperHandle {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, unix))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::process::Stdio;
|
||||
use tokio::time::Duration;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_reaper_state() {
|
||||
let config = ReaperConfig {
|
||||
verbose: true,
|
||||
..Default::default()
|
||||
};
|
||||
let mut state = ReaperState::new(config);
|
||||
|
||||
// 创建一个长时间运行的进程
|
||||
let child = tokio::process::Command::new("sleep")
|
||||
.arg("10")
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.spawn()
|
||||
.unwrap();
|
||||
|
||||
let pid = child.id().unwrap();
|
||||
state.register_child(child);
|
||||
|
||||
// 立即调用 reap_all,进程应该还在运行
|
||||
state.reap_all();
|
||||
|
||||
// 验证进程仍在列表中(因为还没退出)
|
||||
assert!(state.active_children.contains_key(&pid));
|
||||
assert_eq!(state.reaped_count, 0); // 还没有回收任何进程
|
||||
|
||||
// 清理:杀死进程
|
||||
if let Some(mut child) = state.active_children.remove(&pid) {
|
||||
let _ = child.kill().await;
|
||||
let _ = child.wait().await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_long_running_process() {
|
||||
let config = ReaperConfig::default();
|
||||
let mut state = ReaperState::new(config);
|
||||
|
||||
// 创建一个长时间运行的进程
|
||||
let child = tokio::process::Command::new("sleep")
|
||||
.arg("10")
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.spawn()
|
||||
.unwrap();
|
||||
|
||||
let pid = child.id().unwrap();
|
||||
state.register_child(child);
|
||||
|
||||
// 立即尝试回收,进程应该还在运行
|
||||
state.reap_all();
|
||||
|
||||
// 验证进程仍在列表中
|
||||
assert!(state.active_children.contains_key(&pid));
|
||||
|
||||
// 清理:杀死进程
|
||||
if let Some(mut child) = state.active_children.remove(&pid) {
|
||||
let _ = child.kill().await;
|
||||
let _ = child.wait().await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_start_process_reaper() {
|
||||
let handle = start_process_reaper();
|
||||
|
||||
// 创建几个快速退出的子进程
|
||||
for _ in 0..3 {
|
||||
let _ = tokio::process::Command::new("true")
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.spawn();
|
||||
}
|
||||
|
||||
// 等待回收器处理
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
|
||||
// 回收器应该仍在运行
|
||||
assert!(!handle.is_finished());
|
||||
|
||||
// 取消任务
|
||||
handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_stat_file() {
|
||||
let content = "1 (init) S 0 0 0 0 -1 4194560 667 5569406 8 23660837 1 0 0 0 0 0 0 0 20 0 1 0 3642608 1340 18446744073709551615 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0";
|
||||
let info = parse_stat_file(1, content).unwrap();
|
||||
|
||||
assert_eq!(info.pid, 1);
|
||||
assert_eq!(info.ppid, 0);
|
||||
assert_eq!(info.comm, "init");
|
||||
assert_eq!(info.state, 'S');
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_stat_file_with_parentheses_in_comm() {
|
||||
// 进程名包含括号的情况
|
||||
let content = "1234 (test(a)b)) Z 1 1234 1234 0 -1 4194560 0 0 0 0 0 0 0 0 20 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0";
|
||||
let info = parse_stat_file(1234, content).unwrap();
|
||||
|
||||
assert_eq!(info.pid, 1234);
|
||||
assert_eq!(info.ppid, 1);
|
||||
assert_eq!(info.comm, "test(a)b)");
|
||||
assert_eq!(info.state, 'Z'); // 僵尸进程
|
||||
}
|
||||
}
|
||||
111
qiming-rcoder/crates/agent_runner/src/profiler.rs
Normal file
111
qiming-rcoder/crates/agent_runner/src/profiler.rs
Normal file
@@ -0,0 +1,111 @@
|
||||
//! Pyroscope Profiler 集成模块
|
||||
//!
|
||||
//! 提供连续性能分析功能,包括:
|
||||
//! - CPU Profiling
|
||||
//! - Memory Profiling (Heap/Allocs)
|
||||
//! - Stack trace sampling
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use pyroscope::PyroscopeAgent;
|
||||
use pyroscope_pprofrs::{PprofConfig, pprof_backend};
|
||||
use tracing::{debug, info};
|
||||
|
||||
/// Profiler 配置
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProfilerConfig {
|
||||
/// Pyroscope Server 地址
|
||||
pub server_url: String,
|
||||
/// 应用名称(支持标签格式)
|
||||
pub application_name: String,
|
||||
/// 采样频率(Hz)
|
||||
pub sample_rate: u32,
|
||||
/// 是否启用内存 profiling
|
||||
pub enable_memory: bool,
|
||||
/// 是否启用 CPU profiling
|
||||
pub enable_cpu: bool,
|
||||
}
|
||||
|
||||
impl Default for ProfilerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
server_url: "http://pyroscope:4040".to_string(),
|
||||
application_name: "agent_runner{env=dev}".to_string(),
|
||||
sample_rate: 100,
|
||||
enable_memory: true,
|
||||
enable_cpu: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProfilerConfig {
|
||||
/// 从环境变量加载配置
|
||||
pub fn from_env() -> Self {
|
||||
let mut config = Self::default();
|
||||
|
||||
if let Ok(url) = std::env::var("PYROSCOPE_URL") {
|
||||
config.server_url = url;
|
||||
}
|
||||
|
||||
if let Ok(name) = std::env::var("PYROSCOPE_APP_NAME") {
|
||||
config.application_name = name;
|
||||
}
|
||||
|
||||
if let Ok(project_id) = std::env::var("PROJECT_ID") {
|
||||
config.application_name = format!("agent_runner{{project_id={}}}", project_id);
|
||||
}
|
||||
|
||||
config
|
||||
}
|
||||
}
|
||||
|
||||
/// Profiler Guard
|
||||
///
|
||||
/// 当 dropped 时自动停止 profiler
|
||||
pub struct ProfilerGuard {
|
||||
_agent: Option<pyroscope::PyroscopeAgent<pyroscope::pyroscope::PyroscopeAgentRunning>>,
|
||||
}
|
||||
|
||||
/// 初始化并启动 Pyroscope Profiler
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 如果 profiler 初始化失败,返回错误
|
||||
pub fn init_pyroscope_profiler(config: ProfilerConfig) -> Result<ProfilerGuard> {
|
||||
if !config.enable_cpu && !config.enable_memory {
|
||||
debug!("Pyroscope profiler disabled, skipping initialization");
|
||||
return Ok(ProfilerGuard { _agent: None });
|
||||
}
|
||||
|
||||
info!(
|
||||
"Initializing Pyroscope profiler: {}",
|
||||
config.application_name
|
||||
);
|
||||
info!(" Server URL: {}", config.server_url);
|
||||
info!(" Sample rate: {} Hz", config.sample_rate);
|
||||
info!(" CPU profiling: {}", config.enable_cpu);
|
||||
info!(" Memory profiling: {}", config.enable_memory);
|
||||
|
||||
// 使用正确的 API: builder() -> backend() -> build()
|
||||
let agent = PyroscopeAgent::builder(config.server_url, config.application_name)
|
||||
.backend(pprof_backend(
|
||||
PprofConfig::new().sample_rate(config.sample_rate),
|
||||
))
|
||||
.build()
|
||||
.context("Failed to build Pyroscope agent")?;
|
||||
|
||||
// 启动 profiling,返回 PyroscopeAgent<PyroscopeAgentRunning>
|
||||
let agent_running = agent
|
||||
.start()
|
||||
.context("Failed to start Pyroscope profiler")?;
|
||||
|
||||
info!("Pyroscope profiler started successfully");
|
||||
|
||||
Ok(ProfilerGuard {
|
||||
_agent: Some(agent_running),
|
||||
})
|
||||
}
|
||||
|
||||
/// 便捷函数:使用默认配置初始化 profiler
|
||||
pub fn init_pyroscope_profiler_default() -> Result<ProfilerGuard> {
|
||||
init_pyroscope_profiler(ProfilerConfig::from_env())
|
||||
}
|
||||
618
qiming-rcoder/crates/agent_runner/src/proxy_agent/acp_agent.rs
Normal file
618
qiming-rcoder/crates/agent_runner/src/proxy_agent/acp_agent.rs
Normal file
@@ -0,0 +1,618 @@
|
||||
//! ACP Agent Worker 模块 (SACP 版本)
|
||||
//!
|
||||
//! 负责处理 Agent 请求队列,管理 Agent 会话的创建和复用。
|
||||
//! 使用 AcpSessionManager 进行会话管理。
|
||||
//!
|
||||
//! ## SACP 迁移说明
|
||||
//!
|
||||
//! - 移除了 `spawn_blocking` + `LocalSet` 模式(SACP 支持 Send trait)
|
||||
//! - 使用标准 `tokio::spawn` 进行并发处理
|
||||
//! - 简化了并发模型,提高性能
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use agent_abstraction::session::AcpSessionManager;
|
||||
use anyhow::Result;
|
||||
use chrono::Utc;
|
||||
use shared_types::ModelProviderConfig;
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
use agent_client_protocol::schema::SessionId;
|
||||
|
||||
use crate::{
|
||||
agent_runtime::get_concurrency_limit,
|
||||
model::{AgentStatus, ChatPromptResponse, ProjectAndAgentInfo},
|
||||
proxy_agent::SESSION_REQUEST_CONTEXT,
|
||||
service::{AGENT_REGISTRY, AgentSessionRegistry, StateAwareNotifier},
|
||||
utils::ContentBuilder,
|
||||
};
|
||||
|
||||
/// 🔥 配置标志:是否为无限制模式(HTTP Server 部署)
|
||||
static IS_UNLIMITED_MODE: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
/// 设置运行模式(供 main.rs 调用)
|
||||
pub fn set_unlimited_mode(enabled: bool) {
|
||||
IS_UNLIMITED_MODE.store(enabled, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
// 🔥 OpenTelemetry 追踪
|
||||
#[cfg(feature = "otel")]
|
||||
use crate::otel_tracing::RequestSpan;
|
||||
|
||||
// 🔥 简化版本:如果没有 OpenTelemetry,使用空的 span
|
||||
#[cfg(not(feature = "otel"))]
|
||||
struct RequestSpan;
|
||||
|
||||
#[cfg(not(feature = "otel"))]
|
||||
impl RequestSpan {
|
||||
fn new(_project_id: &str, _request_id: &str, _operation: &str) -> Self {
|
||||
Self
|
||||
}
|
||||
}
|
||||
|
||||
/// Agent 请求结构 (SACP 版本)
|
||||
///
|
||||
/// 不再需要 LocalSet,直接在 tokio::spawn 中运行
|
||||
#[derive(Debug)]
|
||||
pub struct AgentRequest {
|
||||
/// Agent 抽象层的 prompt 消息
|
||||
prompt_message: agent_abstraction::PromptMessage,
|
||||
/// 发送回执消息的通道
|
||||
chat_prompt_tx: oneshot::Sender<ChatPromptResponse>,
|
||||
/// 模型提供商配置
|
||||
model_provider: Option<ModelProviderConfig>,
|
||||
/// 🔥 关联的 service UUID(用于 API 密钥管理)
|
||||
service_uuid: Option<String>,
|
||||
/// 🔥 共享的 API 密钥管理器(用于自动清理)
|
||||
shared_api_key_manager: Option<Arc<DashMap<String, ModelProviderConfig>>>,
|
||||
/// 是否跳过槽位限制(HTTP Server 宿主机部署时为 true)
|
||||
skip_slot_limit: bool,
|
||||
}
|
||||
|
||||
impl AgentRequest {
|
||||
pub fn new(
|
||||
prompt_message: agent_abstraction::PromptMessage,
|
||||
model_provider: Option<ModelProviderConfig>,
|
||||
) -> (Self, oneshot::Receiver<ChatPromptResponse>) {
|
||||
let (chat_prompt_tx, chat_prompt_rx) = oneshot::channel();
|
||||
(
|
||||
Self {
|
||||
prompt_message,
|
||||
chat_prompt_tx,
|
||||
model_provider,
|
||||
service_uuid: None,
|
||||
shared_api_key_manager: None,
|
||||
skip_slot_limit: false,
|
||||
},
|
||||
chat_prompt_rx,
|
||||
)
|
||||
}
|
||||
|
||||
/// 设置 service_uuid
|
||||
pub fn with_service_uuid(mut self, service_uuid: Option<String>) -> Self {
|
||||
self.service_uuid = service_uuid;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 shared_api_key_manager
|
||||
pub fn with_key_manager(
|
||||
mut self,
|
||||
key_manager: Option<Arc<DashMap<String, ModelProviderConfig>>>,
|
||||
) -> Self {
|
||||
self.shared_api_key_manager = key_manager;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置是否跳过槽位限制
|
||||
pub fn with_skip_slot_limit(mut self, skip: bool) -> Self {
|
||||
self.skip_slot_limit = skip;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Agent Worker 任务
|
||||
///
|
||||
/// 使用标准 tokio::spawn 处理 Agent 请求队列。
|
||||
/// SACP 支持 Send trait,无需 LocalSet。
|
||||
pub async fn agent_worker(mut request_rx: mpsc::UnboundedReceiver<AgentRequest>) -> Result<()> {
|
||||
use agent_abstraction::session::{AcpAgentWorker, AgentWorker, WorkerRequest};
|
||||
|
||||
info!("agent_worker started (SACP version), listening for requests...");
|
||||
|
||||
// 创建 AcpSessionManager,注入 AGENT_REGISTRY 作为 SessionRegistry
|
||||
// SACP 版本只需要 2 个泛型参数:N (SessionNotifier) 和 R (SessionRegistry)
|
||||
let session_manager = Arc::new(
|
||||
AcpSessionManager::<StateAwareNotifier, AgentSessionRegistry>::new(
|
||||
Arc::new(StateAwareNotifier::new()),
|
||||
AGENT_REGISTRY.clone(),
|
||||
),
|
||||
);
|
||||
|
||||
// 创建 AcpAgentWorker
|
||||
let worker = AcpAgentWorker::new(session_manager);
|
||||
|
||||
while let Some(request) = request_rx.recv().await {
|
||||
let project_id = request.prompt_message.project_id.clone();
|
||||
let request_id = request.prompt_message.request_id.clone();
|
||||
|
||||
info!(
|
||||
"📨 Received request, project_id: {}, request_id: {}",
|
||||
project_id, request_id
|
||||
);
|
||||
|
||||
// 1. 预处理附件(agent_runner 特有逻辑)
|
||||
let attachment_blocks = if !request.prompt_message.attachments.is_empty() {
|
||||
match ContentBuilder::attachments_to_content_blocks(
|
||||
&request.prompt_message.attachments,
|
||||
&request.prompt_message.project_path,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(blocks) => Some(blocks),
|
||||
Err(e) => {
|
||||
error!("Attachment processing failed: {:?}", e);
|
||||
|
||||
if let Err(send_err) = request.chat_prompt_tx.send(ChatPromptResponse {
|
||||
project_id: project_id.clone(),
|
||||
session_id: String::new(),
|
||||
code: shared_types::error_codes::ERR_AGENT_ERROR.to_string(),
|
||||
error: Some(format!(
|
||||
"{}: {:?}",
|
||||
shared_types::error_codes::get_i18n_message_default("error.attachment_processing_failed"),
|
||||
e
|
||||
)),
|
||||
request_id: Some(request_id),
|
||||
service_type: request.prompt_message.service_type.clone(),
|
||||
}) {
|
||||
error!(
|
||||
"Failed to send error response (receiver closed): {:?}",
|
||||
send_err
|
||||
);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// 2. 创建 WorkerRequest
|
||||
let worker_request = WorkerRequest {
|
||||
prompt_message: request.prompt_message.clone(),
|
||||
model_provider: request.model_provider.clone(),
|
||||
attachment_blocks,
|
||||
service_uuid: request.service_uuid.clone(),
|
||||
shared_api_key_manager: request.shared_api_key_manager.clone(),
|
||||
};
|
||||
|
||||
// 3. 调用 AcpAgentWorker 处理(核心业务逻辑)
|
||||
let worker_response = match worker.process_request(worker_request).await {
|
||||
Ok(response) => response,
|
||||
Err(e) => {
|
||||
error!("Worker processing failed: {:?}", e);
|
||||
|
||||
if let Err(send_err) = request.chat_prompt_tx.send(ChatPromptResponse {
|
||||
project_id: project_id.clone(),
|
||||
session_id: String::new(),
|
||||
code: shared_types::error_codes::ERR_AGENT_ERROR.to_string(),
|
||||
error: Some(format!(
|
||||
"{}: {:?}",
|
||||
shared_types::error_codes::get_i18n_message_default("error.processing_failed"),
|
||||
e
|
||||
)),
|
||||
request_id: Some(request_id.clone()),
|
||||
service_type: request.prompt_message.service_type.clone(),
|
||||
}) {
|
||||
error!(
|
||||
"Failed to send error response (receiver closed): {:?}",
|
||||
send_err
|
||||
);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// 4. 更新全局状态(使用统一的 AGENT_REGISTRY)
|
||||
if worker_response.is_new_session {
|
||||
if let Some(handles) = &worker_response.session_handles {
|
||||
debug!("🆕 New session, registering in AGENT_REGISTRY");
|
||||
|
||||
let project_and_agent_info = ProjectAndAgentInfo {
|
||||
project_id: project_id.clone(),
|
||||
session_id: SessionId::new(Arc::from(worker_response.session_id.as_str())),
|
||||
prompt_tx: handles.prompt_tx.clone(),
|
||||
cancel_tx: handles.cancel_tx.clone(),
|
||||
model_provider: request.model_provider.clone(),
|
||||
request_id: Some(request_id.clone()),
|
||||
status: AgentStatus::Active, // 🆕 修复:Worker 处理中应为 Active,而非 Idle
|
||||
last_activity: Utc::now(),
|
||||
created_at: Utc::now(),
|
||||
stop_handle: handles.lifecycle_handle.clone(),
|
||||
};
|
||||
|
||||
// 使用统一的 AGENT_REGISTRY 注册(自动处理所有映射)
|
||||
AGENT_REGISTRY.register(
|
||||
&project_id,
|
||||
&worker_response.session_id,
|
||||
project_and_agent_info,
|
||||
);
|
||||
|
||||
info!(
|
||||
"🔗 Agent 已注册到 AGENT_REGISTRY: project_id={}, session_id={}",
|
||||
project_id, worker_response.session_id
|
||||
);
|
||||
}
|
||||
} else {
|
||||
debug!("♻️ Reusing session, no global Registry update needed");
|
||||
}
|
||||
|
||||
// 5. 更新 SESSION_REQUEST_CONTEXT(请求追踪)
|
||||
SESSION_REQUEST_CONTEXT.insert(project_id, request_id.clone());
|
||||
|
||||
// 6. 转换并发送回执
|
||||
let chat_prompt_response = ChatPromptResponse {
|
||||
project_id: worker_response.project_id,
|
||||
session_id: worker_response.session_id,
|
||||
code: if worker_response.error.is_none() {
|
||||
shared_types::error_codes::SUCCESS.to_string()
|
||||
} else {
|
||||
shared_types::error_codes::ERR_AGENT_ERROR.to_string()
|
||||
},
|
||||
error: worker_response.error,
|
||||
request_id: worker_response.request_id,
|
||||
service_type: worker_response.service_type,
|
||||
};
|
||||
|
||||
if let Err(e) = request.chat_prompt_tx.send(chat_prompt_response) {
|
||||
error!("Failed to send acknowledgment: {:?}", e);
|
||||
}
|
||||
}
|
||||
|
||||
info!("🛑 agent_worker stopped");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 带心跳的 Agent Worker (SACP 版本) - 新架构
|
||||
///
|
||||
/// 使用标准 tokio::spawn 进行并发处理(无需 LocalSet)
|
||||
/// SACP 的 Component<L> trait 要求 Send + 'static,因此可以安全地在多线程环境中使用
|
||||
///
|
||||
/// ## 参数变化
|
||||
///
|
||||
/// - `request_rx`: 使用有界 channel 代替 unbounded
|
||||
/// - `state`: 使用 `Arc<AtomicState>` 代替 `WorkerHandle`
|
||||
/// - `last_heartbeat_ts`: 🔥 P1 修复: 使用 `Arc<AtomicI64>` 代替 `Arc<Mutex<Option<DateTime>>>`
|
||||
/// - `active_requests`: 直接访问,用于请求追踪
|
||||
pub async fn agent_worker_with_heartbeat(
|
||||
mut request_rx: mpsc::Receiver<AgentRequest>,
|
||||
state: Arc<crate::agent_runtime::AtomicState>,
|
||||
last_heartbeat_ts: Arc<std::sync::atomic::AtomicI64>,
|
||||
active_requests: Arc<tokio::sync::Mutex<HashMap<String, chrono::DateTime<chrono::Utc>>>>,
|
||||
) -> Result<()> {
|
||||
info!("agent_worker started (SACP version with heartbeat), listening for requests...");
|
||||
|
||||
use agent_abstraction::session::{AcpAgentWorker, AgentWorker, WorkerRequest};
|
||||
use tokio::time::{Duration, interval};
|
||||
|
||||
// 创建 AcpSessionManager,注入 AGENT_REGISTRY 作为 SessionRegistry
|
||||
// SACP 版本只需要 2 个泛型参数:N (SessionNotifier) 和 R (SessionRegistry)
|
||||
let session_manager = Arc::new(
|
||||
AcpSessionManager::<StateAwareNotifier, AgentSessionRegistry>::new(
|
||||
Arc::new(StateAwareNotifier::new()),
|
||||
AGENT_REGISTRY.clone(),
|
||||
),
|
||||
);
|
||||
|
||||
// 创建 AcpAgentWorker
|
||||
let worker = AcpAgentWorker::new(session_manager);
|
||||
|
||||
// 设置状态为 Running(就绪信号)
|
||||
state.set(crate::agent_runtime::WorkerState::Running);
|
||||
info!("[Worker] SACP Worker initialized, state set to Running");
|
||||
|
||||
// 启动心跳任务 - 🔥 P1 修复: 使用原子操作直接更新 last_heartbeat_ts
|
||||
let last_heartbeat_ts_clone = last_heartbeat_ts.clone();
|
||||
let heartbeat_task = tokio::spawn(async move {
|
||||
let mut heartbeat_interval = interval(Duration::from_secs(5));
|
||||
loop {
|
||||
heartbeat_interval.tick().await;
|
||||
let timestamp = Utc::now();
|
||||
|
||||
// 📊 打印当前 Worker 占用情况
|
||||
if IS_UNLIMITED_MODE.load(Ordering::SeqCst) {
|
||||
// 无限制模式(HTTP Server 部署)- 不显示具体数量,避免误解
|
||||
info!("💓 [Worker] Heartbeat - active sessions: (unlimited)");
|
||||
} else {
|
||||
// 限制模式(Docker 容器部署)
|
||||
let active = AGENT_REGISTRY.stats().agent_count;
|
||||
let limit = get_concurrency_limit();
|
||||
info!(
|
||||
"💓 [Worker] Heartbeat - active sessions: {}/{}",
|
||||
active, limit
|
||||
);
|
||||
}
|
||||
|
||||
// 🔥 P1 修复: 使用原子操作直接更新时间戳(无锁)
|
||||
last_heartbeat_ts_clone.store(
|
||||
timestamp.timestamp_millis(),
|
||||
std::sync::atomic::Ordering::Release,
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
// 主处理循环 - SACP 版本:使用标准 tokio::spawn 进行并发处理
|
||||
while let Some(request) = request_rx.recv().await {
|
||||
let project_id = request.prompt_message.project_id.clone();
|
||||
let request_id = request.prompt_message.request_id.clone();
|
||||
|
||||
info!(
|
||||
"📨 Received request, project_id: {}, request_id: {} - SACP concurrent processing",
|
||||
project_id, request_id
|
||||
);
|
||||
|
||||
// 克隆需要的变量,用于 spawn 任务
|
||||
let worker_clone = worker.clone();
|
||||
|
||||
// 🚀 SACP 版本:直接使用 tokio::spawn(无需 spawn_blocking + LocalSet)
|
||||
// SACP 的 Component<L> trait 要求 Send + 'static,因此可以安全地在多线程环境中使用
|
||||
tokio::spawn(async move {
|
||||
info!(
|
||||
"🔵 [SACP] 开始处理请求 project_id={}, request_id={}",
|
||||
project_id, request_id
|
||||
);
|
||||
|
||||
// 🔥 OpenTelemetry 追踪: 创建请求 span
|
||||
let _otel_span = RequestSpan::new(&project_id, &request_id, "process_agent_request");
|
||||
|
||||
// 1. 预处理附件(agent_runner 特有逻辑)
|
||||
let attachment_blocks = if !request.prompt_message.attachments.is_empty() {
|
||||
match ContentBuilder::attachments_to_content_blocks(
|
||||
&request.prompt_message.attachments,
|
||||
&request.prompt_message.project_path,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(blocks) => Some(blocks),
|
||||
Err(e) => {
|
||||
error!("Attachment processing failed: {:?}", e);
|
||||
|
||||
// 🔥 DeferGuard 自动清理,无需手动调用 clear_pending_if_exists
|
||||
|
||||
if let Err(send_err) = request.chat_prompt_tx.send(ChatPromptResponse {
|
||||
project_id: project_id.clone(),
|
||||
session_id: String::new(),
|
||||
code: shared_types::error_codes::ERR_AGENT_ERROR.to_string(),
|
||||
error: Some(format!(
|
||||
"{}: {:?}",
|
||||
shared_types::error_codes::get_i18n_message_default("error.attachment_processing_failed"),
|
||||
e
|
||||
)),
|
||||
request_id: Some(request_id.clone()),
|
||||
service_type: request.prompt_message.service_type.clone(),
|
||||
}) {
|
||||
error!(
|
||||
"Failed to send error response (receiver closed): {:?}",
|
||||
send_err
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// 2. 创建 WorkerRequest
|
||||
let worker_request = WorkerRequest {
|
||||
prompt_message: request.prompt_message.clone(),
|
||||
model_provider: request.model_provider.clone(),
|
||||
attachment_blocks,
|
||||
service_uuid: request.service_uuid.clone(),
|
||||
shared_api_key_manager: request.shared_api_key_manager.clone(),
|
||||
};
|
||||
|
||||
// 3. 调用 AcpAgentWorker 处理(核心业务逻辑)
|
||||
let worker_response = match worker_clone.process_request(worker_request).await {
|
||||
Ok(response) => response,
|
||||
Err(e) => {
|
||||
error!("Worker processing failed: {:?}", e);
|
||||
|
||||
// 🔥 DeferGuard 自动清理,无需手动调用 clear_pending_if_exists
|
||||
|
||||
if let Err(send_err) = request.chat_prompt_tx.send(ChatPromptResponse {
|
||||
project_id: project_id.clone(),
|
||||
session_id: String::new(),
|
||||
code: shared_types::error_codes::ERR_AGENT_ERROR.to_string(),
|
||||
error: Some(format!(
|
||||
"{}: {:?}",
|
||||
shared_types::error_codes::get_i18n_message_default("error.processing_failed"),
|
||||
e
|
||||
)),
|
||||
request_id: Some(request_id.clone()),
|
||||
service_type: request.prompt_message.service_type.clone(),
|
||||
}) {
|
||||
error!(
|
||||
"Failed to send error response (receiver closed): {:?}",
|
||||
send_err
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// 4. 提取 session_handles(在移动 worker_response 之前)
|
||||
// 关键:需要保存 lifecycle_handle 用于后续等待会话结束
|
||||
let session_handles = worker_response.session_handles.clone();
|
||||
let is_new_session = worker_response.is_new_session;
|
||||
let response_session_id = worker_response.session_id.clone();
|
||||
|
||||
// 5. 更新全局状态(使用统一的 AGENT_REGISTRY)
|
||||
if is_new_session {
|
||||
// 🔥 修复:槽位对应 Agent 生命周期,只在创建新 Agent 时获取槽位
|
||||
// HTTP Server 部署模式跳过槽位限制
|
||||
if request.skip_slot_limit {
|
||||
info!(
|
||||
"⏭️ [原子槽位] 跳过限制(无限制模式): project_id={}",
|
||||
project_id
|
||||
);
|
||||
} else if !AGENT_REGISTRY.try_acquire_session_slot() {
|
||||
let limit = get_concurrency_limit();
|
||||
error!(
|
||||
"🛡️ [原子并发限制] Agent 会话槽位已满 ({}/{}), 拒绝新请求 - project_id={}, request_id={}",
|
||||
AGENT_REGISTRY.active_sessions_count(),
|
||||
limit,
|
||||
project_id,
|
||||
request_id
|
||||
);
|
||||
|
||||
// 清理 Pending 状态(如果已设置)
|
||||
AGENT_REGISTRY.clear_pending_if_exists(&project_id);
|
||||
|
||||
if let Err(send_err) = request.chat_prompt_tx.send(ChatPromptResponse {
|
||||
project_id: project_id.clone(),
|
||||
session_id: String::new(),
|
||||
code: shared_types::error_codes::ERR_TOO_MANY_REQUESTS.to_string(),
|
||||
error: Some(format!(
|
||||
"{}",
|
||||
shared_types::error_codes::get_i18n_message_default("error.system_busy")
|
||||
).replace("{}", &limit.to_string())),
|
||||
request_id: Some(request_id.clone()),
|
||||
service_type: request.prompt_message.service_type.clone(),
|
||||
}) {
|
||||
error!(
|
||||
"Failed to send reject response (receiver closed): {:?}",
|
||||
send_err
|
||||
);
|
||||
}
|
||||
return;
|
||||
} else {
|
||||
info!(
|
||||
"✅ [原子槽位] 成功获取槽位: {}/{} - project_id={}",
|
||||
AGENT_REGISTRY.active_sessions_count(),
|
||||
get_concurrency_limit(),
|
||||
project_id
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(ref handles) = session_handles {
|
||||
debug!("🆕 New session, registering in AGENT_REGISTRY");
|
||||
|
||||
let project_and_agent_info = ProjectAndAgentInfo {
|
||||
project_id: project_id.clone(),
|
||||
session_id: SessionId::new(Arc::from(response_session_id.as_str())),
|
||||
prompt_tx: handles.prompt_tx.clone(),
|
||||
cancel_tx: handles.cancel_tx.clone(),
|
||||
model_provider: request.model_provider.clone(),
|
||||
request_id: Some(request_id.clone()),
|
||||
status: AgentStatus::Active,
|
||||
last_activity: Utc::now(),
|
||||
created_at: Utc::now(),
|
||||
stop_handle: handles.lifecycle_handle.clone(),
|
||||
};
|
||||
|
||||
AGENT_REGISTRY.register(
|
||||
&project_id,
|
||||
&response_session_id,
|
||||
project_and_agent_info,
|
||||
);
|
||||
|
||||
info!(
|
||||
"🔗 Agent 已注册到 AGENT_REGISTRY: project_id={}, session_id={}",
|
||||
project_id, response_session_id
|
||||
);
|
||||
}
|
||||
} else {
|
||||
debug!("♻️ Reusing session, no new slot needed (Agent already holds slot)");
|
||||
}
|
||||
|
||||
// 6. 更新 SESSION_REQUEST_CONTEXT(请求追踪)
|
||||
SESSION_REQUEST_CONTEXT.insert(project_id.clone(), request_id.clone());
|
||||
|
||||
// 7. 转换并发送回执
|
||||
let chat_prompt_response = ChatPromptResponse {
|
||||
project_id: worker_response.project_id,
|
||||
session_id: worker_response.session_id,
|
||||
code: if worker_response.error.is_none() {
|
||||
shared_types::error_codes::SUCCESS.to_string()
|
||||
} else {
|
||||
shared_types::error_codes::ERR_AGENT_ERROR.to_string()
|
||||
},
|
||||
error: worker_response.error,
|
||||
request_id: worker_response.request_id,
|
||||
service_type: worker_response.service_type,
|
||||
};
|
||||
|
||||
if let Err(e) = request.chat_prompt_tx.send(chat_prompt_response) {
|
||||
error!("Failed to send acknowledgment: {:?}", e);
|
||||
} else {
|
||||
info!(
|
||||
"✅ 回执已发送,project_id: {}",
|
||||
request.prompt_message.project_id
|
||||
);
|
||||
}
|
||||
|
||||
// SACP 版本:生命周期管理简化
|
||||
//
|
||||
// 架构设计:
|
||||
// - 槽位对应 Agent 生命周期,而非每次请求
|
||||
// - 新会话:等待 Agent 生命周期结束,然后清理
|
||||
// - 复用会话:立即退出,不释放槽位(因为 Agent 还在运行)
|
||||
if is_new_session {
|
||||
if let Some(ref handles) = session_handles {
|
||||
if let Some(ref lifecycle) = handles.lifecycle_handle {
|
||||
info!(
|
||||
"🔄 [SACP] 新会话:等待 Agent 生命周期 - project_id={}, session_id={}",
|
||||
project_id, response_session_id
|
||||
);
|
||||
|
||||
// 等待以下任一事件:
|
||||
// 1. 用户调用 stop_agent → lifecycle.cancel()
|
||||
// 2. 清理任务停止闲置 Agent(5分钟)→ lifecycle.graceful_stop()
|
||||
// 3. Agent 进程异常退出
|
||||
lifecycle.cancellation_token().cancelled().await;
|
||||
|
||||
// 🔥 关键修复:lifecycle 结束后,主动清理 Agent 并释放槽位
|
||||
// 使用 session-aware 移除,避免旧 session 的 cleanup 误删新 session 的 registry 条目。
|
||||
// 场景:用户快速发送多条消息时,旧 session 被取消,新 session 注册。
|
||||
// 旧 session 的 spawned task 退出时,registry 中已是新 session 的条目。
|
||||
// 如果用 remove_by_project 会误删新 session,导致新请求超时。
|
||||
|
||||
AGENT_REGISTRY.remove_by_project_if_session_matches(&project_id, &response_session_id);
|
||||
|
||||
info!(
|
||||
"🛑 [SACP] Agent 生命周期结束,已清理 Registry - project_id={}, session_id={}",
|
||||
project_id, response_session_id
|
||||
);
|
||||
} else {
|
||||
warn!(
|
||||
"⚠️ [SACP] 新会话缺少 lifecycle_handle - project_id={}",
|
||||
project_id
|
||||
);
|
||||
// 缺少 lifecycle_handle,立即清理
|
||||
AGENT_REGISTRY.remove_by_project_if_session_matches(&project_id, &response_session_id);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
info!(
|
||||
"🔵 [SACP] 复用会话:请求处理完成 - project_id={}, session_id={}",
|
||||
project_id, response_session_id
|
||||
);
|
||||
// 复用会话时不释放槽位,因为槽位对应 Agent 生命周期
|
||||
// 而不是每次请求。Agent 还在运行,槽位应保持占用状态
|
||||
}
|
||||
});
|
||||
|
||||
// 立即继续循环,接收下一个请求 - 不等待上面的 spawn 完成
|
||||
}
|
||||
|
||||
// 清理心跳任务
|
||||
heartbeat_task.abort();
|
||||
|
||||
info!("🛑 agent_worker stopped");
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
//! 定期清理闲置agent的任务
|
||||
//!
|
||||
//! 基于AgentLifecycleGuard的RAII原则,简化清理逻辑:
|
||||
//! 1. 定时扫描识别闲置的agent
|
||||
//! 2. 从PROJECT_AND_AGENT_INFO_MAP中移除
|
||||
//! 3. AgentLifecycleGuard自动drop并清理资源
|
||||
|
||||
use anyhow::Result;
|
||||
use chrono::{DateTime, Utc};
|
||||
use std::time::Duration;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::model::AgentStatus;
|
||||
use crate::service::{AGENT_REGISTRY, SESSION_CACHE};
|
||||
|
||||
/// 清理配置
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CleanupConfig {
|
||||
/// 闲置超时时间(默认30分钟)
|
||||
pub idle_timeout: Duration,
|
||||
/// 清理检查间隔(默认5分钟)
|
||||
pub cleanup_interval: Duration,
|
||||
// 注意:force_terminate_timeout 字段已移除
|
||||
// 因为采用RAII模式,AgentLifecycleGuard会自动处理资源清理
|
||||
// 不需要强制终止超时机制
|
||||
}
|
||||
|
||||
impl Default for CleanupConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
idle_timeout: Duration::from_secs(3 * 60), // 🆕 默认3分钟
|
||||
cleanup_interval: Duration::from_secs(30), // 默认30秒
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 清理任务统计信息
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct CleanupStats {
|
||||
/// 总共清理的agent数量
|
||||
pub total_cleaned: u64,
|
||||
/// 成功清理的agent数量
|
||||
pub success_cleaned: u64,
|
||||
/// 清理失败的agent数量
|
||||
pub failed_cleaned: u64,
|
||||
/// 清理的孤立session数量
|
||||
pub orphaned_sessions_cleaned: u64,
|
||||
/// 清理的SSE消息数量
|
||||
pub sse_messages_cleaned: u64,
|
||||
/// 最后清理时间
|
||||
pub last_cleanup: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
/// Agent清理器 - 基于RAII的简化版本
|
||||
pub struct AgentCleaner {
|
||||
config: CleanupConfig,
|
||||
stats: CleanupStats,
|
||||
}
|
||||
|
||||
impl AgentCleaner {
|
||||
/// 创建新的清理器
|
||||
pub fn new(config: CleanupConfig) -> Self {
|
||||
Self {
|
||||
config,
|
||||
stats: CleanupStats::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查agent是否闲置超时
|
||||
fn is_agent_idle_timeout(
|
||||
&self,
|
||||
last_activity: DateTime<Utc>,
|
||||
current_time: DateTime<Utc>,
|
||||
) -> bool {
|
||||
let duration = current_time.signed_duration_since(last_activity);
|
||||
duration.num_seconds() > 0
|
||||
&& duration.num_seconds() as u64 > self.config.idle_timeout.as_secs()
|
||||
}
|
||||
|
||||
/// 清理孤立的SSE消息数据
|
||||
/// 清理没有project_id引用的session和长期未活跃的session
|
||||
async fn cleanup_orphaned_sse_sessions(&mut self) -> (u64, u64) {
|
||||
let mut orphaned_count = 0;
|
||||
let mut messages_cleared = 0;
|
||||
|
||||
// 使用统一 Registry 收集所有活跃的 session_id
|
||||
let active_session_ids: std::collections::HashSet<String> = AGENT_REGISTRY
|
||||
.iter_agents()
|
||||
.map(|entry| entry.value().session_id.to_string())
|
||||
.collect();
|
||||
|
||||
// 检查SESSION_CACHE中的所有session
|
||||
let mut sessions_to_remove = Vec::new();
|
||||
|
||||
let session_ids: Vec<String> = SESSION_CACHE
|
||||
.iter()
|
||||
.map(|entry| entry.key().clone())
|
||||
.collect();
|
||||
|
||||
for session_id in session_ids {
|
||||
// 如果session_id不在活跃映射中,则为孤立session
|
||||
if !active_session_ids.contains(&session_id) {
|
||||
// 检查session中是否有消息
|
||||
if let Some(session_data_ref) = SESSION_CACHE.get(&session_id) {
|
||||
let session_data = session_data_ref.clone();
|
||||
drop(session_data_ref);
|
||||
|
||||
let message_count = session_data.message_count().await;
|
||||
|
||||
if message_count > 0 {
|
||||
info!(
|
||||
"Found orphaned session: session_id={}, message_count={}",
|
||||
session_id, message_count
|
||||
);
|
||||
|
||||
// 清理这个session的消息 - 直接移除条目
|
||||
if SESSION_CACHE.remove(&session_id).is_some() {
|
||||
messages_cleared += 1;
|
||||
}
|
||||
|
||||
// 如果清理后session为空,标记为待删除
|
||||
if session_data.message_count().await == 0 {
|
||||
sessions_to_remove.push(session_id.clone());
|
||||
}
|
||||
|
||||
orphaned_count += 1;
|
||||
} else {
|
||||
// 没有消息的空session,直接标记删除
|
||||
sessions_to_remove.push(session_id.clone());
|
||||
}
|
||||
} else {
|
||||
// session_data不存在,也标记删除
|
||||
sessions_to_remove.push(session_id.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 删除空的孤立session
|
||||
for session_id in sessions_to_remove {
|
||||
if let Some((_, _)) = SESSION_CACHE.remove(&session_id) {
|
||||
debug!("Removed empty session: {}", session_id);
|
||||
}
|
||||
}
|
||||
|
||||
if orphaned_count > 0 {
|
||||
info!(
|
||||
"Orphaned SSE session cleanup completed: session_count={}, message_count={}",
|
||||
orphaned_count, messages_cleared
|
||||
);
|
||||
}
|
||||
|
||||
(orphaned_count, messages_cleared)
|
||||
}
|
||||
|
||||
/// 执行一次清理操作 - 基于RAII的简化版
|
||||
/// 只需要从 MAP 中移除闲置agent,AgentLifecycleGuard 会自动清理资源
|
||||
async fn cleanup_idle_agents(&mut self) -> Result<CleanupStats> {
|
||||
let current_time = Utc::now();
|
||||
let mut cleaned_count = 0;
|
||||
let mut success_count = 0;
|
||||
let mut failed_count = 0;
|
||||
|
||||
// 使用统一 Registry 获取统计信息
|
||||
let registry_stats = AGENT_REGISTRY.stats();
|
||||
let total_agents = registry_stats.agent_count;
|
||||
|
||||
info!(
|
||||
"Starting cleanup for idle agents and SSE messages, current_time: {}, active_agent_count: {}",
|
||||
current_time, total_agents
|
||||
);
|
||||
|
||||
// 先清理孤立的SSE消息数据
|
||||
let (orphaned_sessions, sse_messages) = self.cleanup_orphaned_sse_sessions().await;
|
||||
|
||||
// 收集需要清理的agent ID(使用统一 Registry 遍历)
|
||||
let mut agents_to_remove = Vec::new();
|
||||
|
||||
for entry in AGENT_REGISTRY.iter_agents() {
|
||||
let project_id = entry.key();
|
||||
let agent_info = entry.value();
|
||||
|
||||
// 只清理Idle状态的agent,避免中断正在执行的任务
|
||||
if agent_info.status == AgentStatus::Idle
|
||||
&& self.is_agent_idle_timeout(agent_info.last_activity, current_time)
|
||||
{
|
||||
let idle_duration = (current_time - agent_info.last_activity).num_seconds();
|
||||
info!(
|
||||
"Found idle agent: project_id={}, status={:?}, last_activity: {}, idle_duration_seconds: {}, created_at: {}",
|
||||
project_id,
|
||||
agent_info.status,
|
||||
agent_info.last_activity,
|
||||
idle_duration,
|
||||
agent_info.created_at
|
||||
);
|
||||
agents_to_remove.push(project_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// 执行清理 - RAII版:直接从 MAP 中移除,AgentLifecycleGuard 会自动清理
|
||||
for project_id in agents_to_remove {
|
||||
match self.cleanup_agent_raii(&project_id) {
|
||||
Ok(_) => {
|
||||
success_count += 1;
|
||||
info!("Agent cleaned successfully: {}", project_id);
|
||||
}
|
||||
Err(e) => {
|
||||
failed_count += 1;
|
||||
warn!("Failed to clean agent: {} - {}", project_id, e);
|
||||
}
|
||||
}
|
||||
cleaned_count += 1;
|
||||
}
|
||||
|
||||
// 更新统计信息
|
||||
self.stats.total_cleaned += cleaned_count;
|
||||
self.stats.success_cleaned += success_count;
|
||||
self.stats.failed_cleaned += failed_count;
|
||||
self.stats.orphaned_sessions_cleaned += orphaned_sessions;
|
||||
self.stats.sse_messages_cleaned += sse_messages;
|
||||
self.stats.last_cleanup = Some(current_time);
|
||||
|
||||
// 清理完成后的统计(使用统一 Registry)
|
||||
let final_stats = AGENT_REGISTRY.stats();
|
||||
let remaining_agents = final_stats.agent_count;
|
||||
let active_sessions = final_stats.session_count;
|
||||
let cached_sessions = SESSION_CACHE.len();
|
||||
|
||||
info!(
|
||||
"Cleanup completed: agent(total={}, success={}, failed={}, remaining={}) | session(active={}, cached={}) | sse_messages(cleared={})",
|
||||
cleaned_count,
|
||||
success_count,
|
||||
failed_count,
|
||||
remaining_agents,
|
||||
active_sessions,
|
||||
cached_sessions,
|
||||
sse_messages
|
||||
);
|
||||
|
||||
Ok(CleanupStats {
|
||||
total_cleaned: cleaned_count,
|
||||
success_cleaned: success_count,
|
||||
failed_cleaned: failed_count,
|
||||
orphaned_sessions_cleaned: orphaned_sessions,
|
||||
sse_messages_cleaned: sse_messages,
|
||||
last_cleanup: Some(current_time),
|
||||
})
|
||||
}
|
||||
|
||||
/// 基于RAII的简化清理方法
|
||||
/// 只需要从MAP中移除agent,AgentLifecycleGuard会自动清理所有资源
|
||||
fn cleanup_agent_raii(&self, project_id: &str) -> Result<()> {
|
||||
debug!("Starting RAII cleanup for agent: {}", project_id);
|
||||
|
||||
// 使用统一 Registry 检查并移除(内部自动同步清理所有映射)
|
||||
if AGENT_REGISTRY.contains_project(project_id) {
|
||||
// 通过统一 Registry 移除,自动清理:
|
||||
// - agent_info_map
|
||||
// - project_to_session 映射
|
||||
// - session_to_project 反向映射
|
||||
let removed = AGENT_REGISTRY.remove_by_project(project_id);
|
||||
|
||||
// 同步清理 SESSION_REQUEST_CONTEXT 中的 request_id
|
||||
crate::proxy_agent::SESSION_REQUEST_CONTEXT.remove(project_id);
|
||||
debug!(
|
||||
"🧼 [cleanup] Cleared project_id from SESSION_REQUEST_CONTEXT: {}",
|
||||
project_id
|
||||
);
|
||||
|
||||
if removed.is_some() {
|
||||
info!(
|
||||
"Agent removed from Registry; AgentLifecycleGuard will clean up resources automatically: {}",
|
||||
project_id
|
||||
);
|
||||
} else {
|
||||
warn!("Tried to remove agent but it was not found: {}", project_id);
|
||||
}
|
||||
} else {
|
||||
warn!("Agent not found in Registry: {}", project_id);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 运行清理任务 - 简化版,只做定时清理
|
||||
pub async fn run(&mut self) {
|
||||
info!("Cleanup task started, config: {:?}", self.config);
|
||||
|
||||
let mut interval = tokio::time::interval(self.config.cleanup_interval);
|
||||
|
||||
loop {
|
||||
interval.tick().await;
|
||||
|
||||
match self.cleanup_idle_agents().await {
|
||||
Ok(stats) => debug!("Periodic cleanup completed: {:?}", stats),
|
||||
Err(e) => warn!("Periodic cleanup failed: {}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取统计信息
|
||||
pub fn get_stats(&self) -> &CleanupStats {
|
||||
&self.stats
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动清理任务 - 普通异步版本
|
||||
///
|
||||
/// 清理任务只操作 Send 数据结构,可以在普通异步线程中运行
|
||||
pub fn start_cleanup_task(config: CleanupConfig) -> tokio::task::JoinHandle<()> {
|
||||
let mut cleaner = AgentCleaner::new(config);
|
||||
|
||||
tokio::task::spawn(async move {
|
||||
cleaner.run().await;
|
||||
})
|
||||
}
|
||||
115
qiming-rcoder/crates/agent_runner/src/proxy_agent/mod.rs
Normal file
115
qiming-rcoder/crates/agent_runner/src/proxy_agent/mod.rs
Normal file
@@ -0,0 +1,115 @@
|
||||
mod acp_agent;
|
||||
pub mod cleanup_task;
|
||||
|
||||
use crate::CancelNotificationRequestWrapper;
|
||||
// 导出 agent_worker 相关类型和函数
|
||||
pub use acp_agent::{AgentRequest, agent_worker_with_heartbeat, set_unlimited_mode};
|
||||
use shared_types::AgentLifecycleGuard;
|
||||
// SACP 类型导入
|
||||
#[cfg(feature = "proxy")]
|
||||
use crate::config::ProxyConfig;
|
||||
use dashmap::DashMap;
|
||||
#[cfg(feature = "proxy")]
|
||||
use rcoder_proxy::{PingoraServerManager, ProxyConfig as PingoraProxyConfig};
|
||||
use agent_client_protocol::schema::{PromptRequest, SessionId};
|
||||
use std::sync::{Arc, LazyLock};
|
||||
use tokio::sync::mpsc;
|
||||
#[cfg(feature = "proxy")]
|
||||
use tracing::{error, info};
|
||||
|
||||
/// Pingora 启动结果
|
||||
///
|
||||
/// 持有关闭信号的发送端,`stop()` 时直接发送信号,无需 Mutex 锁。
|
||||
#[cfg(feature = "proxy")]
|
||||
pub struct PingoraStartResult {
|
||||
/// 关闭信号发送端
|
||||
shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "proxy")]
|
||||
impl PingoraStartResult {
|
||||
/// 停止 Pingora 服务器
|
||||
pub async fn stop(&mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动 Pingora 代理服务
|
||||
///
|
||||
/// 封装 Pingora 的创建和启动逻辑,供 main.rs 和 http_server/start.rs 复用。
|
||||
/// shutdown 通道在外部创建,`stop()` 直接发送信号,不经过 Mutex,消除死锁风险。
|
||||
#[cfg(feature = "proxy")]
|
||||
#[must_use]
|
||||
pub fn start_pingora(
|
||||
proxy_config: &ProxyConfig,
|
||||
shared_api_key_manager: Arc<dashmap::DashMap<String, shared_types::ModelProviderConfig>>,
|
||||
) -> PingoraStartResult {
|
||||
info!(
|
||||
"Starting Pingora reverse proxy service, listening on port: {}",
|
||||
proxy_config.listen_port
|
||||
);
|
||||
info!(
|
||||
"Proxy route format: /proxy/{{port}}{{/path}} - e.g.: /proxy/{}/health",
|
||||
proxy_config.default_backend_port
|
||||
);
|
||||
|
||||
let pingora_config = PingoraProxyConfig {
|
||||
listen_port: proxy_config.listen_port,
|
||||
default_backend_port: proxy_config.default_backend_port,
|
||||
backend_host: proxy_config.backend_host.clone(),
|
||||
port_param: proxy_config.port_param.clone(),
|
||||
config_file: None,
|
||||
verbose: false,
|
||||
};
|
||||
|
||||
// 创建 Pingora 服务器管理器
|
||||
let mut server_manager =
|
||||
PingoraServerManager::new(pingora_config).with_api_key_manager(shared_api_key_manager);
|
||||
|
||||
let pingora_service = server_manager.service();
|
||||
|
||||
// 启动健康检查循环(按配置)
|
||||
if proxy_config.health_check.enabled {
|
||||
let hc = &proxy_config.health_check;
|
||||
pingora_service.start_health_check_loop(hc.interval_seconds, hc.timeout_seconds * 1000);
|
||||
}
|
||||
|
||||
// 在外部创建 shutdown 通道,避免通过 Mutex 发送信号导致死锁
|
||||
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
|
||||
|
||||
// 在后台任务中启动 Pingora(直接 move server_manager,无需 Arc<Mutex<>>)
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = server_manager.start(shutdown_rx).await {
|
||||
error!("Failed to start Pingora proxy server: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
info!(
|
||||
"✅ Pingora 代理服务已启动在端口 {}",
|
||||
proxy_config.listen_port
|
||||
);
|
||||
|
||||
PingoraStartResult {
|
||||
shutdown_tx: Some(shutdown_tx),
|
||||
}
|
||||
}
|
||||
|
||||
/// 会话级别的 request_id 上下文映射(project_id -> request_id)
|
||||
/// 用于在 session_notification 回调中获取当前请求的 request_id
|
||||
/// 避免使用 PROJECT_AND_AGENT_INFO_MAP 导致的锁竞争问题
|
||||
/// 注意:使用 project_id 而非 session_id,确保同一项目的多次请求能自动覆盖为最新值
|
||||
pub static SESSION_REQUEST_CONTEXT: LazyLock<DashMap<String, String>> = LazyLock::new(DashMap::new);
|
||||
|
||||
/// ACP协议的连接信息
|
||||
pub struct AcpConnectionInfo {
|
||||
/// 会话ID
|
||||
pub session_id: SessionId,
|
||||
/// 用于发送 Prompt 的通道
|
||||
pub prompt_tx: mpsc::UnboundedSender<PromptRequest>,
|
||||
/// 用于发送取消通知的通道(使用新类型)
|
||||
pub cancel_tx: mpsc::UnboundedSender<CancelNotificationRequestWrapper>,
|
||||
/// Agent停止句柄(将被包装为守卫并放入 ProjectAndAgentInfo)
|
||||
pub stop_handle: Option<Arc<AgentLifecycleGuard>>,
|
||||
}
|
||||
166
qiming-rcoder/crates/agent_runner/src/router.rs
Normal file
166
qiming-rcoder/crates/agent_runner/src/router.rs
Normal file
@@ -0,0 +1,166 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::agent_runtime::AgentRuntime;
|
||||
use crate::{config::AppConfig, handler};
|
||||
use axum::{Router, response::IntoResponse, routing::get};
|
||||
use dashmap::DashMap;
|
||||
use rcoder_telemetry::{HttpMetricsLayer, TelemetryGuard};
|
||||
use serde::Serialize;
|
||||
use utoipa::OpenApi;
|
||||
use utoipa_swagger_ui::SwaggerUi;
|
||||
|
||||
/// 会话信息结构
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct SessionInfo {
|
||||
pub session_id: String,
|
||||
pub user_id: String,
|
||||
pub project_id: Option<String>,
|
||||
pub created_at: chrono::DateTime<chrono::Utc>,
|
||||
pub last_activity: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
/// 应用状态
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
/// 活跃的会话映射, project_id -> SessionInfo
|
||||
pub sessions: Arc<DashMap<String, SessionInfo>>,
|
||||
/// 应用配置
|
||||
pub config: AppConfig,
|
||||
|
||||
/// 🆕 Agent Runtime(新架构)
|
||||
///
|
||||
/// **推荐**: 使用 `agent_runtime.send().await` 发送任务
|
||||
/// - 支持自动重启
|
||||
/// - 提供健康状态检查
|
||||
/// - 线程安全的原子操作
|
||||
pub local_task_sender: Arc<AgentRuntime>,
|
||||
|
||||
/// 🆕 Agent Runtime(用于健康检查和状态监控)
|
||||
pub agent_runtime: Arc<AgentRuntime>,
|
||||
|
||||
/// Pingora 代理服务引用(用于读取真实指标)
|
||||
#[cfg(feature = "proxy")]
|
||||
pub pingora_service: Option<Arc<rcoder_proxy::PingoraProxyService>>,
|
||||
|
||||
/// 🔒 API 密钥管理器(用于 Pingora 代理注入真实密钥)
|
||||
/// 注意:此现在是 shared_api_key_manager 的包装器
|
||||
pub api_key_manager: Arc<crate::api_key_manager::ApiKeyManager>,
|
||||
|
||||
/// 🔒 共享的 API 密钥 DashMap(直接与 Pingora 共享)
|
||||
/// 存储格式:<UUID> -> ModelProviderConfig
|
||||
pub shared_api_key_manager: Arc<DashMap<String, shared_types::ModelProviderConfig>>,
|
||||
|
||||
/// 🔒 project_id -> service_uuid 映射(用于清理时查找对应的配置)
|
||||
pub project_uuid_map: Arc<DashMap<String, String>>,
|
||||
}
|
||||
|
||||
/// 创建 Axum 路由
|
||||
pub fn create_router(state: Arc<AppState>, telemetry: Option<Arc<TelemetryGuard>>) -> Router {
|
||||
let api_routes = Router::new()
|
||||
.route("/health", get(handler::health_check))
|
||||
.with_state(state.clone());
|
||||
|
||||
let mut router = Router::new().merge(api_routes).merge(create_swagger_ui());
|
||||
|
||||
// 添加 /metrics 端点(如果启用了 Prometheus)
|
||||
if let Some(ref guard) = telemetry {
|
||||
let guard_clone = Arc::clone(guard);
|
||||
router = router.route(
|
||||
"/metrics",
|
||||
get(move || {
|
||||
let guard = Arc::clone(&guard_clone);
|
||||
async move { metrics_handler(guard).await }
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// 应用 HTTP 指标中间件
|
||||
router.layer(HttpMetricsLayer::new())
|
||||
}
|
||||
|
||||
/// Prometheus 指标处理器
|
||||
async fn metrics_handler(telemetry: Arc<TelemetryGuard>) -> impl IntoResponse {
|
||||
match telemetry.render_metrics() {
|
||||
Some(metrics) => (
|
||||
axum::http::StatusCode::OK,
|
||||
[(
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
"text/plain; charset=utf-8",
|
||||
)],
|
||||
metrics,
|
||||
),
|
||||
None => (
|
||||
axum::http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
[(
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
"text/plain; charset=utf-8",
|
||||
)],
|
||||
"Prometheus metrics not enabled".to_string(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// OpenAPI 文档结构
|
||||
#[derive(OpenApi)]
|
||||
#[openapi(
|
||||
paths(
|
||||
handler::health_check,
|
||||
),
|
||||
components(
|
||||
schemas(
|
||||
// 响应结构体
|
||||
handler::HealthResponse,
|
||||
)
|
||||
),
|
||||
tags(
|
||||
(name = "system", description = "系统健康检查和状态监控接口"),
|
||||
),
|
||||
info(
|
||||
description = r#"
|
||||
RCoder AI 服务 API
|
||||
|
||||
基于 ACP (Agent Client Protocol) 的 AI 驱动开发平台,提供完整的 AI 代理集成解决方案。
|
||||
|
||||
## 主要功能
|
||||
|
||||
- **智能对话**: 支持文本、图像、音频、文档等多媒体内容的 AI 交互
|
||||
- **实时通知**: 通过 SSE 协议提供 AI 代理执行进度的实时推送
|
||||
- **会话管理**: 完整的会话生命周期管理,支持任务取消
|
||||
- **项目隔离**: 每个对话在独立的项目工作空间中进行,确保安全性
|
||||
|
||||
## 技术架构
|
||||
|
||||
- **协议**: ACP (Agent Client Protocol) v0.4
|
||||
- **gRPC 通信**: rcoder 通过 gRPC 与 agent_runner 通信
|
||||
- **健康检查**: 提供服务健康状态查询
|
||||
|
||||
## gRPC 接口
|
||||
|
||||
agent_runner 主要通过 gRPC 提供服务:
|
||||
- `Chat` - 聊天对话
|
||||
- `SubscribeProgress` - 订阅进度事件(Server Streaming)
|
||||
- `CancelSession` - 取消会话
|
||||
- `GetStatus` - 查询状态
|
||||
"#,
|
||||
title = "RCoder AI API",
|
||||
version = "2.0.0",
|
||||
license(name = "MIT OR Apache-2.0", url = "https://opensource.org/licenses/MIT"),
|
||||
contact(
|
||||
name = "RCoder Team",
|
||||
email = "team@rcoder.com",
|
||||
url = "https://github.com/rcoder/rcoder"
|
||||
)
|
||||
),
|
||||
servers(
|
||||
(url = "http://localhost:50051", description = "gRPC 服务 (agent_runner)"),
|
||||
(url = "http://localhost:8087", description = "HTTP API 服务 (rcoder)"),
|
||||
)
|
||||
)]
|
||||
pub struct ApiDoc;
|
||||
|
||||
/// 创建 Swagger UI 路由
|
||||
pub fn create_swagger_ui() -> SwaggerUi {
|
||||
SwaggerUi::new("/api/docs")
|
||||
.url("/api/docs/openapi.json", ApiDoc::openapi())
|
||||
.config(utoipa_swagger_ui::Config::new(["/api/docs/openapi.json"]))
|
||||
}
|
||||
1001
qiming-rcoder/crates/agent_runner/src/service/agent_registry.rs
Normal file
1001
qiming-rcoder/crates/agent_runner/src/service/agent_registry.rs
Normal file
File diff suppressed because it is too large
Load Diff
582
qiming-rcoder/crates/agent_runner/src/service/chat_handler.rs
Normal file
582
qiming-rcoder/crates/agent_runner/src/service/chat_handler.rs
Normal file
@@ -0,0 +1,582 @@
|
||||
//! Chat Handler 共享逻辑
|
||||
//!
|
||||
//! 封装 chat 请求处理的核心业务逻辑,供 gRPC 和 HTTP 复用。
|
||||
//!
|
||||
//! ## 设计原则
|
||||
//!
|
||||
//! - **RAII 状态管理**: 使用 PendingGuard 自动管理 Pending 状态
|
||||
//! - **DashMap Entry API**: 避免读写锁竞态条件
|
||||
//! - **Fail Fast**: 尽早暴露问题,便于定位修复
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
use shared_types::{
|
||||
Attachment, CancelNotificationRequestWrapper, CancelResult, ChatAgentConfig, ChatPromptBuilder,
|
||||
ModelProviderConfig, ServiceType, error_codes,
|
||||
};
|
||||
use agent_client_protocol::schema::{CancelNotification, SessionId};
|
||||
use tokio::time::Duration;
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
use crate::AgentRuntime;
|
||||
use crate::agent_runtime::WorkerState;
|
||||
use crate::proxy_agent::AgentRequest;
|
||||
use crate::service::{AGENT_REGISTRY, PendingGuard, SESSION_CACHE};
|
||||
|
||||
/// Chat Handler 输入参数
|
||||
///
|
||||
/// 包含处理 chat 请求所需的所有参数,与协议无关(gRPC/HTTP)
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChatHandlerInput {
|
||||
/// 项目 ID
|
||||
pub project_id: String,
|
||||
/// 项目工作目录(由调用方根据环境决定)
|
||||
pub project_dir: PathBuf,
|
||||
/// 会话 ID(可选,用于复用会话)
|
||||
pub session_id: Option<String>,
|
||||
/// 用户提示词
|
||||
pub prompt: String,
|
||||
/// 请求 ID(用于追踪)
|
||||
pub request_id: String,
|
||||
/// 附件列表
|
||||
pub attachments: Vec<Attachment>,
|
||||
/// 数据源附件列表
|
||||
pub data_source_attachments: Vec<String>,
|
||||
/// 模型配置(可选)
|
||||
pub model_config: Option<ModelProviderConfig>,
|
||||
/// 服务类型
|
||||
pub service_type: ServiceType,
|
||||
/// Agent 配置覆盖(可选)
|
||||
pub agent_config_override: Option<ChatAgentConfig>,
|
||||
/// 系统提示覆盖(可选)
|
||||
pub system_prompt_override: Option<String>,
|
||||
/// 用户提示模板覆盖(可选)
|
||||
pub user_prompt_template_override: Option<String>,
|
||||
/// 是否跳过槽位限制(HTTP Server 宿主机部署时为 true)
|
||||
pub skip_slot_limit: bool,
|
||||
}
|
||||
|
||||
/// Chat Handler 输出结果
|
||||
///
|
||||
/// 统一的响应结构,可转换为 gRPC 或 HTTP 响应
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChatHandlerOutput {
|
||||
/// 项目 ID
|
||||
pub project_id: String,
|
||||
/// 会话 ID
|
||||
pub session_id: String,
|
||||
/// 是否成功
|
||||
pub success: bool,
|
||||
/// 错误消息(可选)
|
||||
pub error: Option<String>,
|
||||
/// 错误码(可选)
|
||||
pub error_code: Option<String>,
|
||||
/// 请求 ID(可选)
|
||||
pub request_id: Option<String>,
|
||||
/// 是否需要降级处理
|
||||
pub need_fallback: bool,
|
||||
/// 降级原因(可选)
|
||||
pub fallback_reason: Option<String>,
|
||||
}
|
||||
|
||||
impl ChatHandlerOutput {
|
||||
/// 创建错误响应
|
||||
pub fn error(
|
||||
project_id: String,
|
||||
session_id: String,
|
||||
error_msg: String,
|
||||
error_code: String,
|
||||
) -> Self {
|
||||
Self {
|
||||
project_id,
|
||||
session_id,
|
||||
success: false,
|
||||
error: Some(error_msg),
|
||||
error_code: Some(error_code),
|
||||
request_id: None,
|
||||
need_fallback: false,
|
||||
fallback_reason: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建 Agent Busy 错误响应
|
||||
pub fn agent_busy(project_id: String, session_id: Option<String>) -> Self {
|
||||
Self {
|
||||
project_id,
|
||||
session_id: session_id.unwrap_or_default(),
|
||||
success: false,
|
||||
error: Some(error_codes::get_i18n_message_default("error.agent_busy")),
|
||||
error_code: Some(error_codes::ERR_AGENT_BUSY.to_string()),
|
||||
request_id: None,
|
||||
need_fallback: false,
|
||||
fallback_reason: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Chat Handler 依赖上下文
|
||||
///
|
||||
/// 包含处理 chat 请求所需的运行时依赖
|
||||
pub struct ChatHandlerContext {
|
||||
/// Agent 运行时
|
||||
pub agent_runtime: Arc<AgentRuntime>,
|
||||
/// 共享的 API 密钥管理器
|
||||
pub shared_api_key_manager: Arc<DashMap<String, ModelProviderConfig>>,
|
||||
/// project_id -> UUID 映射
|
||||
pub project_uuid_map: Arc<DashMap<String, String>>,
|
||||
}
|
||||
|
||||
/// 取消当前正在执行的 Agent 任务
|
||||
///
|
||||
/// 发送取消通知并等待取消完成,超时时间为 10 秒
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `cancel_tx` - 取消通知发送通道
|
||||
/// * `session_id` - 当前会话 ID
|
||||
/// * `project_id` - 项目 ID
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Ok(())` - 取消成功,Agent 状态已恢复为 Idle
|
||||
/// * `Err(ChatHandlerOutput)` - 取消失败,包含错误响应
|
||||
async fn cancel_current_task(
|
||||
cancel_tx: &tokio::sync::mpsc::Sender<CancelNotificationRequestWrapper>,
|
||||
session_id: &str,
|
||||
project_id: &str,
|
||||
) -> Result<(), ChatHandlerOutput> {
|
||||
info!(
|
||||
"[ChatHandler] Cancelling current task: project_id={}, session_id={}",
|
||||
project_id, session_id
|
||||
);
|
||||
|
||||
// 1. 检查 cancel_tx 是否有效
|
||||
if cancel_tx.is_closed() {
|
||||
error!(
|
||||
"[ChatHandler] Cancel channel closed: project_id={}, session_id={}",
|
||||
project_id, session_id
|
||||
);
|
||||
return Err(ChatHandlerOutput::error(
|
||||
project_id.to_string(),
|
||||
session_id.to_string(),
|
||||
error_codes::get_i18n_message_default("error.cancel_channel_closed"),
|
||||
error_codes::ERR_SERVICE_UNAVAILABLE.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// 2. 创建 oneshot channel 等待取消结果
|
||||
let (result_tx, result_rx) = tokio::sync::oneshot::channel::<CancelResult>();
|
||||
let cancel_notification = CancelNotification::new(SessionId::new(Arc::from(session_id)));
|
||||
let cancel_request = CancelNotificationRequestWrapper {
|
||||
cancel_notification,
|
||||
result_tx,
|
||||
};
|
||||
|
||||
// 3. 发送取消通知
|
||||
if let Err(e) = cancel_tx.send(cancel_request).await {
|
||||
error!(
|
||||
"[ChatHandler] Failed to send cancel notification: project_id={}, error={}",
|
||||
project_id, e
|
||||
);
|
||||
return Err(ChatHandlerOutput::error(
|
||||
project_id.to_string(),
|
||||
session_id.to_string(),
|
||||
format!(
|
||||
"{}: {}",
|
||||
error_codes::get_i18n_message_default("error.cancel_failed"),
|
||||
e
|
||||
),
|
||||
error_codes::ERR_INTERNAL_SERVER_ERROR.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// 4. 等待取消结果(超时 10 秒)
|
||||
match tokio::time::timeout(Duration::from_secs(10), result_rx).await {
|
||||
Ok(Ok(cancel_result)) => {
|
||||
if cancel_result.is_success() {
|
||||
info!(
|
||||
"[ChatHandler] Cancel notification sent successfully, proceeding with new request: project_id={}, session_id={}",
|
||||
project_id, session_id
|
||||
);
|
||||
|
||||
// 🎯 关键设计:cancel 后立即返回,不等待 session 移除
|
||||
//
|
||||
// 上下文连续性保证:
|
||||
// - 不等待 session 移除 → session 保持在 Registry 中
|
||||
// - get_or_create_session → is_channel_closed()=false → 复用同一 session
|
||||
// - 新 prompt 发送到同一 session 的 prompt_tx → 同一 Agent 子进程处理
|
||||
// - Agent 子进程保持存活 → 内存中的对话上下文连续
|
||||
//
|
||||
// 时序:
|
||||
// 1. CancelResult::Success → cancel 通知已发送给 Agent
|
||||
// 2. SACP inner loop 收到 cancel → is_cancelled=true → 等待 Agent 响应或超时
|
||||
// 3. inner loop 退出 → outer loop 继续等待 prompt_rx
|
||||
// 4. 新请求的 prompt 到达 → session_cancelled 重置 → 处理新 prompt
|
||||
// 5. 同一 Agent 子进程处理新 prompt → 上下文连续
|
||||
//
|
||||
// 最坏情况延迟:inner cancel timeout (10s) — Agent 不响应 cancel 时
|
||||
Ok(())
|
||||
} else {
|
||||
let error_msg = cancel_result.error_message().unwrap_or("Unknown error");
|
||||
error!(
|
||||
"[ChatHandler] Cancel failed: project_id={}, error={}",
|
||||
project_id, error_msg
|
||||
);
|
||||
Err(ChatHandlerOutput::error(
|
||||
project_id.to_string(),
|
||||
session_id.to_string(),
|
||||
format!(
|
||||
"{}: {}",
|
||||
error_codes::get_i18n_message_default("error.cancel_failed"),
|
||||
error_msg
|
||||
),
|
||||
error_codes::ERR_AGENT_ERROR.to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
Ok(Err(_)) => {
|
||||
error!(
|
||||
"[ChatHandler] Cancel result channel dropped: project_id={}",
|
||||
project_id
|
||||
);
|
||||
Err(ChatHandlerOutput::error(
|
||||
project_id.to_string(),
|
||||
session_id.to_string(),
|
||||
error_codes::get_i18n_message_default("error.cancel_channel_dropped"),
|
||||
error_codes::ERR_INTERNAL_SERVER_ERROR.to_string(),
|
||||
))
|
||||
}
|
||||
Err(_) => {
|
||||
error!(
|
||||
"[ChatHandler] Cancel timeout (10s): project_id={}",
|
||||
project_id
|
||||
);
|
||||
Err(ChatHandlerOutput::error(
|
||||
project_id.to_string(),
|
||||
session_id.to_string(),
|
||||
error_codes::get_i18n_message_default("error.cancel_timeout"),
|
||||
error_codes::ERR_CANCEL_FAILED.to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行 Chat 请求的核心逻辑
|
||||
///
|
||||
/// 封装了 chat 请求的完整处理流程:
|
||||
/// 1. Agent Busy 检查
|
||||
/// 2. PendingGuard RAII 状态管理
|
||||
/// 3. Session 清理逻辑
|
||||
/// 4. 目录创建
|
||||
/// 5. ChatPrompt 构建
|
||||
/// 6. API Key 管理
|
||||
/// 7. AgentRequest 创建和发送
|
||||
/// 8. 等待响应
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `input` - Chat 请求输入参数
|
||||
/// * `context` - 运行时上下文依赖
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// 返回统一的 `ChatHandlerOutput` 结果
|
||||
pub async fn handle_chat_core(
|
||||
input: ChatHandlerInput,
|
||||
context: &ChatHandlerContext,
|
||||
) -> ChatHandlerOutput {
|
||||
let project_id = input.project_id.clone();
|
||||
let session_id = input.session_id.clone();
|
||||
let request_id = input.request_id.clone();
|
||||
|
||||
info!(
|
||||
"[ChatHandler] Starting to process request: project_id={}, session_id={:?}, prompt_len={}, has_model_config={}",
|
||||
project_id,
|
||||
session_id,
|
||||
input.prompt.len(),
|
||||
input.model_config.is_some()
|
||||
);
|
||||
|
||||
// ========== 步骤1: 查询现有 Agent 状态 ==========
|
||||
// 优先通过 session_id 查找,回退到 project_id 查找
|
||||
let agent_info_ref = if let Some(ref sid) = session_id {
|
||||
info!(
|
||||
"[ChatHandler] Looking up Agent by session_id: session_id={}",
|
||||
sid
|
||||
);
|
||||
AGENT_REGISTRY.get_agent_info_by_session(sid)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let agent_info_ref = agent_info_ref.or_else(|| {
|
||||
info!(
|
||||
"[ChatHandler] Looking up Agent by project_id: project_id={}",
|
||||
project_id
|
||||
);
|
||||
AGENT_REGISTRY.get_agent_info(&project_id)
|
||||
});
|
||||
|
||||
// ========== 步骤2: 检查 Agent Busy 状态,如果忙则取消当前任务 ==========
|
||||
use crate::model::AgentStatus;
|
||||
if let Some(agent_info) = agent_info_ref {
|
||||
if agent_info.status == AgentStatus::Active || agent_info.status == AgentStatus::Pending {
|
||||
info!(
|
||||
"[ChatHandler] Agent Busy, cancelling current task: project_id={}, status={:?}, session_id={:?}",
|
||||
project_id, agent_info.status, session_id
|
||||
);
|
||||
|
||||
// 获取 cancel_tx 和 session_id,并释放 DashMap 读锁(防死锁)
|
||||
let cancel_tx = agent_info.cancel_tx.clone();
|
||||
// 优先使用请求中的 session_id,如果为空则从 agent_info 中获取
|
||||
let actual_session_id = session_id
|
||||
.clone()
|
||||
.filter(|s| !s.is_empty())
|
||||
.unwrap_or_else(|| agent_info.session_id.to_string());
|
||||
drop(agent_info);
|
||||
|
||||
// 取消当前任务
|
||||
if let Err(cancel_error) =
|
||||
cancel_current_task(&cancel_tx, &actual_session_id, &project_id).await
|
||||
{
|
||||
// 取消失败,返回错误
|
||||
error!(
|
||||
"[ChatHandler] Failed to cancel current task: project_id={}, error={:?}",
|
||||
project_id, cancel_error
|
||||
);
|
||||
return cancel_error;
|
||||
}
|
||||
|
||||
info!(
|
||||
"[ChatHandler] Current task cancelled, proceeding with new request: project_id={}",
|
||||
project_id
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ========== 步骤3: 创建 PendingGuard(RAII)==========
|
||||
// 自动在作用域结束时清理,避免状态泄漏
|
||||
let pending_guard = PendingGuard::new(&AGENT_REGISTRY, &project_id);
|
||||
info!(
|
||||
"[ChatHandler] Created PendingGuard: project_id={}",
|
||||
project_id
|
||||
);
|
||||
|
||||
// ========== 步骤4: 清理无效 session ==========
|
||||
// 只在 session 不存在时才清理无效的 session_id
|
||||
if let Some(ref sid) = session_id {
|
||||
let session_exists = AGENT_REGISTRY.contains_session(sid);
|
||||
|
||||
if !session_exists && SESSION_CACHE.remove(sid).is_some() {
|
||||
info!(
|
||||
"[ChatHandler] session does not exist, removing invalid session: session_id={}",
|
||||
sid
|
||||
);
|
||||
} else if session_exists {
|
||||
info!("[ChatHandler] Reusing existing session: session_id={}", sid);
|
||||
}
|
||||
}
|
||||
|
||||
// ========== 步骤5: 获取项目工作目录 ==========
|
||||
let project_dir = input.project_dir.clone();
|
||||
info!(
|
||||
"[ChatHandler] Project working directory: {:?}, service_type={:?}",
|
||||
project_dir, input.service_type
|
||||
);
|
||||
|
||||
// 确保目录存在
|
||||
if !project_dir.exists() {
|
||||
if let Err(e) = tokio::fs::create_dir_all(&project_dir).await {
|
||||
error!("[ChatHandler] Failed to create project directory: {}", e);
|
||||
return ChatHandlerOutput::error(
|
||||
project_id,
|
||||
session_id.unwrap_or_default(),
|
||||
format!(
|
||||
"{}: {}",
|
||||
error_codes::get_i18n_message_default("error.create_project_dir_failed"),
|
||||
e
|
||||
),
|
||||
error_codes::ERR_INTERNAL_SERVER_ERROR.to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ========== 步骤6: 构建 ChatPrompt 和 PromptMessage ==========
|
||||
let chat_prompt = match ChatPromptBuilder::default()
|
||||
.project_id(project_id.clone())
|
||||
.project_path(project_dir)
|
||||
.session_id(session_id.clone())
|
||||
.prompt(input.prompt)
|
||||
.attachments(input.attachments)
|
||||
.data_source_attachments(input.data_source_attachments)
|
||||
.service_type(input.service_type)
|
||||
.request_id(request_id.clone())
|
||||
.system_prompt_override(input.system_prompt_override)
|
||||
.user_prompt_template_override(input.user_prompt_template_override)
|
||||
.agent_config_override(input.agent_config_override)
|
||||
.build()
|
||||
{
|
||||
Ok(prompt) => prompt,
|
||||
Err(e) => {
|
||||
error!("[ChatHandler] Failed to build ChatPrompt: {}", e);
|
||||
return ChatHandlerOutput::error(
|
||||
project_id,
|
||||
session_id.unwrap_or_default(),
|
||||
format!(
|
||||
"{}: {}",
|
||||
error_codes::get_i18n_message_default("error.build_chat_prompt_failed"),
|
||||
e
|
||||
),
|
||||
error_codes::ERR_INTERNAL_SERVER_ERROR.to_string(),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// 转换为 PromptMessage
|
||||
let prompt_message = agent_abstraction::PromptMessage::from(chat_prompt);
|
||||
|
||||
// ========== 步骤7: 管理 API 密钥配置 ==========
|
||||
let model_provider = input.model_config;
|
||||
|
||||
// 生成唯一的 service UUID(用于 API 密钥管理)
|
||||
let service_uuid = if model_provider.is_some() {
|
||||
Some(uuid::Uuid::new_v4().to_string())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// 存储 API 配置到共享 DashMap
|
||||
if let (Some(ref provider), Some(ref service_uuid_ref)) =
|
||||
(model_provider.as_ref(), service_uuid.as_ref())
|
||||
{
|
||||
debug!(
|
||||
"[ChatHandler] 使用模型配置: provider={}, model={}, base_url={}, api_protocol={:?}, requires_openai_auth={}, service_uuid={}",
|
||||
provider.name,
|
||||
provider.default_model,
|
||||
provider.base_url,
|
||||
provider.api_protocol,
|
||||
provider.requires_openai_auth,
|
||||
service_uuid_ref
|
||||
);
|
||||
|
||||
// 存储 ModelProviderConfig 到共享 DashMap(使用 UUID 作为 key)
|
||||
context
|
||||
.shared_api_key_manager
|
||||
.insert(service_uuid_ref.to_string(), (*provider).clone());
|
||||
|
||||
// 存储 project_id -> UUID 映射(用于后续清理时查找)
|
||||
context
|
||||
.project_uuid_map
|
||||
.insert(project_id.clone(), service_uuid_ref.to_string());
|
||||
|
||||
info!(
|
||||
"[ChatHandler] Stored API config: service_uuid={}, provider_name={}, base_url={}",
|
||||
service_uuid_ref,
|
||||
provider.name,
|
||||
shared_types::mask_url(&provider.base_url)
|
||||
);
|
||||
} else {
|
||||
warn!("[ChatHandler] No model config provided; falling back to env vars or defaults");
|
||||
}
|
||||
|
||||
// ========== 步骤8: 检查 Worker 状态 ==========
|
||||
match context.agent_runtime.state() {
|
||||
WorkerState::Running => {
|
||||
// 正常操作,继续处理
|
||||
}
|
||||
WorkerState::Starting => {
|
||||
warn!("[ChatHandler] Agent Worker is starting; request may be delayed");
|
||||
}
|
||||
WorkerState::Stopping | WorkerState::Stopped => {
|
||||
// PendingGuard 自动清理(在 drop 时)
|
||||
error!("[ChatHandler] Agent Worker unavailable");
|
||||
return ChatHandlerOutput::error(
|
||||
project_id,
|
||||
session_id.unwrap_or_default(),
|
||||
error_codes::get_i18n_message_default("error.agent_worker_unavailable"),
|
||||
error_codes::ERR_SERVICE_UNAVAILABLE.to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ========== 步骤9: 发送任务到 Agent Worker ==========
|
||||
// 创建请求并设置 UUID 和密钥管理器
|
||||
let (agent_request, chat_prompt_rx) = AgentRequest::new(prompt_message, model_provider);
|
||||
let agent_request = agent_request
|
||||
.with_service_uuid(service_uuid)
|
||||
.with_key_manager(Some(context.shared_api_key_manager.clone()))
|
||||
.with_skip_slot_limit(input.skip_slot_limit);
|
||||
|
||||
if let Err(e) = context.agent_runtime.send(agent_request).await {
|
||||
// PendingGuard 自动清理(在 drop 时)
|
||||
error!("[ChatHandler] Failed to send task: {}", e);
|
||||
return ChatHandlerOutput::error(
|
||||
project_id,
|
||||
session_id.unwrap_or_default(),
|
||||
format!(
|
||||
"{}: {}",
|
||||
error_codes::get_i18n_message_default("error.send_task_failed"),
|
||||
e
|
||||
),
|
||||
error_codes::ERR_INTERNAL_SERVER_ERROR.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
// ========== 步骤10: 等待响应(5 分钟超时)==========
|
||||
match tokio::time::timeout(std::time::Duration::from_secs(300), chat_prompt_rx).await {
|
||||
Ok(Ok(response)) => {
|
||||
let output = ChatHandlerOutput {
|
||||
project_id: response.project_id,
|
||||
session_id: response.session_id,
|
||||
success: response.error.is_none(),
|
||||
error: response.error,
|
||||
error_code: if response.code != error_codes::SUCCESS {
|
||||
Some(response.code)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
request_id: Some(request_id),
|
||||
need_fallback: false,
|
||||
fallback_reason: None,
|
||||
};
|
||||
|
||||
info!(
|
||||
"[ChatHandler] Chat completed: success={}, session_id={}",
|
||||
output.success, output.session_id
|
||||
);
|
||||
|
||||
// 只有请求成功时才提交 PendingGuard 保留 Pending 状态
|
||||
// 失败时 PendingGuard 自动 drop 清理,允许下次请求重新创建 Agent
|
||||
if output.success {
|
||||
pending_guard.commit_success();
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
// PendingGuard 自动清理(在 drop 时)
|
||||
error!("[ChatHandler] Chat response channel dropped: {}", e);
|
||||
ChatHandlerOutput::error(
|
||||
project_id,
|
||||
session_id.unwrap_or_default(),
|
||||
format!(
|
||||
"{}: {}",
|
||||
error_codes::get_i18n_message_default("error.request_processing_failed"),
|
||||
e
|
||||
),
|
||||
error_codes::ERR_INTERNAL_SERVER_ERROR.to_string(),
|
||||
)
|
||||
}
|
||||
Err(_elapsed) => {
|
||||
// PendingGuard 自动清理(在 drop 时)
|
||||
error!("[ChatHandler] ⏰ Chat request timeout (300s): project_id={}", project_id);
|
||||
ChatHandlerOutput::error(
|
||||
project_id,
|
||||
session_id.unwrap_or_default(),
|
||||
error_codes::get_i18n_message_default("error.request_processing_failed")
|
||||
+ ": request timeout (300s)",
|
||||
error_codes::ERR_INTERNAL_SERVER_ERROR.to_string(),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,351 @@
|
||||
//! Agent Runner 本地服务实现
|
||||
//!
|
||||
//! 实现 `AgentHttpService` trait,直接调用本地服务(AGENT_REGISTRY, SESSION_CACHE, handle_chat_core)
|
||||
//! 适用于 Agent Runner 的 HTTP Server 模式
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use dashmap::DashMap;
|
||||
use agent_client_protocol::schema::{CancelNotification, SessionId};
|
||||
use shared_types::{
|
||||
agent_http_service::AgentHttpService,
|
||||
rcoder_agent_types::{
|
||||
RcoderAgentCancelRequest, RcoderAgentCancelResponse, RcoderAgentStopRequest,
|
||||
RcoderAgentStopResponse,
|
||||
},
|
||||
AgentStatusResponse, ChatResponse, HttpResult, RcoderChatRequest, ServiceType,
|
||||
};
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::AgentRuntime;
|
||||
use crate::service::{
|
||||
chat_handler::{ChatHandlerContext, ChatHandlerInput},
|
||||
handle_chat_core, AGENT_REGISTRY, SESSION_CACHE, SessionData,
|
||||
};
|
||||
|
||||
/// Agent Runner 本地服务实现
|
||||
///
|
||||
/// 直接调用本地 AGENT_REGISTRY、SESSION_CACHE、handle_chat_core()
|
||||
/// 不需要 gRPC 转发
|
||||
pub struct LocalAgentHttpService {
|
||||
/// Agent 运行时
|
||||
pub agent_runtime: Arc<AgentRuntime>,
|
||||
/// 共享 API Key 管理器
|
||||
pub shared_api_key_manager: Arc<DashMap<String, shared_types::ModelProviderConfig>>,
|
||||
/// project_id -> UUID 映射
|
||||
pub project_uuid_map: Arc<DashMap<String, String>>,
|
||||
/// 项目工作目录根路径
|
||||
pub projects_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl LocalAgentHttpService {
|
||||
/// 创建新的 LocalAgentHttpService
|
||||
pub fn new(
|
||||
agent_runtime: Arc<AgentRuntime>,
|
||||
shared_api_key_manager: Arc<DashMap<String, shared_types::ModelProviderConfig>>,
|
||||
project_uuid_map: Arc<DashMap<String, String>>,
|
||||
projects_dir: PathBuf,
|
||||
) -> Self {
|
||||
Self {
|
||||
agent_runtime,
|
||||
shared_api_key_manager,
|
||||
project_uuid_map,
|
||||
projects_dir,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHttpService for LocalAgentHttpService {
|
||||
/// Chat 对话请求
|
||||
async fn chat(&self, request: RcoderChatRequest) -> HttpResult<ChatResponse> {
|
||||
// 0. 验证 prompt 不为空
|
||||
if request.prompt.is_empty() {
|
||||
warn!("[LocalAgent] Empty prompt received");
|
||||
return HttpResult::error(
|
||||
shared_types::error_codes::ERR_VALIDATION,
|
||||
"Prompt cannot be empty",
|
||||
);
|
||||
}
|
||||
|
||||
info!(
|
||||
"📨 [LocalAgent] Chat request: prompt_len={}, project_id={:?}, session_id={:?}",
|
||||
request.prompt.len(),
|
||||
request.project_id,
|
||||
request.session_id
|
||||
);
|
||||
|
||||
// 1. 生成 project_id(如果未提供)
|
||||
let project_id = request
|
||||
.project_id
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string().replace('-', ""));
|
||||
|
||||
// 2. 自动查找现有 session_id(如果未提供)
|
||||
let session_id = request.session_id.or_else(|| {
|
||||
AGENT_REGISTRY
|
||||
.get_agent_info(&project_id)
|
||||
.map(|info| info.session_id.to_string())
|
||||
});
|
||||
|
||||
// 3. 创建项目工作目录
|
||||
let project_dir = self.projects_dir.join(&project_id);
|
||||
if let Err(e) = tokio::fs::create_dir_all(&project_dir).await {
|
||||
let error_msg = format!("Failed to create project dir: {}", e);
|
||||
warn!("[LocalAgent] {}", error_msg);
|
||||
return HttpResult::error(shared_types::error_codes::ERR_INTERNAL_SERVER_ERROR, &error_msg);
|
||||
}
|
||||
|
||||
// 4. 生成 request_id(如果未提供)
|
||||
let request_id = request
|
||||
.request_id
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
||||
|
||||
// 5. 构建 ChatHandlerInput
|
||||
let input = ChatHandlerInput {
|
||||
project_id: project_id.clone(),
|
||||
project_dir,
|
||||
session_id,
|
||||
prompt: request.prompt,
|
||||
request_id: request_id.clone(),
|
||||
attachments: request.attachments,
|
||||
data_source_attachments: request.data_source_attachments,
|
||||
model_config: request.model_provider,
|
||||
service_type: ServiceType::RCoder,
|
||||
agent_config_override: request.agent_config,
|
||||
system_prompt_override: request.system_prompt,
|
||||
user_prompt_template_override: request.user_prompt,
|
||||
skip_slot_limit: true, // HTTP Server 部署,跳过槽位限制
|
||||
};
|
||||
|
||||
// 6. 构建 ChatHandlerContext
|
||||
let context = ChatHandlerContext {
|
||||
agent_runtime: self.agent_runtime.clone(),
|
||||
shared_api_key_manager: self.shared_api_key_manager.clone(),
|
||||
project_uuid_map: self.project_uuid_map.clone(),
|
||||
};
|
||||
|
||||
// 7. 调用 handle_chat_core
|
||||
let output = handle_chat_core(input, &context).await;
|
||||
|
||||
// 8. 将 session 写入 SESSION_CACHE(SSE 进度流需要)
|
||||
let session_id_str = output.session_id.clone();
|
||||
if SESSION_CACHE.get(&session_id_str).is_none() {
|
||||
let session_data = SessionData::new(1000);
|
||||
SESSION_CACHE.insert(session_id_str, session_data);
|
||||
}
|
||||
|
||||
// 9. 构建响应
|
||||
if output.error.is_some() || !output.success {
|
||||
let error_code = output.error_code.as_deref().unwrap_or(shared_types::error_codes::ERR_INTERNAL_SERVER_ERROR);
|
||||
let error_msg = output.error.unwrap_or_else(|| "Unknown error".to_string());
|
||||
HttpResult::error(error_code, &error_msg)
|
||||
} else {
|
||||
HttpResult::success(ChatResponse {
|
||||
project_id: output.project_id,
|
||||
session_id: output.session_id,
|
||||
error: output.error,
|
||||
request_id: Some(request_id),
|
||||
need_fallback: None,
|
||||
fallback_reason: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 查询 Agent 状态
|
||||
async fn get_status(&self, project_id: &str) -> HttpResult<AgentStatusResponse> {
|
||||
info!(
|
||||
"🔍 [LocalAgent] Status query: project_id={}",
|
||||
project_id
|
||||
);
|
||||
|
||||
if let Some(info) = AGENT_REGISTRY.get_agent_info(project_id) {
|
||||
// Agent 存在且活跃
|
||||
let response = AgentStatusResponse {
|
||||
project_id: project_id.to_string(),
|
||||
is_alive: true,
|
||||
session_id: Some(info.session_id.to_string()),
|
||||
status: Some(info.status.clone()),
|
||||
last_activity: Some(info.last_activity),
|
||||
created_at: Some(info.created_at),
|
||||
model_provider: None, // AgentRegistry 不存储 model_provider
|
||||
};
|
||||
|
||||
info!(
|
||||
"✅ [LocalAgent] Agent status: project_id={}, is_alive=true, status={:?}",
|
||||
project_id,
|
||||
info.status
|
||||
);
|
||||
|
||||
HttpResult::success(response)
|
||||
} else {
|
||||
// Agent 不存在
|
||||
info!(
|
||||
"📭 [LocalAgent] Agent not found: project_id={}",
|
||||
project_id
|
||||
);
|
||||
|
||||
let response = AgentStatusResponse {
|
||||
project_id: project_id.to_string(),
|
||||
is_alive: false,
|
||||
session_id: None,
|
||||
status: None,
|
||||
last_activity: None,
|
||||
created_at: None,
|
||||
model_provider: None,
|
||||
};
|
||||
|
||||
HttpResult::success(response)
|
||||
}
|
||||
}
|
||||
|
||||
/// 停止 Agent
|
||||
async fn stop(&self, request: RcoderAgentStopRequest) -> HttpResult<RcoderAgentStopResponse> {
|
||||
info!(
|
||||
"🛑 [LocalAgent] Stop request: project_id={}",
|
||||
request.project_id
|
||||
);
|
||||
|
||||
// 1. 获取 Agent 信息
|
||||
if let Some(agent_info) = AGENT_REGISTRY.get_agent_info(&request.project_id) {
|
||||
let session_id = agent_info.session_id.to_string();
|
||||
let cancel_tx = agent_info.cancel_tx.clone();
|
||||
|
||||
// 释放读锁
|
||||
drop(agent_info);
|
||||
|
||||
// 2. 发送取消信号(如果 channel 仍然打开)
|
||||
if !cancel_tx.is_closed() {
|
||||
let session_id_obj = SessionId::new(Arc::from(session_id.as_str()));
|
||||
let cancel_notification = CancelNotification::new(session_id_obj);
|
||||
|
||||
let (result_tx, _result_rx) = oneshot::channel();
|
||||
let cancel_request = shared_types::CancelNotificationRequestWrapper {
|
||||
cancel_notification,
|
||||
result_tx,
|
||||
};
|
||||
|
||||
match cancel_tx.send(cancel_request).await {
|
||||
Ok(_) => {
|
||||
info!("[LocalAgent] Cancel signal sent: session_id={}", session_id);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [LocalAgent] Failed to send cancel signal: session_id={}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 从 AGENT_REGISTRY 移除 Agent
|
||||
AGENT_REGISTRY.remove_by_project(&request.project_id);
|
||||
|
||||
info!(
|
||||
"✅ [LocalAgent] Agent stopped: project_id={}",
|
||||
request.project_id
|
||||
);
|
||||
|
||||
HttpResult::success(RcoderAgentStopResponse {
|
||||
success: true,
|
||||
project_id: request.project_id,
|
||||
session_id: Some(session_id),
|
||||
message: "Agent stopped successfully".to_string(),
|
||||
})
|
||||
} else {
|
||||
// Agent 不存在,幂等返回成功
|
||||
info!(
|
||||
"ℹ️ [LocalAgent] Agent not found, returning success idempotently: project_id={}",
|
||||
request.project_id
|
||||
);
|
||||
|
||||
HttpResult::success(RcoderAgentStopResponse {
|
||||
success: true,
|
||||
project_id: request.project_id,
|
||||
session_id: None,
|
||||
message: "Agent not found".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 取消正在执行的任务
|
||||
async fn cancel(&self, request: RcoderAgentCancelRequest) -> HttpResult<RcoderAgentCancelResponse> {
|
||||
info!(
|
||||
"🚫 [LocalAgent] Cancel request: project_id={}, session_id={:?}",
|
||||
request.project_id, request.session_id
|
||||
);
|
||||
|
||||
// 1. 查找 session_id(如果未提供,从 AGENT_REGISTRY 获取)
|
||||
let session_id = if let Some(sid) = request.session_id {
|
||||
sid
|
||||
} else {
|
||||
match AGENT_REGISTRY.get_agent_info(&request.project_id) {
|
||||
Some(info) => info.session_id.to_string(),
|
||||
None => {
|
||||
// Agent 不存在,幂等返回成功
|
||||
info!(
|
||||
"ℹ️ [LocalAgent] Agent not found, returning success idempotently: project_id={}",
|
||||
request.project_id
|
||||
);
|
||||
return HttpResult::success(RcoderAgentCancelResponse {
|
||||
success: true,
|
||||
session_id: String::new(),
|
||||
});
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// 2. 获取 Agent 信息并发送取消信号
|
||||
if let Some(agent_info) = AGENT_REGISTRY.get_agent_info(&request.project_id) {
|
||||
let cancel_tx = agent_info.cancel_tx.clone();
|
||||
|
||||
// 释放读锁
|
||||
drop(agent_info);
|
||||
|
||||
// 检查是否已经空闲或停止中(幂等性)
|
||||
if cancel_tx.is_closed() {
|
||||
info!(
|
||||
"ℹ️ [LocalAgent] Agent stopped, cancel channel is closed: session_id={}",
|
||||
session_id
|
||||
);
|
||||
} else {
|
||||
// 创建取消通知
|
||||
let session_id_obj = SessionId::new(Arc::from(session_id.as_str()));
|
||||
let cancel_notification = CancelNotification::new(session_id_obj);
|
||||
|
||||
// 创建 oneshot channel 接收取消结果(HTTP 不等待结果,直接丢弃)
|
||||
let (result_tx, _result_rx) = oneshot::channel();
|
||||
let cancel_request = shared_types::CancelNotificationRequestWrapper {
|
||||
cancel_notification,
|
||||
result_tx,
|
||||
};
|
||||
|
||||
// 发送取消信号(异步)
|
||||
match cancel_tx.send(cancel_request).await {
|
||||
Ok(_) => {
|
||||
info!("[LocalAgent] Cancel signal sent: session_id={}", session_id);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [LocalAgent] Failed to send cancel signal: session_id={}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Session 不存在,幂等返回成功
|
||||
info!(
|
||||
"ℹ️ [LocalAgent] Agent not found, returning success idempotently: session_id={}",
|
||||
session_id
|
||||
);
|
||||
}
|
||||
|
||||
HttpResult::success(RcoderAgentCancelResponse {
|
||||
success: true,
|
||||
session_id,
|
||||
})
|
||||
}
|
||||
}
|
||||
11
qiming-rcoder/crates/agent_runner/src/service/mod.rs
Normal file
11
qiming-rcoder/crates/agent_runner/src/service/mod.rs
Normal file
@@ -0,0 +1,11 @@
|
||||
pub mod agent_registry;
|
||||
pub mod chat_handler;
|
||||
pub mod local_agent_service;
|
||||
mod session_cache;
|
||||
mod session_notifier;
|
||||
mod state_aware_notifier;
|
||||
|
||||
pub use agent_registry::{AGENT_REGISTRY, AgentSessionRegistry, PendingGuard};
|
||||
pub use chat_handler::{ChatHandlerContext, ChatHandlerInput, ChatHandlerOutput, handle_chat_core};
|
||||
pub use session_cache::{SESSION_CACHE, SessionData, push_session_update_with_project};
|
||||
pub use state_aware_notifier::StateAwareNotifier;
|
||||
463
qiming-rcoder/crates/agent_runner/src/service/session_cache.rs
Normal file
463
qiming-rcoder/crates/agent_runner/src/service/session_cache.rs
Normal file
@@ -0,0 +1,463 @@
|
||||
//! 全局Session缓存模块
|
||||
//!
|
||||
//! 使用LazyLock初始化全局DashMap,按session_id分组缓存统一会话消息到ringbuf循环缓冲区
|
||||
|
||||
use crate::{SessionNotify, UnifiedSessionMessage};
|
||||
use anyhow::Result;
|
||||
use dashmap::DashMap;
|
||||
use ringbuf::HeapRb;
|
||||
use ringbuf::traits::{Consumer, Observer, Producer, Split};
|
||||
use std::sync::{Arc, LazyLock};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use super::AGENT_REGISTRY;
|
||||
|
||||
/// 日志截取的最大长度默认值
|
||||
const MAX_LOG_TRUNCATE_LEN: usize = 50;
|
||||
|
||||
/// 截取消息内容用于日志打印(防止日志膨胀)
|
||||
///
|
||||
/// 社区常见做法:
|
||||
/// - tracing: 通过 Subscriber 配置限制字段大小
|
||||
/// - serde: 自定义 Serialize 实现截断
|
||||
/// - 简单场景: chars().take() + 长度检查(本实现)
|
||||
fn truncate_message_for_log(data: &serde_json::Value, max_len: usize) -> String {
|
||||
// 边界检查:max_len 为 0 时返回空,避免无效计算
|
||||
if max_len == 0 {
|
||||
return String::new();
|
||||
}
|
||||
|
||||
// 优先提取原始字符串内容,避免 JSON 序列化后的引号包裹
|
||||
let s = match data.as_str() {
|
||||
Some(inner) => inner.to_string(),
|
||||
None => data.to_string(),
|
||||
};
|
||||
|
||||
// chars().count() 是一次性遍历,可以与 take() 合并优化但牺牲可读性
|
||||
// 当前实现清晰且性能可接受(短字符串场景)
|
||||
let char_count = s.chars().count();
|
||||
if char_count <= max_len {
|
||||
return s;
|
||||
}
|
||||
|
||||
// 使用 chars() 安全截取 UTF-8 字符边界
|
||||
let truncated: String = s.chars().take(max_len).collect();
|
||||
format!("{}... (truncated)", truncated)
|
||||
}
|
||||
|
||||
/// 全局Session缓存 - LazyLock初始化
|
||||
pub static SESSION_CACHE: LazyLock<DashMap<String, Arc<SessionData>>> = LazyLock::new(DashMap::new);
|
||||
|
||||
/// Session数据包装 - 极简版本,专注消息传输
|
||||
pub struct SessionData {
|
||||
command_tx: mpsc::UnboundedSender<SessionCommand>,
|
||||
// 🎯 极简优化:直接存储当前连接,无需命令传递
|
||||
current_sender: Arc<tokio::sync::Mutex<Option<mpsc::Sender<UnifiedSessionMessage>>>>,
|
||||
current_cancel: Arc<tokio::sync::Mutex<Option<CancellationToken>>>,
|
||||
}
|
||||
|
||||
impl SessionData {
|
||||
pub fn new(max_size: usize) -> Arc<Self> {
|
||||
let start_time = std::time::Instant::now();
|
||||
debug!(
|
||||
"⏱️ [SessionData::new] Starting creation, max_size={}",
|
||||
max_size
|
||||
);
|
||||
|
||||
let channel_start = std::time::Instant::now();
|
||||
let (command_tx, command_rx) = mpsc::unbounded_channel();
|
||||
debug!(
|
||||
"⏱️ [SessionData::new] Channel creation took: {:?}",
|
||||
channel_start.elapsed()
|
||||
);
|
||||
|
||||
let arc_start = std::time::Instant::now();
|
||||
let session = Arc::new(SessionData {
|
||||
command_tx,
|
||||
current_sender: Arc::new(tokio::sync::Mutex::new(None)),
|
||||
current_cancel: Arc::new(tokio::sync::Mutex::new(None)),
|
||||
});
|
||||
debug!(
|
||||
"⏱️ [SessionData::new] Arc creation took: {:?}",
|
||||
arc_start.elapsed()
|
||||
);
|
||||
|
||||
let spawn_start = std::time::Instant::now();
|
||||
SessionWorker::spawn(
|
||||
max_size,
|
||||
command_rx,
|
||||
session.current_sender.clone(),
|
||||
session.current_cancel.clone(),
|
||||
);
|
||||
debug!(
|
||||
"⏱️ [SessionData::new] SessionWorker::spawn took: {:?}",
|
||||
spawn_start.elapsed()
|
||||
);
|
||||
|
||||
debug!(
|
||||
"⏱️ [SessionData::new] Total creation took: {:?}",
|
||||
start_time.elapsed()
|
||||
);
|
||||
session
|
||||
}
|
||||
|
||||
pub async fn message_count(&self) -> usize {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
if self
|
||||
.command_tx
|
||||
.send(SessionCommand::MessageCount { ack: tx })
|
||||
.is_err()
|
||||
{
|
||||
warn!("Failed to send message_count command; worker has exited");
|
||||
return 0;
|
||||
}
|
||||
rx.await.unwrap_or(0)
|
||||
}
|
||||
|
||||
pub async fn create_new_connection(
|
||||
&self,
|
||||
buffer_size: usize,
|
||||
) -> Result<(mpsc::Receiver<UnifiedSessionMessage>, CancellationToken)> {
|
||||
let start_time = std::time::Instant::now();
|
||||
debug!(
|
||||
"⏱️ [create_new_connection] Starting connection creation, buffer_size={}",
|
||||
buffer_size
|
||||
);
|
||||
|
||||
let token_start = std::time::Instant::now();
|
||||
let cancellation_token = CancellationToken::new();
|
||||
debug!(
|
||||
"⏱️ [create_new_connection] CancellationToken creation took: {:?}",
|
||||
token_start.elapsed()
|
||||
);
|
||||
|
||||
let channel_start = std::time::Instant::now();
|
||||
let (tx, rx) = mpsc::channel(buffer_size);
|
||||
debug!(
|
||||
"⏱️ [create_new_connection] mpsc channel creation took: {:?}",
|
||||
channel_start.elapsed()
|
||||
);
|
||||
|
||||
let setup_start = std::time::Instant::now();
|
||||
// 🛡️ 关键修复:使用 lock() 而非 try_lock(),确保连接一定被设置
|
||||
// try_lock() 可能失败导致 current_sender 未设置,造成消息丢失
|
||||
{
|
||||
// 取消之前的连接
|
||||
let mut current_cancel_guard = self.current_cancel.lock().await;
|
||||
if let Some(token) = current_cancel_guard.take() {
|
||||
token.cancel();
|
||||
}
|
||||
// 设置新的取消令牌
|
||||
*current_cancel_guard = Some(cancellation_token.clone());
|
||||
|
||||
// 设置新的发送器
|
||||
let mut current_sender_guard = self.current_sender.lock().await;
|
||||
*current_sender_guard = Some(tx);
|
||||
}
|
||||
debug!(
|
||||
"⏱️ [create_new_connection] Connection state setup took: {:?}",
|
||||
setup_start.elapsed()
|
||||
);
|
||||
|
||||
debug!(
|
||||
"⏱️ [create_new_connection] Total connection creation took: {:?}",
|
||||
start_time.elapsed()
|
||||
);
|
||||
Ok((rx, cancellation_token))
|
||||
}
|
||||
|
||||
pub fn push_message(&self, message: UnifiedSessionMessage) {
|
||||
if self
|
||||
.command_tx
|
||||
.send(SessionCommand::Push { message })
|
||||
.is_err()
|
||||
{
|
||||
warn!("Failed to push message; worker has exited");
|
||||
}
|
||||
}
|
||||
|
||||
/// 主动关闭当前 SSE 连接
|
||||
///
|
||||
/// 当用户取消任务时,需要主动关闭 SSE 连接,而不是让客户端一直等待
|
||||
///
|
||||
/// 关闭机制:
|
||||
/// 1. 触发 CancellationToken,让 SSE 流立即退出循环
|
||||
/// 2. 显式关闭 channel 发送端,让 rx.recv() 立即返回 None
|
||||
/// 3. 清空连接状态,防止新的消息被发送
|
||||
pub async fn close_current_connection(&self) {
|
||||
// 🎯 主动触发取消令牌,关闭 SSE 连接
|
||||
let mut current_cancel_guard = self.current_cancel.lock().await;
|
||||
if let Some(token) = current_cancel_guard.take() {
|
||||
info!("🔌 [SessionData] Triggering CancellationToken to close SSE connection");
|
||||
token.cancel();
|
||||
}
|
||||
drop(current_cancel_guard);
|
||||
|
||||
// 🎯 显式关闭 channel 发送端,让接收端立即感知到连接关闭
|
||||
let mut current_sender_guard = self.current_sender.lock().await;
|
||||
if current_sender_guard.take().is_some() {
|
||||
info!(
|
||||
"🔌 [SessionData] Explicitly closed channel sender; receiver disconnects immediately"
|
||||
);
|
||||
// 当 Sender 被 drop 时,Receiver 的 recv() 会返回 None
|
||||
// 这里通过 take() 将 sender 从 Option 中移除,触发 drop
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct SessionWorker {
|
||||
max_size: usize,
|
||||
command_rx: mpsc::UnboundedReceiver<SessionCommand>,
|
||||
// 🎯 极简优化:直接共享连接状态,无需命令传递
|
||||
current_sender: Arc<tokio::sync::Mutex<Option<mpsc::Sender<UnifiedSessionMessage>>>>,
|
||||
current_cancel: Arc<tokio::sync::Mutex<Option<CancellationToken>>>,
|
||||
}
|
||||
|
||||
impl SessionWorker {
|
||||
fn spawn(
|
||||
max_size: usize,
|
||||
command_rx: mpsc::UnboundedReceiver<SessionCommand>,
|
||||
current_sender: Arc<tokio::sync::Mutex<Option<mpsc::Sender<UnifiedSessionMessage>>>>,
|
||||
current_cancel: Arc<tokio::sync::Mutex<Option<CancellationToken>>>,
|
||||
) {
|
||||
let start_time = std::time::Instant::now();
|
||||
debug!(
|
||||
"⏱️ [SessionWorker::spawn] Starting SessionWorker creation, max_size={}",
|
||||
max_size
|
||||
);
|
||||
|
||||
let worker = SessionWorker {
|
||||
max_size,
|
||||
command_rx,
|
||||
current_sender,
|
||||
current_cancel,
|
||||
};
|
||||
|
||||
let spawn_start = std::time::Instant::now();
|
||||
tokio::spawn(worker.run());
|
||||
debug!(
|
||||
"⏱️ [SessionWorker::spawn] tokio::spawn took: {:?}",
|
||||
spawn_start.elapsed()
|
||||
);
|
||||
debug!(
|
||||
"⏱️ [SessionWorker::spawn] Total spawn took: {:?}",
|
||||
start_time.elapsed()
|
||||
);
|
||||
}
|
||||
|
||||
async fn run(mut self) {
|
||||
let (mut producer, mut consumer) = HeapRb::new(self.max_size).split();
|
||||
let mut buffered_len = 0usize;
|
||||
|
||||
while let Some(cmd) = self.command_rx.recv().await {
|
||||
match cmd {
|
||||
SessionCommand::Push { message } => {
|
||||
let should_buffer = !matches!(
|
||||
message.message_type,
|
||||
crate::model::SessionMessageType::Heartbeat
|
||||
);
|
||||
|
||||
if should_buffer {
|
||||
if producer.is_full() {
|
||||
let _ = consumer.try_pop();
|
||||
buffered_len = buffered_len.saturating_sub(1);
|
||||
}
|
||||
if producer.try_push(message.clone()).is_ok() {
|
||||
buffered_len += 1;
|
||||
} else {
|
||||
warn!("Ring buffer push failed; real-time delivery only");
|
||||
}
|
||||
}
|
||||
|
||||
// 🛡️ 关键修复:使用 lock().await 确保消息一定被发送
|
||||
// try_lock() 可能失败导致消息丢失,造成 SSE 卡死
|
||||
let mut current_sender_guard = self.current_sender.lock().await;
|
||||
if let Some(sender) = current_sender_guard.as_mut() {
|
||||
if sender.try_send(message.clone()).is_err() {
|
||||
// 如果发送失败,可能是缓冲区满了或连接已关闭
|
||||
// 注意:truncate 在锁内执行,但 50 字符开销可忽略
|
||||
warn!(
|
||||
"⚠️ SSE sender send failed, disabling real-time delivery: message_type={:?}, sub_type={}, data={}",
|
||||
message.message_type,
|
||||
message.sub_type,
|
||||
truncate_message_for_log(&message.data, MAX_LOG_TRUNCATE_LEN)
|
||||
);
|
||||
*current_sender_guard = None;
|
||||
}
|
||||
} else {
|
||||
// 连接不存在,跳过实时推送(记录为 info 级别,便于排查问题)
|
||||
info!(
|
||||
"📭 SSE sender missing, skipping real-time delivery (message buffered in ring buffer): message_type={:?}, sub_type={}, data={}",
|
||||
message.message_type,
|
||||
message.sub_type,
|
||||
truncate_message_for_log(&message.data, MAX_LOG_TRUNCATE_LEN)
|
||||
);
|
||||
}
|
||||
}
|
||||
SessionCommand::Clear { ack } => {
|
||||
let mut cleared = 0usize;
|
||||
while consumer.try_pop().is_some() {
|
||||
cleared += 1;
|
||||
}
|
||||
buffered_len = 0;
|
||||
let _ = ack.send(cleared);
|
||||
}
|
||||
SessionCommand::MessageCount { ack } => {
|
||||
let _ = ack.send(buffered_len);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
debug!("🔚 SessionWorker stopped");
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum SessionCommand {
|
||||
Push { message: UnifiedSessionMessage },
|
||||
Clear { ack: oneshot::Sender<usize> },
|
||||
MessageCount { ack: oneshot::Sender<usize> },
|
||||
}
|
||||
|
||||
/// 便捷函数:添加SessionNotify消息(自动转换为统一格式)
|
||||
///
|
||||
/// 如果 SESSION_CACHE 中不存在该 session_id 的条目,会自动创建。
|
||||
/// 这解决了 Agent 开始推送消息时 SESSION_CACHE 条目尚未由 HTTP 处理器创建的竞态问题。
|
||||
pub async fn push_session_update(session_id: &str, notify: SessionNotify) -> Result<()> {
|
||||
use dashmap::mapref::entry::Entry;
|
||||
|
||||
// 🛡️ 关键修复:使用 entry API 原子性地获取或创建 SessionData
|
||||
// 之前直接用 get(),如果 session 不存在就丢弃消息。
|
||||
// 竞态场景:Agent 开始推送消息 → push_session_update 查找 SESSION_CACHE → 不存在(因为
|
||||
// handle_chat_core 尚未返回,computer_chat.rs 还未创建条目)→ 消息被丢弃
|
||||
let session_data = {
|
||||
match SESSION_CACHE.entry(session_id.to_string()) {
|
||||
Entry::Occupied(entry) => entry.get().clone(),
|
||||
Entry::Vacant(entry) => {
|
||||
let data = SessionData::new(1000);
|
||||
info!(
|
||||
"📦 [push_session_update] SESSION_CACHE auto-created: session_id={}",
|
||||
session_id
|
||||
);
|
||||
entry.insert(data.clone());
|
||||
data
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let unified_message = notify.to_unified_message();
|
||||
|
||||
debug!(
|
||||
"📥 Pushing message to cache: session_id={}, message_type={:?}, sub_type={}",
|
||||
session_id, unified_message.message_type, unified_message.sub_type
|
||||
);
|
||||
|
||||
session_data.push_message(unified_message);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 便捷函数:添加SessionNotify消息并管理Project-Session映射
|
||||
///
|
||||
/// 这个函数会自动确保project_id只对应一个活跃的session_id
|
||||
///
|
||||
/// 这个函数会自动确保project_id只对应一个活跃的session_id
|
||||
/// 当检测到session_id变化时,会自动清理旧session的数据
|
||||
pub async fn push_session_update_with_project(
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
notify: SessionNotify,
|
||||
) -> Result<()> {
|
||||
// 确保project_id对应正确的session_id,如果变化则清理旧数据
|
||||
let cleared_count = ensure_project_session(project_id, session_id).await;
|
||||
|
||||
if cleared_count > 0 {
|
||||
info!(
|
||||
"📝 [push_session_update_with_project] Session changed, cleaned {} old messages: project_id={}, new_session_id={}",
|
||||
cleared_count, project_id, session_id
|
||||
);
|
||||
}
|
||||
|
||||
// 推送消息到新的session
|
||||
push_session_update(session_id, notify).await
|
||||
}
|
||||
|
||||
/// 确保project_id对应正确的session_id
|
||||
///
|
||||
/// 使用统一的 AGENT_REGISTRY 管理 project-session 映射
|
||||
/// 如果project_id对应的session_id发生变化,会自动清理旧session的数据
|
||||
/// 如果session_id相同,则不做任何操作
|
||||
///
|
||||
/// 参数:
|
||||
/// - project_id: 项目ID
|
||||
/// - session_id: 当前会话ID
|
||||
///
|
||||
/// 返回值: 如果清理了旧数据则返回清理的消息数量,否则返回0
|
||||
pub async fn ensure_project_session(project_id: &str, session_id: &str) -> usize {
|
||||
// 使用统一 Registry 检查当前映射
|
||||
let mapped_session_id = AGENT_REGISTRY.get_session_by_project(project_id);
|
||||
|
||||
match mapped_session_id {
|
||||
Some(mapped_sid) if mapped_sid == session_id => {
|
||||
// session_id 相同,不需要做任何操作
|
||||
debug!(
|
||||
"📋 Project session mapping unchanged: project_id={}, session_id={}",
|
||||
project_id, session_id
|
||||
);
|
||||
0
|
||||
}
|
||||
Some(old_session_id) => {
|
||||
// session_id 发生变化,需要清理旧 session 的数据
|
||||
info!(
|
||||
"🔄 Detected project session change: project_id={}, old_session_id={}, new_session_id={}",
|
||||
project_id, old_session_id, session_id
|
||||
);
|
||||
|
||||
// 🛡️ 关键修复:先主动关闭旧 session 的 SSE 连接,再移除缓存
|
||||
// 之前直接 remove 导致旧 SSE 连接的心跳流继续发送但不再收到业务消息,
|
||||
// 前端如果没有及时关闭旧连接,会看到孤立的心跳流
|
||||
let cleared_count = if let Some((_, old_session_data)) =
|
||||
SESSION_CACHE.remove(&old_session_id)
|
||||
{
|
||||
old_session_data.close_current_connection().await;
|
||||
info!(
|
||||
"🔌 [ensure_project_session] Closed old session SSE connection: old_session_id={}",
|
||||
old_session_id
|
||||
);
|
||||
1 // 移除了1个session
|
||||
} else {
|
||||
0 // session不存在
|
||||
};
|
||||
|
||||
// 更新 AGENT_REGISTRY 中的映射关系
|
||||
let _ = AGENT_REGISTRY.update_session(project_id, session_id);
|
||||
|
||||
if cleared_count > 0 {
|
||||
info!(
|
||||
"🧹 Cleared old session data and updated mapping: project_id={}, old_session_id={}, new_session_id={}, cleared_count={}",
|
||||
project_id, old_session_id, session_id, cleared_count
|
||||
);
|
||||
} else {
|
||||
info!(
|
||||
"📝 Updated project session mapping: project_id={}, old_session_id={}, new_session_id={}",
|
||||
project_id, old_session_id, session_id
|
||||
);
|
||||
}
|
||||
|
||||
cleared_count
|
||||
}
|
||||
None => {
|
||||
// 第一次建立映射关系(无旧映射)
|
||||
// 注意:此时 AGENT_REGISTRY 中可能还没有这个 project 的记录
|
||||
// 这种情况下不需要调用 update_session,因为 agent 注册时会调用 register
|
||||
info!(
|
||||
"🆕 Project session first seen: project_id={}, session_id={}",
|
||||
project_id, session_id
|
||||
);
|
||||
0
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
//! SessionNotifier 实现
|
||||
//!
|
||||
//! 实现 agent_abstraction 定义的 SessionNotifier trait,
|
||||
//! 用于推送 SSE 消息到前端。
|
||||
|
||||
use agent_abstraction::SessionNotifier;
|
||||
use async_trait::async_trait;
|
||||
use shared_types::{
|
||||
AgentSessionUpdate, SessionNotify, SessionPromptEnd, SessionPromptError, SessionPromptStart,
|
||||
};
|
||||
|
||||
use super::push_session_update_with_project;
|
||||
|
||||
/// SSE 消息推送器
|
||||
///
|
||||
/// 实现 SessionNotifier trait,将会话消息推送到 SSE 连接。
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct SseSessionNotifier;
|
||||
|
||||
impl SseSessionNotifier {
|
||||
/// 创建新的 SSE 消息推送器
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 anyhow::Error 转换为 Box<dyn std::error::Error + Send + Sync>
|
||||
fn convert_error(e: anyhow::Error) -> Box<dyn std::error::Error + Send + Sync> {
|
||||
Box::new(std::io::Error::other(e.to_string()))
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SessionNotifier for SseSessionNotifier {
|
||||
async fn notify_prompt_start(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
let notify = SessionNotify::SessionPromptStart(SessionPromptStart {
|
||||
session_id: session_id.to_string(),
|
||||
request_id,
|
||||
});
|
||||
|
||||
push_session_update_with_project(project_id, session_id, notify)
|
||||
.await
|
||||
.map_err(convert_error)
|
||||
}
|
||||
|
||||
async fn notify_prompt_end(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
stop_reason: agent_client_protocol::schema::StopReason,
|
||||
error_message: Option<String>,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
let notify = SessionNotify::SessionPromptEnd(SessionPromptEnd {
|
||||
session_id: session_id.to_string(),
|
||||
stop_reason,
|
||||
error_message,
|
||||
request_id,
|
||||
});
|
||||
|
||||
push_session_update_with_project(project_id, session_id, notify)
|
||||
.await
|
||||
.map_err(convert_error)
|
||||
}
|
||||
|
||||
async fn notify_prompt_error(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
error: agent_client_protocol::schema::Error,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
let notify = SessionNotify::SessionPromptError(SessionPromptError {
|
||||
session_id: session_id.to_string(),
|
||||
error,
|
||||
request_id,
|
||||
});
|
||||
|
||||
push_session_update_with_project(project_id, session_id, notify)
|
||||
.await
|
||||
.map_err(convert_error)
|
||||
}
|
||||
|
||||
async fn notify_session_update(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
session_update: agent_client_protocol::schema::SessionUpdate,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
let notify = SessionNotify::AgentSessionUpdate(Box::new(AgentSessionUpdate {
|
||||
session_id: session_id.to_string(),
|
||||
session_update,
|
||||
request_id,
|
||||
}));
|
||||
|
||||
push_session_update_with_project(project_id, session_id, notify)
|
||||
.await
|
||||
.map_err(convert_error)
|
||||
}
|
||||
|
||||
async fn notify(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
notify: SessionNotify,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
push_session_update_with_project(project_id, session_id, notify)
|
||||
.await
|
||||
.map_err(convert_error)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
//! StateAwareNotifier 实现
|
||||
//!
|
||||
//! 状态感知的 SessionNotifier 包装器,在推送 SSE 消息的同时同步更新 Agent 状态。
|
||||
//!
|
||||
//! ## 核心职责
|
||||
//! 1. 委托给 SseSessionNotifier 推送 SSE 消息
|
||||
//! 2. 同步更新 AGENT_REGISTRY 状态
|
||||
//! 3. 保持状态转换的原子性和一致性
|
||||
|
||||
use agent_abstraction::SessionNotifier;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
use super::AGENT_REGISTRY;
|
||||
use super::session_notifier::SseSessionNotifier;
|
||||
use shared_types::{AgentStatus, SessionNotify};
|
||||
|
||||
/// 状态感知的 SessionNotifier 包装器
|
||||
///
|
||||
/// 通过委托模式包装 SseSessionNotifier,在推送 SSE 消息的同时同步更新
|
||||
/// AGENT_REGISTRY 中的 Agent 状态。
|
||||
///
|
||||
/// # 设计特点
|
||||
/// - **委托模式**:所有 SSE 推送操作委托给内部的 SseSessionNotifier
|
||||
/// - **原子性状态更新**:使用 AGENT_REGISTRY 确保状态更新的原子性
|
||||
/// - **状态同步顺序**:
|
||||
/// - PromptStart: 先更新状态为 Active,再推送 SSE
|
||||
/// - PromptEnd: 先推送 SSE,再恢复状态为 Idle
|
||||
/// - PromptError: 先推送错误消息,再恢复状态为 Idle
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StateAwareNotifier {
|
||||
/// 内部 SSE 推送器
|
||||
inner: Arc<SseSessionNotifier>,
|
||||
}
|
||||
|
||||
impl StateAwareNotifier {
|
||||
/// 创建新的 StateAwareNotifier 实例
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
inner: Arc::new(SseSessionNotifier::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 🔥 P1 修复: 更新 Agent 状态(原子操作)
|
||||
///
|
||||
/// 使用 `try_update_agent_info` 方法实现原子性状态更新,
|
||||
/// 避免 TOCTOU 竞态条件:读锁释放 → 时间窗口 → 写锁更新。
|
||||
///
|
||||
/// # 参数
|
||||
/// - `project_id`: 项目 ID
|
||||
/// - `status`: 新的 Agent 状态
|
||||
fn update_agent_status(&self, project_id: &str, status: AgentStatus) {
|
||||
AGENT_REGISTRY.try_update_agent_info(project_id, |info| {
|
||||
let old_status = info.status;
|
||||
if old_status != status {
|
||||
info.status = status;
|
||||
info.last_activity = chrono::Utc::now();
|
||||
debug!(
|
||||
"🔄 [atomic_status] Project[{}] status: {:?} -> {:?}",
|
||||
project_id, old_status, status
|
||||
);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for StateAwareNotifier {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SessionNotifier for StateAwareNotifier {
|
||||
/// 推送会话开始通知
|
||||
///
|
||||
/// 顺序:
|
||||
/// 1. 更新 Agent 状态为 Active
|
||||
/// 2. 推送 SessionPromptStart 到 SSE
|
||||
async fn notify_prompt_start(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
info!(
|
||||
"📨 Project[{}] sending SessionPromptStart notification, session_id={}, request_id={:?}",
|
||||
project_id, session_id, request_id
|
||||
);
|
||||
|
||||
// 1. 更新状态为 Active(在推送 SSE 之前)
|
||||
self.update_agent_status(project_id, AgentStatus::Active);
|
||||
|
||||
// 2. 推送 SSE 消息
|
||||
self.inner
|
||||
.notify_prompt_start(project_id, session_id, request_id)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 推送会话结束通知
|
||||
///
|
||||
/// 顺序:
|
||||
/// 1. 推送 SessionPromptEnd 到 SSE
|
||||
/// 2. 恢复 Agent 状态为 Idle
|
||||
async fn notify_prompt_end(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
stop_reason: agent_client_protocol::schema::StopReason,
|
||||
error_message: Option<String>,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
info!(
|
||||
"✅ Project[{}] sending SessionPromptEnd notification, session_id={}, stop_reason={:?}, request_id={:?}",
|
||||
project_id, session_id, stop_reason, request_id
|
||||
);
|
||||
|
||||
// 1. 推送 SSE 消息(在恢复状态之前)
|
||||
let result = self
|
||||
.inner
|
||||
.notify_prompt_end(
|
||||
project_id,
|
||||
session_id,
|
||||
stop_reason,
|
||||
error_message,
|
||||
request_id,
|
||||
)
|
||||
.await;
|
||||
|
||||
// 2. 恢复状态为 Idle
|
||||
self.update_agent_status(project_id, AgentStatus::Idle);
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// 推送会话错误通知
|
||||
///
|
||||
/// 顺序:
|
||||
/// 1. 推送 SessionPromptError 到 SSE
|
||||
/// 2. 恢复 Agent 状态为 Idle
|
||||
async fn notify_prompt_error(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
error: agent_client_protocol::schema::Error,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
error!(
|
||||
"❌ Project[{}] sending SessionPromptError notification, session_id={}, error_code={}, error_message={}, request_id={:?}",
|
||||
project_id, session_id, error.code, error.message, request_id
|
||||
);
|
||||
|
||||
// 1. 推送错误消息(在恢复状态之前)
|
||||
let result = self
|
||||
.inner
|
||||
.notify_prompt_error(project_id, session_id, error, request_id)
|
||||
.await;
|
||||
|
||||
// 2. 恢复状态为 Idle
|
||||
self.update_agent_status(project_id, AgentStatus::Idle);
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// 推送 Agent 会话更新通知
|
||||
///
|
||||
/// 不更新 Agent 状态,仅推送 SSE 消息。
|
||||
async fn notify_session_update(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
session_update: agent_client_protocol::schema::SessionUpdate,
|
||||
request_id: Option<String>,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
debug!(
|
||||
"🔄 Project[{}] sending SessionUpdate notification, session_id={}",
|
||||
project_id, session_id
|
||||
);
|
||||
|
||||
// 委托给内部 notifier,不更新状态
|
||||
self.inner
|
||||
.notify_session_update(project_id, session_id, session_update, request_id)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 推送通用会话通知
|
||||
///
|
||||
/// 不更新 Agent 状态,仅推送 SSE 消息。
|
||||
async fn notify(
|
||||
&self,
|
||||
project_id: &str,
|
||||
session_id: &str,
|
||||
notify: SessionNotify,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
debug!(
|
||||
"📢 Project[{}] sending generic session notification, session_id={}",
|
||||
project_id, session_id
|
||||
);
|
||||
|
||||
// 委托给内部 notifier
|
||||
self.inner.notify(project_id, session_id, notify).await
|
||||
}
|
||||
}
|
||||
126
qiming-rcoder/crates/agent_runner/src/testing/blocking.rs
Normal file
126
qiming-rcoder/crates/agent_runner/src/testing/blocking.rs
Normal file
@@ -0,0 +1,126 @@
|
||||
//! 测试用阻塞注入
|
||||
//!
|
||||
//! 用于模拟极端场景 (Agent 阻塞、超时等)
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
/// 阻塞配置
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BlockingConfig {
|
||||
/// 是否阻塞 new_session
|
||||
pub block_new_session: bool,
|
||||
/// 是否阻塞 prompt
|
||||
pub block_prompt: bool,
|
||||
/// 阻塞持续时间
|
||||
pub block_duration: Duration,
|
||||
}
|
||||
|
||||
impl Default for BlockingConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
block_new_session: false,
|
||||
block_prompt: false,
|
||||
block_duration: Duration::from_secs(30),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 全局阻塞配置 (线程安全)
|
||||
pub static BLOCKING_CONFIG: std::sync::LazyLock<Arc<std::sync::RwLock<BlockingConfig>>> =
|
||||
std::sync::LazyLock::new(|| Arc::new(std::sync::RwLock::new(BlockingConfig::default())));
|
||||
|
||||
/// 注入阻塞 (用于测试)
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use agent_runner::testing::blocking::{BlockingConfig, inject_blocking};
|
||||
/// use std::time::Duration;
|
||||
///
|
||||
/// // 注入 30 秒阻塞
|
||||
/// inject_blocking(BlockingConfig {
|
||||
/// block_prompt: true,
|
||||
/// block_duration: Duration::from_secs(30),
|
||||
/// ..Default::default()
|
||||
/// });
|
||||
/// ```
|
||||
pub fn inject_blocking(config: BlockingConfig) {
|
||||
tracing::warn!("🧪 [TEST] Blocking config updated: {:?}", config);
|
||||
let mut global = BLOCKING_CONFIG.write().unwrap();
|
||||
*global = config;
|
||||
}
|
||||
|
||||
/// 检查并执行阻塞 (在 agent_worker_with_heartbeat 中调用)
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `blocking_type` - 阻塞类型: "new_session" 或 "prompt"
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// # use agent_runner::testing::blocking::maybe_block;
|
||||
/// # async fn test() {
|
||||
/// // 在 agent_worker_with_heartbeat 中调用
|
||||
/// maybe_block("prompt").await;
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn maybe_block(blocking_type: &str) {
|
||||
let config = BLOCKING_CONFIG.read().unwrap();
|
||||
|
||||
let should_block = match blocking_type {
|
||||
"new_session" => config.block_new_session,
|
||||
"prompt" => config.block_prompt,
|
||||
_ => false,
|
||||
};
|
||||
|
||||
if should_block {
|
||||
tracing::warn!(
|
||||
"🧪 [TEST] Injected blocking: {}, duration: {:?}",
|
||||
blocking_type,
|
||||
config.block_duration
|
||||
);
|
||||
tokio::time::sleep(config.block_duration).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 重置阻塞配置 (用于测试清理)
|
||||
pub fn reset_blocking() {
|
||||
inject_blocking(BlockingConfig::default());
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_blocking_config_default() {
|
||||
let config = BlockingConfig::default();
|
||||
assert!(!config.block_new_session);
|
||||
assert!(!config.block_prompt);
|
||||
assert_eq!(config.block_duration, Duration::from_secs(30));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inject_and_reset_blocking() {
|
||||
// 注入配置
|
||||
inject_blocking(BlockingConfig {
|
||||
block_prompt: true,
|
||||
block_duration: Duration::from_secs(10),
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
let config = BLOCKING_CONFIG.read().unwrap();
|
||||
assert!(config.block_prompt);
|
||||
assert_eq!(config.block_duration, Duration::from_secs(10));
|
||||
|
||||
// 重置配置
|
||||
drop(config);
|
||||
reset_blocking();
|
||||
|
||||
let config = BLOCKING_CONFIG.read().unwrap();
|
||||
assert!(!config.block_prompt);
|
||||
assert_eq!(config.block_duration, Duration::from_secs(30));
|
||||
}
|
||||
}
|
||||
226
qiming-rcoder/crates/agent_runner/src/testing/fixtures.rs
Normal file
226
qiming-rcoder/crates/agent_runner/src/testing/fixtures.rs
Normal file
@@ -0,0 +1,226 @@
|
||||
//! 测试 Fixtures
|
||||
//!
|
||||
//! 提供测试用的构造器和辅助工具
|
||||
|
||||
use crate::proxy_agent::AgentRequest;
|
||||
use agent_abstraction::PromptMessage;
|
||||
use shared_types::{Attachment, ServiceType};
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// 测试请求构造器
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use agent_runner::testing::fixtures::TestRequestBuilder;
|
||||
///
|
||||
/// // 创建基础测试请求
|
||||
/// let (req, resp_rx) = TestRequestBuilder::new()
|
||||
/// .project_id("test-project")
|
||||
/// .content("Hello, Agent!")
|
||||
/// .build();
|
||||
/// ```
|
||||
pub struct TestRequestBuilder {
|
||||
project_id: String,
|
||||
content: String,
|
||||
request_id: String,
|
||||
attachments: Vec<Attachment>,
|
||||
session_id: Option<String>,
|
||||
project_path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl Default for TestRequestBuilder {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl TestRequestBuilder {
|
||||
/// 创建新的测试请求构造器
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
project_id: "test-project".to_string(),
|
||||
content: "Hello, Agent!".to_string(),
|
||||
request_id: uuid::Uuid::new_v4().to_string().replace("-", ""),
|
||||
attachments: vec![],
|
||||
session_id: None,
|
||||
project_path: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 设置 project_id
|
||||
pub fn project_id(mut self, id: &str) -> Self {
|
||||
self.project_id = id.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置请求内容
|
||||
pub fn content(mut self, content: &str) -> Self {
|
||||
self.content = content.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 request_id
|
||||
pub fn request_id(mut self, request_id: &str) -> Self {
|
||||
self.request_id = request_id.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 session_id
|
||||
pub fn session_id(mut self, session_id: &str) -> Self {
|
||||
self.session_id = Some(session_id.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置项目路径
|
||||
pub fn project_path(mut self, path: &str) -> Self {
|
||||
self.project_path = Some(PathBuf::from(path));
|
||||
self
|
||||
}
|
||||
|
||||
/// 添加附件
|
||||
pub fn add_attachment(mut self, attachment: Attachment) -> Self {
|
||||
self.attachments.push(attachment);
|
||||
self
|
||||
}
|
||||
|
||||
/// 构建 PromptMessage(用于测试验证)
|
||||
///
|
||||
/// 返回构建的 PromptMessage,可以直接访问字段进行断言验证
|
||||
pub fn build_prompt_message(&self) -> PromptMessage {
|
||||
PromptMessage {
|
||||
content: self.content.clone(),
|
||||
project_id: self.project_id.clone(),
|
||||
project_path: self
|
||||
.project_path
|
||||
.clone()
|
||||
.unwrap_or_else(|| PathBuf::from("/tmp/test")),
|
||||
session_id: self.session_id.clone(),
|
||||
request_id: self.request_id.clone(),
|
||||
attachments: self.attachments.clone(),
|
||||
data_source_attachments: vec![],
|
||||
service_type: ServiceType::RCoder,
|
||||
system_prompt_override: None,
|
||||
user_prompt_template_override: None,
|
||||
agent_config_override: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 构建测试请求
|
||||
pub fn build(
|
||||
self,
|
||||
) -> (
|
||||
AgentRequest,
|
||||
tokio::sync::oneshot::Receiver<crate::model::ChatPromptResponse>,
|
||||
) {
|
||||
let prompt_message = self.build_prompt_message();
|
||||
AgentRequest::new(prompt_message, None)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_test_request_builder_default() {
|
||||
let builder = TestRequestBuilder::new();
|
||||
let prompt = builder.build_prompt_message();
|
||||
|
||||
// 验证默认值
|
||||
assert_eq!(prompt.project_id, "test-project");
|
||||
assert_eq!(prompt.content, "Hello, Agent!");
|
||||
assert!(!prompt.request_id.is_empty(), "request_id 应该自动生成");
|
||||
assert_eq!(prompt.project_path, PathBuf::from("/tmp/test"));
|
||||
assert!(prompt.session_id.is_none());
|
||||
assert!(prompt.attachments.is_empty());
|
||||
assert_eq!(prompt.service_type, ServiceType::RCoder);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_test_request_builder_with_project_id() {
|
||||
let prompt = TestRequestBuilder::new()
|
||||
.project_id("custom-project")
|
||||
.build_prompt_message();
|
||||
|
||||
assert_eq!(prompt.project_id, "custom-project");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_test_request_builder_with_content() {
|
||||
let prompt = TestRequestBuilder::new()
|
||||
.content("Custom content")
|
||||
.build_prompt_message();
|
||||
|
||||
assert_eq!(prompt.content, "Custom content");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_test_request_builder_with_session_id() {
|
||||
let prompt = TestRequestBuilder::new()
|
||||
.session_id("test-session-123")
|
||||
.build_prompt_message();
|
||||
|
||||
assert_eq!(prompt.session_id, Some("test-session-123".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_test_request_builder_with_request_id() {
|
||||
let prompt = TestRequestBuilder::new()
|
||||
.request_id("custom-request-id")
|
||||
.build_prompt_message();
|
||||
|
||||
assert_eq!(prompt.request_id, "custom-request-id");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_test_request_builder_with_project_path() {
|
||||
let prompt = TestRequestBuilder::new()
|
||||
.project_path("/home/user/project")
|
||||
.build_prompt_message();
|
||||
|
||||
assert_eq!(prompt.project_path, PathBuf::from("/home/user/project"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_test_request_builder_chain() {
|
||||
// 验证链式调用和所有字段设置
|
||||
let prompt = TestRequestBuilder::new()
|
||||
.project_id("my-project")
|
||||
.content("Test content")
|
||||
.request_id("custom-request-id")
|
||||
.session_id("session-456")
|
||||
.project_path("/home/user/project")
|
||||
.build_prompt_message();
|
||||
|
||||
assert_eq!(prompt.project_id, "my-project");
|
||||
assert_eq!(prompt.content, "Test content");
|
||||
assert_eq!(prompt.request_id, "custom-request-id");
|
||||
assert_eq!(prompt.session_id, Some("session-456".to_string()));
|
||||
assert_eq!(prompt.project_path, PathBuf::from("/home/user/project"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_test_request_builder_build_returns_receiver() {
|
||||
// 验证 build() 返回的接收器可用
|
||||
let (_req, mut rx) = TestRequestBuilder::new().build();
|
||||
|
||||
// 验证接收器尚未收到消息
|
||||
assert!(rx.try_recv().is_err(), "接收器应该为空");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_test_request_builder_request_id_auto_generated() {
|
||||
// 验证每次构建生成不同的 request_id
|
||||
let builder1 = TestRequestBuilder::new();
|
||||
let builder2 = TestRequestBuilder::new();
|
||||
|
||||
let prompt1 = builder1.build_prompt_message();
|
||||
let prompt2 = builder2.build_prompt_message();
|
||||
|
||||
assert_ne!(
|
||||
prompt1.request_id, prompt2.request_id,
|
||||
"每次构建应该生成不同的 request_id"
|
||||
);
|
||||
}
|
||||
}
|
||||
14
qiming-rcoder/crates/agent_runner/src/testing/mod.rs
Normal file
14
qiming-rcoder/crates/agent_runner/src/testing/mod.rs
Normal file
@@ -0,0 +1,14 @@
|
||||
//! 测试辅助模块
|
||||
//!
|
||||
//! 仅在 `testing` feature 启用时编译
|
||||
|
||||
// 阻塞注入模块 (用于极端场景测试)
|
||||
#[cfg(feature = "test-blocking")]
|
||||
pub mod blocking;
|
||||
|
||||
#[cfg(feature = "test-blocking")]
|
||||
pub use blocking::{BlockingConfig, inject_blocking};
|
||||
|
||||
// 测试 Fixtures (通用测试辅助工具)
|
||||
pub mod fixtures;
|
||||
pub use fixtures::TestRequestBuilder;
|
||||
428
qiming-rcoder/crates/agent_runner/src/utils/content_builder.rs
Normal file
428
qiming-rcoder/crates/agent_runner/src/utils/content_builder.rs
Normal file
@@ -0,0 +1,428 @@
|
||||
//! Content Builder 工具
|
||||
//!
|
||||
//! 将 Attachment 转换为 ACP 协议的 ContentBlock
|
||||
|
||||
use anyhow::Result;
|
||||
use base64::{Engine as _, engine::general_purpose};
|
||||
use agent_client_protocol::schema::{
|
||||
AudioContent, BlobResourceContents, ContentBlock, EmbeddedResource, EmbeddedResourceResource,
|
||||
ImageContent, ResourceLink, TextContent, TextResourceContents,
|
||||
};
|
||||
use std::path::Path;
|
||||
|
||||
use crate::model::{Attachment, AttachmentSource};
|
||||
|
||||
/// Content Builder - 将附件转换为 ACP ContentBlock
|
||||
pub struct ContentBuilder;
|
||||
|
||||
impl ContentBuilder {
|
||||
/// 将单个附件转换为 ContentBlock
|
||||
pub async fn attachment_to_content_block(
|
||||
attachment: &Attachment,
|
||||
project_path: &std::path::Path,
|
||||
) -> Result<Option<ContentBlock>> {
|
||||
match attachment {
|
||||
Attachment::Text(text_attachment) => {
|
||||
Self::text_to_content_block(text_attachment, project_path).await
|
||||
}
|
||||
Attachment::Image(image_attachment) => {
|
||||
Self::image_to_content_block(image_attachment, project_path).await
|
||||
}
|
||||
Attachment::Audio(audio_attachment) => {
|
||||
Self::audio_to_content_block(audio_attachment, project_path).await
|
||||
}
|
||||
Attachment::Document(document_attachment) => {
|
||||
Self::document_to_content_block(document_attachment, project_path).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 将多个附件转换为 ContentBlock 列表
|
||||
/// 文件不存在或无法读取的附件会被静默忽略
|
||||
pub async fn attachments_to_content_blocks(
|
||||
attachments: &[Attachment],
|
||||
project_path: &std::path::Path,
|
||||
) -> Result<Vec<ContentBlock>> {
|
||||
let mut content_blocks = Vec::new();
|
||||
|
||||
for attachment in attachments {
|
||||
match Self::attachment_to_content_block(attachment, project_path).await {
|
||||
Ok(Some(content_block)) => {
|
||||
content_blocks.push(content_block);
|
||||
}
|
||||
Ok(None) => {
|
||||
// 文件不存在或无法读取,静默忽略
|
||||
tracing::warn!(
|
||||
"Attachment could not be loaded and was ignored: attachment_id={}",
|
||||
attachment.id()
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
// 其他错误(如网络错误),记录警告但继续处理
|
||||
tracing::warn!(
|
||||
"⚠️ Attachment conversion failed and was ignored: attachment_id={}, error={:?}",
|
||||
attachment.id(),
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(content_blocks)
|
||||
}
|
||||
|
||||
/// 文本附件转换为 ContentBlock
|
||||
/// 如果文件不存在或是压缩文件,返回 Ok(None)
|
||||
async fn text_to_content_block(
|
||||
text_attachment: &crate::model::TextAttachment,
|
||||
project_path: &std::path::Path,
|
||||
) -> Result<Option<ContentBlock>> {
|
||||
let text_content = match &text_attachment.source {
|
||||
AttachmentSource::FilePath { path } => {
|
||||
let file_path = project_path.join(path);
|
||||
|
||||
// 检查文件是否存在
|
||||
if !file_path.exists() {
|
||||
tracing::warn!("Attachment file not found, ignored: {:?}", file_path);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// 检查是否为二进制文件(常见压缩格式)
|
||||
if let Some(ext) = file_path.extension() {
|
||||
let ext_str = ext.to_string_lossy().to_lowercase();
|
||||
if matches!(
|
||||
ext_str.as_str(),
|
||||
"gz" | "zip" | "tar" | "bz2" | "xz" | "7z" | "rar"
|
||||
) {
|
||||
tracing::warn!(
|
||||
"⚠️ Compressed files are not supported as Text attachments, ignored: {:?} (extension: {})",
|
||||
file_path,
|
||||
ext_str
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
|
||||
// 尝试读取文本文件
|
||||
let text = match tokio::fs::read_to_string(&file_path).await {
|
||||
Ok(text) => text,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"⚠️ Failed to read text file, ignored: {:?}, error: {} (may be binary or not UTF-8 encoded)",
|
||||
file_path,
|
||||
e
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
TextContent::new(text)
|
||||
}
|
||||
AttachmentSource::Base64 { data, mime_type: _ } => {
|
||||
let text = String::from_utf8(general_purpose::STANDARD.decode(data)?)?;
|
||||
|
||||
TextContent::new(text)
|
||||
}
|
||||
AttachmentSource::Url { url } => {
|
||||
let client = reqwest::Client::new();
|
||||
let response = client.get(url).send().await?;
|
||||
let text = response.text().await?;
|
||||
|
||||
TextContent::new(text)
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Some(ContentBlock::Text(text_content)))
|
||||
}
|
||||
|
||||
/// 图像附件转换为 ContentBlock
|
||||
/// 如果文件不存在,返回 Ok(None)
|
||||
async fn image_to_content_block(
|
||||
image_attachment: &crate::model::ImageAttachment,
|
||||
project_path: &std::path::Path,
|
||||
) -> Result<Option<ContentBlock>> {
|
||||
let (data, uri) = match &image_attachment.source {
|
||||
AttachmentSource::FilePath { path } => {
|
||||
let file_path = project_path.join(path);
|
||||
|
||||
// 检查文件是否存在
|
||||
if !file_path.exists() {
|
||||
tracing::warn!("Image file not found, ignored: {:?}", file_path);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// 尝试读取文件
|
||||
let data = match tokio::fs::read(&file_path).await {
|
||||
Ok(data) => data,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to read image file, ignored: {:?}, error: {}",
|
||||
file_path,
|
||||
e
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
let base64_data = general_purpose::STANDARD.encode(data);
|
||||
let uri = file_path.to_string_lossy().to_string();
|
||||
(base64_data, Some(uri))
|
||||
}
|
||||
AttachmentSource::Base64 { data, .. } => (data.clone(), None),
|
||||
AttachmentSource::Url { url } => {
|
||||
let client = reqwest::Client::new();
|
||||
let response = client.get(url).send().await?;
|
||||
let data = response.bytes().await?;
|
||||
let base64_data = general_purpose::STANDARD.encode(data);
|
||||
(base64_data, Some(url.clone()))
|
||||
}
|
||||
};
|
||||
|
||||
let mut image_content = ImageContent::new(data, image_attachment.mime_type.clone());
|
||||
if let Some(u) = uri {
|
||||
image_content = image_content.uri(u);
|
||||
}
|
||||
Ok(Some(ContentBlock::Image(image_content)))
|
||||
}
|
||||
|
||||
/// 音频附件转换为 ContentBlock
|
||||
/// 如果文件不存在,返回 Ok(None)
|
||||
async fn audio_to_content_block(
|
||||
audio_attachment: &crate::model::AudioAttachment,
|
||||
project_path: &std::path::Path,
|
||||
) -> Result<Option<ContentBlock>> {
|
||||
let data = match &audio_attachment.source {
|
||||
AttachmentSource::FilePath { path } => {
|
||||
let file_path = project_path.join(path);
|
||||
|
||||
// 检查文件是否存在
|
||||
if !file_path.exists() {
|
||||
tracing::warn!("Audio file not found, ignored: {:?}", file_path);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// 尝试读取文件
|
||||
let data = match tokio::fs::read(&file_path).await {
|
||||
Ok(data) => data,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to read audio file, ignored: {:?}, error: {}",
|
||||
file_path,
|
||||
e
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
general_purpose::STANDARD.encode(data)
|
||||
}
|
||||
AttachmentSource::Base64 { data, .. } => data.clone(),
|
||||
AttachmentSource::Url { url } => {
|
||||
let client = reqwest::Client::new();
|
||||
let response = client.get(url).send().await?;
|
||||
let data = response.bytes().await?;
|
||||
general_purpose::STANDARD.encode(data)
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Some(ContentBlock::Audio(AudioContent::new(
|
||||
data,
|
||||
audio_attachment.mime_type.clone(),
|
||||
))))
|
||||
}
|
||||
|
||||
/// 文档附件转换为 ContentBlock
|
||||
/// 如果文件不存在,返回 Ok(None)
|
||||
async fn document_to_content_block(
|
||||
document_attachment: &crate::model::DocumentAttachment,
|
||||
project_path: &std::path::Path,
|
||||
) -> Result<Option<ContentBlock>> {
|
||||
match &document_attachment.source {
|
||||
AttachmentSource::FilePath { path } => {
|
||||
let file_path = project_path.join(path);
|
||||
|
||||
// 检查文件是否存在
|
||||
if !file_path.exists() {
|
||||
tracing::warn!("Document file not found, ignored: {:?}", file_path);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let uri = file_path.to_string_lossy().to_string();
|
||||
|
||||
// 尝试读取为文本,如果失败则作为二进制处理
|
||||
if document_attachment.mime_type.starts_with("text/") {
|
||||
match tokio::fs::read_to_string(&file_path).await {
|
||||
Ok(text) => {
|
||||
let mut text_contents = TextResourceContents::new(text, uri);
|
||||
if let Some(mime_type) = Some(document_attachment.mime_type.clone()) {
|
||||
text_contents = text_contents.mime_type(mime_type);
|
||||
}
|
||||
Ok(Some(ContentBlock::Resource(EmbeddedResource::new(
|
||||
EmbeddedResourceResource::TextResourceContents(text_contents),
|
||||
))))
|
||||
}
|
||||
Err(_) => {
|
||||
// 文本读取失败,尝试作为二进制处理
|
||||
Self::handle_binary_document(
|
||||
&file_path,
|
||||
&document_attachment.mime_type,
|
||||
&uri,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Self::handle_binary_document(&file_path, &document_attachment.mime_type, &uri)
|
||||
.await
|
||||
}
|
||||
}
|
||||
AttachmentSource::Base64 { data, mime_type } => {
|
||||
let uri = format!("base64://{}", document_attachment.id);
|
||||
|
||||
if mime_type.starts_with("text/") {
|
||||
match String::from_utf8(general_purpose::STANDARD.decode(data)?) {
|
||||
Ok(text) => {
|
||||
let mut text_contents = TextResourceContents::new(text, uri);
|
||||
text_contents = text_contents.mime_type(mime_type.clone());
|
||||
Ok(Some(ContentBlock::Resource(EmbeddedResource::new(
|
||||
EmbeddedResourceResource::TextResourceContents(text_contents),
|
||||
))))
|
||||
}
|
||||
Err(_) => {
|
||||
// 文本解码失败,作为二进制处理
|
||||
let mut blob_contents = BlobResourceContents::new(data.clone(), uri);
|
||||
blob_contents = blob_contents.mime_type(mime_type.clone());
|
||||
Ok(Some(ContentBlock::Resource(EmbeddedResource::new(
|
||||
EmbeddedResourceResource::BlobResourceContents(blob_contents),
|
||||
))))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let mut blob_contents = BlobResourceContents::new(data.clone(), uri);
|
||||
blob_contents = blob_contents.mime_type(mime_type.clone());
|
||||
Ok(Some(ContentBlock::Resource(EmbeddedResource::new(
|
||||
EmbeddedResourceResource::BlobResourceContents(blob_contents),
|
||||
))))
|
||||
}
|
||||
}
|
||||
AttachmentSource::Url { url } => {
|
||||
// URL 资源作为 ResourceLink 处理
|
||||
let name = document_attachment
|
||||
.filename
|
||||
.clone()
|
||||
.unwrap_or_else(|| "document".to_string());
|
||||
let mut resource_link = ResourceLink::new(name, url.clone());
|
||||
resource_link = resource_link.mime_type(document_attachment.mime_type.clone());
|
||||
|
||||
// 只有当值存在时才设置
|
||||
if let Some(description) = &document_attachment.description {
|
||||
resource_link = resource_link.description(description.clone());
|
||||
}
|
||||
if let Some(size) = document_attachment.size {
|
||||
resource_link = resource_link.size(size as i64);
|
||||
}
|
||||
if let Some(filename) = &document_attachment.filename {
|
||||
resource_link = resource_link.title(filename.clone());
|
||||
}
|
||||
Ok(Some(ContentBlock::ResourceLink(resource_link)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 处理二进制文档
|
||||
/// 如果文件无法读取,返回 Ok(None)
|
||||
async fn handle_binary_document(
|
||||
file_path: &Path,
|
||||
mime_type: &str,
|
||||
uri: &str,
|
||||
) -> Result<Option<ContentBlock>> {
|
||||
let data = match tokio::fs::read(file_path).await {
|
||||
Ok(data) => data,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"⚠️ Cannot read binary document, ignored: {:?}, error: {}",
|
||||
file_path,
|
||||
e
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
let blob = general_purpose::STANDARD.encode(data);
|
||||
|
||||
let mut blob_contents = BlobResourceContents::new(blob, uri.to_string());
|
||||
blob_contents = blob_contents.mime_type(mime_type.to_string());
|
||||
Ok(Some(ContentBlock::Resource(EmbeddedResource::new(
|
||||
EmbeddedResourceResource::BlobResourceContents(blob_contents),
|
||||
))))
|
||||
}
|
||||
|
||||
/// 从文件扩展名推断 MIME 类型
|
||||
pub fn infer_mime_type_from_extension(filename: &str) -> &'static str {
|
||||
let path = std::path::Path::new(filename);
|
||||
match path.extension().and_then(|ext| ext.to_str()) {
|
||||
Some("txt") => "text/plain",
|
||||
Some("md") => "text/markdown",
|
||||
Some("json") => "application/json",
|
||||
Some("xml") => "application/xml",
|
||||
Some("html") => "text/html",
|
||||
Some("css") => "text/css",
|
||||
Some("js") => "application/javascript",
|
||||
Some("ts") => "application/typescript",
|
||||
Some("pdf") => "application/pdf",
|
||||
Some("jpg") | Some("jpeg") => "image/jpeg",
|
||||
Some("png") => "image/png",
|
||||
Some("gif") => "image/gif",
|
||||
Some("svg") => "image/svg+xml",
|
||||
Some("webp") => "image/webp",
|
||||
Some("mp3") => "audio/mpeg",
|
||||
Some("wav") => "audio/wav",
|
||||
Some("ogg") => "audio/ogg",
|
||||
Some("m4a") => "audio/mp4",
|
||||
Some("mp4") => "video/mp4",
|
||||
Some("webm") => "video/webm",
|
||||
Some("zip") => "application/zip",
|
||||
Some("tar") => "application/x-tar",
|
||||
Some("gz") => "application/gzip",
|
||||
Some("rar") => "application/x-rar-compressed",
|
||||
Some("7z") => "application/x-7z-compressed",
|
||||
Some("doc") => "application/msword",
|
||||
Some("docx") => {
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
}
|
||||
Some("xls") => "application/vnd.ms-excel",
|
||||
Some("xlsx") => "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
Some("ppt") => "application/vnd.ms-powerpoint",
|
||||
Some("pptx") => {
|
||||
"application/vnd.openxmlformats-officedocument.presentationml.presentation"
|
||||
}
|
||||
_ => "application/octet-stream",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_infer_mime_type() {
|
||||
assert_eq!(
|
||||
ContentBuilder::infer_mime_type_from_extension("test.txt"),
|
||||
"text/plain"
|
||||
);
|
||||
assert_eq!(
|
||||
ContentBuilder::infer_mime_type_from_extension("image.jpg"),
|
||||
"image/jpeg"
|
||||
);
|
||||
assert_eq!(
|
||||
ContentBuilder::infer_mime_type_from_extension("audio.mp3"),
|
||||
"audio/mpeg"
|
||||
);
|
||||
assert_eq!(
|
||||
ContentBuilder::infer_mime_type_from_extension("document.pdf"),
|
||||
"application/pdf"
|
||||
);
|
||||
assert_eq!(
|
||||
ContentBuilder::infer_mime_type_from_extension("unknown.xyz"),
|
||||
"application/octet-stream"
|
||||
);
|
||||
}
|
||||
}
|
||||
370
qiming-rcoder/crates/agent_runner/src/utils/file_utils.rs
Normal file
370
qiming-rcoder/crates/agent_runner/src/utils/file_utils.rs
Normal file
@@ -0,0 +1,370 @@
|
||||
//! 文件处理工具函数
|
||||
//!
|
||||
//! 提供文件读取、验证、转换等实用功能
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use base64::{Engine as _, engine::general_purpose};
|
||||
use std::path::{Path, PathBuf};
|
||||
use tokio::fs;
|
||||
use tracing::warn;
|
||||
|
||||
use super::ContentBuilder;
|
||||
use shared_types::{Attachment, AttachmentError, AttachmentSource};
|
||||
|
||||
/// 文件处理配置
|
||||
pub struct FileConfig {
|
||||
/// 最大文件大小(字节)
|
||||
pub max_file_size: u64,
|
||||
/// 允许的文件扩展名
|
||||
pub allowed_extensions: Vec<String>,
|
||||
/// 临时文件目录
|
||||
pub temp_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl Default for FileConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_file_size: 50 * 1024 * 1024, // 50MB
|
||||
allowed_extensions: vec![
|
||||
"txt", "md", "json", "xml", "html", "css", "js", "ts", // 文本文件
|
||||
"jpg", "jpeg", "png", "gif", "svg", "webp", // 图像文件
|
||||
"mp3", "wav", "ogg", "m4a", // 音频文件
|
||||
"pdf", "doc", "docx", "xls", "xlsx", "ppt", "pptx", // 文档文件
|
||||
"zip", "tar", "gz", "rar", "7z", // 压缩文件
|
||||
]
|
||||
.iter()
|
||||
.map(|s| s.to_string())
|
||||
.collect(),
|
||||
temp_dir: std::env::temp_dir(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 文件处理工具
|
||||
pub struct FileUtils {
|
||||
config: FileConfig,
|
||||
}
|
||||
|
||||
impl FileUtils {
|
||||
/// 创建新的文件处理工具实例
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
config: FileConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用自定义配置创建文件处理工具
|
||||
pub fn with_config(config: FileConfig) -> Self {
|
||||
Self { config }
|
||||
}
|
||||
|
||||
/// 验证文件扩展名是否允许
|
||||
pub fn validate_extension(&self, filename: &str) -> Result<()> {
|
||||
let path = Path::new(filename);
|
||||
let extension = path
|
||||
.extension()
|
||||
.and_then(|ext| ext.to_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("File has no extension"))?;
|
||||
|
||||
if !self
|
||||
.config
|
||||
.allowed_extensions
|
||||
.contains(&extension.to_lowercase())
|
||||
{
|
||||
return Err(anyhow::anyhow!("Unsupported file extension: {}", extension));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 验证文件大小
|
||||
pub async fn validate_file_size(&self, file_path: &Path) -> Result<()> {
|
||||
let metadata = fs::metadata(file_path)
|
||||
.await
|
||||
.context("unable to get file metadata")?;
|
||||
|
||||
if metadata.len() > self.config.max_file_size {
|
||||
return Err(AttachmentError::FileSizeExceeded(metadata.len()).into());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 读取文件内容为 base64
|
||||
pub async fn read_file_as_base64(&self, file_path: &Path) -> Result<String> {
|
||||
self.validate_file_size(file_path).await?;
|
||||
|
||||
let content = fs::read(file_path).await.context("Failed to read file")?;
|
||||
|
||||
Ok(general_purpose::STANDARD.encode(content))
|
||||
}
|
||||
|
||||
/// 读取文本文件
|
||||
pub async fn read_text_file(&self, file_path: &Path) -> Result<String> {
|
||||
self.validate_file_size(file_path).await?;
|
||||
|
||||
let content = fs::read_to_string(file_path)
|
||||
.await
|
||||
.context("Failed to read text file")?;
|
||||
|
||||
Ok(content)
|
||||
}
|
||||
|
||||
/// 从文件路径创建附件
|
||||
pub async fn create_attachment_from_file(
|
||||
&self,
|
||||
file_path: &Path,
|
||||
project_path: &Path,
|
||||
description: Option<String>,
|
||||
) -> Result<Attachment> {
|
||||
// 验证文件扩展名
|
||||
let filename = file_path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("Invalid filename"))?;
|
||||
|
||||
self.validate_extension(filename)?;
|
||||
|
||||
// 获取相对于项目路径的路径
|
||||
let relative_path = file_path
|
||||
.strip_prefix(project_path)
|
||||
.unwrap_or(file_path)
|
||||
.to_string_lossy()
|
||||
.to_string();
|
||||
|
||||
// 获取 MIME 类型
|
||||
let mime_type = ContentBuilder::infer_mime_type_from_extension(filename);
|
||||
|
||||
// 根据文件类型创建相应的附件
|
||||
let attachment = if mime_type.starts_with("image/") {
|
||||
Attachment::new_image(
|
||||
AttachmentSource::FilePath {
|
||||
path: relative_path,
|
||||
},
|
||||
mime_type.to_string(),
|
||||
)
|
||||
} else if mime_type.starts_with("audio/") {
|
||||
Attachment::new_audio(
|
||||
AttachmentSource::FilePath {
|
||||
path: relative_path,
|
||||
},
|
||||
mime_type.to_string(),
|
||||
)
|
||||
} else if mime_type.starts_with("text/") || mime_type == "application/json" {
|
||||
let mut attachment = Attachment::new_text(AttachmentSource::FilePath {
|
||||
path: relative_path,
|
||||
});
|
||||
if let Attachment::Text(ref mut text_attachment) = attachment {
|
||||
text_attachment.description = description;
|
||||
}
|
||||
attachment
|
||||
} else {
|
||||
Attachment::new_document(
|
||||
AttachmentSource::FilePath {
|
||||
path: relative_path,
|
||||
},
|
||||
mime_type.to_string(),
|
||||
)
|
||||
};
|
||||
|
||||
Ok(attachment)
|
||||
}
|
||||
|
||||
/// 从 base64 数据创建附件
|
||||
pub fn create_attachment_from_base64(
|
||||
&self,
|
||||
data: String,
|
||||
mime_type: String,
|
||||
filename: Option<String>,
|
||||
description: Option<String>,
|
||||
) -> Result<Attachment> {
|
||||
// 验证数据大小
|
||||
let decoded_size = general_purpose::STANDARD.decode(&data)?.len() as u64;
|
||||
if decoded_size > self.config.max_file_size {
|
||||
return Err(AttachmentError::FileSizeExceeded(decoded_size).into());
|
||||
}
|
||||
|
||||
let source = AttachmentSource::Base64 {
|
||||
data,
|
||||
mime_type: mime_type.clone(),
|
||||
};
|
||||
|
||||
let attachment = if mime_type.starts_with("image/") {
|
||||
let mut attachment = Attachment::new_image(source, mime_type);
|
||||
if let Attachment::Image(ref mut image_attachment) = attachment {
|
||||
image_attachment.filename = filename;
|
||||
image_attachment.description = description;
|
||||
}
|
||||
attachment
|
||||
} else if mime_type.starts_with("audio/") {
|
||||
let mut attachment = Attachment::new_audio(source, mime_type);
|
||||
if let Attachment::Audio(ref mut audio_attachment) = attachment {
|
||||
audio_attachment.filename = filename;
|
||||
audio_attachment.description = description;
|
||||
}
|
||||
attachment
|
||||
} else if mime_type.starts_with("text/") || mime_type == "application/json" {
|
||||
let mut attachment = Attachment::new_text(source);
|
||||
if let Attachment::Text(ref mut text_attachment) = attachment {
|
||||
text_attachment.filename = filename;
|
||||
text_attachment.description = description;
|
||||
}
|
||||
attachment
|
||||
} else {
|
||||
let mut attachment = Attachment::new_document(source, mime_type);
|
||||
if let Attachment::Document(ref mut doc_attachment) = attachment {
|
||||
doc_attachment.filename = filename;
|
||||
doc_attachment.description = description;
|
||||
doc_attachment.size = Some(decoded_size);
|
||||
}
|
||||
attachment
|
||||
};
|
||||
|
||||
Ok(attachment)
|
||||
}
|
||||
|
||||
/// 从 URL 创建附件
|
||||
pub async fn create_attachment_from_url(
|
||||
&self,
|
||||
url: &str,
|
||||
filename: Option<String>,
|
||||
description: Option<String>,
|
||||
) -> Result<Attachment> {
|
||||
// 获取 URL 头信息以确定文件类型和大小
|
||||
let client = reqwest::Client::new();
|
||||
let response = client.head(url).send().await?;
|
||||
|
||||
// 检查文件大小
|
||||
if let Some(content_length) = response.headers().get(reqwest::header::CONTENT_LENGTH) {
|
||||
let length: u64 = content_length.to_str()?.parse()?;
|
||||
if length > self.config.max_file_size {
|
||||
return Err(AttachmentError::FileSizeExceeded(length).into());
|
||||
}
|
||||
}
|
||||
|
||||
// 获取 MIME 类型
|
||||
let mime_type = response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or("application/octet-stream");
|
||||
|
||||
let source = AttachmentSource::Url {
|
||||
url: url.to_string(),
|
||||
};
|
||||
|
||||
let attachment = if mime_type.starts_with("image/") {
|
||||
let mut attachment = Attachment::new_image(source, mime_type.to_string());
|
||||
if let Attachment::Image(ref mut image_attachment) = attachment {
|
||||
image_attachment.filename = filename;
|
||||
image_attachment.description = description;
|
||||
}
|
||||
attachment
|
||||
} else if mime_type.starts_with("audio/") {
|
||||
let mut attachment = Attachment::new_audio(source, mime_type.to_string());
|
||||
if let Attachment::Audio(ref mut audio_attachment) = attachment {
|
||||
audio_attachment.filename = filename;
|
||||
audio_attachment.description = description;
|
||||
}
|
||||
attachment
|
||||
} else if mime_type.starts_with("text/") || mime_type == "application/json" {
|
||||
let mut attachment = Attachment::new_text(source);
|
||||
if let Attachment::Text(ref mut text_attachment) = attachment {
|
||||
text_attachment.filename = filename;
|
||||
text_attachment.description = description;
|
||||
}
|
||||
attachment
|
||||
} else {
|
||||
let mut attachment = Attachment::new_document(source, mime_type.to_string());
|
||||
if let Attachment::Document(ref mut doc_attachment) = attachment {
|
||||
doc_attachment.filename = filename;
|
||||
doc_attachment.description = description;
|
||||
}
|
||||
attachment
|
||||
};
|
||||
|
||||
Ok(attachment)
|
||||
}
|
||||
|
||||
/// 批量处理文件创建附件
|
||||
pub async fn create_attachments_from_files(
|
||||
&self,
|
||||
file_paths: &[PathBuf],
|
||||
project_path: &Path,
|
||||
) -> Result<Vec<Attachment>> {
|
||||
let mut attachments = Vec::new();
|
||||
|
||||
for file_path in file_paths {
|
||||
match self
|
||||
.create_attachment_from_file(file_path, project_path, None)
|
||||
.await
|
||||
{
|
||||
Ok(attachment) => attachments.push(attachment),
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"Skipping file {:?}; failed to create attachment: {}",
|
||||
file_path, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(attachments)
|
||||
}
|
||||
|
||||
/// 清理临时文件
|
||||
pub async fn cleanup_temp_files(&self, temp_files: &[PathBuf]) -> Result<()> {
|
||||
for temp_file in temp_files {
|
||||
if temp_file.exists() {
|
||||
fs::remove_file(temp_file)
|
||||
.await
|
||||
.with_context(|| format!("failed to delete temp file: {:?}", temp_file))?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FileUtils {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_validate_extension() {
|
||||
let file_utils = FileUtils::new();
|
||||
|
||||
// 测试允许的扩展名
|
||||
assert!(file_utils.validate_extension("test.txt").is_ok());
|
||||
assert!(file_utils.validate_extension("image.jpg").is_ok());
|
||||
assert!(file_utils.validate_extension("audio.mp3").is_ok());
|
||||
|
||||
// 测试不允许的扩展名
|
||||
assert!(file_utils.validate_extension("test.exe").is_err());
|
||||
assert!(file_utils.validate_extension("test").is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_file_as_base64() {
|
||||
let file_utils = FileUtils::new();
|
||||
|
||||
// 创建临时文件
|
||||
let mut temp_file = NamedTempFile::new().unwrap();
|
||||
writeln!(temp_file, "Hello, World!").unwrap();
|
||||
let temp_path = temp_file.path();
|
||||
|
||||
let base64_content = file_utils.read_file_as_base64(temp_path).await.unwrap();
|
||||
|
||||
// 验证 base64 编码是否正确
|
||||
let decoded = general_purpose::STANDARD.decode(base64_content).unwrap();
|
||||
let decoded_text = String::from_utf8(decoded).unwrap();
|
||||
assert_eq!(decoded_text.trim(), "Hello, World!");
|
||||
}
|
||||
}
|
||||
6
qiming-rcoder/crates/agent_runner/src/utils/mod.rs
Normal file
6
qiming-rcoder/crates/agent_runner/src/utils/mod.rs
Normal file
@@ -0,0 +1,6 @@
|
||||
//! 工具函数模块
|
||||
|
||||
mod content_builder;
|
||||
mod file_utils;
|
||||
|
||||
pub use content_builder::*;
|
||||
903
qiming-rcoder/crates/agent_runner/tests/concurrency_raii_test.rs
Normal file
903
qiming-rcoder/crates/agent_runner/tests/concurrency_raii_test.rs
Normal file
@@ -0,0 +1,903 @@
|
||||
//! 并发和 RAII 设计测试
|
||||
//!
|
||||
//! 验证以下功能:
|
||||
//! 1. 并发独立启动 agent,agent 之间互不影响
|
||||
//! 2. agent 销毁的正确性
|
||||
//! 3. RAII 设计(PendingGuard)是否可以正常快速销毁 agent
|
||||
//! 4. 原子计数器的并发安全性
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Barrier;
|
||||
|
||||
// 重新导出必要的类型
|
||||
use agent_runner::agent_runtime::{get_concurrency_limit, init_concurrency_limit};
|
||||
use agent_runner::service::AgentSessionRegistry;
|
||||
use agent_runner::service::PendingGuard;
|
||||
use agent_client_protocol::schema::SessionId;
|
||||
use shared_types::{AgentStatus, ProjectAndAgentInfo, SessionEntry};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
// ============================================================================
|
||||
// 1. PendingGuard RAII 测试
|
||||
// ============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_pending_guard_auto_cleanup_on_drop() {
|
||||
let registry = AgentSessionRegistry::new();
|
||||
|
||||
// 设置 Pending 状态
|
||||
{
|
||||
let _guard = PendingGuard::new(®istry, "test-project");
|
||||
// 验证 Pending 状态已设置
|
||||
assert!(registry.contains_project("test-project"));
|
||||
let info = registry.get_agent_info("test-project").unwrap();
|
||||
assert_eq!(format!("{:?}", info.status), "Pending");
|
||||
}
|
||||
// guard 已 drop,Pending 状态应该被清理
|
||||
assert!(!registry.contains_project("test-project"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pending_guard_commit_success_prevents_cleanup() {
|
||||
let registry = AgentSessionRegistry::new();
|
||||
|
||||
{
|
||||
let guard = PendingGuard::new(®istry, "test-project");
|
||||
|
||||
// 验证 Pending 状态已设置
|
||||
assert!(registry.contains_project("test-project"));
|
||||
|
||||
// 提交成功,防止清理
|
||||
guard.commit_success();
|
||||
}
|
||||
|
||||
// Pending 状态应该保留
|
||||
assert!(registry.contains_project("test-project"));
|
||||
let info = registry.get_agent_info("test-project").unwrap();
|
||||
assert_eq!(format!("{:?}", info.status), "Pending");
|
||||
|
||||
// 清理
|
||||
registry.remove_by_project("test-project");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pending_guard_early_return_cleanup() {
|
||||
let registry = AgentSessionRegistry::new();
|
||||
|
||||
// 模拟早期返回场景(使用 return 代替 panic,因为 DashMap 不支持 catch_unwind)
|
||||
let early_return = || {
|
||||
let _guard = PendingGuard::new(®istry, "test-project");
|
||||
// 早期返回(模拟错误场景)
|
||||
return false;
|
||||
};
|
||||
|
||||
// 调用后,guard 已经被 drop,应该被清理
|
||||
early_return();
|
||||
|
||||
assert!(!registry.contains_project("test-project"));
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 2. 原子计数器并发安全性测试
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_atomic_slot_counter_concurrent_acquisition() {
|
||||
// 重置并发限制为默认值,防止其他测试影响
|
||||
init_concurrency_limit(10);
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
let limit = get_concurrency_limit();
|
||||
let num_tasks = limit * 2;
|
||||
let barrier = Arc::new(Barrier::new(num_tasks));
|
||||
let successful_count = Arc::new(AtomicUsize::new(0));
|
||||
let failed_count = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
let mut handles = vec![];
|
||||
|
||||
// 启动并发任务尝试获取槽位
|
||||
for i in 0..num_tasks {
|
||||
let registry_clone = registry.clone();
|
||||
let barrier_clone = barrier.clone();
|
||||
let successful_count_clone = successful_count.clone();
|
||||
let failed_count_clone = failed_count.clone();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
// 等待所有任务就绪
|
||||
barrier_clone.wait().await;
|
||||
|
||||
// 尝试获取槽位
|
||||
if registry_clone.try_acquire_session_slot() {
|
||||
successful_count_clone.fetch_add(1, Ordering::Relaxed);
|
||||
// 模拟工作
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
registry_clone.release_session_slot();
|
||||
} else {
|
||||
failed_count_clone.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
});
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
// 等待所有任务完成
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
}
|
||||
|
||||
// 验证: 只有 WORKER_THREAD_POOL_SIZE 个任务成功
|
||||
let successful = successful_count.load(Ordering::Relaxed);
|
||||
let failed = failed_count.load(Ordering::Relaxed);
|
||||
let limit = get_concurrency_limit();
|
||||
|
||||
assert_eq!(successful, limit, "应该有 {} 个任务成功获取槽位", limit);
|
||||
assert_eq!(failed, limit, "应该有 {} 个任务失败(槽位已满)", limit);
|
||||
|
||||
// 验证: 计数器最终应该回到 0
|
||||
assert_eq!(registry.active_sessions_count(), 0, "所有槽位应该被释放");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_atomic_slot_counter_stress_test() {
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
let num_iterations = 1000;
|
||||
let successful_count = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
let mut handles = vec![];
|
||||
|
||||
// 启动大量并发任务
|
||||
for _ in 0..50 {
|
||||
let registry_clone = registry.clone();
|
||||
let successful_count_clone = successful_count.clone();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
for j in 0..num_iterations {
|
||||
if registry_clone.try_acquire_session_slot() {
|
||||
successful_count_clone.fetch_add(1, Ordering::Relaxed);
|
||||
// 使用简单的随机性(不依赖 rand)
|
||||
let delay = (j % 10) as u64;
|
||||
tokio::time::sleep(Duration::from_micros(delay * 10)).await;
|
||||
registry_clone.release_session_slot();
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
// 等待所有任务完成
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
}
|
||||
|
||||
// 验证: 计数器最终应该回到 0
|
||||
assert_eq!(registry.active_sessions_count(), 0, "所有槽位应该被释放");
|
||||
|
||||
println!(
|
||||
"压力测试完成: {} 次成功获取槽位",
|
||||
successful_count.load(Ordering::Relaxed)
|
||||
);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 3. Agent 并发独立性测试
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_concurrent_agents_independence() {
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
let num_agents = 10;
|
||||
let barrier = Arc::new(Barrier::new(num_agents));
|
||||
let mut handles = vec![];
|
||||
|
||||
// 并发创建多个 agent
|
||||
for i in 0..num_agents {
|
||||
let registry_clone = registry.clone();
|
||||
let barrier_clone = barrier.clone();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
let project_id = format!("project-{}", i);
|
||||
let session_id = format!("session-{}", i);
|
||||
|
||||
// 创建 AgentInfo
|
||||
let (prompt_tx, _) = mpsc::channel(100);
|
||||
let (cancel_tx, _) = mpsc::channel(100);
|
||||
|
||||
let agent_info = ProjectAndAgentInfo {
|
||||
project_id: project_id.clone(),
|
||||
session_id: SessionId::new(Arc::from(session_id.as_str())),
|
||||
prompt_tx,
|
||||
cancel_tx,
|
||||
model_provider: None,
|
||||
request_id: None,
|
||||
status: AgentStatus::Active,
|
||||
last_activity: chrono::Utc::now(),
|
||||
created_at: chrono::Utc::now(),
|
||||
stop_handle: None,
|
||||
};
|
||||
|
||||
// 注册 agent
|
||||
registry_clone.register(&project_id, &session_id, agent_info);
|
||||
|
||||
// 等待所有 agent 就绪
|
||||
barrier_clone.wait().await;
|
||||
|
||||
// 验证: 当前 agent 存在
|
||||
assert!(registry_clone.contains_project(&project_id));
|
||||
|
||||
// 验证: 其他 agent 也存在(不会相互覆盖)
|
||||
let stats = registry_clone.stats();
|
||||
assert_eq!(stats.agent_count, num_agents);
|
||||
|
||||
// 模拟工作
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
|
||||
// 清理
|
||||
registry_clone.remove_by_project(&project_id);
|
||||
});
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
// 等待所有任务完成
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
}
|
||||
|
||||
// 验证: 所有 agent 已被清理
|
||||
assert_eq!(registry.stats().agent_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_concurrent_agent_state_updates() {
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
let project_id = "test-project";
|
||||
let session_id = "test-session";
|
||||
|
||||
// 创建初始 agent
|
||||
let (prompt_tx, _) = mpsc::channel(100);
|
||||
let (cancel_tx, _) = mpsc::channel(100);
|
||||
|
||||
let agent_info = ProjectAndAgentInfo {
|
||||
project_id: project_id.to_string(),
|
||||
session_id: SessionId::new(Arc::from(session_id)),
|
||||
prompt_tx,
|
||||
cancel_tx,
|
||||
model_provider: None,
|
||||
request_id: None,
|
||||
status: AgentStatus::Idle,
|
||||
last_activity: chrono::Utc::now(),
|
||||
created_at: chrono::Utc::now(),
|
||||
stop_handle: None,
|
||||
};
|
||||
|
||||
registry.register(project_id, session_id, agent_info);
|
||||
|
||||
// 并发更新状态
|
||||
let num_updates = 100;
|
||||
let mut handles = vec![];
|
||||
|
||||
for i in 0..10 {
|
||||
let registry_clone = registry.clone();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
for j in 0..num_updates {
|
||||
// 使用原子性更新
|
||||
registry_clone.try_update_agent_info(project_id, |info| {
|
||||
// 模拟状态切换
|
||||
if j % 2 == 0 {
|
||||
info.status = AgentStatus::Active;
|
||||
} else {
|
||||
info.status = AgentStatus::Idle;
|
||||
}
|
||||
info.last_activity = chrono::Utc::now();
|
||||
true
|
||||
});
|
||||
tokio::time::sleep(Duration::from_micros(10)).await;
|
||||
}
|
||||
});
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
// 等待所有更新完成
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
}
|
||||
|
||||
// 验证: agent 仍然存在,没有数据损坏
|
||||
assert!(registry.contains_project(project_id));
|
||||
let info = registry.get_agent_info(project_id).unwrap();
|
||||
assert!(matches!(
|
||||
info.status,
|
||||
AgentStatus::Active | AgentStatus::Idle
|
||||
));
|
||||
|
||||
// 清理
|
||||
registry.remove_by_project(project_id);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 4. Agent 销毁测试
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_lifecycle_cleanup() {
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
|
||||
// 创建多个 agent
|
||||
let num_agents = 5;
|
||||
for i in 0..num_agents {
|
||||
let project_id = format!("project-{}", i);
|
||||
let session_id = format!("session-{}", i);
|
||||
|
||||
let (prompt_tx, _) = mpsc::channel(100);
|
||||
let (cancel_tx, _) = mpsc::channel(100);
|
||||
|
||||
let agent_info = ProjectAndAgentInfo {
|
||||
project_id: project_id.clone(),
|
||||
session_id: SessionId::new(Arc::from(session_id.as_str())),
|
||||
prompt_tx,
|
||||
cancel_tx,
|
||||
model_provider: None,
|
||||
request_id: None,
|
||||
status: AgentStatus::Active,
|
||||
last_activity: chrono::Utc::now(),
|
||||
created_at: chrono::Utc::now(),
|
||||
stop_handle: None,
|
||||
};
|
||||
|
||||
registry.register(&project_id, &session_id, agent_info);
|
||||
}
|
||||
|
||||
// 验证: 所有 agent 已注册
|
||||
assert_eq!(registry.stats().agent_count, num_agents);
|
||||
|
||||
// 销毁所有 agent
|
||||
for i in 0..num_agents {
|
||||
let project_id = format!("project-{}", i);
|
||||
let removed = registry.remove_by_project(&project_id);
|
||||
assert!(removed.is_some(), "应该能移除 agent");
|
||||
}
|
||||
|
||||
// 验证: 所有 agent 已被清理
|
||||
assert_eq!(registry.stats().agent_count, 0);
|
||||
|
||||
// 验证: 映射关系也被清理
|
||||
for i in 0..num_agents {
|
||||
let session_id = format!("session-{}", i);
|
||||
assert!(!registry.contains_session(&session_id));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_concurrent_removal() {
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
|
||||
// 创建大量 agent
|
||||
let num_agents = 100;
|
||||
for i in 0..num_agents {
|
||||
let project_id = format!("project-{}", i);
|
||||
let session_id = format!("session-{}", i);
|
||||
|
||||
let (prompt_tx, _) = mpsc::channel(100);
|
||||
let (cancel_tx, _) = mpsc::channel(100);
|
||||
|
||||
let agent_info = ProjectAndAgentInfo {
|
||||
project_id: project_id.clone(),
|
||||
session_id: SessionId::new(Arc::from(session_id.as_str())),
|
||||
prompt_tx,
|
||||
cancel_tx,
|
||||
model_provider: None,
|
||||
request_id: None,
|
||||
status: AgentStatus::Active,
|
||||
last_activity: chrono::Utc::now(),
|
||||
created_at: chrono::Utc::now(),
|
||||
stop_handle: None,
|
||||
};
|
||||
|
||||
registry.register(&project_id, &session_id, agent_info);
|
||||
}
|
||||
|
||||
// 并发移除所有 agent
|
||||
let mut handles = vec![];
|
||||
for i in 0..num_agents {
|
||||
let registry_clone = registry.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
let project_id = format!("project-{}", i);
|
||||
registry_clone.remove_by_project(&project_id);
|
||||
});
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
// 等待所有移除完成
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
}
|
||||
|
||||
// 验证: 所有 agent 已被清理
|
||||
assert_eq!(registry.stats().agent_count, 0);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 5. RAII 快速销毁测试
|
||||
// ============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_raii_fast_destruction() {
|
||||
let registry = AgentSessionRegistry::new();
|
||||
let num_guards = 1000;
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// 创建大量 guard
|
||||
for i in 0..num_guards {
|
||||
let project_id = format!("project-{}", i);
|
||||
let _guard = PendingGuard::new(®istry, &project_id);
|
||||
// guard 立即被 drop
|
||||
}
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
// 验证: 销毁应该很快(< 10ms)
|
||||
assert!(
|
||||
elapsed.as_millis() < 10,
|
||||
"RAII 销毁应该快速完成,实际耗时: {:?}",
|
||||
elapsed
|
||||
);
|
||||
|
||||
// 验证: 所有项目都被清理
|
||||
assert_eq!(registry.stats().agent_count, 0);
|
||||
|
||||
println!(
|
||||
"RAII 快速销毁测试: {} 个 guard 在 {:?} 内销毁",
|
||||
num_guards, elapsed
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pending_guard_with_tokio_spawn() {
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
let num_tasks = 50;
|
||||
|
||||
let mut handles = vec![];
|
||||
|
||||
// 并发创建 guard
|
||||
for i in 0..num_tasks {
|
||||
let registry_clone = registry.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
let project_id = format!("project-{}", i);
|
||||
let _guard = PendingGuard::new(®istry_clone, &project_id);
|
||||
// 模拟异步工作
|
||||
tokio::time::sleep(Duration::from_millis(1)).await;
|
||||
// guard 在这里被 drop
|
||||
});
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
// 等待所有任务完成
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
}
|
||||
|
||||
// 验证: 所有项目都被清理
|
||||
assert_eq!(registry.stats().agent_count, 0);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 6. 边界条件测试
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_slot_counter_underflow_protection() {
|
||||
let registry = AgentSessionRegistry::new();
|
||||
|
||||
// 尝试释放从未获取的槽位
|
||||
for _ in 0..10 {
|
||||
registry.release_session_slot();
|
||||
}
|
||||
|
||||
// 验证: 计数器不会下溢(使用 saturating_sub)
|
||||
let count = registry.active_sessions_count();
|
||||
assert_eq!(count, 0, "计数器应该保持为 0,不会下溢");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_pending_guards_same_project() {
|
||||
let registry = AgentSessionRegistry::new();
|
||||
let project_id = "test-project";
|
||||
|
||||
// 创建多个 guard(模拟并发请求)
|
||||
{
|
||||
let _guard1 = PendingGuard::new(®istry, project_id);
|
||||
// 第二个 guard 会更新现有项目为 Pending(已经是 Pending,无操作)
|
||||
let _guard2 = PendingGuard::new(®istry, project_id);
|
||||
|
||||
let info = registry.get_agent_info(project_id).unwrap();
|
||||
assert_eq!(format!("{:?}", info.status), "Pending");
|
||||
}
|
||||
|
||||
// 所有 guard 都 drop,应该被清理
|
||||
assert!(!registry.contains_project(project_id));
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 7. 压力测试:高并发场景
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_high_concurrency_stress() {
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
let num_requests = 1000;
|
||||
let barrier = Arc::new(Barrier::new(num_requests));
|
||||
let success_count = Arc::new(AtomicUsize::new(0));
|
||||
let fail_count = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
let mut handles = vec![];
|
||||
|
||||
// 模拟高并发请求
|
||||
for i in 0..num_requests {
|
||||
let registry_clone = registry.clone();
|
||||
let barrier_clone = barrier.clone();
|
||||
let success_count_clone = success_count.clone();
|
||||
let fail_count_clone = fail_count.clone();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
let project_id = format!("project-{}", i);
|
||||
|
||||
// 等待所有任务就绪
|
||||
barrier_clone.wait().await;
|
||||
|
||||
// 尝试获取槽位
|
||||
if registry_clone.try_acquire_session_slot() {
|
||||
success_count_clone.fetch_add(1, Ordering::Relaxed);
|
||||
|
||||
// 使用 PendingGuard
|
||||
let _guard = PendingGuard::new(®istry_clone, &project_id);
|
||||
|
||||
// 模拟工作
|
||||
let delay = (i % 10) as u64;
|
||||
tokio::time::sleep(Duration::from_millis(delay)).await;
|
||||
|
||||
// 正常流程:禁用 guard,手动释放
|
||||
drop(_guard);
|
||||
registry_clone.release_session_slot();
|
||||
registry_clone.clear_pending_if_exists(&project_id);
|
||||
} else {
|
||||
fail_count_clone.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
});
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
// 等待所有任务完成
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
}
|
||||
|
||||
// 验证: 只有 WORKER_THREAD_POOL_SIZE 个请求成功
|
||||
let success = success_count.load(Ordering::Relaxed);
|
||||
let fail = fail_count.load(Ordering::Relaxed);
|
||||
|
||||
assert_eq!(success + fail, num_requests, "所有请求都应该被处理");
|
||||
|
||||
assert_eq!(registry.active_sessions_count(), 0, "所有槽位应该被释放");
|
||||
|
||||
println!(
|
||||
"High-concurrency stress test: {} success, {} failed",
|
||||
success, fail
|
||||
);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 8. PendingGuard 与 SessionManager 竞态条件修复测试
|
||||
// ============================================================================
|
||||
|
||||
/// 测试场景:PendingGuard 创建的占位符应该被真实会话替换
|
||||
///
|
||||
/// 这是修复的核心场景:
|
||||
/// 1. PendingGuard 创建 pending 占位符
|
||||
/// 2. 模拟 SessionManager 检测到 pending 并创建真实会话
|
||||
/// 3. 验证 pending 占位符被正确替换
|
||||
#[tokio::test]
|
||||
async fn test_pending_placeholder_replaced_by_real_session() {
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
let project_id = "test-pending-replace";
|
||||
|
||||
// 第一阶段:创建 PendingGuard(模拟 gRPC 层的行为)
|
||||
let guard = PendingGuard::new(®istry, project_id);
|
||||
|
||||
// 验证:pending 占位符已创建
|
||||
assert!(registry.contains_project(project_id));
|
||||
{
|
||||
let pending_info = registry.get_agent_info(project_id).unwrap();
|
||||
assert_eq!(format!("{:?}", pending_info.status), "Pending");
|
||||
assert_eq!(pending_info.session_id.to_string(), "pending");
|
||||
} // 释放 Ref 锁
|
||||
|
||||
// 第二阶段:模拟 SessionManager 检测到 Pending 状态并创建真实会话
|
||||
// 检查状态(模拟 session_manager.rs 中的逻辑)
|
||||
let should_replace = {
|
||||
let info = registry.get_agent_info(project_id).unwrap();
|
||||
*info.status() == AgentStatus::Pending
|
||||
}; // 释放 Ref 锁
|
||||
assert!(should_replace, "应该检测到 Pending 占位符");
|
||||
|
||||
// 创建真实会话
|
||||
let real_session_id = "real-session-123";
|
||||
let (prompt_tx, _prompt_rx) = mpsc::channel(100);
|
||||
let (cancel_tx, _cancel_rx) = mpsc::channel(100);
|
||||
|
||||
let real_session = ProjectAndAgentInfo {
|
||||
project_id: project_id.to_string(),
|
||||
session_id: SessionId::new(Arc::from(real_session_id)),
|
||||
prompt_tx,
|
||||
cancel_tx,
|
||||
model_provider: None,
|
||||
request_id: None,
|
||||
status: AgentStatus::Idle,
|
||||
last_activity: chrono::Utc::now(),
|
||||
created_at: chrono::Utc::now(),
|
||||
stop_handle: None,
|
||||
};
|
||||
|
||||
// 第三阶段:原子性替换(模拟 session_manager.rs 中的 Entry API 逻辑)
|
||||
// 使用 DashMap 的 entry API 进行原子性替换
|
||||
use dashmap::mapref::entry::Entry;
|
||||
match registry.as_ref().inner_mut().entry(project_id.to_string()) {
|
||||
Entry::Vacant(entry) => {
|
||||
entry.insert(real_session.clone());
|
||||
}
|
||||
Entry::Occupied(mut entry) => {
|
||||
// 检查仍然是 Pending(防止其他线程已经插入了真实会话)
|
||||
// 提取状态值,避免借用冲突
|
||||
let is_pending = {
|
||||
let existing = entry.get();
|
||||
*existing.status() == AgentStatus::Pending
|
||||
};
|
||||
if is_pending {
|
||||
entry.insert(real_session.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 验证:pending 占位符已被替换为真实会话
|
||||
assert!(registry.contains_project(project_id));
|
||||
let final_info = registry.get_agent_info(project_id).unwrap();
|
||||
assert_eq!(format!("{:?}", final_info.status), "Idle");
|
||||
assert_eq!(final_info.session_id.to_string(), real_session_id);
|
||||
assert!(!final_info.prompt_tx.is_closed());
|
||||
|
||||
// PendingGuard 不需要 commit(因为 pending 已被替换)
|
||||
drop(guard);
|
||||
|
||||
// 验证:真实会话仍然存在(没有被 PendingGuard 清理)
|
||||
assert!(registry.contains_project(project_id));
|
||||
|
||||
// 清理
|
||||
registry.remove_by_project(project_id);
|
||||
}
|
||||
|
||||
/// 测试场景:并发创建时,只有一个真实会话被保留
|
||||
///
|
||||
/// 验证修复的并发安全性:
|
||||
/// 1. PendingGuard 创建 pending 占位符
|
||||
/// 2. 多个线程尝试替换 pending
|
||||
/// 3. 只有一个真实会话被保留
|
||||
#[tokio::test]
|
||||
async fn test_concurrent_pending_replacement() {
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
let project_id = "test-concurrent-replace";
|
||||
let num_threads = 5;
|
||||
let barrier = Arc::new(Barrier::new(num_threads));
|
||||
let success_count = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
// 第一阶段:创建 PendingGuard
|
||||
let _guard = PendingGuard::new(®istry, project_id);
|
||||
|
||||
// 验证 pending 占位符
|
||||
assert!(registry.contains_project(project_id));
|
||||
{
|
||||
let pending_info = registry.get_agent_info(project_id).unwrap();
|
||||
assert_eq!(format!("{:?}", pending_info.status), "Pending");
|
||||
} // 释放 Ref 锁
|
||||
|
||||
// 第二阶段:多个线程并发尝试替换 pending
|
||||
let mut handles = vec![];
|
||||
for i in 0..num_threads {
|
||||
let registry_clone = registry.clone();
|
||||
let barrier_clone = barrier.clone();
|
||||
let success_count_clone = success_count.clone();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
// 等待所有线程就绪
|
||||
barrier_clone.wait().await;
|
||||
|
||||
// 每个线程创建一个"真实会话"
|
||||
let session_id = format!("session-{}", i);
|
||||
let (prompt_tx, _prompt_rx) = mpsc::channel(100);
|
||||
let (cancel_tx, _cancel_rx) = mpsc::channel(100);
|
||||
|
||||
let real_session = ProjectAndAgentInfo {
|
||||
project_id: project_id.to_string(),
|
||||
session_id: SessionId::new(Arc::from(session_id.clone())),
|
||||
prompt_tx,
|
||||
cancel_tx,
|
||||
model_provider: None,
|
||||
request_id: None,
|
||||
status: AgentStatus::Idle,
|
||||
last_activity: chrono::Utc::now(),
|
||||
created_at: chrono::Utc::now(),
|
||||
stop_handle: None,
|
||||
};
|
||||
|
||||
// 尝试原子性替换
|
||||
use dashmap::mapref::entry::Entry;
|
||||
let replaced = match registry_clone
|
||||
.as_ref()
|
||||
.inner_mut()
|
||||
.entry(project_id.to_string())
|
||||
{
|
||||
Entry::Occupied(mut entry) => {
|
||||
// 只有 pending 才替换
|
||||
// 提取状态值,避免借用冲突
|
||||
let is_pending = {
|
||||
let existing = entry.get();
|
||||
*existing.status() == AgentStatus::Pending
|
||||
};
|
||||
if is_pending {
|
||||
entry.insert(real_session.clone());
|
||||
success_count_clone.fetch_add(1, Ordering::Relaxed);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
Entry::Vacant(_) => false,
|
||||
};
|
||||
|
||||
replaced
|
||||
});
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
// 等待所有线程完成
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
}
|
||||
|
||||
// 验证:只有一个线程成功替换
|
||||
let success = success_count.load(Ordering::Relaxed);
|
||||
assert_eq!(success, 1, "应该只有一个线程成功替换 pending");
|
||||
|
||||
// 验证:最终只有一个会话存在
|
||||
assert!(registry.contains_project(project_id));
|
||||
let final_info = registry.get_agent_info(project_id).unwrap();
|
||||
assert_eq!(format!("{:?}", final_info.status), "Idle");
|
||||
|
||||
// 清理
|
||||
registry.remove_by_project(project_id);
|
||||
}
|
||||
|
||||
/// 测试场景:模拟真实的 session_manager.rs 逻辑流程
|
||||
///
|
||||
/// 这是一个端到端测试,模拟完整的修复流程:
|
||||
/// 1. PendingGuard 创建占位符
|
||||
/// 2. 检测到 Pending 状态
|
||||
/// 3. 释放锁,创建真实会话
|
||||
/// 4. 原子性插入/替换
|
||||
#[tokio::test]
|
||||
async fn test_session_manager_pending_replacement_flow() {
|
||||
let registry = Arc::new(AgentSessionRegistry::new());
|
||||
let project_id = "test-e2e-flow";
|
||||
|
||||
// ========== 第一阶段:PendingGuard 创建占位符 ==========
|
||||
{
|
||||
let _guard = PendingGuard::new(®istry, project_id);
|
||||
|
||||
// 验证占位符
|
||||
{
|
||||
let info = registry.get_agent_info(project_id).unwrap();
|
||||
assert_eq!(format!("{:?}", info.status), "Pending");
|
||||
} // 释放 Ref 锁
|
||||
|
||||
// ========== 第二阶段:模拟 SessionManager 的 get_or_create_session ==========
|
||||
|
||||
// 2.1 快速检查:发现 entry 存在
|
||||
let entry_exists = registry.contains_project(project_id);
|
||||
assert!(entry_exists);
|
||||
|
||||
// 2.2 显式检查 Pending 状态
|
||||
let should_replace = {
|
||||
let info = registry.get_agent_info(project_id).unwrap();
|
||||
*info.status() == AgentStatus::Pending
|
||||
}; // 释放 Ref 锁
|
||||
assert!(should_replace, "应该检测到 Pending 状态");
|
||||
|
||||
// 2.3 创建真实会话(不持有锁)
|
||||
let real_session_id = "ses_real_12345";
|
||||
let (prompt_tx, _prompt_rx) = mpsc::channel(100);
|
||||
let (cancel_tx, _cancel_rx) = mpsc::channel(100);
|
||||
|
||||
let real_session = ProjectAndAgentInfo {
|
||||
project_id: project_id.to_string(),
|
||||
session_id: SessionId::new(Arc::from(real_session_id)),
|
||||
prompt_tx,
|
||||
cancel_tx,
|
||||
model_provider: None,
|
||||
request_id: None,
|
||||
status: AgentStatus::Idle,
|
||||
last_activity: chrono::Utc::now(),
|
||||
created_at: chrono::Utc::now(),
|
||||
stop_handle: None,
|
||||
};
|
||||
|
||||
// ========== 第三阶段:原子性替换 ==========
|
||||
|
||||
// 使用 DashMap entry API 进行原子性操作
|
||||
use dashmap::mapref::entry::Entry;
|
||||
let was_pending = match registry.as_ref().inner_mut().entry(project_id.to_string()) {
|
||||
Entry::Occupied(mut entry) => {
|
||||
// 提取状态值,避免借用冲突
|
||||
let is_pending = {
|
||||
let existing = entry.get();
|
||||
*existing.status() == AgentStatus::Pending
|
||||
};
|
||||
if is_pending {
|
||||
entry.insert(real_session.clone());
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
Entry::Vacant(_) => false,
|
||||
};
|
||||
|
||||
assert!(was_pending, "应该成功替换 pending 占位符");
|
||||
|
||||
// 验证:真实会话已插入
|
||||
let final_info = registry.get_agent_info(project_id).unwrap();
|
||||
assert_eq!(final_info.session_id.to_string(), real_session_id);
|
||||
assert_eq!(format!("{:?}", final_info.status), "Idle");
|
||||
assert!(!final_info.prompt_tx.is_closed());
|
||||
}
|
||||
|
||||
// PendingGuard 已 drop,但真实会话应该保留
|
||||
assert!(registry.contains_project(project_id));
|
||||
let final_info = registry.get_agent_info(project_id).unwrap();
|
||||
assert_eq!(format!("{:?}", final_info.status), "Idle");
|
||||
|
||||
// 清理
|
||||
registry.remove_by_project(project_id);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 调试测试 - 定位挂起问题
|
||||
// ============================================================================
|
||||
|
||||
#[test]
|
||||
fn debug_simple_registry_operations() {
|
||||
use agent_runner::service::AgentSessionRegistry;
|
||||
use shared_types::SessionEntry;
|
||||
|
||||
let registry = AgentSessionRegistry::new();
|
||||
registry.set_pending("test-1");
|
||||
|
||||
// 测试 contains_project
|
||||
assert!(registry.contains_project("test-1"));
|
||||
|
||||
// 测试 get_agent_info
|
||||
if let Some(info) = registry.get_agent_info("test-1") {
|
||||
// 直接访问字段而不是通过 trait 方法
|
||||
let status = &info.status;
|
||||
assert_eq!(format!("{:?}", status), "Pending");
|
||||
} else {
|
||||
panic!("get_agent_info returned None");
|
||||
}
|
||||
|
||||
registry.remove_by_project("test-1");
|
||||
assert!(!registry.contains_project("test-1"));
|
||||
}
|
||||
@@ -0,0 +1,300 @@
|
||||
//! 超时保护测试
|
||||
//!
|
||||
//! 验证超时保护机制的正确性:
|
||||
//! - 使用 Tokio paused_time 加速测试
|
||||
//! - 验证超时后返回错误
|
||||
//! - 验证超时后清理资源
|
||||
|
||||
use std::time::Duration;
|
||||
use tokio::time::timeout;
|
||||
|
||||
/// 测试基础超时机制
|
||||
#[tokio::test]
|
||||
async fn test_basic_timeout_mechanism() {
|
||||
// 使用 timeout 包装一个长时间运行的操作
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let result = timeout(Duration::from_millis(100), async {
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
"completed"
|
||||
})
|
||||
.await;
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
// 验证: 应该在 100ms 左右超时,而不是等待 5 秒
|
||||
assert!(result.is_err(), "应该在 100ms 时超时");
|
||||
assert!(elapsed.as_millis() < 200, "超时测试应该快速完成");
|
||||
}
|
||||
|
||||
/// 测试 paused_time 加速超时测试
|
||||
///
|
||||
/// 这是 Tokio 的最佳实践: 使用 paused_time 可以让长时间的 sleep 立即完成
|
||||
///
|
||||
/// 注意:在 paused_time 模式下:
|
||||
/// - tokio::time::Instant 会跟随虚拟时间
|
||||
/// - std::time::Instant 保持真实墙钟时间
|
||||
/// - 我们使用 std::time::Instant 来验证测试确实快速完成
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn test_timeout_with_paused_time() {
|
||||
// 使用 std::time::Instant 来测量真实墙钟时间
|
||||
let wall_clock_start = std::time::Instant::now();
|
||||
|
||||
// 注入 200 秒的阻塞 (但在 paused_time 下会立即完成)
|
||||
let result = timeout(Duration::from_secs(100), async {
|
||||
tokio::time::sleep(Duration::from_secs(200)).await;
|
||||
"completed"
|
||||
})
|
||||
.await;
|
||||
|
||||
let wall_clock_elapsed = wall_clock_start.elapsed();
|
||||
|
||||
// 验证: 应该在 100 "虚拟秒" 后超时
|
||||
assert!(result.is_err(), "应该在 100 秒时超时");
|
||||
|
||||
// 验证: 真实墙钟时间应该非常短 (< 100ms)
|
||||
assert!(
|
||||
wall_clock_elapsed.as_millis() < 100,
|
||||
"paused_time 测试应该瞬间完成,实际耗时: {:?}",
|
||||
wall_clock_elapsed
|
||||
);
|
||||
}
|
||||
|
||||
/// 测试超时前的成功完成
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn test_success_before_timeout() {
|
||||
let wall_clock_start = std::time::Instant::now();
|
||||
|
||||
// 任务在 50 秒完成,超时时间是 100 秒
|
||||
let result = timeout(Duration::from_secs(100), async {
|
||||
tokio::time::sleep(Duration::from_secs(50)).await;
|
||||
"completed"
|
||||
})
|
||||
.await;
|
||||
|
||||
let wall_clock_elapsed = wall_clock_start.elapsed();
|
||||
|
||||
// 验证: 应该成功完成
|
||||
assert!(result.is_ok(), "应该在 50 秒时成功完成");
|
||||
assert_eq!(result.unwrap(), "completed");
|
||||
|
||||
// 验证: 真实墙钟时间应该非常短
|
||||
assert!(
|
||||
wall_clock_elapsed.as_millis() < 100,
|
||||
"paused_time 测试应该快速完成"
|
||||
);
|
||||
}
|
||||
|
||||
/// 测试分级超时警告
|
||||
#[tokio::test]
|
||||
async fn test_graduated_timeout_warnings() {
|
||||
// 测试不同时长的请求应该触发不同级别的警告
|
||||
let test_cases = vec![
|
||||
(30, false, false), // 30 秒: 无警告
|
||||
(65, true, false), // 65 秒: 黄色警告 (> 60s)
|
||||
(125, true, true), // 125 秒: 红色警告 (> 120s)
|
||||
];
|
||||
|
||||
for (duration_seconds, should_yellow_warn, should_red_warn) in test_cases {
|
||||
let has_yellow = duration_seconds > 60;
|
||||
let has_red = duration_seconds > 120;
|
||||
|
||||
assert_eq!(
|
||||
has_yellow, should_yellow_warn,
|
||||
"{}秒请求的黄色警告判断错误",
|
||||
duration_seconds
|
||||
);
|
||||
assert_eq!(
|
||||
has_red, should_red_warn,
|
||||
"{}秒请求的红色警告判断错误",
|
||||
duration_seconds
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 测试超时后资源清理(使用 RAII Guard 模式)
|
||||
#[tokio::test]
|
||||
async fn test_resource_cleanup_after_timeout() {
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
let active_requests = Arc::new(std::sync::Mutex::new(HashMap::new()));
|
||||
let request_id = "timeout-request".to_string();
|
||||
|
||||
// 使用 RAII Guard 模式测试资源清理
|
||||
struct RequestGuard {
|
||||
active_requests: Arc<std::sync::Mutex<HashMap<String, chrono::DateTime<chrono::Utc>>>>,
|
||||
request_id: String,
|
||||
}
|
||||
|
||||
impl RequestGuard {
|
||||
fn new(
|
||||
active_requests: Arc<std::sync::Mutex<HashMap<String, chrono::DateTime<chrono::Utc>>>>,
|
||||
request_id: String,
|
||||
) -> Self {
|
||||
// 在单独的作用域中获取锁,确保锁在 move 之前释放
|
||||
{
|
||||
let mut reqs = active_requests.lock().unwrap();
|
||||
reqs.insert(request_id.clone(), chrono::Utc::now());
|
||||
}
|
||||
Self {
|
||||
active_requests,
|
||||
request_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for RequestGuard {
|
||||
fn drop(&mut self) {
|
||||
let mut reqs = self.active_requests.lock().unwrap();
|
||||
reqs.remove(&self.request_id);
|
||||
}
|
||||
}
|
||||
|
||||
// 创建 guard
|
||||
{
|
||||
let _guard = RequestGuard::new(active_requests.clone(), request_id.clone());
|
||||
assert_eq!(active_requests.lock().unwrap().len(), 1);
|
||||
// guard 在这里被 drop
|
||||
}
|
||||
|
||||
// 验证: 资源已被 Drop 自动清理
|
||||
assert_eq!(
|
||||
active_requests.lock().unwrap().len(),
|
||||
0,
|
||||
"RAII Guard 应该在 drop 时自动清理资源"
|
||||
);
|
||||
}
|
||||
|
||||
/// 测试多个请求的超时检测
|
||||
#[tokio::test]
|
||||
async fn test_multiple_requests_timeout_detection() {
|
||||
use std::collections::HashMap;
|
||||
|
||||
let mut active_requests: HashMap<String, chrono::DateTime<chrono::Utc>> = HashMap::new();
|
||||
let now = chrono::Utc::now();
|
||||
|
||||
// 添加不同时长的请求
|
||||
active_requests.insert("req-1".to_string(), now - chrono::Duration::seconds(30)); // 30 秒 - 正常
|
||||
active_requests.insert("req-2".to_string(), now - chrono::Duration::seconds(70)); // 70 秒 - 黄色警告
|
||||
active_requests.insert("req-3".to_string(), now - chrono::Duration::seconds(130)); // 130 秒 - 红色警告
|
||||
|
||||
let mut normal_count = 0;
|
||||
let mut yellow_count = 0;
|
||||
let mut red_count = 0;
|
||||
|
||||
for start_time in active_requests.values() {
|
||||
let duration = (now - *start_time).num_seconds();
|
||||
|
||||
if duration > 120 {
|
||||
red_count += 1;
|
||||
} else if duration > 60 {
|
||||
yellow_count += 1;
|
||||
} else {
|
||||
normal_count += 1;
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(normal_count, 1, "应该有 1 个正常请求");
|
||||
assert_eq!(yellow_count, 1, "应该有 1 个黄色警告");
|
||||
assert_eq!(red_count, 1, "应该有 1 个红色警告");
|
||||
}
|
||||
|
||||
/// 测试超时阈值配置验证
|
||||
#[tokio::test]
|
||||
async fn test_timeout_threshold_configuration() {
|
||||
// 测试不同的超时阈值配置
|
||||
struct TimeoutConfig {
|
||||
name: &'static str,
|
||||
threshold_seconds: u64,
|
||||
description: &'static str,
|
||||
}
|
||||
|
||||
let configs = vec![
|
||||
TimeoutConfig {
|
||||
name: "new_session",
|
||||
threshold_seconds: 100,
|
||||
description: "MCP 服务器启动可能较慢",
|
||||
},
|
||||
TimeoutConfig {
|
||||
name: "monitor_warn",
|
||||
threshold_seconds: 60,
|
||||
description: "监控黄色警告阈值",
|
||||
},
|
||||
TimeoutConfig {
|
||||
name: "monitor_error",
|
||||
threshold_seconds: 120,
|
||||
description: "监控红色警告阈值",
|
||||
},
|
||||
TimeoutConfig {
|
||||
name: "monitor_restart",
|
||||
threshold_seconds: 180,
|
||||
description: "触发重启阈值",
|
||||
},
|
||||
];
|
||||
|
||||
for config in configs {
|
||||
assert!(
|
||||
config.threshold_seconds > 0,
|
||||
"{} 超时阈值应该大于 0",
|
||||
config.name
|
||||
);
|
||||
println!(
|
||||
"配置: {} = {}秒 ({})",
|
||||
config.name, config.threshold_seconds, config.description
|
||||
);
|
||||
}
|
||||
|
||||
// 验证阈值递增关系
|
||||
assert!(60 < 120, "warn 阈值应小于 error 阈值");
|
||||
assert!(120 < 180, "error 阈值应小于 restart 阈值");
|
||||
}
|
||||
|
||||
/// 测试超时不会导致死锁
|
||||
#[tokio::test]
|
||||
async fn test_timeout_does_not_cause_deadlock() {
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
let task_started = Arc::new(AtomicBool::new(false));
|
||||
let task_started_clone = task_started.clone();
|
||||
|
||||
// 启动一个任务
|
||||
let handle = tokio::spawn(async move {
|
||||
task_started_clone.store(true, Ordering::SeqCst);
|
||||
// 模拟长时间操作
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
"completed"
|
||||
});
|
||||
|
||||
// 等待任务启动
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
|
||||
// 验证任务已启动
|
||||
assert!(task_started.load(Ordering::SeqCst), "任务应该已启动");
|
||||
|
||||
// 使用短超时等待
|
||||
let result = timeout(Duration::from_millis(100), handle).await;
|
||||
|
||||
// 验证: 超时发生,但没有死锁
|
||||
assert!(result.is_err(), "应该超时");
|
||||
}
|
||||
|
||||
/// 测试 tokio::time::Instant 在 paused_time 模式下的行为
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn test_tokio_instant_with_paused_time() {
|
||||
// tokio::time::Instant 跟随虚拟时间
|
||||
let tokio_start = tokio::time::Instant::now();
|
||||
|
||||
// 前进 10 秒虚拟时间
|
||||
tokio::time::sleep(Duration::from_secs(10)).await;
|
||||
|
||||
let tokio_elapsed = tokio_start.elapsed();
|
||||
|
||||
// tokio::time::Instant 应该显示 10 秒
|
||||
assert!(
|
||||
tokio_elapsed >= Duration::from_secs(10),
|
||||
"tokio::time::Instant 应该跟随虚拟时间: {:?}",
|
||||
tokio_elapsed
|
||||
);
|
||||
}
|
||||
14
qiming-rcoder/crates/container-runtime-api/Cargo.toml
Normal file
14
qiming-rcoder/crates/container-runtime-api/Cargo.toml
Normal file
@@ -0,0 +1,14 @@
|
||||
[package]
|
||||
name = "container-runtime-api"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
shared_types = { path = "../shared_types" }
|
||||
async-trait = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
chrono = { workspace = true, features = ["serde"] }
|
||||
8
qiming-rcoder/crates/container-runtime-api/src/lib.rs
Normal file
8
qiming-rcoder/crates/container-runtime-api/src/lib.rs
Normal file
@@ -0,0 +1,8 @@
|
||||
//! Container Runtime API
|
||||
//!
|
||||
//! This crate provides the `ContainerRuntime` trait abstraction for different
|
||||
//! container runtimes (Docker, Kubernetes, etc.).
|
||||
|
||||
pub mod runtime_trait;
|
||||
|
||||
pub use runtime_trait::*;
|
||||
305
qiming-rcoder/crates/container-runtime-api/src/runtime_trait.rs
Normal file
305
qiming-rcoder/crates/container-runtime-api/src/runtime_trait.rs
Normal file
@@ -0,0 +1,305 @@
|
||||
//! Container runtime abstraction trait
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use shared_types::{ContainerBasicInfo, ServiceResourceLimits, ServiceType};
|
||||
use thiserror::Error;
|
||||
|
||||
/// Container runtime errors
|
||||
#[derive(Error, Debug)]
|
||||
pub enum ContainerRuntimeError {
|
||||
#[error("Connection error: {0}")]
|
||||
ConnectionError(String),
|
||||
|
||||
#[error("Container creation failed: {0}")]
|
||||
ContainerCreationError(String),
|
||||
|
||||
#[error("Container start failed: {0}")]
|
||||
ContainerStartError(String),
|
||||
|
||||
#[error("Container stop failed: {0}")]
|
||||
ContainerStopError(String),
|
||||
|
||||
#[error("Container not found: {0}")]
|
||||
ContainerNotFound(String),
|
||||
|
||||
#[error("Configuration error: {0}")]
|
||||
ConfigurationError(String),
|
||||
|
||||
#[error("Timeout: {0}")]
|
||||
Timeout(String),
|
||||
|
||||
#[error("Kubernetes error: {0}")]
|
||||
K8sError(String),
|
||||
|
||||
#[error("Docker error: {0}")]
|
||||
DockerError(String),
|
||||
}
|
||||
|
||||
/// Result type for container operations
|
||||
pub type ContainerRuntimeResult<T> = Result<T, ContainerRuntimeError>;
|
||||
|
||||
/// Container runtime status
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ContainerRuntimeStatus {
|
||||
Pending,
|
||||
Running,
|
||||
Succeeded,
|
||||
Failed,
|
||||
Unknown(String),
|
||||
}
|
||||
|
||||
/// Basic container info returned by runtime
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RuntimeContainerInfo {
|
||||
pub container_id: String,
|
||||
pub container_name: String,
|
||||
pub container_ip: String,
|
||||
pub status: ContainerRuntimeStatus,
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 已被移除的容器信息(用于清理关联资源)
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RemovedContainerInfo {
|
||||
/// 容器名称(稳定标识符)
|
||||
pub container_name: String,
|
||||
/// 容器 IP
|
||||
pub container_ip: String,
|
||||
/// 容器标识符(project_id 或 user_id)
|
||||
pub identifier: String,
|
||||
/// 服务类型
|
||||
pub service_type: ServiceType,
|
||||
}
|
||||
|
||||
impl From<ContainerRuntimeStatus> for String {
|
||||
fn from(status: ContainerRuntimeStatus) -> Self {
|
||||
match status {
|
||||
ContainerRuntimeStatus::Pending => "pending".to_string(),
|
||||
ContainerRuntimeStatus::Running => "running".to_string(),
|
||||
ContainerRuntimeStatus::Succeeded => "succeeded".to_string(),
|
||||
ContainerRuntimeStatus::Failed => "failed".to_string(),
|
||||
ContainerRuntimeStatus::Unknown(s) => s,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parameters for creating a container
|
||||
///
|
||||
/// Bundles all parameters needed for container creation to avoid
|
||||
/// long parameter lists that hurt code readability and maintainability.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContainerCreateParams {
|
||||
/// Project identifier (used as container name base for RCoder service)
|
||||
pub project_id: Option<String>,
|
||||
/// User identifier (used as container name base for ComputerAgentRunner)
|
||||
pub user_id: Option<String>,
|
||||
/// Workspace path on host
|
||||
pub host_workspace_path: String,
|
||||
/// Service type determining container purpose
|
||||
pub service_type: ServiceType,
|
||||
/// Optional resource constraints
|
||||
pub resource_limits: Option<ServiceResourceLimits>,
|
||||
/// Pod identifier for container reuse (for multi-tenant scenarios)
|
||||
pub pod_id: Option<String>,
|
||||
/// Isolation type: tenant|space|project (for multi-tenant scenarios)
|
||||
pub isolation_type: Option<String>,
|
||||
/// Tenant identifier (for multi-tenant scenarios)
|
||||
pub tenant_id: Option<String>,
|
||||
/// Space identifier (for multi-tenant scenarios)
|
||||
pub space_id: Option<String>,
|
||||
}
|
||||
|
||||
impl ContainerCreateParams {
|
||||
/// Create a new builder for container create params
|
||||
pub fn builder() -> ContainerCreateParamsBuilder {
|
||||
ContainerCreateParamsBuilder::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ContainerCreateParamsBuilder {
|
||||
project_id: Option<String>,
|
||||
user_id: Option<String>,
|
||||
host_workspace_path: Option<String>,
|
||||
service_type: Option<ServiceType>,
|
||||
resource_limits: Option<ServiceResourceLimits>,
|
||||
pod_id: Option<String>,
|
||||
isolation_type: Option<String>,
|
||||
tenant_id: Option<String>,
|
||||
space_id: Option<String>,
|
||||
}
|
||||
|
||||
impl ContainerCreateParamsBuilder {
|
||||
pub fn project_id(mut self, project_id: impl Into<String>) -> Self {
|
||||
self.project_id = Some(project_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn user_id(mut self, user_id: impl Into<String>) -> Self {
|
||||
self.user_id = Some(user_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn host_workspace_path(mut self, host_workspace_path: impl Into<String>) -> Self {
|
||||
self.host_workspace_path = Some(host_workspace_path.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn service_type(mut self, service_type: ServiceType) -> Self {
|
||||
self.service_type = Some(service_type);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn resource_limits(mut self, resource_limits: ServiceResourceLimits) -> Self {
|
||||
self.resource_limits = Some(resource_limits);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn pod_id(mut self, pod_id: impl Into<String>) -> Self {
|
||||
self.pod_id = Some(pod_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn isolation_type(mut self, isolation_type: impl Into<String>) -> Self {
|
||||
self.isolation_type = Some(isolation_type.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn tenant_id(mut self, tenant_id: impl Into<String>) -> Self {
|
||||
self.tenant_id = Some(tenant_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn space_id(mut self, space_id: impl Into<String>) -> Self {
|
||||
self.space_id = Some(space_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn build(self) -> ContainerCreateParams {
|
||||
ContainerCreateParams {
|
||||
project_id: self.project_id,
|
||||
user_id: self.user_id,
|
||||
host_workspace_path: self
|
||||
.host_workspace_path
|
||||
.unwrap_or_else(|| String::new()),
|
||||
service_type: self.service_type.unwrap_or(ServiceType::RCoder),
|
||||
resource_limits: self.resource_limits,
|
||||
pod_id: self.pod_id,
|
||||
isolation_type: self.isolation_type,
|
||||
tenant_id: self.tenant_id,
|
||||
space_id: self.space_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Abstraction trait for container runtimes (Docker, Kubernetes, etc.)
|
||||
///
|
||||
/// This trait follows the Interface Segregation Principle - it provides
|
||||
/// a lean interface with only the methods that callers actually need.
|
||||
#[async_trait]
|
||||
pub trait ContainerRuntime: Send + Sync {
|
||||
/// Create and start a container
|
||||
async fn create_container(
|
||||
&self,
|
||||
params: ContainerCreateParams,
|
||||
) -> ContainerRuntimeResult<ContainerBasicInfo>;
|
||||
|
||||
/// Get container information by project_id
|
||||
async fn get_container_info(
|
||||
&self,
|
||||
project_id: &str,
|
||||
) -> ContainerRuntimeResult<Option<ContainerBasicInfo>>;
|
||||
|
||||
/// Get container information by identifier + service type.
|
||||
///
|
||||
/// `identifier` means:
|
||||
/// - RCoder: `project_id`
|
||||
/// - ComputerAgentRunner: `user_id`
|
||||
async fn get_container_info_by_identifier(
|
||||
&self,
|
||||
identifier: &str,
|
||||
service_type: &ServiceType,
|
||||
) -> ContainerRuntimeResult<Option<ContainerBasicInfo>> {
|
||||
if matches!(service_type, ServiceType::RCoder) {
|
||||
return self.get_container_info(identifier).await;
|
||||
}
|
||||
|
||||
let info = self.find_container(identifier, service_type).await?;
|
||||
Ok(info.map(|pod| ContainerBasicInfo {
|
||||
container_id: pod.container_id,
|
||||
container_name: pod.container_name,
|
||||
container_ip: pod.container_ip.clone(),
|
||||
internal_port: shared_types::GRPC_DEFAULT_PORT,
|
||||
external_port: 0,
|
||||
project_id: identifier.to_string(),
|
||||
status: String::from(pod.status),
|
||||
created_at: pod.created_at,
|
||||
service_url: format!(
|
||||
"http://{}:{}",
|
||||
pod.container_ip,
|
||||
shared_types::GRPC_DEFAULT_PORT
|
||||
),
|
||||
}))
|
||||
}
|
||||
|
||||
/// Find container by project_id (returns None if not running)
|
||||
async fn find_container(
|
||||
&self,
|
||||
project_id: &str,
|
||||
service_type: &ServiceType,
|
||||
) -> ContainerRuntimeResult<Option<RuntimeContainerInfo>>;
|
||||
|
||||
/// Stop and remove container
|
||||
async fn stop_container(&self, project_id: &str) -> ContainerRuntimeResult<()>;
|
||||
|
||||
/// Stop and remove container by identifier + service type.
|
||||
///
|
||||
/// `identifier` means:
|
||||
/// - RCoder: `project_id`
|
||||
/// - ComputerAgentRunner: `user_id`
|
||||
async fn stop_container_by_identifier(
|
||||
&self,
|
||||
identifier: &str,
|
||||
_service_type: &ServiceType,
|
||||
) -> ContainerRuntimeResult<()> {
|
||||
self.stop_container(identifier).await
|
||||
}
|
||||
|
||||
/// Get container status
|
||||
async fn is_container_running(&self, project_id: &str) -> ContainerRuntimeResult<bool>;
|
||||
|
||||
/// Get container status by identifier + service type.
|
||||
async fn is_container_running_by_identifier(
|
||||
&self,
|
||||
identifier: &str,
|
||||
service_type: &ServiceType,
|
||||
) -> ContainerRuntimeResult<bool> {
|
||||
Ok(self
|
||||
.find_container(identifier, service_type)
|
||||
.await?
|
||||
.map(|c| c.status == ContainerRuntimeStatus::Running)
|
||||
.unwrap_or(false))
|
||||
}
|
||||
|
||||
/// List all containers managed by this runtime
|
||||
async fn list_containers(&self) -> ContainerRuntimeResult<Vec<RuntimeContainerInfo>>;
|
||||
|
||||
/// 同步缓存状态,清理失效的容器记录
|
||||
///
|
||||
/// 对于 Docker:遍历 ContainerStateActor 缓存,通过 Docker API 验证容器是否仍存在
|
||||
/// 对于 K8s:遍历 pod_cache,通过 K8s API 验证 Pod 是否仍存在
|
||||
///
|
||||
/// # Returns
|
||||
/// 返回元组 (已检查数量, 已移除容器信息列表)
|
||||
async fn sync_states(&self) -> ContainerRuntimeResult<(u32, Vec<RemovedContainerInfo>)> {
|
||||
// 默认实现:不做任何事(向后兼容)
|
||||
Ok((0, Vec::new()))
|
||||
}
|
||||
|
||||
/// Cleanup all containers (used on shutdown)
|
||||
async fn cleanup_all(&self) -> ContainerRuntimeResult<()>;
|
||||
|
||||
/// Health check - verify runtime is accessible
|
||||
async fn health_check(&self) -> ContainerRuntimeResult<()>;
|
||||
}
|
||||
61
qiming-rcoder/crates/docker_manager/Cargo.toml
Normal file
61
qiming-rcoder/crates/docker_manager/Cargo.toml
Normal file
@@ -0,0 +1,61 @@
|
||||
[package]
|
||||
name = "docker_manager"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
# Docker API
|
||||
bollard = { workspace = true }
|
||||
|
||||
# Async runtime
|
||||
tokio = { workspace = true }
|
||||
futures-util = { workspace = true }
|
||||
|
||||
# HTTP client (for health checks)
|
||||
reqwest = { workspace = true }
|
||||
|
||||
# Serialization
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
|
||||
# Error handling
|
||||
anyhow = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
# Logging
|
||||
tracing = { workspace = true }
|
||||
|
||||
# UUID
|
||||
uuid = { workspace = true }
|
||||
|
||||
# Date/time
|
||||
chrono = { workspace = true }
|
||||
|
||||
# Utilities
|
||||
dashmap = { workspace = true }
|
||||
moka = { version = "0.12", features = ["future"] }
|
||||
|
||||
# Shared types
|
||||
shared_types = { path = "../shared_types" }
|
||||
|
||||
# Runtime API (always needed)
|
||||
container-runtime-api = { path = "../container-runtime-api" }
|
||||
|
||||
# Async runtime
|
||||
async-trait = { workspace = true }
|
||||
|
||||
# Kubernetes runtime support (feature-gated)
|
||||
kube = { version = "0.98", features = ["runtime", "client", "kube-runtime", "config"], optional = true }
|
||||
k8s-openapi = { version = "0.24", features = ["v1_30"], optional = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
# Kubernetes runtime support
|
||||
kubernetes = ["kube", "k8s-openapi"]
|
||||
# 🔧 eBPF 调试模式:启用容器特权,允许使用 eBPF 诊断工具
|
||||
# 警告:仅在开发调试时使用,生产环境请勿启用
|
||||
ebpf-debug = []
|
||||
234
qiming-rcoder/crates/docker_manager/README.md
Normal file
234
qiming-rcoder/crates/docker_manager/README.md
Normal file
@@ -0,0 +1,234 @@
|
||||
# Docker Manager
|
||||
|
||||
基于 bollard 库的 Docker 容器动态管理模块,用于 RCoder 项目中的 Docker Agent 管理。
|
||||
|
||||
## 功能特性
|
||||
|
||||
- ✅ 动态创建和销毁 Docker 容器
|
||||
- ✅ 支持项目工作目录挂载
|
||||
- ✅ 容器状态监控和日志获取
|
||||
- ✅ 支持环境变量和端口映射配置
|
||||
- ✅ 自动镜像拉取和资源限制
|
||||
- ✅ 容器生命周期管理
|
||||
|
||||
## 核心组件
|
||||
|
||||
### 1. DockerManager
|
||||
主要的 Docker 管理器,提供容器的完整生命周期管理。
|
||||
|
||||
```rust
|
||||
use docker_manager::{DockerManager, DockerManagerConfig};
|
||||
|
||||
// 创建 Docker 管理器
|
||||
let config = DockerManagerConfig::default();
|
||||
let docker_manager = DockerManager::new(config).await?;
|
||||
|
||||
// 创建容器
|
||||
let container_info = docker_manager.create_container(config).await?;
|
||||
|
||||
// 停止容器
|
||||
docker_manager.stop_container("project_id").await?;
|
||||
```
|
||||
|
||||
### 2. DockerContainerConfig
|
||||
容器配置结构体,定义容器的各种参数。
|
||||
|
||||
```rust
|
||||
use docker_manager::DockerContainerConfig;
|
||||
|
||||
let config = DockerContainerConfig {
|
||||
project_id: "my_project".to_string(),
|
||||
image: "registry.yichamao.com/rcoder:latest".to_string(),
|
||||
host_path: "/path/to/project".to_string(),
|
||||
container_path: "/app/workspace".to_string(),
|
||||
env_vars: env_map,
|
||||
port_bindings: port_map,
|
||||
..Default::default()
|
||||
};
|
||||
```
|
||||
|
||||
### 3. DockerUtils
|
||||
工具函数集合,简化常见操作。
|
||||
|
||||
```rust
|
||||
use docker_manager::DockerUtils;
|
||||
|
||||
// 根据项目ID创建配置
|
||||
let config = DockerUtils::create_config_from_project_id(
|
||||
"project_123",
|
||||
"./project_workspace",
|
||||
Some("custom:image".to_string()),
|
||||
);
|
||||
```
|
||||
|
||||
## 使用场景
|
||||
|
||||
### 1. RCoder Docker Agent
|
||||
|
||||
在 RCoder 项目中,Docker Agent 用于为每个项目创建独立的运行环境:
|
||||
|
||||
```rust
|
||||
use docker_manager::{DockerAgentManager, DockerUtils};
|
||||
use rcoder::model::{AgentType, ChatPrompt};
|
||||
|
||||
// 创建 Docker Agent 管理器
|
||||
let docker_agent_manager = DockerAgentManager::new().await?;
|
||||
|
||||
// 为项目创建 Docker Agent
|
||||
let chat_prompt = ChatPrompt {
|
||||
project_id: "project_123".to_string(),
|
||||
agent_type: AgentType::Docker,
|
||||
// ... 其他字段
|
||||
};
|
||||
|
||||
let docker_agent = docker_agent_manager.create_docker_agent(
|
||||
&chat_prompt,
|
||||
AgentType::Docker
|
||||
).await?;
|
||||
```
|
||||
|
||||
### 2. 项目隔离
|
||||
|
||||
每个项目在独立的 Docker 容器中运行,确保环境隔离:
|
||||
|
||||
- 挂载路径: `./project_workspace/{project_id}` → `/app/workspace`
|
||||
- 镜像: `registry.yichamao.com/rcoder:latest`
|
||||
- 网络: 使用 host 模式以获得更好的性能
|
||||
|
||||
### 3. 环境变量配置
|
||||
|
||||
Docker Agent 支持丰富的环境变量配置:
|
||||
|
||||
```rust
|
||||
let mut env_vars = HashMap::new();
|
||||
env_vars.insert("ANTHROPIC_AUTH_TOKEN".to_string(), "your_token".to_string());
|
||||
env_vars.insert("PROJECT_ID".to_string(), "project_123".to_string());
|
||||
env_vars.insert("DOCKER_AGENT_TYPE".to_string(), "claude".to_string());
|
||||
```
|
||||
|
||||
## 配置选项
|
||||
|
||||
### DockerManagerConfig
|
||||
|
||||
```rust
|
||||
pub struct DockerManagerConfig {
|
||||
pub docker_host: Option<String>, // Docker 守护进程地址
|
||||
pub default_image: String, // 默认镜像
|
||||
pub default_network_mode: String, // 默认网络模式
|
||||
pub default_work_dir: String, // 默认工作目录
|
||||
pub auto_cleanup: bool, // 是否启用自动清理
|
||||
pub container_ttl_seconds: Option<u64>, // 容器存活时间
|
||||
}
|
||||
```
|
||||
|
||||
### 环境变量配置
|
||||
|
||||
可以通过环境变量配置 Docker 管理器:
|
||||
|
||||
- `DOCKER_HOST`: Docker 守护进程地址
|
||||
- `DEFAULT_DOCKER_IMAGE`: 默认镜像
|
||||
- `DOCKER_NETWORK_MODE`: 网络模式
|
||||
- `DOCKER_WORK_DIR`: 工作目录
|
||||
- `DOCKER_AUTO_CLEANUP`: 自动清理
|
||||
- `DOCKER_CONTAINER_TTL`: 容器TTL
|
||||
|
||||
## 集成到 RCoder
|
||||
|
||||
### 1. 启用 Docker Agent
|
||||
|
||||
设置环境变量启用 Docker Agent:
|
||||
|
||||
```bash
|
||||
export USE_DOCKER_AGENT=true
|
||||
```
|
||||
|
||||
### 2. 项目目录结构
|
||||
|
||||
```
|
||||
./project_workspace/
|
||||
├── project_123/ # 项目 123 的工作目录
|
||||
│ ├── src/
|
||||
│ ├── Cargo.toml
|
||||
│ └── ...
|
||||
├── project_456/ # 项目 456 的工作目录
|
||||
│ └── ...
|
||||
```
|
||||
|
||||
### 3. Docker 镜像
|
||||
|
||||
使用预构建的 RCoder Docker 镜像:
|
||||
- 镜像地址: `registry.yichamao.com/rcoder:latest`
|
||||
- 包含完整的 AI 开发工具链
|
||||
- 支持 Claude Code 和 Codex Agent
|
||||
|
||||
## 错误处理
|
||||
|
||||
所有操作都返回 `DockerResult<T>`,包含详细的错误信息:
|
||||
|
||||
```rust
|
||||
match docker_manager.create_container(config).await {
|
||||
Ok(container_info) => {
|
||||
println!("容器创建成功: {}", container_info.container_name);
|
||||
}
|
||||
Err(DockerError::ContainerCreationError(msg)) => {
|
||||
eprintln!("容器创建失败: {}", msg);
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("其他错误: {}", e);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## 日志和监控
|
||||
|
||||
### 获取容器日志
|
||||
|
||||
```rust
|
||||
// 获取最后 50 行日志
|
||||
let logs = docker_manager.get_container_logs("project_id", 50).await?;
|
||||
println!("容器日志:\n{}", logs);
|
||||
```
|
||||
|
||||
### 监控容器状态
|
||||
|
||||
```rust
|
||||
// 检查容器状态
|
||||
let status = docker_manager.update_container_status("project_id").await?;
|
||||
if let Some(status) = status {
|
||||
println!("容器状态: {:?}", status);
|
||||
}
|
||||
```
|
||||
|
||||
## 最佳实践
|
||||
|
||||
1. **资源管理**: 及时停止不用的容器以释放资源
|
||||
2. **错误处理**: 始终检查操作结果并处理错误
|
||||
3. **日志监控**: 定期检查容器日志以了解运行状态
|
||||
4. **环境隔离**: 为每个项目使用独立的容器
|
||||
5. **镜像管理**: 使用固定版本的镜像以避免意外更新
|
||||
|
||||
## 故障排除
|
||||
|
||||
### 常见问题
|
||||
|
||||
1. **Docker 连接失败**
|
||||
- 检查 Docker 守护进程是否运行
|
||||
- 验证 Docker socket 权限
|
||||
|
||||
2. **镜像拉取失败**
|
||||
- 检查网络连接
|
||||
- 验证镜像地址和认证信息
|
||||
|
||||
3. **容器启动失败**
|
||||
- 检查资源限制(内存、CPU)
|
||||
- 验证挂载路径权限
|
||||
- 查看容器日志了解详细错误
|
||||
|
||||
### 调试模式
|
||||
|
||||
启用详细日志进行调试:
|
||||
|
||||
```bash
|
||||
export RUST_LOG=debug
|
||||
cargo run --example basic_usage
|
||||
```
|
||||
@@ -0,0 +1,350 @@
|
||||
//! 容器配置构建器
|
||||
//!
|
||||
//! 使用 Builder 模式构建 DockerContainerConfig
|
||||
|
||||
use crate::{DockerContainerConfig, DockerResult, MountPoint, ResourceLimits};
|
||||
use std::collections::HashMap;
|
||||
use tracing::debug;
|
||||
|
||||
/// 容器配置构建器
|
||||
///
|
||||
/// 使用 Builder 模式提供灵活的容器配置构建接口
|
||||
///
|
||||
/// # Examples
|
||||
/// ```no_run
|
||||
/// use docker_manager::container_builder::ContainerConfigBuilder;
|
||||
///
|
||||
/// # async fn example() -> docker_manager::DockerResult<()> {
|
||||
/// let config = ContainerConfigBuilder::new("project-123")
|
||||
/// .image("registry.example.com/agent-runner:latest")
|
||||
/// .host_path("/host/path")
|
||||
/// .container_path("/app/project_workspace/project-123")
|
||||
/// .env("PROJECT_ID", "project-123")
|
||||
/// .network_name("rcoder_agent-network")
|
||||
/// .build()?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub struct ContainerConfigBuilder {
|
||||
project_id: String,
|
||||
image: Option<String>,
|
||||
name_prefix: Option<String>,
|
||||
host_path: Option<String>,
|
||||
container_path: Option<String>,
|
||||
work_dir: Option<String>,
|
||||
env_vars: HashMap<String, String>,
|
||||
port_bindings: HashMap<String, String>,
|
||||
network_mode: Option<String>,
|
||||
auto_remove: bool,
|
||||
resource_limits: Option<ResourceLimits>,
|
||||
extra_mounts: Vec<MountPoint>,
|
||||
command: Option<Vec<String>>,
|
||||
entrypoint: Option<Vec<String>>,
|
||||
network_name: Option<String>,
|
||||
// 新增字段(隔离类型支持)
|
||||
pod_id: Option<String>,
|
||||
tenant_id: Option<String>,
|
||||
space_id: Option<String>,
|
||||
isolation_type: Option<String>,
|
||||
}
|
||||
|
||||
impl ContainerConfigBuilder {
|
||||
/// 创建新的容器配置构建器
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `project_id` - 项目ID(必需)
|
||||
pub fn new(project_id: impl Into<String>) -> Self {
|
||||
Self {
|
||||
project_id: project_id.into(),
|
||||
image: None,
|
||||
name_prefix: None,
|
||||
host_path: None,
|
||||
container_path: None,
|
||||
work_dir: None,
|
||||
env_vars: HashMap::new(),
|
||||
port_bindings: HashMap::new(),
|
||||
network_mode: None,
|
||||
auto_remove: false,
|
||||
resource_limits: None,
|
||||
extra_mounts: Vec::new(),
|
||||
command: None,
|
||||
entrypoint: None,
|
||||
network_name: None,
|
||||
// 新增字段(隔离类型支持)
|
||||
pod_id: None,
|
||||
tenant_id: None,
|
||||
space_id: None,
|
||||
isolation_type: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 设置 Docker 镜像
|
||||
pub fn image(mut self, image: impl Into<String>) -> Self {
|
||||
self.image = Some(image.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置容器名称前缀
|
||||
pub fn name_prefix(mut self, prefix: impl Into<String>) -> Self {
|
||||
self.name_prefix = Some(prefix.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置宿主机路径
|
||||
pub fn host_path(mut self, path: impl Into<String>) -> Self {
|
||||
self.host_path = Some(path.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置容器内路径
|
||||
pub fn container_path(mut self, path: impl Into<String>) -> Self {
|
||||
self.container_path = Some(path.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置工作目录
|
||||
pub fn work_dir(mut self, dir: impl Into<String>) -> Self {
|
||||
self.work_dir = Some(dir.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 添加单个环境变量
|
||||
pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
|
||||
self.env_vars.insert(key.into(), value.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 批量添加环境变量
|
||||
pub fn envs(mut self, vars: HashMap<String, String>) -> Self {
|
||||
self.env_vars.extend(vars);
|
||||
self
|
||||
}
|
||||
|
||||
/// 添加端口映射
|
||||
pub fn port_binding(
|
||||
mut self,
|
||||
container_port: impl Into<String>,
|
||||
host_port: impl Into<String>,
|
||||
) -> Self {
|
||||
self.port_bindings
|
||||
.insert(container_port.into(), host_port.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 批量添加端口映射
|
||||
pub fn port_bindings(mut self, bindings: HashMap<String, String>) -> Self {
|
||||
self.port_bindings.extend(bindings);
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置网络模式
|
||||
pub fn network_mode(mut self, mode: impl Into<String>) -> Self {
|
||||
self.network_mode = Some(mode.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置自动删除标志
|
||||
pub fn auto_remove(mut self, enabled: bool) -> Self {
|
||||
self.auto_remove = enabled;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置资源限制
|
||||
pub fn resource_limits(mut self, limits: ResourceLimits) -> Self {
|
||||
self.resource_limits = Some(limits);
|
||||
self
|
||||
}
|
||||
|
||||
/// 添加单个挂载点
|
||||
pub fn add_mount(mut self, mount: MountPoint) -> Self {
|
||||
self.extra_mounts.push(mount);
|
||||
self
|
||||
}
|
||||
|
||||
/// 批量添加挂载点
|
||||
pub fn add_mounts(mut self, mounts: Vec<MountPoint>) -> Self {
|
||||
self.extra_mounts.extend(mounts);
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置启动命令
|
||||
pub fn command(mut self, command: Vec<String>) -> Self {
|
||||
self.command = Some(command);
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置入口点
|
||||
pub fn entrypoint(mut self, entrypoint: Vec<String>) -> Self {
|
||||
self.entrypoint = Some(entrypoint);
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置网络名称
|
||||
pub fn network_name(mut self, name: impl Into<String>) -> Self {
|
||||
self.network_name = Some(name.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 pod_id(容器唯一标识,用于容器复用)
|
||||
pub fn pod_id(mut self, pod_id: impl Into<String>) -> Self {
|
||||
self.pod_id = Some(pod_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 tenant_id(租户 ID)
|
||||
pub fn tenant_id(mut self, tenant_id: impl Into<String>) -> Self {
|
||||
self.tenant_id = Some(tenant_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 space_id(空间 ID)
|
||||
pub fn space_id(mut self, space_id: impl Into<String>) -> Self {
|
||||
self.space_id = Some(space_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 isolation_type(隔离类型)
|
||||
pub fn isolation_type(mut self, isolation_type: impl Into<String>) -> Self {
|
||||
self.isolation_type = Some(isolation_type.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 构建 DockerContainerConfig
|
||||
///
|
||||
/// # Returns
|
||||
/// * `DockerResult<DockerContainerConfig>` - 构建的配置或错误
|
||||
pub fn build(self) -> DockerResult<DockerContainerConfig> {
|
||||
debug!("builtcontainerconfig, projectID: {}", self.project_id);
|
||||
|
||||
// 使用默认值或提供的值
|
||||
let image = self.image.unwrap_or_else(crate::default_docker_image);
|
||||
let name_prefix = self
|
||||
.name_prefix
|
||||
.unwrap_or_else(|| "rcoder-agent".to_string());
|
||||
let host_path = self.host_path.unwrap_or_default();
|
||||
let container_path = self
|
||||
.container_path
|
||||
.unwrap_or_else(|| crate::DEFAULT_WORK_DIR.to_string());
|
||||
let work_dir = self
|
||||
.work_dir
|
||||
.unwrap_or_else(|| crate::DEFAULT_WORK_DIR.to_string());
|
||||
let network_mode = self
|
||||
.network_mode
|
||||
.unwrap_or_else(|| crate::DEFAULT_NETWORK_MODE.to_string());
|
||||
|
||||
let config = DockerContainerConfig {
|
||||
project_id: self.project_id,
|
||||
image,
|
||||
name_prefix,
|
||||
host_path,
|
||||
container_path,
|
||||
work_dir,
|
||||
env_vars: self.env_vars,
|
||||
port_bindings: self.port_bindings,
|
||||
network_mode,
|
||||
auto_remove: self.auto_remove,
|
||||
resource_limits: self.resource_limits,
|
||||
extra_mounts: self.extra_mounts,
|
||||
command: self.command,
|
||||
entrypoint: self.entrypoint,
|
||||
network_name: self.network_name,
|
||||
// 新增字段(隔离类型支持)
|
||||
pod_id: self.pod_id,
|
||||
tenant_id: self.tenant_id,
|
||||
space_id: self.space_id,
|
||||
isolation_type: self.isolation_type,
|
||||
};
|
||||
|
||||
debug!(
|
||||
"Container config built: image={}, network={:?}, mounts={}",
|
||||
config.image,
|
||||
config.network_name,
|
||||
config.extra_mounts.len()
|
||||
);
|
||||
|
||||
Ok(config)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_builder_minimal() {
|
||||
let config = ContainerConfigBuilder::new("test-project").build().unwrap();
|
||||
|
||||
assert_eq!(config.project_id, "test-project");
|
||||
assert_eq!(config.name_prefix, "rcoder-agent");
|
||||
assert!(!config.auto_remove);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_full() {
|
||||
let config = ContainerConfigBuilder::new("test-project")
|
||||
.image("custom-image:latest")
|
||||
.name_prefix("custom-prefix")
|
||||
.host_path("/host/path")
|
||||
.container_path("/container/path")
|
||||
.work_dir("/work")
|
||||
.env("KEY1", "value1")
|
||||
.env("KEY2", "value2")
|
||||
.port_binding("8080", "8080")
|
||||
.network_mode("bridge")
|
||||
.auto_remove(true)
|
||||
.network_name("test-network")
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(config.project_id, "test-project");
|
||||
assert_eq!(config.image, "custom-image:latest");
|
||||
assert_eq!(config.name_prefix, "custom-prefix");
|
||||
assert_eq!(config.host_path, "/host/path");
|
||||
assert_eq!(config.container_path, "/container/path");
|
||||
assert_eq!(config.work_dir, "/work");
|
||||
assert_eq!(config.env_vars.len(), 2);
|
||||
assert_eq!(config.port_bindings.len(), 1);
|
||||
assert_eq!(config.network_mode, "bridge");
|
||||
assert!(config.auto_remove);
|
||||
assert_eq!(config.network_name, Some("test-network".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_with_mounts() {
|
||||
let mount = MountPoint {
|
||||
host_path: "/host/mount".to_string(),
|
||||
container_path: "/container/mount".to_string(),
|
||||
read_only: false,
|
||||
};
|
||||
|
||||
let config = ContainerConfigBuilder::new("test-project")
|
||||
.add_mount(mount)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(config.extra_mounts.len(), 1);
|
||||
assert_eq!(config.extra_mounts[0].host_path, "/host/mount");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_with_resource_limits() {
|
||||
let limits = ResourceLimits {
|
||||
memory_limit: Some((512 * 1024 * 1024) as f64), // 512MB
|
||||
cpu_limit: Some(1.0),
|
||||
swap_limit: None,
|
||||
};
|
||||
|
||||
let config = ContainerConfigBuilder::new("test-project")
|
||||
.resource_limits(limits)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
assert!(config.resource_limits.is_some());
|
||||
let resource_limits = config.resource_limits.unwrap();
|
||||
assert_eq!(
|
||||
resource_limits.memory_limit,
|
||||
Some((512 * 1024 * 1024) as f64)
|
||||
);
|
||||
assert_eq!(resource_limits.cpu_limit, Some(1.0));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
//! 容器构建器模块
|
||||
//!
|
||||
//! 提供 Builder 模式的容器配置构建和挂载点处理功能
|
||||
|
||||
pub mod config_builder;
|
||||
pub mod mount_processor;
|
||||
|
||||
pub use config_builder::*;
|
||||
pub use mount_processor::*;
|
||||
@@ -0,0 +1,222 @@
|
||||
//! 挂载点处理器
|
||||
//!
|
||||
//! 处理容器挂载点的路径解析和变量替换
|
||||
|
||||
use crate::path::HostPathResolver;
|
||||
use crate::{DockerResult, MountPoint};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use tracing::{debug, info};
|
||||
|
||||
/// 挂载点处理器
|
||||
///
|
||||
/// 提供挂载点路径解析、变量替换等功能
|
||||
pub struct MountProcessor {
|
||||
resolver: HostPathResolver,
|
||||
}
|
||||
|
||||
impl MountProcessor {
|
||||
/// 创建新的挂载点处理器
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `resolver` - 路径解析器
|
||||
pub fn new(resolver: HostPathResolver) -> Self {
|
||||
Self { resolver }
|
||||
}
|
||||
|
||||
/// 创建新的挂载点处理器(异步,自动创建路径解析器)
|
||||
///
|
||||
/// # Returns
|
||||
/// * `DockerResult<Self>` - 挂载点处理器或错误
|
||||
pub async fn new_async() -> DockerResult<Self> {
|
||||
let resolver = HostPathResolver::new().await?;
|
||||
Ok(Self { resolver })
|
||||
}
|
||||
|
||||
/// 使用指定的 Docker socket 创建挂载点处理器
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `docker_socket_path` - Docker socket 路径
|
||||
///
|
||||
/// # Returns
|
||||
/// * `DockerResult<Self>` - 挂载点处理器或错误
|
||||
pub async fn new_with_docker_socket(docker_socket_path: Option<String>) -> DockerResult<Self> {
|
||||
let resolver = HostPathResolver::new_with_docker_socket(docker_socket_path).await?;
|
||||
Ok(Self { resolver })
|
||||
}
|
||||
|
||||
/// 处理单个挂载点
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `container_path` - 容器内路径
|
||||
/// * `host_path` - 宿主机路径(可能包含变量或相对路径)
|
||||
/// * `read_only` - 是否只读
|
||||
/// * `variables` - 变量映射表(可选)
|
||||
///
|
||||
/// # Returns
|
||||
/// * `DockerResult<MountPoint>` - 处理后的挂载点或错误
|
||||
pub fn process_mount(
|
||||
&self,
|
||||
container_path: impl AsRef<str>,
|
||||
host_path: impl AsRef<str>,
|
||||
read_only: bool,
|
||||
variables: Option<&HashMap<String, String>>,
|
||||
) -> DockerResult<MountPoint> {
|
||||
let container_path = container_path.as_ref();
|
||||
let mut host_path = host_path.as_ref().to_string();
|
||||
|
||||
debug!(
|
||||
"Processing mount point: {} -> {} (read_only: {})",
|
||||
container_path, host_path, read_only
|
||||
);
|
||||
|
||||
// 变量替换
|
||||
if let Some(vars) = variables {
|
||||
for (key, value) in vars {
|
||||
let pattern = format!("{{{}}}", key);
|
||||
if host_path.contains(&pattern) {
|
||||
host_path = host_path.replace(&pattern, value);
|
||||
debug!(" substituted: {} -> {}", pattern, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 路径解析
|
||||
let normalized_host_path = self.resolve_path(&host_path)?;
|
||||
|
||||
info!(
|
||||
"Mount point processing completed: {} -> {}",
|
||||
container_path, normalized_host_path
|
||||
);
|
||||
|
||||
Ok(MountPoint {
|
||||
container_path: container_path.to_string(),
|
||||
host_path: normalized_host_path,
|
||||
read_only,
|
||||
})
|
||||
}
|
||||
|
||||
/// 批量处理挂载点
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `mounts` - 挂载点列表 (container_path, host_path, read_only)
|
||||
/// * `variables` - 变量映射表(可选)
|
||||
///
|
||||
/// # Returns
|
||||
/// * `DockerResult<Vec<MountPoint>>` - 处理后的挂载点列表或错误
|
||||
pub fn process_mounts(
|
||||
&self,
|
||||
mounts: Vec<(String, String, bool)>,
|
||||
variables: Option<&HashMap<String, String>>,
|
||||
) -> DockerResult<Vec<MountPoint>> {
|
||||
debug!("Processing {} mounts", mounts.len());
|
||||
|
||||
let processed: DockerResult<Vec<MountPoint>> = mounts
|
||||
.into_iter()
|
||||
.map(|(container_path, host_path, read_only)| {
|
||||
self.process_mount(&container_path, &host_path, read_only, variables)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let processed = processed?;
|
||||
info!(
|
||||
"Mount processing completed: {} mounts",
|
||||
processed.len()
|
||||
);
|
||||
|
||||
Ok(processed)
|
||||
}
|
||||
|
||||
/// 解析路径(处理相对路径和容器内路径)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `path` - 待解析的路径
|
||||
///
|
||||
/// # Returns
|
||||
/// * `DockerResult<String>` - 解析后的宿主机绝对路径
|
||||
fn resolve_path(&self, path: &str) -> DockerResult<String> {
|
||||
let path_obj = Path::new(path);
|
||||
|
||||
// 处理相对路径:转换为容器内绝对路径
|
||||
let container_absolute_path = if path_obj.is_relative() {
|
||||
let current_dir =
|
||||
std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("/app"));
|
||||
current_dir.join(path_obj)
|
||||
} else {
|
||||
path_obj.to_path_buf()
|
||||
};
|
||||
|
||||
// 检查是否为容器内路径
|
||||
if container_absolute_path.starts_with("/app") {
|
||||
// 容器内路径:转换为宿主机路径
|
||||
debug!(
|
||||
"Detected container path, resolving to host path: {}",
|
||||
container_absolute_path.display()
|
||||
);
|
||||
let host_abs_path = self
|
||||
.resolver
|
||||
.resolve_to_host_path(&container_absolute_path)?;
|
||||
Ok(host_abs_path.to_string_lossy().to_string())
|
||||
} else {
|
||||
// 可能已经是宿主机路径,直接使用
|
||||
debug!(
|
||||
"Using potential host path: {}",
|
||||
container_absolute_path.display()
|
||||
);
|
||||
Ok(container_absolute_path.to_string_lossy().to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取路径解析器的引用
|
||||
pub fn resolver(&self) -> &HostPathResolver {
|
||||
&self.resolver
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_mount_processor_variable_substitution() {
|
||||
// 注意:此测试需要在容器环境中运行才能完整验证路径解析
|
||||
// 这里仅测试变量替换逻辑
|
||||
|
||||
let mut variables = HashMap::new();
|
||||
variables.insert("project_id".to_string(), "test-123".to_string());
|
||||
|
||||
let host_path = "/path/to/{project_id}/data";
|
||||
let mut result = host_path.to_string();
|
||||
|
||||
for (key, value) in &variables {
|
||||
let pattern = format!("{{{}}}", key);
|
||||
result = result.replace(&pattern, value);
|
||||
}
|
||||
|
||||
assert_eq!(result, "/path/to/test-123/data");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mount_point_structure() {
|
||||
let mount = MountPoint {
|
||||
container_path: "/app/data".to_string(),
|
||||
host_path: "/host/data".to_string(),
|
||||
read_only: false,
|
||||
};
|
||||
|
||||
assert_eq!(mount.container_path, "/app/data");
|
||||
assert_eq!(mount.host_path, "/host/data");
|
||||
assert!(!mount.read_only);
|
||||
}
|
||||
|
||||
// 注意:以下测试需要在 Docker 容器环境中运行
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn test_mount_processor_creation() {
|
||||
// 此测试仅在容器内有效
|
||||
if std::env::var("HOSTNAME").is_ok() {
|
||||
let result = MountProcessor::new_async().await;
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,302 @@
|
||||
//! 容器自检测器
|
||||
//!
|
||||
//! 用于在容器内部通过 Docker API 检测自己的挂载信息,
|
||||
//! 获取容器内路径对应的宿主机绝对路径
|
||||
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
use bollard::{API_DEFAULT_VERSION, Docker};
|
||||
use tokio::fs;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
/// 容器自检测器
|
||||
///
|
||||
/// 用于检测当前容器的挂载信息,获取容器内路径对应的宿主机路径
|
||||
pub struct ContainerSelfInspector {
|
||||
/// Docker 客户端
|
||||
docker_client: Docker,
|
||||
/// 当前容器ID
|
||||
container_id: String,
|
||||
}
|
||||
|
||||
impl ContainerSelfInspector {
|
||||
/// 创建新的容器自检测器
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `docker_socket_path` - Docker socket 路径
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Result<Self>` - 检测器实例或错误
|
||||
///
|
||||
/// # Examples
|
||||
/// ```rust,no_run
|
||||
/// use docker_manager::ContainerSelfInspector;
|
||||
///
|
||||
/// # async fn example() -> anyhow::Result<()> {
|
||||
/// let inspector = ContainerSelfInspector::new("/var/run/docker.sock").await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn new(docker_socket_path: &str) -> Result<Self> {
|
||||
info!(
|
||||
"Initializing container self inspector, Docker socket: {}",
|
||||
docker_socket_path
|
||||
);
|
||||
|
||||
// 创建 Docker 客户端
|
||||
let docker_client =
|
||||
Docker::connect_with_socket(docker_socket_path, 120, API_DEFAULT_VERSION)
|
||||
.context("Failed to connect to Docker socket")?;
|
||||
|
||||
// 测试 Docker 连接
|
||||
docker_client
|
||||
.ping()
|
||||
.await
|
||||
.context("Failed to test Docker connection, please check socket path and permissions")?;
|
||||
|
||||
info!("Docker connectionsucceeded");
|
||||
|
||||
// 获取当前容器ID
|
||||
let container_id = Self::get_current_container_id()
|
||||
.await
|
||||
.context("Failed to get current container ID")?;
|
||||
|
||||
info!("Detected container ID: {}", container_id);
|
||||
|
||||
Ok(Self {
|
||||
docker_client,
|
||||
container_id,
|
||||
})
|
||||
}
|
||||
|
||||
/// 检测容器内路径对应的宿主机路径
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `container_path` - 容器内路径(如 "/app/project_workspace")
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Result<String>` - 宿主机绝对路径或错误
|
||||
///
|
||||
/// # Examples
|
||||
/// ```rust,no_run
|
||||
/// use docker_manager::ContainerSelfInspector;
|
||||
///
|
||||
/// # async fn example() -> anyhow::Result<()> {
|
||||
/// # let inspector = ContainerSelfInspector::new("/var/run/docker.sock").await?;
|
||||
/// let host_path = inspector.detect_host_path_for_container_dir("/app/project_workspace").await?;
|
||||
/// println!(" message path: {:?}", host_path);
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn detect_host_path_for_container_dir(&self, container_path: &str) -> Result<String> {
|
||||
info!("Detecting path {} host path", container_path);
|
||||
|
||||
// 获取容器详细信息
|
||||
let inspect_result = self
|
||||
.docker_client
|
||||
.inspect_container(
|
||||
&self.container_id,
|
||||
None::<bollard::query_parameters::InspectContainerOptions>,
|
||||
)
|
||||
.await
|
||||
.context("Failed to call Docker inspect API")?;
|
||||
|
||||
debug!(
|
||||
"Container inspect result: {:?}",
|
||||
serde_json::to_string_pretty(&inspect_result)?
|
||||
);
|
||||
|
||||
// 解析挂载信息
|
||||
if let Some(mounts) = inspect_result.mounts {
|
||||
debug!("Container has {} mounts", mounts.len());
|
||||
|
||||
for (index, mount) in mounts.iter().enumerate() {
|
||||
let mount_destination = mount
|
||||
.destination
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow!("mount {} has no destination field", index))?
|
||||
.clone();
|
||||
|
||||
debug!(
|
||||
"Mount point {}: {} -> {}",
|
||||
index,
|
||||
mount_destination,
|
||||
mount.source.as_ref().unwrap_or(&String::new()).clone()
|
||||
);
|
||||
|
||||
// 检查是否是我们要找的路径
|
||||
if mount_destination == container_path {
|
||||
let host_path = mount
|
||||
.source
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow!("mount {} has no source field", index))?
|
||||
.clone();
|
||||
|
||||
info!(
|
||||
" mount: {} -> {}",
|
||||
container_path, host_path
|
||||
);
|
||||
return Ok(host_path);
|
||||
}
|
||||
}
|
||||
|
||||
// 如果没找到,列出所有挂载点供调试
|
||||
warn!(
|
||||
"not found path {} in mount, mount info:",
|
||||
container_path
|
||||
);
|
||||
for (index, mount) in mounts.iter().enumerate() {
|
||||
if let (Some(dest), Some(source)) = (&mount.destination, &mount.source) {
|
||||
warn!(" {}: {} -> {}", index, dest, source);
|
||||
}
|
||||
}
|
||||
|
||||
bail!("mount info for path {} not found", container_path);
|
||||
} else {
|
||||
bail!("container has no mount info (mounts field is empty)");
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取当前容器ID
|
||||
///
|
||||
/// 通过读取 `/proc/self/cgroup` 文件解析容器ID
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Result<String>` - 容器ID或错误
|
||||
async fn get_current_container_id() -> Result<String> {
|
||||
debug!("starting get containerID");
|
||||
|
||||
let cgroup_content = fs::read_to_string("/proc/self/cgroup")
|
||||
.await
|
||||
.with_context(|| "Failed to read /proc/self/cgroup file")?;
|
||||
|
||||
debug!("cgroup file: {}", cgroup_content);
|
||||
|
||||
// 解析 cgroup 文件获取容器ID
|
||||
// 格式示例: 12:perf_event:/docker/abc123def456...
|
||||
for line in cgroup_content.lines() {
|
||||
debug!(" cgroup line: {}", line);
|
||||
|
||||
let parts: Vec<&str> = line.split(':').collect();
|
||||
if parts.len() >= 3 {
|
||||
let cgroup_path = parts[2];
|
||||
|
||||
// 检查是否是 Docker 容器
|
||||
if cgroup_path.contains("/docker/") || cgroup_path.contains(".scope") {
|
||||
debug!("Found Docker cgroup: {}", cgroup_path);
|
||||
|
||||
// 提取容器ID
|
||||
let container_id = if cgroup_path.contains("/docker/") {
|
||||
// 格式: /docker/abc123def456...
|
||||
let id_parts: Vec<&str> = cgroup_path.split('/').collect();
|
||||
if id_parts.len() >= 3 {
|
||||
id_parts[2].to_string()
|
||||
} else {
|
||||
continue;
|
||||
}
|
||||
} else if cgroup_path.contains(".scope") {
|
||||
// 格式: /system.slice/docker-abc123def456...scope
|
||||
let scope_name = cgroup_path.split('/').next_back().unwrap_or("");
|
||||
if scope_name.starts_with("docker-") && scope_name.ends_with(".scope") {
|
||||
// 移除 "docker-" 前缀和 ".scope" 后缀
|
||||
let id = &scope_name[7..scope_name.len() - 6];
|
||||
id.to_string()
|
||||
} else {
|
||||
continue;
|
||||
}
|
||||
} else {
|
||||
continue;
|
||||
};
|
||||
|
||||
// 验证容器ID格式(应该是64个字符的十六进制字符串)
|
||||
if container_id.len() == 64
|
||||
&& container_id.chars().all(|c| c.is_ascii_hexdigit())
|
||||
{
|
||||
info!("Detected container ID: {}", container_id);
|
||||
return Ok(container_id);
|
||||
} else {
|
||||
debug!("Skipping container ID: {}", container_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 如果 cgroup 方法失败,尝试其他方法
|
||||
warn!("Unable to get container ID from cgroup");
|
||||
|
||||
// 方法2:尝试读取 /proc/1/cgroup(主进程)
|
||||
if let Ok(cgroup_content) = fs::read_to_string("/proc/1/cgroup").await {
|
||||
debug!(" reading /proc/1/cgroup");
|
||||
for line in cgroup_content.lines() {
|
||||
if line.contains("/docker/") || line.contains(".scope") {
|
||||
debug!(" /proc/1/cgroup line: {}", line);
|
||||
// 类似的解析逻辑...
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 方法3:尝试读取主机名(某些环境容器ID会作为主机名)
|
||||
if let Ok(hostname) = std::env::var("HOSTNAME") {
|
||||
debug!("check HOSTNAME: {}", hostname);
|
||||
if hostname.len() == 12 && hostname.chars().all(|c| c.is_ascii_hexdigit()) {
|
||||
// 可能是短格式的容器ID(前12位)
|
||||
info!(
|
||||
" HOSTNAME get containerID: {}",
|
||||
hostname
|
||||
);
|
||||
return Ok(hostname);
|
||||
}
|
||||
}
|
||||
|
||||
bail!("unable to get current container ID, please ensure container has sufficient permissions to access /proc/self/cgroup");
|
||||
}
|
||||
|
||||
/// 验证 Docker socket 连接
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Result<()>` - 连接成功或错误
|
||||
pub async fn verify_docker_connection(&self) -> Result<()> {
|
||||
self.docker_client
|
||||
.ping()
|
||||
.await
|
||||
.context("Docker socket connection test failed")?;
|
||||
info!("Docker socket connection succeeded");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 获取容器所有挂载点信息(用于调试)
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Result<Vec<(String, String)>>` - 挂载点列表(容器路径 -> 宿主机路径)
|
||||
pub async fn get_all_mounts(&self) -> Result<Vec<(String, String)>> {
|
||||
let inspect_result = self
|
||||
.docker_client
|
||||
.inspect_container(
|
||||
&self.container_id,
|
||||
None::<bollard::query_parameters::InspectContainerOptions>,
|
||||
)
|
||||
.await
|
||||
.context("Failed to call Docker inspect API")?;
|
||||
|
||||
let mut mounts = Vec::new();
|
||||
|
||||
if let Some(mount_infos) = inspect_result.mounts {
|
||||
for mount in mount_infos {
|
||||
if let (Some(dest), Some(source)) = (&mount.destination, &mount.source) {
|
||||
mounts.push((dest.clone(), source.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(mounts)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_container_id_parsing() {
|
||||
// 这里应该模拟 cgroup 文件内容进行测试
|
||||
// 由于测试环境不在容器内,这个测试可能需要跳过
|
||||
}
|
||||
}
|
||||
364
qiming-rcoder/crates/docker_manager/src/container_state_actor.rs
Normal file
364
qiming-rcoder/crates/docker_manager/src/container_state_actor.rs
Normal file
@@ -0,0 +1,364 @@
|
||||
//! Container State Actor
|
||||
//!
|
||||
//! 使用 Actor 模式管理容器状态,避免 DashMap 跨 await 持有锁导致的死锁问题。
|
||||
//!
|
||||
//! # 架构
|
||||
//! - `ContainerStateActor`: 独占 HashMap,在独立 task 中运行,处理所有状态操作
|
||||
//! - `ContainerStateHandle`: 可克隆的句柄,提供 async 方法与 Actor 通信
|
||||
//!
|
||||
//! # 优点
|
||||
//! - 完全无锁,不会死锁
|
||||
//! - 状态变更顺序化,更易调试
|
||||
//! - 符合 Rust async 最佳实践
|
||||
|
||||
use crate::DockerContainerInfo;
|
||||
use std::collections::HashMap;
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use tracing::{debug, error, warn};
|
||||
|
||||
/// Actor 通道缓冲区大小
|
||||
const CHANNEL_BUFFER_SIZE: usize = 256;
|
||||
|
||||
/// 容器状态操作命令
|
||||
#[derive(Debug)]
|
||||
pub enum ContainerStateCommand {
|
||||
/// 获取容器信息
|
||||
Get {
|
||||
key: String,
|
||||
reply: oneshot::Sender<Option<DockerContainerInfo>>,
|
||||
},
|
||||
/// 插入/更新容器信息
|
||||
Insert {
|
||||
key: String,
|
||||
info: DockerContainerInfo,
|
||||
},
|
||||
/// 移除容器信息
|
||||
Remove {
|
||||
key: String,
|
||||
reply: oneshot::Sender<Option<DockerContainerInfo>>,
|
||||
},
|
||||
/// 获取所有容器列表
|
||||
List {
|
||||
reply: oneshot::Sender<Vec<DockerContainerInfo>>,
|
||||
},
|
||||
/// 获取所有 key 列表
|
||||
Keys { reply: oneshot::Sender<Vec<String>> },
|
||||
/// 获取容器数量
|
||||
Len { reply: oneshot::Sender<usize> },
|
||||
/// 检查 key 是否存在
|
||||
Contains {
|
||||
key: String,
|
||||
reply: oneshot::Sender<bool>,
|
||||
},
|
||||
/// 通过回调更新容器信息(用于原地更新)
|
||||
UpdateWith {
|
||||
key: String,
|
||||
/// 更新后的信息(如果 key 存在)
|
||||
updated_info: DockerContainerInfo,
|
||||
reply: oneshot::Sender<bool>,
|
||||
},
|
||||
/// 条件移除:只有当 container_id 匹配时才移除
|
||||
RemoveIfContainerId {
|
||||
key: String,
|
||||
container_id: String,
|
||||
reply: oneshot::Sender<Option<DockerContainerInfo>>,
|
||||
},
|
||||
}
|
||||
|
||||
/// 容器状态 Actor
|
||||
///
|
||||
/// 独占 HashMap,在独立 task 中运行
|
||||
pub struct ContainerStateActor {
|
||||
containers: HashMap<String, DockerContainerInfo>,
|
||||
receiver: mpsc::Receiver<ContainerStateCommand>,
|
||||
}
|
||||
|
||||
impl ContainerStateActor {
|
||||
/// 创建新的 Actor 和 Handle
|
||||
pub fn new() -> (Self, ContainerStateHandle) {
|
||||
let (sender, receiver) = mpsc::channel(CHANNEL_BUFFER_SIZE);
|
||||
let actor = Self {
|
||||
containers: HashMap::new(),
|
||||
receiver,
|
||||
};
|
||||
let handle = ContainerStateHandle { sender };
|
||||
(actor, handle)
|
||||
}
|
||||
|
||||
/// 运行 Actor 事件循环
|
||||
///
|
||||
/// 这个方法应该在 `tokio::spawn` 中调用
|
||||
pub async fn run(mut self) {
|
||||
debug!("🚀 [ACTOR] ContainerStateActor started");
|
||||
|
||||
while let Some(cmd) = self.receiver.recv().await {
|
||||
self.handle_command(cmd);
|
||||
}
|
||||
|
||||
debug!("🛑 [ACTOR] ContainerStateActor stopped (all senders dropped)");
|
||||
}
|
||||
|
||||
/// 处理单个命令
|
||||
fn handle_command(&mut self, cmd: ContainerStateCommand) {
|
||||
match cmd {
|
||||
ContainerStateCommand::Get { key, reply } => {
|
||||
let result = self.containers.get(&key).cloned();
|
||||
if reply.send(result).is_err() {
|
||||
warn!("[ACTOR] Get reply channel closed");
|
||||
}
|
||||
}
|
||||
ContainerStateCommand::Insert { key, info } => {
|
||||
self.containers.insert(key, info);
|
||||
}
|
||||
ContainerStateCommand::Remove { key, reply } => {
|
||||
let result = self.containers.remove(&key);
|
||||
if reply.send(result).is_err() {
|
||||
warn!("[ACTOR] Remove reply channel closed");
|
||||
}
|
||||
}
|
||||
ContainerStateCommand::List { reply } => {
|
||||
let result: Vec<_> = self.containers.values().cloned().collect();
|
||||
if reply.send(result).is_err() {
|
||||
warn!("[ACTOR] List reply channel closed");
|
||||
}
|
||||
}
|
||||
ContainerStateCommand::Keys { reply } => {
|
||||
let result: Vec<_> = self.containers.keys().cloned().collect();
|
||||
if reply.send(result).is_err() {
|
||||
warn!("[ACTOR] Keys reply channel closed");
|
||||
}
|
||||
}
|
||||
ContainerStateCommand::Len { reply } => {
|
||||
if reply.send(self.containers.len()).is_err() {
|
||||
warn!("[ACTOR] Len reply channel closed");
|
||||
}
|
||||
}
|
||||
ContainerStateCommand::Contains { key, reply } => {
|
||||
if reply.send(self.containers.contains_key(&key)).is_err() {
|
||||
warn!("[ACTOR] Contains reply channel closed");
|
||||
}
|
||||
}
|
||||
ContainerStateCommand::UpdateWith {
|
||||
key,
|
||||
updated_info,
|
||||
reply,
|
||||
} => {
|
||||
let existed = if let std::collections::hash_map::Entry::Occupied(mut e) =
|
||||
self.containers.entry(key)
|
||||
{
|
||||
e.insert(updated_info);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
};
|
||||
if reply.send(existed).is_err() {
|
||||
warn!("[ACTOR] UpdateWith reply channel closed");
|
||||
}
|
||||
}
|
||||
ContainerStateCommand::RemoveIfContainerId {
|
||||
key,
|
||||
container_id,
|
||||
reply,
|
||||
} => {
|
||||
let should_remove = if let Some(info) = self.containers.get(&key) {
|
||||
info.container_id == container_id
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
let result = if should_remove {
|
||||
self.containers.remove(&key)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
if reply.send(result).is_err() {
|
||||
warn!("[ACTOR] RemoveIfContainerId reply channel closed");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 容器状态句柄
|
||||
///
|
||||
/// 可克隆,用于与 Actor 通信
|
||||
#[derive(Clone)]
|
||||
pub struct ContainerStateHandle {
|
||||
sender: mpsc::Sender<ContainerStateCommand>,
|
||||
}
|
||||
|
||||
impl ContainerStateHandle {
|
||||
/// 获取容器信息
|
||||
pub async fn get(&self, key: &str) -> Option<DockerContainerInfo> {
|
||||
let (reply, rx) = oneshot::channel();
|
||||
if self
|
||||
.sender
|
||||
.send(ContainerStateCommand::Get {
|
||||
key: key.to_string(),
|
||||
reply,
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
error!("[HANDLE] Failed to send Get command - actor stopped");
|
||||
return None;
|
||||
}
|
||||
rx.await.unwrap_or(None)
|
||||
}
|
||||
|
||||
/// 插入/更新容器信息
|
||||
pub async fn insert(&self, key: String, info: DockerContainerInfo) {
|
||||
if self
|
||||
.sender
|
||||
.send(ContainerStateCommand::Insert { key, info })
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
error!("[HANDLE] Failed to send Insert command - actor stopped");
|
||||
}
|
||||
}
|
||||
|
||||
/// 移除容器信息
|
||||
pub async fn remove(&self, key: &str) -> Option<DockerContainerInfo> {
|
||||
let (reply, rx) = oneshot::channel();
|
||||
if self
|
||||
.sender
|
||||
.send(ContainerStateCommand::Remove {
|
||||
key: key.to_string(),
|
||||
reply,
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
error!("[HANDLE] Failed to send Remove command - actor stopped");
|
||||
return None;
|
||||
}
|
||||
rx.await.unwrap_or(None)
|
||||
}
|
||||
|
||||
/// 获取所有容器列表
|
||||
pub async fn list(&self) -> Vec<DockerContainerInfo> {
|
||||
let (reply, rx) = oneshot::channel();
|
||||
if self
|
||||
.sender
|
||||
.send(ContainerStateCommand::List { reply })
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
error!("[HANDLE] Failed to send List command - actor stopped");
|
||||
return Vec::new();
|
||||
}
|
||||
rx.await.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// 获取所有 key 列表
|
||||
pub async fn keys(&self) -> Vec<String> {
|
||||
let (reply, rx) = oneshot::channel();
|
||||
if self
|
||||
.sender
|
||||
.send(ContainerStateCommand::Keys { reply })
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
error!("[HANDLE] Failed to send Keys command - actor stopped");
|
||||
return Vec::new();
|
||||
}
|
||||
rx.await.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// 获取容器数量
|
||||
pub async fn len(&self) -> usize {
|
||||
let (reply, rx) = oneshot::channel();
|
||||
if self
|
||||
.sender
|
||||
.send(ContainerStateCommand::Len { reply })
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
error!("[HANDLE] Failed to send Len command - actor stopped");
|
||||
return 0;
|
||||
}
|
||||
rx.await.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// 检查是否为空
|
||||
pub async fn is_empty(&self) -> bool {
|
||||
self.len().await == 0
|
||||
}
|
||||
|
||||
/// 检查 key 是否存在
|
||||
pub async fn contains_key(&self, key: &str) -> bool {
|
||||
let (reply, rx) = oneshot::channel();
|
||||
if self
|
||||
.sender
|
||||
.send(ContainerStateCommand::Contains {
|
||||
key: key.to_string(),
|
||||
reply,
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
error!("[HANDLE] Failed to send Contains command - actor stopped");
|
||||
return false;
|
||||
}
|
||||
rx.await.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// 条件更新:如果 key 存在则更新
|
||||
///
|
||||
/// 返回 true 表示Update succeeded,false 表示 key 不存在
|
||||
pub async fn update_if_exists(&self, key: &str, info: DockerContainerInfo) -> bool {
|
||||
let (reply, rx) = oneshot::channel();
|
||||
if self
|
||||
.sender
|
||||
.send(ContainerStateCommand::UpdateWith {
|
||||
key: key.to_string(),
|
||||
updated_info: info,
|
||||
reply,
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
error!("[HANDLE] Failed to send UpdateWith command - actor stopped");
|
||||
return false;
|
||||
}
|
||||
rx.await.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// 条件移除:只有当 container_id 匹配时才移除
|
||||
///
|
||||
/// 防止在清理时误删刚重启的容器(CAS 操作)
|
||||
pub async fn remove_if_container_id(
|
||||
&self,
|
||||
key: &str,
|
||||
container_id: &str,
|
||||
) -> Option<DockerContainerInfo> {
|
||||
let (reply, rx) = oneshot::channel();
|
||||
if self
|
||||
.sender
|
||||
.send(ContainerStateCommand::RemoveIfContainerId {
|
||||
key: key.to_string(),
|
||||
container_id: container_id.to_string(),
|
||||
reply,
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
error!("[HANDLE] Failed to send RemoveIfContainerId command - actor stopped");
|
||||
return None;
|
||||
}
|
||||
rx.await.unwrap_or(None)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ContainerStateHandle {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ContainerStateHandle")
|
||||
.field("sender_closed", &self.sender.is_closed())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
// Tests removed intentionally.
|
||||
// Integration tests should be implemented in `tests/` directory if needed.
|
||||
600
qiming-rcoder/crates/docker_manager/src/container_stop.rs
Normal file
600
qiming-rcoder/crates/docker_manager/src/container_stop.rs
Normal file
@@ -0,0 +1,600 @@
|
||||
//! 统一的容器停止模块
|
||||
//!
|
||||
//! 提供两种容器停止策略:
|
||||
//! 1. 启动时清理(startup_cleanup):用于服务启动时清理遗留容器
|
||||
//! - 使用5秒超时
|
||||
//! - 过滤409冲突错误(容器已在删除中)
|
||||
//! - 不阻塞服务启动
|
||||
//!
|
||||
//! 2. 运行时清理(runtime_cleanup):用于运行时快速清理容器
|
||||
//! - 使用3秒优雅停止超时
|
||||
//! - 超时后立即强制停止
|
||||
//! - 快速释放资源
|
||||
//!
|
||||
//! # 使用示例
|
||||
//!
|
||||
//! ```rust,no_run
|
||||
//! use docker_manager::container_stop;
|
||||
//! use docker_manager::DockerManager;
|
||||
//! use std::sync::Arc;
|
||||
//!
|
||||
//! # async fn example() -> anyhow::Result<()> {
|
||||
//! # let docker_manager = Arc::new(DockerManager::new(Default::default()).await?);
|
||||
//!
|
||||
//! // 启动时清理
|
||||
//! let result = container_stop::startup_cleanup_containers(
|
||||
//! &docker_manager,
|
||||
//! "rcoder-agent-*"
|
||||
//! ).await?;
|
||||
//!
|
||||
//! // 运行时清理单个容器
|
||||
//! container_stop::runtime_cleanup_container(
|
||||
//! &docker_manager,
|
||||
//! "container_id_123"
|
||||
//! ).await?;
|
||||
//! # Ok(())
|
||||
//! # }
|
||||
//! ```
|
||||
|
||||
use crate::{CleanupResult, ContainerRemovalFailure, DockerError, DockerManager, DockerResult};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// 启动清理超时时间(秒)
|
||||
///
|
||||
/// 启动时使用较短的超时时间,快速清理遗留容器
|
||||
const STARTUP_CLEANUP_TIMEOUT_SECONDS: u64 = 5;
|
||||
|
||||
/// 运行时清理超时时间(秒)
|
||||
///
|
||||
/// 运行时给容器3秒优雅退出时间,然后强制停止
|
||||
const RUNTIME_CLEANUP_TIMEOUT_SECONDS: u64 = 3;
|
||||
|
||||
/// 容器停止后的等待时间(毫秒)
|
||||
///
|
||||
/// 给Docker一些时间完成清理操作
|
||||
const POST_STOP_WAIT_MS: u64 = 100;
|
||||
|
||||
/// 启动时容器清理策略
|
||||
///
|
||||
/// 用于服务启动时清理遗留的容器。此函数会:
|
||||
/// - 查找匹配指定模式的所有容器
|
||||
/// - 🚀 并发停止所有容器(提高清理速度)
|
||||
/// - 使用5秒超时停止每个容器
|
||||
/// - 过滤409冲突错误(容器已在删除中)
|
||||
/// - 返回详细的清理统计信息
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `docker_manager` - Docker管理器实例
|
||||
/// * `pattern` - 容器名称匹配模式(如 "rcoder-agent-*")
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// 返回 `CleanupResult` 包含清理统计信息
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```rust,no_run
|
||||
/// use docker_manager::container_stop;
|
||||
/// use docker_manager::DockerManager;
|
||||
/// use std::sync::Arc;
|
||||
///
|
||||
/// # async fn example() -> anyhow::Result<()> {
|
||||
/// # let docker_manager = Arc::new(DockerManager::new(Default::default()).await?);
|
||||
///
|
||||
/// let result = container_stop::startup_cleanup_containers(
|
||||
/// &docker_manager,
|
||||
/// "rcoder-agent-*"
|
||||
/// ).await?;
|
||||
///
|
||||
/// println!("cleanup message {} message container", result.successfully_removed);
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn startup_cleanup_containers(
|
||||
docker_manager: &Arc<DockerManager>,
|
||||
pattern: &str,
|
||||
) -> DockerResult<CleanupResult> {
|
||||
info!(
|
||||
"[STARTUP_CLEANUP] Starting cleanup container: pattern={}",
|
||||
pattern
|
||||
);
|
||||
let start_time = Instant::now();
|
||||
|
||||
// 查找匹配模式的容器
|
||||
let matched_containers = docker_manager.list_containers_with_pattern(pattern).await?;
|
||||
|
||||
let total_found = matched_containers.len();
|
||||
info!(
|
||||
"[STARTUP_CLEANUP] Found {} containers",
|
||||
total_found
|
||||
);
|
||||
|
||||
if total_found == 0 {
|
||||
return Ok(CleanupResult {
|
||||
total_found: 0,
|
||||
successfully_removed: 0,
|
||||
failed_removals: 0,
|
||||
skipped_running: 0,
|
||||
removed_container_ids: Vec::new(),
|
||||
failed_removals_details: Vec::new(),
|
||||
duration_ms: start_time.elapsed().as_millis() as u64,
|
||||
});
|
||||
}
|
||||
|
||||
let mut successfully_removed = 0;
|
||||
let mut failed_removals = 0;
|
||||
let mut removed_container_ids = Vec::new();
|
||||
let mut failed_removals_details = Vec::new();
|
||||
|
||||
// 🚀 并发停止所有容器
|
||||
let mut tasks = Vec::new();
|
||||
for container in &matched_containers {
|
||||
if let Some(container_id) = &container.id {
|
||||
let container_name = container
|
||||
.names
|
||||
.as_ref()
|
||||
.and_then(|names| names.first())
|
||||
.map(|n| n.trim_start_matches('/'))
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
let docker_manager_clone = Arc::clone(docker_manager);
|
||||
let container_id_clone = container_id.clone();
|
||||
let task = tokio::spawn(async move {
|
||||
let result =
|
||||
stop_container_startup_mode(&docker_manager_clone, &container_id_clone).await;
|
||||
(container_id_clone, container_name, result)
|
||||
});
|
||||
tasks.push(task);
|
||||
}
|
||||
}
|
||||
|
||||
// 等待所有任务完成
|
||||
for task in tasks {
|
||||
if let Ok((container_id, container_name, result)) = task.await {
|
||||
match result {
|
||||
Ok(_) => {
|
||||
successfully_removed += 1;
|
||||
removed_container_ids.push(container_id.clone());
|
||||
info!(
|
||||
"[STARTUP_CLEANUP] Container cleanup succeeded: container_id={}, name={}",
|
||||
container_id, container_name
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
// 检查是否为409冲突错误
|
||||
if is_409_conflict_error(&e) {
|
||||
info!(
|
||||
"[STARTUP_CLEANUP] Container already being removed, skipping: container_id={}, name={}",
|
||||
container_id, container_name
|
||||
);
|
||||
// 409错误不计入失败统计
|
||||
successfully_removed += 1;
|
||||
removed_container_ids.push(container_id);
|
||||
} else {
|
||||
failed_removals += 1;
|
||||
warn!(
|
||||
"[STARTUP_CLEANUP] Container cleanup failed: container_id={}, name={}, error={}",
|
||||
container_id, container_name, e
|
||||
);
|
||||
failed_removals_details.push(ContainerRemovalFailure {
|
||||
container_id,
|
||||
container_name,
|
||||
error_message: e.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let duration_ms = start_time.elapsed().as_millis() as u64;
|
||||
|
||||
info!(
|
||||
"[STARTUP_CLEANUP] Cleanup completed: total={}, success={}, failed={}, duration={}ms",
|
||||
total_found, successfully_removed, failed_removals, duration_ms
|
||||
);
|
||||
|
||||
Ok(CleanupResult {
|
||||
total_found,
|
||||
successfully_removed,
|
||||
failed_removals,
|
||||
skipped_running: 0, // 启动清理不跳过运行中的容器
|
||||
removed_container_ids,
|
||||
failed_removals_details,
|
||||
duration_ms,
|
||||
})
|
||||
}
|
||||
|
||||
/// 停止单个容器(启动模式)
|
||||
///
|
||||
/// 使用启动清理的超时设置停止容器
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `docker_manager` - Docker管理器实例
|
||||
/// * `container_id` - 容器ID
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// 成功返回 `Ok(())`,失败返回 `DockerError`
|
||||
async fn stop_container_startup_mode(
|
||||
docker_manager: &Arc<DockerManager>,
|
||||
container_id: &str,
|
||||
) -> DockerResult<()> {
|
||||
docker_manager
|
||||
.stop_container_by_id_with_timeout(container_id, STARTUP_CLEANUP_TIMEOUT_SECONDS)
|
||||
.await?;
|
||||
|
||||
// 给Docker一些时间完成清理
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(POST_STOP_WAIT_MS)).await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 检查是否为409冲突错误
|
||||
///
|
||||
/// 409错误表示容器已经在删除过程中,这在启动清理时是正常情况
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `error` - Docker错误
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// 如果是409冲突错误返回 `true`,否则返回 `false`
|
||||
fn is_409_conflict_error(error: &DockerError) -> bool {
|
||||
// 检查是否是 Docker 409 冲突错误(容器删除已在进行中)
|
||||
matches!(
|
||||
error,
|
||||
DockerError::BollardError(bollard::errors::Error::DockerResponseServerError {
|
||||
status_code: 409,
|
||||
..
|
||||
})
|
||||
)
|
||||
}
|
||||
|
||||
/// 运行时容器清理策略(单个容器)
|
||||
///
|
||||
/// 用于运行时快速清理单个容器。此函数会:
|
||||
/// - 使用3秒优雅停止超时
|
||||
/// - 超时后立即强制停止
|
||||
/// - 快速释放资源
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `docker_manager` - Docker管理器实例
|
||||
/// * `container_id` - 容器ID
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// 成功返回 `Ok(())`,失败返回 `DockerError`
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```rust,no_run
|
||||
/// use docker_manager::container_stop;
|
||||
/// use docker_manager::DockerManager;
|
||||
/// use std::sync::Arc;
|
||||
///
|
||||
/// # async fn example() -> anyhow::Result<()> {
|
||||
/// # let docker_manager = Arc::new(DockerManager::new(Default::default()).await?);
|
||||
///
|
||||
/// container_stop::runtime_cleanup_container(
|
||||
/// &docker_manager,
|
||||
/// "container_id_123"
|
||||
/// ).await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn runtime_cleanup_container(
|
||||
docker_manager: &Arc<DockerManager>,
|
||||
container_id: &str,
|
||||
) -> DockerResult<()> {
|
||||
info!(
|
||||
"[RUNTIME_CLEANUP] Starting to stop container: container_id={}",
|
||||
container_id
|
||||
);
|
||||
|
||||
match stop_container_runtime_mode(docker_manager, container_id).await {
|
||||
Ok(_) => {
|
||||
info!(
|
||||
"[RUNTIME_CLEANUP] Container stop succeeded: container_id={}",
|
||||
container_id
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"[RUNTIME_CLEANUP] Container stop failed: container_id={}, error={}",
|
||||
container_id, e
|
||||
);
|
||||
Err(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 运行时容器清理策略(批量)
|
||||
///
|
||||
/// 用于运行时批量清理多个容器。此函数会:
|
||||
/// - 🚀 并发停止所有容器(提高清理速度)
|
||||
/// - 使用3秒优雅停止超时
|
||||
/// - 超时后立即强制停止
|
||||
/// - 返回详细的清理统计信息
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `docker_manager` - Docker管理器实例
|
||||
/// * `container_ids` - 容器ID列表
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// 返回 `CleanupResult` 包含清理统计信息
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```rust,no_run
|
||||
/// use docker_manager::container_stop;
|
||||
/// use docker_manager::DockerManager;
|
||||
/// use std::sync::Arc;
|
||||
///
|
||||
/// # async fn example() -> anyhow::Result<()> {
|
||||
/// # let docker_manager = Arc::new(DockerManager::new(Default::default()).await?);
|
||||
///
|
||||
/// let container_ids = vec!["id1".to_string(), "id2".to_string()];
|
||||
/// let result = container_stop::runtime_cleanup_containers(
|
||||
/// &docker_manager,
|
||||
/// container_ids
|
||||
/// ).await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn runtime_cleanup_containers(
|
||||
docker_manager: &Arc<DockerManager>,
|
||||
container_ids: Vec<String>,
|
||||
) -> DockerResult<CleanupResult> {
|
||||
info!(
|
||||
"[RUNTIME_CLEANUP] Starting batch container cleanup: count={}",
|
||||
container_ids.len()
|
||||
);
|
||||
let start_time = Instant::now();
|
||||
|
||||
let total_found = container_ids.len();
|
||||
let mut successfully_removed = 0;
|
||||
let mut failed_removals = 0;
|
||||
let mut removed_container_ids = Vec::new();
|
||||
let mut failed_removals_details = Vec::new();
|
||||
|
||||
// 🚀 并发停止所有容器
|
||||
let mut tasks = Vec::new();
|
||||
for container_id in &container_ids {
|
||||
let docker_manager_clone = Arc::clone(docker_manager);
|
||||
let container_id_clone = container_id.clone();
|
||||
let task = tokio::spawn(async move {
|
||||
let result =
|
||||
stop_container_runtime_mode(&docker_manager_clone, &container_id_clone).await;
|
||||
(container_id_clone, result)
|
||||
});
|
||||
tasks.push(task);
|
||||
}
|
||||
|
||||
// 等待所有任务完成
|
||||
for task in tasks {
|
||||
if let Ok((container_id, result)) = task.await {
|
||||
match result {
|
||||
Ok(_) => {
|
||||
successfully_removed += 1;
|
||||
removed_container_ids.push(container_id.clone());
|
||||
info!(
|
||||
"[RUNTIME_CLEANUP] Container cleanup succeeded: container_id={}",
|
||||
container_id
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
failed_removals += 1;
|
||||
warn!(
|
||||
"[RUNTIME_CLEANUP] Container cleanup failed: container_id={}, error={}",
|
||||
container_id, e
|
||||
);
|
||||
failed_removals_details.push(ContainerRemovalFailure {
|
||||
container_id: container_id.clone(),
|
||||
container_name: container_id.clone(), // 批量清理时可能不知道名称
|
||||
error_message: e.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let duration_ms = start_time.elapsed().as_millis() as u64;
|
||||
|
||||
info!(
|
||||
"[RUNTIME_CLEANUP] Batch cleanup completed: total={}, success={}, failed={}, duration={}ms",
|
||||
total_found, successfully_removed, failed_removals, duration_ms
|
||||
);
|
||||
|
||||
Ok(CleanupResult {
|
||||
total_found,
|
||||
successfully_removed,
|
||||
failed_removals,
|
||||
skipped_running: 0, // 运行时清理不跳过运行中的容器
|
||||
removed_container_ids,
|
||||
failed_removals_details,
|
||||
duration_ms,
|
||||
})
|
||||
}
|
||||
|
||||
/// 停止单个容器(运行时模式)
|
||||
///
|
||||
/// 使用运行时清理的超时设置停止容器
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `docker_manager` - Docker管理器实例
|
||||
/// * `container_id` - 容器ID
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// 成功返回 `Ok(())`,失败返回 `DockerError`
|
||||
async fn stop_container_runtime_mode(
|
||||
docker_manager: &Arc<DockerManager>,
|
||||
container_id: &str,
|
||||
) -> DockerResult<()> {
|
||||
docker_manager
|
||||
.stop_container_by_id_with_timeout(container_id, RUNTIME_CLEANUP_TIMEOUT_SECONDS)
|
||||
.await?;
|
||||
|
||||
// 给Docker一些时间完成清理
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(POST_STOP_WAIT_MS)).await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 为所有启用的服务构建容器清理模式列表
|
||||
///
|
||||
/// 从多镜像配置中获取所有启用的服务类型,并生成对应的容器名称模式。
|
||||
/// 使用 ServiceImageConfig.container_prefix() 获取配置的前缀,确保与容器创建时使用的前缀一致。
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `multi_image_config` - 多镜像配置
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// 返回容器名称模式列表,如 `["rcoder-agent-*", "rcoder-computer-agent-runner-*"]`
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```rust,no_run
|
||||
/// use docker_manager::container_stop;
|
||||
/// use shared_types;
|
||||
///
|
||||
/// let config = shared_types::create_default_multi_image_config();
|
||||
/// let patterns = container_stop::get_container_patterns_for_enabled_services(&config);
|
||||
/// println!("container message : {:?}", patterns);
|
||||
/// ```
|
||||
pub fn get_container_patterns_for_enabled_services(
|
||||
multi_image_config: &shared_types::MultiImageConfig,
|
||||
) -> Vec<String> {
|
||||
// 直接遍历 services,获取启用的服务配置并使用其 container_prefix()
|
||||
multi_image_config
|
||||
.services
|
||||
.values()
|
||||
.filter(|config| config.enabled)
|
||||
.map(|config| {
|
||||
let prefix = config.container_prefix();
|
||||
let pattern = format!("{}-*", prefix);
|
||||
tracing::debug!(
|
||||
"🔍 [CLEANUP_PATTERN] Service type: {:?}, using prefix: {}, pattern: {}",
|
||||
config.service_type,
|
||||
prefix,
|
||||
pattern
|
||||
);
|
||||
pattern
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 清理所有启用服务的容器(启动时清理)
|
||||
///
|
||||
/// 自动从配置中获取所有启用的服务,并清理对应的容器。此函数会:
|
||||
/// - 从配置中读取启用的服务类型
|
||||
/// - 为每个服务类型生成容器模式
|
||||
/// - 并行清理多个服务类型的容器
|
||||
/// - 聚合所有清理结果
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `docker_manager` - Docker管理器实例
|
||||
/// * `multi_image_config` - 多镜像配置
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// 返回聚合的 `CleanupResult` 包含所有服务的清理统计信息
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```rust,no_run
|
||||
/// use docker_manager::container_stop;
|
||||
/// use docker_manager::DockerManager;
|
||||
/// use shared_types;
|
||||
/// use std::sync::Arc;
|
||||
///
|
||||
/// # async fn example() -> anyhow::Result<()> {
|
||||
/// # let docker_manager = Arc::new(DockerManager::new(Default::default()).await?);
|
||||
/// let config = shared_types::create_default_multi_image_config();
|
||||
///
|
||||
/// let result = container_stop::startup_cleanup_all_enabled_services(
|
||||
/// &docker_manager,
|
||||
/// &config
|
||||
/// ).await?;
|
||||
///
|
||||
/// println!("cleanup message {} message container", result.successfully_removed);
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn startup_cleanup_all_enabled_services(
|
||||
docker_manager: &Arc<DockerManager>,
|
||||
multi_image_config: &shared_types::MultiImageConfig,
|
||||
) -> DockerResult<CleanupResult> {
|
||||
let patterns = get_container_patterns_for_enabled_services(multi_image_config);
|
||||
|
||||
if patterns.is_empty() {
|
||||
warn!("No patterns, skip container cleanup");
|
||||
return Ok(CleanupResult::default());
|
||||
}
|
||||
|
||||
info!("🧹 Starting cleanup container: {:?}", patterns);
|
||||
let start_time = Instant::now();
|
||||
|
||||
// 并行清理多个服务类型的容器
|
||||
let cleanup_tasks: Vec<_> = patterns
|
||||
.into_iter()
|
||||
.map(|pattern| {
|
||||
let docker_manager = docker_manager.clone();
|
||||
let pattern_clone = pattern.clone();
|
||||
tokio::spawn(async move {
|
||||
startup_cleanup_containers(&docker_manager, &pattern_clone).await
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 聚合所有清理结果
|
||||
let mut aggregated_result = CleanupResult::default();
|
||||
for task in cleanup_tasks {
|
||||
match task.await {
|
||||
Ok(Ok(result)) => {
|
||||
aggregated_result.total_found += result.total_found;
|
||||
aggregated_result.successfully_removed += result.successfully_removed;
|
||||
aggregated_result.failed_removals += result.failed_removals;
|
||||
aggregated_result.skipped_running += result.skipped_running;
|
||||
aggregated_result
|
||||
.removed_container_ids
|
||||
.extend(result.removed_container_ids);
|
||||
aggregated_result
|
||||
.failed_removals_details
|
||||
.extend(result.failed_removals_details);
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
warn!("cleanup container failed: {}", e);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("cleanup failed: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
aggregated_result.duration_ms = start_time.elapsed().as_millis() as u64;
|
||||
|
||||
info!(
|
||||
"[MULTI_SERVICE_CLEANUP] Multi-service cleanup completed: total={}, success={}, failed={}, duration={}ms",
|
||||
aggregated_result.total_found,
|
||||
aggregated_result.successfully_removed,
|
||||
aggregated_result.failed_removals,
|
||||
aggregated_result.duration_ms
|
||||
);
|
||||
|
||||
Ok(aggregated_result)
|
||||
}
|
||||
125
qiming-rcoder/crates/docker_manager/src/health/http_health.rs
Normal file
125
qiming-rcoder/crates/docker_manager/src/health/http_health.rs
Normal file
@@ -0,0 +1,125 @@
|
||||
//! HTTP 健康检查
|
||||
//!
|
||||
//! 从 docker_container_agent.rs 迁移
|
||||
|
||||
use crate::{DockerError, DockerResult};
|
||||
use reqwest::Client;
|
||||
use std::time::Duration;
|
||||
use tokio::time::timeout;
|
||||
use tracing::{debug, info};
|
||||
|
||||
/// HTTP 健康检查器
|
||||
pub struct HttpHealthChecker {
|
||||
client: Client,
|
||||
max_attempts: u32,
|
||||
timeout_seconds: u64,
|
||||
}
|
||||
|
||||
impl HttpHealthChecker {
|
||||
/// 创建新的健康检查器
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `max_attempts` - 最大尝试次数
|
||||
/// * `timeout_seconds` - 每次尝试的超时时间(秒)
|
||||
pub fn new(max_attempts: u32, timeout_seconds: u64) -> Self {
|
||||
Self {
|
||||
client: Client::new(),
|
||||
max_attempts,
|
||||
timeout_seconds,
|
||||
}
|
||||
}
|
||||
|
||||
/// 默认配置的健康检查器(60次,每次2秒,总计约180秒)
|
||||
/// 容器启动包含 MCP Proxy 等服务,可能需要 60-90 秒
|
||||
pub fn default_checker() -> Self {
|
||||
Self::new(60, 2)
|
||||
}
|
||||
|
||||
/// 等待服务就绪
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `base_url` - 服务基础URL (如: "http://172.17.0.2:8086")
|
||||
/// * `health_path` - 健康检查路径 (默认: "/health")
|
||||
///
|
||||
/// # Returns
|
||||
/// * `DockerResult<()>` - 成功或超时错误
|
||||
pub async fn wait_for_ready(
|
||||
&self,
|
||||
base_url: &str,
|
||||
health_path: Option<&str>,
|
||||
) -> DockerResult<()> {
|
||||
let health_url = format!(
|
||||
"{}/{}",
|
||||
base_url.trim_end_matches('/'),
|
||||
health_path.unwrap_or("health").trim_start_matches('/')
|
||||
);
|
||||
|
||||
info!("Health check started: {}", health_url);
|
||||
|
||||
for attempt in 0..self.max_attempts {
|
||||
match timeout(
|
||||
Duration::from_secs(self.timeout_seconds),
|
||||
self.client.get(&health_url).send(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(response)) if response.status().is_success() => {
|
||||
info!("health check passed");
|
||||
return Ok(());
|
||||
}
|
||||
Ok(Ok(response)) => {
|
||||
debug!(
|
||||
"Service returned non-success status: {}, waiting... ({}/{})",
|
||||
response.status(),
|
||||
attempt + 1,
|
||||
self.max_attempts
|
||||
);
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
debug!(
|
||||
"Connection failed: {}, continuing to wait... ({}/{})",
|
||||
e,
|
||||
attempt + 1,
|
||||
self.max_attempts
|
||||
);
|
||||
}
|
||||
Err(_) => {
|
||||
debug!(
|
||||
"Connection timeout, continuing to wait... ({}/{})",
|
||||
attempt + 1,
|
||||
self.max_attempts
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// 每 10 次尝试输出一次 info 日志
|
||||
if (attempt + 1) % 10 == 0 {
|
||||
info!(
|
||||
"Still waiting for service to start... ({}/{})",
|
||||
attempt + 1,
|
||||
self.max_attempts
|
||||
);
|
||||
}
|
||||
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
Err(DockerError::ContainerStartError(format!(
|
||||
"Wait for service startup timeout: {} (attempted {} times)",
|
||||
health_url, self.max_attempts
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
/// 便捷函数: 等待服务就绪(使用默认配置)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `base_url` - 服务基础URL
|
||||
///
|
||||
/// # Returns
|
||||
/// * `DockerResult<()>` - 成功或超时错误
|
||||
pub async fn wait_for_service_ready(base_url: &str) -> DockerResult<()> {
|
||||
HttpHealthChecker::default_checker()
|
||||
.wait_for_ready(base_url, None)
|
||||
.await
|
||||
}
|
||||
7
qiming-rcoder/crates/docker_manager/src/health/mod.rs
Normal file
7
qiming-rcoder/crates/docker_manager/src/health/mod.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
//! 容器健康检查模块
|
||||
|
||||
pub mod http_health;
|
||||
pub mod service_health;
|
||||
|
||||
pub use http_health::*;
|
||||
pub use service_health::*;
|
||||
264
qiming-rcoder/crates/docker_manager/src/health/service_health.rs
Normal file
264
qiming-rcoder/crates/docker_manager/src/health/service_health.rs
Normal file
@@ -0,0 +1,264 @@
|
||||
//! 服务健康检查模块
|
||||
//!
|
||||
//! 提供服务层面的健康检查功能,包括 HTTP 健康端点和 gRPC 连接检查。
|
||||
//! 这是对 Docker 容器状态检查的补充,用于确认容器内服务是否真正可用。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::time::Duration;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::time::timeout;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
/// 默认 HTTP 健康检查端口
|
||||
pub const DEFAULT_HTTP_HEALTH_PORT: u16 = 8086;
|
||||
|
||||
/// 默认 gRPC 服务端口
|
||||
pub const DEFAULT_GRPC_PORT: u16 = 50051;
|
||||
|
||||
/// 服务健康检查超时时间(秒)
|
||||
const HEALTH_CHECK_TIMEOUT_SECS: u64 = 3;
|
||||
|
||||
/// 服务健康状态
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ServiceHealthStatus {
|
||||
/// HTTP 健康端点是否可达
|
||||
pub http_healthy: bool,
|
||||
/// gRPC 端口是否可连接
|
||||
pub grpc_healthy: bool,
|
||||
/// 上次检查时间
|
||||
pub last_check_time: DateTime<Utc>,
|
||||
/// 连续失败次数(HTTP 和 gRPC 都失败算一次)
|
||||
/// TODO: 用于后续自动重启功能
|
||||
pub consecutive_failures: u32,
|
||||
}
|
||||
|
||||
impl ServiceHealthStatus {
|
||||
/// 创建新的健康状态(初始化为未知状态)
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
http_healthy: false,
|
||||
grpc_healthy: false,
|
||||
last_check_time: Utc::now(),
|
||||
consecutive_failures: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查服务是否完全健康(HTTP 和 gRPC 都正常)
|
||||
pub fn is_fully_healthy(&self) -> bool {
|
||||
self.http_healthy && self.grpc_healthy
|
||||
}
|
||||
|
||||
/// 检查服务是否部分健康(至少有一个正常)
|
||||
pub fn is_partially_healthy(&self) -> bool {
|
||||
self.http_healthy || self.grpc_healthy
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ServiceHealthStatus {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// 服务健康检查器
|
||||
pub struct ServiceHealthChecker {
|
||||
client: Client,
|
||||
http_port: u16,
|
||||
grpc_port: u16,
|
||||
timeout_secs: u64,
|
||||
}
|
||||
|
||||
impl ServiceHealthChecker {
|
||||
/// 创建新的健康检查器
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
client: Client::builder()
|
||||
.timeout(Duration::from_secs(HEALTH_CHECK_TIMEOUT_SECS))
|
||||
.build()
|
||||
.unwrap_or_else(|_| Client::new()),
|
||||
http_port: DEFAULT_HTTP_HEALTH_PORT,
|
||||
grpc_port: DEFAULT_GRPC_PORT,
|
||||
timeout_secs: HEALTH_CHECK_TIMEOUT_SECS,
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用自定义端口创建检查器
|
||||
pub fn with_ports(http_port: u16, grpc_port: u16) -> Self {
|
||||
Self {
|
||||
http_port,
|
||||
grpc_port,
|
||||
..Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查 HTTP 健康端点
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `ip` - 容器 IP 地址
|
||||
///
|
||||
/// # Returns
|
||||
/// * `true` - 健康端点返回成功状态码
|
||||
/// * `false` - 连接失败或返回错误状态码
|
||||
pub async fn check_http_health(&self, ip: &str) -> bool {
|
||||
let url = format!("http://{}:{}/health", ip, self.http_port);
|
||||
|
||||
match timeout(
|
||||
Duration::from_secs(self.timeout_secs),
|
||||
self.client.get(&url).send(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(response)) => {
|
||||
let is_healthy = response.status().is_success();
|
||||
debug!(
|
||||
"HTTP health check {}: status={}, healthy={}",
|
||||
url,
|
||||
response.status(),
|
||||
is_healthy
|
||||
);
|
||||
is_healthy
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
debug!("HTTP health check failed {}: {}", url, e);
|
||||
false
|
||||
}
|
||||
Err(_) => {
|
||||
debug!("HTTP health check timeout {}", url);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查 gRPC 端口连通性
|
||||
///
|
||||
/// 通过 TCP 连接测试 gRPC 端口是否可达。
|
||||
/// 注意:这只是连接测试,不执行实际的 gRPC 健康检查协议。
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `ip` - 容器 IP 地址
|
||||
///
|
||||
/// # Returns
|
||||
/// * `true` - 能够建立 TCP 连接
|
||||
/// * `false` - 连接失败或超时
|
||||
pub async fn check_grpc_connectivity(&self, ip: &str) -> bool {
|
||||
let addr = format!("{}:{}", ip, self.grpc_port);
|
||||
|
||||
match timeout(
|
||||
Duration::from_secs(self.timeout_secs),
|
||||
TcpStream::connect(&addr),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(_stream)) => {
|
||||
debug!("gRPC port connection: {}", addr);
|
||||
true
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
debug!("gRPC portconnectionfailed {}: {}", addr, e);
|
||||
false
|
||||
}
|
||||
Err(_) => {
|
||||
debug!("gRPC portconnectiontimeout {}", addr);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行完整的服务健康检查
|
||||
///
|
||||
/// 同时检查 HTTP 健康端点和 gRPC 端口连通性。
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `container_ip` - 容器 IP 地址
|
||||
/// * `previous_failures` - 之前的连续失败次数(用于累加)
|
||||
///
|
||||
/// # Returns
|
||||
/// * `ServiceHealthStatus` - 包含所有检查结果的健康状态
|
||||
pub async fn check_service(
|
||||
&self,
|
||||
container_ip: &str,
|
||||
previous_failures: u32,
|
||||
) -> ServiceHealthStatus {
|
||||
// 并行执行两个检查
|
||||
let (http_healthy, grpc_healthy) = tokio::join!(
|
||||
self.check_http_health(container_ip),
|
||||
self.check_grpc_connectivity(container_ip)
|
||||
);
|
||||
|
||||
let is_fully_healthy = http_healthy && grpc_healthy;
|
||||
|
||||
// 更新连续失败次数
|
||||
let consecutive_failures = if is_fully_healthy {
|
||||
0 // 完全健康,重置计数
|
||||
} else {
|
||||
previous_failures + 1
|
||||
};
|
||||
|
||||
if !is_fully_healthy {
|
||||
warn!(
|
||||
"Service health check: IP={}, HTTP={}, gRPC={}, consecutive_failures={}",
|
||||
container_ip, http_healthy, grpc_healthy, consecutive_failures
|
||||
);
|
||||
}
|
||||
|
||||
ServiceHealthStatus {
|
||||
http_healthy,
|
||||
grpc_healthy,
|
||||
last_check_time: Utc::now(),
|
||||
consecutive_failures,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ServiceHealthChecker {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// 便捷函数:使用默认配置执行服务健康检查
|
||||
pub async fn check_service_health(container_ip: &str) -> ServiceHealthStatus {
|
||||
ServiceHealthChecker::new()
|
||||
.check_service(container_ip, 0)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_service_health_status_new() {
|
||||
let status = ServiceHealthStatus::new();
|
||||
assert!(!status.http_healthy);
|
||||
assert!(!status.grpc_healthy);
|
||||
assert_eq!(status.consecutive_failures, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_fully_healthy() {
|
||||
let mut status = ServiceHealthStatus::new();
|
||||
assert!(!status.is_fully_healthy());
|
||||
|
||||
status.http_healthy = true;
|
||||
assert!(!status.is_fully_healthy());
|
||||
|
||||
status.grpc_healthy = true;
|
||||
assert!(status.is_fully_healthy());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_partially_healthy() {
|
||||
let mut status = ServiceHealthStatus::new();
|
||||
assert!(!status.is_partially_healthy());
|
||||
|
||||
status.http_healthy = true;
|
||||
assert!(status.is_partially_healthy());
|
||||
|
||||
status.http_healthy = false;
|
||||
status.grpc_healthy = true;
|
||||
assert!(status.is_partially_healthy());
|
||||
}
|
||||
}
|
||||
155
qiming-rcoder/crates/docker_manager/src/image_selector.rs
Normal file
155
qiming-rcoder/crates/docker_manager/src/image_selector.rs
Normal file
@@ -0,0 +1,155 @@
|
||||
//! 镜像选择器
|
||||
//!
|
||||
//! 根据服务类型选择合适的 Docker 镜像。
|
||||
//! 简化版本:针对只有2种镜像的静态映射场景进行了优化。
|
||||
|
||||
use crate::utils::DockerUtils;
|
||||
use crate::{DockerError, DockerResult};
|
||||
use shared_types::{MultiImageConfig, ProjectImageOverrides, ServiceType};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
/// 简化的镜像选择器
|
||||
///
|
||||
/// 根据服务类型选择合适的 Docker 镜像。
|
||||
/// 针对2种镜像的静态映射场景进行了优化,移除了不必要的缓存。
|
||||
/// 强制要求明确指定服务类型,不支持默认值。
|
||||
pub struct ImageSelector {
|
||||
/// 多镜像配置
|
||||
config: MultiImageConfig,
|
||||
/// 当前平台
|
||||
platform: String,
|
||||
}
|
||||
|
||||
impl ImageSelector {
|
||||
/// 创建新的镜像选择器
|
||||
pub fn new(config: MultiImageConfig) -> Self {
|
||||
let platform = DockerUtils::get_optimal_platform();
|
||||
debug!("created selector for: {}", platform);
|
||||
|
||||
Self { config, platform }
|
||||
}
|
||||
|
||||
/// 根据服务类型和项目配置选择镜像
|
||||
///
|
||||
/// 注意:service_type 不能为空,必须明确指定。
|
||||
/// 会自动验证服务是否已启用。
|
||||
/// 简化版本:直接计算镜像名称,无缓存
|
||||
pub async fn select_image(
|
||||
&self,
|
||||
service_type: &ServiceType,
|
||||
project_overrides: Option<&ProjectImageOverrides>,
|
||||
) -> DockerResult<String> {
|
||||
// 强制验证:service_type 必须明确指定并启用
|
||||
if !self.is_service_enabled(service_type) {
|
||||
return Err(DockerError::ConfigurationError(format!(
|
||||
"service type '{}' is not enabled or configuration does not exist",
|
||||
service_type
|
||||
)));
|
||||
}
|
||||
|
||||
// 直接计算镜像名称,无需缓存
|
||||
let image_name = self
|
||||
.select_service_image(service_type, project_overrides)
|
||||
.await?;
|
||||
|
||||
info!(
|
||||
"Selected image: {} (service: {}, platform: {})",
|
||||
image_name, service_type, self.platform
|
||||
);
|
||||
|
||||
Ok(image_name)
|
||||
}
|
||||
|
||||
/// 获取服务配置
|
||||
pub async fn get_service_config(
|
||||
&self,
|
||||
service_type: &ServiceType,
|
||||
) -> DockerResult<shared_types::ServiceImageConfig> {
|
||||
// 强制验证:service_type 必须明确指定并启用
|
||||
if !self.is_service_enabled(service_type) {
|
||||
return Err(DockerError::ConfigurationError(format!(
|
||||
"service type '{}' is not enabled or configuration does not exist",
|
||||
service_type
|
||||
)));
|
||||
}
|
||||
|
||||
// 从配置中获取服务配置
|
||||
let service_key = service_type.to_string();
|
||||
match self.config.services.get(&service_key) {
|
||||
Some(service_config) => {
|
||||
info!("Get config succeeded: {}", service_key);
|
||||
Ok(service_config.clone())
|
||||
}
|
||||
None => Err(DockerError::ConfigurationError(format!(
|
||||
"configuration for service type '{}' does not exist",
|
||||
service_type
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查服务是否已启用和配置
|
||||
pub fn is_service_enabled(&self, service_type: &ServiceType) -> bool {
|
||||
let service_key = service_type.to_string();
|
||||
info!(
|
||||
"[IMAGE_SELECTOR] Checking if service is enabled: service_type={:?}, service_key={}",
|
||||
service_type, service_key
|
||||
);
|
||||
|
||||
if let Some(service_config) = self.config.services.get(&service_key) {
|
||||
info!(
|
||||
"[IMAGE_SELECTOR] Service found: enabled={}, arm64_image={:?}",
|
||||
service_config.enabled, service_config.arm64_image
|
||||
);
|
||||
service_config.enabled
|
||||
} else {
|
||||
warn!(
|
||||
"[IMAGE_SELECTOR] Service type '{}' not found in config, available services: {:?}",
|
||||
service_key,
|
||||
self.config.services.keys().collect::<Vec<_>>()
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// 从服务特定配置选择镜像
|
||||
/// 简化版本:针对2种镜像的静态映射
|
||||
async fn select_service_image(
|
||||
&self,
|
||||
service_type: &ServiceType,
|
||||
_project_overrides: Option<&ProjectImageOverrides>,
|
||||
) -> DockerResult<String> {
|
||||
let service_key = service_type.to_string();
|
||||
|
||||
// 1. 优先使用服务特定配置
|
||||
if let Some(service_config) = self.config.services.get(&service_key) {
|
||||
// 服务级通用镜像(最高优先级)
|
||||
if let Some(image) = &service_config.image {
|
||||
debug!(" using image: {}", image);
|
||||
return Ok(image.clone());
|
||||
}
|
||||
|
||||
// 平台特定镜像
|
||||
if self.platform == "linux/arm64" {
|
||||
if let Some(arm64_image) = &service_config.arm64_image {
|
||||
debug!(" using ARM64 image: {}", arm64_image);
|
||||
return Ok(arm64_image.clone());
|
||||
}
|
||||
} else if let Some(amd64_image) = &service_config.amd64_image {
|
||||
debug!(" using AMD64 image: {}", amd64_image);
|
||||
return Ok(amd64_image.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 使用全局默认配置
|
||||
if let Some(default_image) = &self.config.global_defaults.default_image {
|
||||
debug!(" using default image: {}", default_image);
|
||||
return Ok(default_image.clone());
|
||||
}
|
||||
|
||||
// 3. 配置错误:不应该发生,因为默认配置已经设置了镜像
|
||||
Err(DockerError::ConfigurationError(format!(
|
||||
"Service type '{}' has no available image config, please check the configuration file",
|
||||
service_key
|
||||
)))
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user