byy/qiming
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

添加qiming-rcoder模块

Browse Source
This commit is contained in:
Codex
2026-06-01 13:54:52 +08:00
parent 8092c4b1f8
commit 4b1a580132
539 changed files with 151650 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

903
qiming-rcoder/crates/agent_runner/tests/concurrency_raii_test.rs Normal file
View 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(&registry, "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(&registry, "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(&registry, "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(&registry, &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(&registry_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(&registry, project_id);
// 第二个 guard 会更新现有项目为 Pending（已经是 Pending，无操作）
let _guard2 = PendingGuard::new(&registry, 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(&registry_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(&registry, 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(&registry, 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(&registry, 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"));
}

300
qiming-rcoder/crates/agent_runner/tests/timeout_protection_test.rs Normal file
View File

@@ -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
);
}