From 40c4e82a941f23b17c91b8e7ffac6ef24ddf29f4 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Wed, 1 Jul 2026 21:39:00 +0800 Subject: [PATCH 1/5] fix(active-lease): protect foreground agents from idle cleanup Add a shared active lease registry and renew endpoints for conversations and teams. Wire the registry into app services and skip ACP idle cleanup while a foreground lease is active. Cover the registry, task cleanup selection, service renewal paths, and app HTTP renewal behavior. --- crates/aionui-ai-agent/src/active_lease.rs | 105 +++++++ crates/aionui-ai-agent/src/idle_scanner.rs | 4 +- crates/aionui-ai-agent/src/lib.rs | 2 + crates/aionui-ai-agent/src/task_manager.rs | 90 ++++-- crates/aionui-app/src/router/state.rs | 6 +- crates/aionui-app/src/services.rs | 12 +- crates/aionui-app/tests/active_lease_e2e.rs | 274 ++++++++++++++++++ crates/aionui-conversation/src/routes.rs | 14 + crates/aionui-conversation/src/service.rs | 44 ++- crates/aionui-conversation/src/state.rs | 3 +- .../aionui-conversation/tests/active_lease.rs | 150 ++++++++++ crates/aionui-team/src/routes.rs | 15 + crates/aionui-team/src/service.rs | 48 ++- .../tests/session_service_integration.rs | 92 +++++- 14 files changed, 830 insertions(+), 29 deletions(-) create mode 100644 crates/aionui-ai-agent/src/active_lease.rs create mode 100644 crates/aionui-app/tests/active_lease_e2e.rs create mode 100644 crates/aionui-conversation/tests/active_lease.rs diff --git a/crates/aionui-ai-agent/src/active_lease.rs b/crates/aionui-ai-agent/src/active_lease.rs new file mode 100644 index 000000000..d089089a4 --- /dev/null +++ b/crates/aionui-ai-agent/src/active_lease.rs @@ -0,0 +1,105 @@ +use aionui_common::{TimestampMs, now_ms}; +use dashmap::DashMap; + +pub const ACTIVE_LEASE_TTL_MS: TimestampMs = 90_000; + +#[derive(Debug)] +pub struct ActiveLeaseRegistry { + leases: DashMap, + ttl_ms: TimestampMs, +} + +impl ActiveLeaseRegistry { + pub fn new() -> Self { + Self::with_ttl_ms(ACTIVE_LEASE_TTL_MS) + } + + pub fn with_ttl_ms(ttl_ms: TimestampMs) -> Self { + Self { + leases: DashMap::new(), + ttl_ms, + } + } + + pub fn renew(&self, conversation_id: &str) -> TimestampMs { + let expires_at = now_ms().saturating_add(self.ttl_ms); + self.leases.insert(conversation_id.to_owned(), expires_at); + expires_at + } + + pub fn renew_many<'a>(&self, conversation_ids: impl IntoIterator) -> (usize, TimestampMs) { + let expires_at = now_ms().saturating_add(self.ttl_ms); + let mut count = 0; + for conversation_id in conversation_ids { + self.leases.insert(conversation_id.to_owned(), expires_at); + count += 1; + } + (count, expires_at) + } + + pub fn active_until(&self, conversation_id: &str) -> Option { + let expires_at = *self.leases.get(conversation_id)?; + if expires_at > now_ms() { + Some(expires_at) + } else { + self.leases.remove(conversation_id); + None + } + } + + pub fn is_active(&self, conversation_id: &str) -> bool { + self.active_until(conversation_id).is_some() + } +} + +impl Default for ActiveLeaseRegistry { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + #[test] + fn renew_sets_active_lease() { + let registry = ActiveLeaseRegistry::new(); + + let expires_at = registry.renew("conv-1"); + + assert_eq!(registry.active_until("conv-1"), Some(expires_at)); + assert!(registry.is_active("conv-1")); + } + + #[test] + fn renew_many_sets_all_leases_and_returns_count() { + let registry = ActiveLeaseRegistry::new(); + + let (count, expires_at) = registry.renew_many(["conv-1", "conv-2"]); + + assert_eq!(count, 2); + assert_eq!(registry.active_until("conv-1"), Some(expires_at)); + assert_eq!(registry.active_until("conv-2"), Some(expires_at)); + } + + #[test] + fn active_lookup_returns_none_for_missing_lease() { + let registry = ActiveLeaseRegistry::new(); + + assert_eq!(registry.active_until("missing"), None); + assert!(!registry.is_active("missing")); + } + + #[test] + fn active_lookup_lazily_removes_expired_lease() { + let registry = ActiveLeaseRegistry::with_ttl_ms(1); + registry.renew("conv-1"); + + std::thread::sleep(Duration::from_millis(5)); + + assert_eq!(registry.active_until("conv-1"), None); + assert!(!registry.leases.contains_key("conv-1")); + } +} diff --git a/crates/aionui-ai-agent/src/idle_scanner.rs b/crates/aionui-ai-agent/src/idle_scanner.rs index aba5b0bb1..2ba0c7e75 100644 --- a/crates/aionui-ai-agent/src/idle_scanner.rs +++ b/crates/aionui-ai-agent/src/idle_scanner.rs @@ -15,7 +15,7 @@ const SCAN_INTERVAL_SECS: u64 = 60; /// Start the background idle agent scanner. /// /// Periodically scans active tasks and kills ACP agents that have been -/// idle (finished + no activity) beyond the configured threshold. +/// idle (finished or warmup-only + no activity) beyond the configured threshold. /// /// The scanner runs until the provided `shutdown` signal resolves. pub fn start_idle_scanner( @@ -59,7 +59,7 @@ fn scan_and_cleanup(manager: &Arc, threshold_ms: i64) { let idle_ids = manager.collect_idle(threshold_ms); if idle_ids.is_empty() { - debug!(active = manager.active_count(), "Idle scan: no idle agents found"); + debug!(active_count = manager.active_count(), "Idle scan: no idle agents found"); return; } diff --git a/crates/aionui-ai-agent/src/lib.rs b/crates/aionui-ai-agent/src/lib.rs index 4b0d24cab..d949c8538 100644 --- a/crates/aionui-ai-agent/src/lib.rs +++ b/crates/aionui-ai-agent/src/lib.rs @@ -1,6 +1,7 @@ #![warn(clippy::disallowed_types)] //! AI agent lifecycle, worker task dispatch, and skill management. +pub mod active_lease; pub(crate) mod agent_runtime; pub mod agent_task; pub mod capability; @@ -21,6 +22,7 @@ pub mod shared_kernel; pub mod task_manager; pub mod types; +pub use active_lease::{ACTIVE_LEASE_TTL_MS, ActiveLeaseRegistry}; pub use agent_runtime::AgentRuntime; #[cfg(any(test, feature = "test-support"))] pub use agent_task::IMockAgent; diff --git a/crates/aionui-ai-agent/src/task_manager.rs b/crates/aionui-ai-agent/src/task_manager.rs index c99dfdec4..8063f1ef1 100644 --- a/crates/aionui-ai-agent/src/task_manager.rs +++ b/crates/aionui-ai-agent/src/task_manager.rs @@ -7,8 +7,9 @@ use async_trait::async_trait; use dashmap::DashMap; use futures_util::future::{BoxFuture, join_all}; use tokio::sync::OnceCell; -use tracing::{info, warn}; +use tracing::{debug, info, warn}; +use crate::active_lease::ActiveLeaseRegistry; use crate::agent_task::AgentInstance; use crate::error::AgentError; use crate::types::BuildTaskOptions; @@ -63,8 +64,10 @@ pub trait IWorkerTaskManager: Send + Sync { /// Collect tasks eligible for idle cleanup. /// /// Returns conversation IDs of tasks that: - /// - have `status == Some(Finished)` + /// - are ACP agents + /// - have `status == None` or `status == Some(Finished)` /// - have been idle longer than `idle_threshold_ms` + /// - are not protected by an active foreground lease fn collect_idle(&self, idle_threshold_ms: TimestampMs) -> Vec; } @@ -78,13 +81,19 @@ type TaskSlot = Arc>; pub struct WorkerTaskManagerImpl { tasks: DashMap, factory: AgentFactory, + active_leases: Arc, } impl WorkerTaskManagerImpl { pub fn new(factory: AgentFactory) -> Self { + Self::new_with_active_leases(factory, Arc::new(ActiveLeaseRegistry::new())) + } + + pub fn new_with_active_leases(factory: AgentFactory, active_leases: Arc) -> Self { Self { tasks: DashMap::new(), factory, + active_leases, } } @@ -197,23 +206,37 @@ impl IWorkerTaskManager for WorkerTaskManagerImpl { let last_activity_at = agent.last_activity_at(); let idle_ms = now.saturating_sub(last_activity_at); - let selected = agent_type == AgentType::Acp - && status == Some(ConversationStatus::Finished) - && idle_ms > idle_threshold_ms; - if selected { - info!( + if agent_type != AgentType::Acp + || !matches!(status, None | Some(ConversationStatus::Finished)) + || idle_ms <= idle_threshold_ms + { + return None; + } + + if let Some(expires_at) = self.active_leases.active_until(entry.key()) { + debug!( conversation_id = %entry.key(), - ?agent_type, ?status, idle_ms, - threshold_ms = idle_threshold_ms, - last_activity_at, - "Idle scan: selected idle agent" + lease_expires_in_ms = expires_at.saturating_sub(now), + reason = %"ActiveLease", + "Idle scan: active lease protects idle agent" ); - Some(entry.key().clone()) - } else { - None + return None; } + + let idle_class = if status.is_none() { "WarmupOnly" } else { "Finished" }; + info!( + conversation_id = %entry.key(), + ?agent_type, + ?status, + idle_ms, + threshold_ms = idle_threshold_ms, + idle_class = %idle_class, + reason = %"IdleTimeout", + "Idle scan: selected idle agent" + ); + Some(entry.key().clone()) }) .collect() } @@ -537,7 +560,7 @@ mod tests { } #[test] - fn collect_idle_finds_finished_and_stale_acp_tasks() { + fn collect_idle_finds_finished_and_warmup_only_stale_acp_tasks() { let factory: AgentFactory = Arc::new(|_| async { unreachable!() }.boxed()); let mgr = WorkerTaskManagerImpl::new(factory); @@ -556,6 +579,12 @@ mod tests { ), ); + // ACP + warmup-only + old activity → should be collected + insert( + "conv-warmup-only", + mock_instance(MockAgent::new("conv-warmup-only", None).with_last_activity(now_ms() - 600_000)), + ); + // ACP + Finished + recent activity → should NOT be collected insert( "conv-recent", @@ -584,8 +613,32 @@ mod tests { ); let idle = mgr.collect_idle(300_000); // 5-min threshold - assert_eq!(idle.len(), 1); - assert_eq!(idle[0], "conv-stale"); + assert_eq!(idle.len(), 2); + assert!(idle.contains(&"conv-stale".to_owned())); + assert!(idle.contains(&"conv-warmup-only".to_owned())); + } + + #[test] + fn collect_idle_skips_tasks_with_active_lease() { + let active_leases = Arc::new(crate::ActiveLeaseRegistry::new()); + active_leases.renew("conv-active"); + let factory: AgentFactory = Arc::new(|_| async { unreachable!() }.boxed()); + let mgr = WorkerTaskManagerImpl::new_with_active_leases(factory, active_leases); + + let cell: OnceCell = OnceCell::new(); + cell.set(mock_instance( + MockAgent::new("conv-active", Some(ConversationStatus::Finished)).with_last_activity(now_ms() - 600_000), + )) + .ok(); + mgr.tasks.insert("conv-active".into(), Arc::new(cell)); + + let captured = capture_logs(tracing::Level::DEBUG, || { + let idle = mgr.collect_idle(300_000); + assert!(idle.is_empty()); + }); + + assert!(captured.contains("reason=ActiveLease")); + assert!(captured.contains("lease_expires_in_ms=")); } #[test] @@ -611,7 +664,8 @@ mod tests { assert!(captured.contains("status=Some(Finished)")); assert!(captured.contains("idle_ms=")); assert!(captured.contains("threshold_ms=5000")); - assert!(captured.contains("last_activity_at=")); + assert!(captured.contains("idle_class=Finished")); + assert!(captured.contains("reason=IdleTimeout")); } #[test] diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index 468ee06bd..d18224237 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -374,6 +374,7 @@ pub fn build_conversation_state( ConversationRouterState { service: conversation_service, task_manager: services.worker_task_manager.clone(), + active_leases: services.active_lease_registry.clone(), } } @@ -594,7 +595,10 @@ pub fn build_team_state( cancellation_port, backend_binary_path, ); - TeamRouterState { service } + TeamRouterState { + service, + active_leases: services.active_lease_registry.clone(), + } } /// Build the default `CronRouterState` from application services. diff --git a/crates/aionui-app/src/services.rs b/crates/aionui-app/src/services.rs index 9a8856714..71a9c8889 100644 --- a/crates/aionui-app/src/services.rs +++ b/crates/aionui-app/src/services.rs @@ -5,8 +5,8 @@ use std::sync::Arc; use crate::config::{AppConfig, derive_encryption_key}; use aionui_ai_agent::{ - AcpSessionSyncService, AcpSkillManager, AgentFactoryDeps, AgentRegistry, IWorkerTaskManager, WorkerTaskManagerImpl, - build_agent_factory, + AcpSessionSyncService, AcpSkillManager, ActiveLeaseRegistry, AgentFactoryDeps, AgentRegistry, IWorkerTaskManager, + WorkerTaskManagerImpl, build_agent_factory, }; use aionui_auth::{CookieConfig, JwtService, QrTokenStore, resolve_jwt_secret}; use aionui_common::OnConversationDelete; @@ -29,6 +29,7 @@ pub struct AppServices { pub ws_manager: Arc, pub event_bus: Arc, pub worker_task_manager: Arc, + pub active_lease_registry: Arc, pub conversation_runtime_state: Arc, pub conversation_service: ConversationService, /// Same instance as `worker_task_manager`, exposed through the @@ -164,7 +165,11 @@ impl AppServices { // Agent factory is now wired. Future extension/custom agents // that get written to `agent_metadata` will show up after the // relevant service calls `AgentRegistry::hydrate`. - let task_manager_concrete = Arc::new(WorkerTaskManagerImpl::new(factory)); + let active_lease_registry = Arc::new(ActiveLeaseRegistry::new()); + let task_manager_concrete = Arc::new(WorkerTaskManagerImpl::new_with_active_leases( + factory, + active_lease_registry.clone(), + )); let worker_task_manager: Arc = task_manager_concrete.clone(); let task_manager_delete_hook: Arc = task_manager_concrete; let conversation_runtime_state = Arc::new(ConversationRuntimeStateService::default()); @@ -189,6 +194,7 @@ impl AppServices { ws_manager: Arc::new(WebSocketManager::new()), event_bus, worker_task_manager, + active_lease_registry, conversation_runtime_state, conversation_service, task_manager_delete_hook: Some(task_manager_delete_hook), diff --git a/crates/aionui-app/tests/active_lease_e2e.rs b/crates/aionui-app/tests/active_lease_e2e.rs new file mode 100644 index 000000000..d789a7ab3 --- /dev/null +++ b/crates/aionui-app/tests/active_lease_e2e.rs @@ -0,0 +1,274 @@ +mod common; + +use aionui_common::now_ms; +use aionui_db::models::TeamRow; +use aionui_db::{ITeamRepository, SqliteTeamRepository}; +use aionui_team::{TeamAgent, TeammateRole}; +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use serde_json::json; +use tower::ServiceExt; + +use common::{body_json, build_app, json_with_token, setup_and_login}; + +fn empty_post_with_token(uri: &str, token: &str, csrf: &str) -> Request { + Request::builder() + .method("POST") + .uri(uri) + .header("authorization", format!("Bearer {token}")) + .header("x-csrf-token", csrf) + .header("cookie", format!("aionui-csrf-token={csrf}")) + .body(Body::empty()) + .unwrap() +} + +fn empty_post_with_auth_without_csrf(uri: &str, token: &str) -> Request { + Request::builder() + .method("POST") + .uri(uri) + .header("authorization", format!("Bearer {token}")) + .body(Body::empty()) + .unwrap() +} + +fn empty_post_without_auth(uri: &str, csrf: &str) -> Request { + Request::builder() + .method("POST") + .uri(uri) + .header("x-csrf-token", csrf) + .header("cookie", format!("aionui-csrf-token={csrf}")) + .body(Body::empty()) + .unwrap() +} + +async fn create_conversation(app: &mut axum::Router, token: &str, csrf: &str) -> String { + let req = json_with_token( + "POST", + "/api/conversations", + json!({ + "type": "acp", + "name": "Lease Conversation", + "extra": {} + }), + token, + csrf, + ); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::CREATED); + let body = body_json(resp).await; + body["data"]["id"].as_str().unwrap().to_owned() +} + +fn team_agent(slot_id: &str, name: &str, role: TeammateRole, conversation_id: &str) -> TeamAgent { + TeamAgent { + slot_id: slot_id.to_owned(), + name: name.to_owned(), + role, + conversation_id: conversation_id.to_owned(), + backend: "acp".to_owned(), + model: "claude".to_owned(), + assistant_id: None, + status: None, + conversation_type: None, + cli_path: None, + } +} + +async fn insert_team(services: &aionui_app::AppServices, user_id: &str, team_id: &str, agents: Vec) { + let repo = SqliteTeamRepository::new(services.database.pool().clone()); + repo.create_team(&TeamRow { + id: team_id.to_owned(), + user_id: user_id.to_owned(), + name: "Lease Team".to_owned(), + workspace: String::new(), + workspace_mode: "shared".to_owned(), + agents: serde_json::to_string(&agents).unwrap(), + lead_agent_id: agents.first().map(|agent| agent.slot_id.clone()), + session_mode: None, + agents_version: "1.0.1".to_owned(), + created_at: now_ms(), + updated_at: now_ms(), + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn conversation_active_lease_renews_owned_conversation() { + let (mut app, services) = build_app().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let conversation_id = create_conversation(&mut app, &token, &csrf).await; + + let req = empty_post_with_token( + &format!("/api/conversations/{conversation_id}/active-lease"), + &token, + &csrf, + ); + let resp = app.oneshot(req).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::OK); + assert!(services.active_lease_registry.active_until(&conversation_id).is_some()); +} + +#[tokio::test] +async fn conversation_active_lease_rejects_missing_auth() { + let (mut app, services) = build_app().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let conversation_id = create_conversation(&mut app, &token, &csrf).await; + + let resp = app + .oneshot(empty_post_without_auth( + &format!("/api/conversations/{conversation_id}/active-lease"), + &csrf, + )) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let body = body_json(resp).await; + assert_eq!(body["code"], "UNAUTHORIZED"); +} + +#[tokio::test] +async fn conversation_active_lease_rejects_missing_csrf() { + let (mut app, services) = build_app().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let conversation_id = create_conversation(&mut app, &token, &csrf).await; + + let resp = app + .oneshot(empty_post_with_auth_without_csrf( + &format!("/api/conversations/{conversation_id}/active-lease"), + &token, + )) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + let body = body_json(resp).await; + assert_eq!(body["code"], "CSRF_INVALID"); +} + +#[tokio::test] +async fn conversation_active_lease_rejects_cross_user_without_renewing() { + let (mut app, services) = build_app().await; + let (owner_token, owner_csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let conversation_id = create_conversation(&mut app, &owner_token, &owner_csrf).await; + let (other_token, other_csrf) = setup_and_login(&mut app, &services, "other", "StrongP@ss2").await; + + let req = empty_post_with_token( + &format!("/api/conversations/{conversation_id}/active-lease"), + &other_token, + &other_csrf, + ); + let resp = app.oneshot(req).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + let body = body_json(resp).await; + assert_eq!(body["code"], "NOT_FOUND"); + assert!(services.active_lease_registry.active_until(&conversation_id).is_none()); +} + +#[tokio::test] +async fn team_active_lease_renews_member_conversations() { + let (mut app, services) = build_app().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let owner = services.user_repo.find_by_username("admin").await.unwrap().unwrap(); + insert_team( + &services, + &owner.id, + "team-lease", + vec![ + team_agent("lead", "Lead", TeammateRole::Lead, "team-conv-lead"), + team_agent("worker", "Worker", TeammateRole::Teammate, "team-conv-worker"), + ], + ) + .await; + + let req = empty_post_with_token("/api/teams/team-lease/active-lease", &token, &csrf); + let resp = app.oneshot(req).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::OK); + assert!(services.active_lease_registry.active_until("team-conv-lead").is_some()); + assert!( + services + .active_lease_registry + .active_until("team-conv-worker") + .is_some() + ); +} + +#[tokio::test] +async fn team_active_lease_allows_empty_agents() { + let (mut app, services) = build_app().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let owner = services.user_repo.find_by_username("admin").await.unwrap().unwrap(); + insert_team(&services, &owner.id, "team-empty", vec![]).await; + + let req = empty_post_with_token("/api/teams/team-empty/active-lease", &token, &csrf); + let resp = app.oneshot(req).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::OK); + assert!(services.active_lease_registry.active_until("team-empty").is_none()); + assert!(services.active_lease_registry.active_until("unrelated").is_none()); +} + +#[tokio::test] +async fn team_active_lease_rejects_missing_auth() { + let (mut app, services) = build_app().await; + let (token, _csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let owner = services.user_repo.find_by_username("admin").await.unwrap().unwrap(); + insert_team(&services, &owner.id, "team-auth", vec![]).await; + + let resp = app + .oneshot(empty_post_without_auth("/api/teams/team-auth/active-lease", &_csrf)) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let body = body_json(resp).await; + assert_eq!(body["code"], "UNAUTHORIZED"); + drop(token); +} + +#[tokio::test] +async fn team_active_lease_rejects_missing_csrf() { + let (mut app, services) = build_app().await; + let (token, _csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let owner = services.user_repo.find_by_username("admin").await.unwrap().unwrap(); + insert_team(&services, &owner.id, "team-csrf", vec![]).await; + + let resp = app + .oneshot(empty_post_with_auth_without_csrf( + "/api/teams/team-csrf/active-lease", + &token, + )) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + let body = body_json(resp).await; + assert_eq!(body["code"], "CSRF_INVALID"); +} + +#[tokio::test] +async fn team_active_lease_rejects_cross_user_without_renewing() { + let (mut app, services) = build_app().await; + let (_owner_token, _owner_csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let owner = services.user_repo.find_by_username("admin").await.unwrap().unwrap(); + insert_team( + &services, + &owner.id, + "team-cross-user", + vec![team_agent("lead", "Lead", TeammateRole::Lead, "team-cross-conv")], + ) + .await; + let (other_token, other_csrf) = setup_and_login(&mut app, &services, "other", "StrongP@ss2").await; + + let req = empty_post_with_token("/api/teams/team-cross-user/active-lease", &other_token, &other_csrf); + let resp = app.oneshot(req).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + let body = body_json(resp).await; + assert_eq!(body["code"], "FORBIDDEN"); + assert!(services.active_lease_registry.active_until("team-cross-conv").is_none()); +} diff --git a/crates/aionui-conversation/src/routes.rs b/crates/aionui-conversation/src/routes.rs index 45b38b32d..0f99388e5 100644 --- a/crates/aionui-conversation/src/routes.rs +++ b/crates/aionui-conversation/src/routes.rs @@ -105,6 +105,7 @@ pub fn conversation_routes(state: ConversationRouterState) -> Router { .route("/api/conversations/{id}/artifacts/{artifactId}", patch(update_artifact)) .route("/api/conversations/{id}/cancel", post(cancel)) .route("/api/conversations/{id}/warmup", post(warmup)) + .route("/api/conversations/{id}/active-lease", post(active_lease)) // Confirmation system .route("/api/conversations/{id}/confirmations", get(list_confirmations)) .route("/api/conversations/{id}/confirmations/{callId}/confirm", post(confirm)) @@ -317,6 +318,19 @@ async fn warmup( Ok(Json(ApiResponse::success())) } +async fn active_lease( + State(state): State, + Extension(user): Extension, + Path(id): Path, +) -> Result>, ApiError> { + state + .service + .renew_active_lease(&user.id, &id, &state.active_leases) + .await + .map_err(ApiError::from)?; + Ok(Json(ApiResponse::success())) +} + async fn search_messages( State(state): State, Extension(user): Extension, diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index 3ac0c0fe2..df8df4a60 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -5,7 +5,9 @@ use std::sync::Arc; use aionui_ai_agent::session_context::{AgentSessionContext, AgentSessionKind}; use aionui_ai_agent::types::BuildTaskOptions; -use aionui_ai_agent::{AgentAvailabilityFeedbackPort, AgentError, AgentInstance, AgentSendError, IWorkerTaskManager}; +use aionui_ai_agent::{ + ActiveLeaseRegistry, AgentAvailabilityFeedbackPort, AgentError, AgentInstance, AgentSendError, IWorkerTaskManager, +}; use crate::message_cursor::{decode_message_cursor, encode_message_cursor}; use crate::response_middleware::ICronService; @@ -2861,6 +2863,46 @@ impl ConversationService { }) } + pub async fn renew_active_lease( + &self, + user_id: &str, + conversation_id: &str, + active_leases: &ActiveLeaseRegistry, + ) -> Result<(), ConversationError> { + let row = match self.conversation_repo.get(conversation_id).await { + Ok(row) => row, + Err(error) => { + warn!( + kind = "conversation", + conversation_id, + user_id, + error = %ErrorChain(&error), + "Conversation active lease renew failed" + ); + return Err(error.into()); + } + }; + + let Some(row) = row.filter(|row| row.user_id == user_id) else { + debug!( + kind = "conversation", + conversation_id, user_id, "Conversation active lease renew rejected" + ); + return Err(ConversationError::NotFound { + id: conversation_id.to_owned(), + }); + }; + + let expires_at = active_leases.renew(&row.id); + debug!( + kind = "conversation", + conversation_id = %row.id, + expires_at, + "Conversation active lease renewed" + ); + Ok(()) + } + /// Pre-initialize an agent task for a conversation (warmup). /// /// This builds the agent task without sending a message, so the diff --git a/crates/aionui-conversation/src/state.rs b/crates/aionui-conversation/src/state.rs index 50cbd5cd1..a9ae4f56c 100644 --- a/crates/aionui-conversation/src/state.rs +++ b/crates/aionui-conversation/src/state.rs @@ -1,11 +1,12 @@ use std::sync::Arc; use crate::service::ConversationService; -use aionui_ai_agent::IWorkerTaskManager; +use aionui_ai_agent::{ActiveLeaseRegistry, IWorkerTaskManager}; /// Shared state for conversation route handlers. #[derive(Clone)] pub struct ConversationRouterState { pub service: ConversationService, pub task_manager: Arc, + pub active_leases: Arc, } diff --git a/crates/aionui-conversation/tests/active_lease.rs b/crates/aionui-conversation/tests/active_lease.rs new file mode 100644 index 000000000..6ec12f5d4 --- /dev/null +++ b/crates/aionui-conversation/tests/active_lease.rs @@ -0,0 +1,150 @@ +use std::sync::Arc; + +use aionui_ai_agent::{ActiveLeaseRegistry, AgentError, IWorkerTaskManager}; +use aionui_api_types::{CreateConversationRequest, WebSocketMessage}; +use aionui_common::{AgentKillReason, TimestampMs}; +use aionui_conversation::skill_resolver::SkillResolver; +use aionui_conversation::{ConversationError, ConversationService}; +use aionui_db::{SqliteConversationRepository, init_database_memory}; +use aionui_realtime::EventBroadcaster; +use serde_json::json; + +struct NoopBroadcaster; + +impl EventBroadcaster for NoopBroadcaster { + fn broadcast(&self, _event: WebSocketMessage) {} +} + +struct NoopTaskManager; + +#[async_trait::async_trait] +impl IWorkerTaskManager for NoopTaskManager { + fn get_task(&self, _: &str) -> Option { + None + } + + async fn get_or_build_task( + &self, + _: &str, + _: aionui_ai_agent::types::BuildTaskOptions, + ) -> Result { + Err(AgentError::internal("noop")) + } + + fn kill(&self, _: &str, _: Option) -> Result<(), AgentError> { + Ok(()) + } + + fn kill_and_wait( + &self, + _: &str, + _: Option, + ) -> std::pin::Pin + Send>> { + Box::pin(std::future::ready(())) + } + + async fn clear(&self) {} + + fn active_count(&self) -> usize { + 0 + } + + fn collect_idle(&self, _: TimestampMs) -> Vec { + vec![] + } +} + +struct EmptySkillResolver; + +#[async_trait::async_trait] +impl SkillResolver for EmptySkillResolver { + async fn auto_inject_names(&self) -> Vec { + Vec::new() + } + + async fn resolve_skills(&self, _names: &[String]) -> Vec { + Vec::new() + } + + async fn link_workspace_skills( + &self, + _workspace: &std::path::Path, + _rel_dirs: &[&str], + _skills: &[aionui_extension::ResolvedAgentSkill], + ) -> usize { + 0 + } +} + +const USER_ID: &str = "system_default_user"; +const OTHER_USER_ID: &str = "other-user"; + +async fn setup() -> ConversationService { + let db = init_database_memory().await.unwrap(); + let repo = Arc::new(SqliteConversationRepository::new(db.pool().clone())); + let agent_metadata_repo: Arc = + Arc::new(aionui_db::SqliteAgentMetadataRepository::new(db.pool().clone())); + let acp_session_repo: Arc = + Arc::new(aionui_db::SqliteAcpSessionRepository::new(db.pool().clone())); + + ConversationService::new( + std::env::temp_dir(), + Arc::new(NoopBroadcaster), + Arc::new(EmptySkillResolver), + Arc::new(NoopTaskManager), + repo, + agent_metadata_repo, + acp_session_repo, + ) +} + +fn make_create_req() -> CreateConversationRequest { + serde_json::from_value(json!({ + "type": "acp", + "extra": { "workspace": std::env::temp_dir().to_string_lossy() } + })) + .unwrap() +} + +#[tokio::test] +async fn renew_active_lease_records_owned_conversation() { + let service = setup().await; + let conversation = service.create(USER_ID, make_create_req()).await.unwrap(); + let active_leases = ActiveLeaseRegistry::new(); + + service + .renew_active_lease(USER_ID, &conversation.id, &active_leases) + .await + .unwrap(); + + assert!(active_leases.is_active(&conversation.id)); +} + +#[tokio::test] +async fn renew_active_lease_rejects_other_users_conversation() { + let service = setup().await; + let conversation = service.create(USER_ID, make_create_req()).await.unwrap(); + let active_leases = ActiveLeaseRegistry::new(); + + let err = service + .renew_active_lease(OTHER_USER_ID, &conversation.id, &active_leases) + .await + .unwrap_err(); + + assert!(matches!(err, ConversationError::NotFound { .. })); + assert!(!active_leases.is_active(&conversation.id)); +} + +#[tokio::test] +async fn renew_active_lease_rejects_missing_conversation() { + let service = setup().await; + let active_leases = ActiveLeaseRegistry::new(); + + let err = service + .renew_active_lease(USER_ID, "missing-conversation", &active_leases) + .await + .unwrap_err(); + + assert!(matches!(err, ConversationError::NotFound { .. })); + assert!(!active_leases.is_active("missing-conversation")); +} diff --git a/crates/aionui-team/src/routes.rs b/crates/aionui-team/src/routes.rs index 48a45fca2..1a906ff2b 100644 --- a/crates/aionui-team/src/routes.rs +++ b/crates/aionui-team/src/routes.rs @@ -8,6 +8,7 @@ use axum::extract::{Extension, Json, Path, State}; use axum::http::StatusCode; use axum::routing::{get, post}; +use aionui_ai_agent::ActiveLeaseRegistry; use aionui_api_types::{ AddAgentRequest, ApiResponse, CancelTeamChildTurnRequest, CancelTeamRunRequest, CreateTeamRequest, PauseTeamSlotRequest, RenameAgentRequest, RenameTeamRequest, SendAgentMessageRequest, SendTeamMessageRequest, @@ -23,6 +24,7 @@ use crate::service::TeamSessionService; #[derive(Clone)] pub struct TeamRouterState { pub service: Arc, + pub active_leases: Arc, } fn db_error_to_api_error(err: DbError) -> ApiError { @@ -87,6 +89,7 @@ pub fn team_routes(state: TeamRouterState) -> Router { post(pause_slot_work), ) .route("/api/teams/{id}/session", post(ensure_session).delete(stop_session)) + .route("/api/teams/{id}/active-lease", post(active_lease)) .route("/api/teams/{id}/session-mode", post(set_session_mode)) .with_state(state) } @@ -290,6 +293,18 @@ async fn set_session_mode( Ok(Json(ApiResponse::success())) } +async fn active_lease( + State(state): State, + Extension(user): Extension, + Path(id): Path, +) -> Result>, ApiError> { + state + .service + .renew_active_lease(&user.id, &id, &state.active_leases) + .await?; + Ok(Json(ApiResponse::success())) +} + async fn ensure_session( State(state): State, Extension(user): Extension, diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 895a37d1d..7ba7b4878 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -5,7 +5,7 @@ pub(crate) mod spawn_support; use std::path::PathBuf; use std::sync::{Arc, Weak}; -use aionui_ai_agent::{AgentError, AgentInstance, IWorkerTaskManager}; +use aionui_ai_agent::{ActiveLeaseRegistry, AgentError, AgentInstance, IWorkerTaskManager}; use aionui_api_types::{ AddAgentRequest, CreateTeamRequest, TeamAgentResponse, TeamMcpPhase, TeamMcpStatusPayload, TeamResponse, TeamRunAckResponse, TeamRunStateResponse, TeamRunTargetRole, WebSocketMessage, @@ -18,7 +18,7 @@ use aionui_db::{ }; use aionui_realtime::EventBroadcaster; use dashmap::DashMap; -use tracing::{info, warn}; +use tracing::{debug, info, warn}; use crate::error::TeamError; use crate::event_loop::AgentLoopContext; @@ -133,6 +133,50 @@ impl TeamSessionService { Ok(Team::from_row(&row)?) } + pub async fn renew_active_lease( + &self, + user_id: &str, + team_id: &str, + active_leases: &ActiveLeaseRegistry, + ) -> Result<(), TeamError> { + let team = match self.load_owned_team(user_id, team_id).await { + Ok(team) => team, + Err(error @ (TeamError::TeamNotFound(_) | TeamError::Forbidden(_))) => { + debug!( + kind = "team", + team_id, + user_id, + error = %error, + "Team active lease renew rejected" + ); + return Err(error); + } + Err(error) => { + warn!( + kind = "team", + team_id, + user_id, + error = %error, + "Team active lease renew failed" + ); + return Err(error); + } + }; + + let conversation_ids = team + .agents + .iter() + .map(|agent| agent.conversation_id.as_str()) + .filter(|conversation_id| !conversation_id.trim().is_empty()); + let (covered_count, expires_at) = active_leases.renew_many(conversation_ids); + + debug!( + kind = "team", + team_id, covered_count, expires_at, "Team active lease renewed" + ); + Ok(()) + } + /// Restore sessions for all existing teams. Called once at app startup /// so that MCP servers are available before any user sends a message. pub async fn restore_all_sessions(&self) { diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index ccbeea468..1a1aa3fee 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -9,7 +9,7 @@ use aionui_ai_agent::session_context::{ }; use aionui_ai_agent::task_manager::AgentFactory; use aionui_ai_agent::types::BuildTaskOptions; -use aionui_ai_agent::{AgentError, IWorkerTaskManager, WorkerTaskManagerImpl}; +use aionui_ai_agent::{ActiveLeaseRegistry, AgentError, IWorkerTaskManager, WorkerTaskManagerImpl}; use aionui_api_types::{AcpBuildExtra, AddAgentRequest, CreateTeamRequest, TeamAgentInput, WebSocketMessage}; use aionui_common::{AgentKillReason, AgentType, PaginatedResult, ProviderWithModel}; use aionui_db::models::{ @@ -1725,6 +1725,96 @@ async fn reset_auto_started_session(svc: &Arc, tm: &Arc, team_id: &str, workspace: &str) { repo.update_team( team_id, From 6ad201d84957e0e3bd41e49aa6986ea9b0bb0fdb Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 2 Jul 2026 11:51:51 +0800 Subject: [PATCH 2/5] Add conversation runtime ensure endpoint --- crates/aionui-api-types/src/conversation.rs | 9 ++ crates/aionui-api-types/src/lib.rs | 4 +- .../tests/acp_config_options_e2e.rs | 95 ++++++++++++++++++- crates/aionui-conversation/src/routes.rs | 20 +++- crates/aionui-conversation/src/service.rs | 55 +++++++++-- .../aionui-conversation/src/service_test.rs | 45 +++++++++ 6 files changed, 215 insertions(+), 13 deletions(-) diff --git a/crates/aionui-api-types/src/conversation.rs b/crates/aionui-api-types/src/conversation.rs index d753a5b5e..93fba083f 100644 --- a/crates/aionui-api-types/src/conversation.rs +++ b/crates/aionui-api-types/src/conversation.rs @@ -4,6 +4,8 @@ use aionui_common::{ }; use serde::{Deserialize, Serialize}; +use crate::acp::AcpConfigOptionDto; + /// Per-MCP snapshot status stored in `conversation.extra`. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "snake_case")] @@ -138,6 +140,13 @@ pub struct ConversationRuntimeSummary { pub turn_id: Option, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct EnsureConversationRuntimeResponse { + pub recovered: bool, + pub config_options: Vec, + pub runtime: ConversationRuntimeSummary, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct ConversationAssistantIdentityResponse { pub id: String, diff --git a/crates/aionui-api-types/src/lib.rs b/crates/aionui-api-types/src/lib.rs index fbcd5b660..1617f1c18 100644 --- a/crates/aionui-api-types/src/lib.rs +++ b/crates/aionui-api-types/src/lib.rs @@ -78,8 +78,8 @@ pub use conversation::{ ConversationArtifactListResponse, ConversationArtifactResponse, ConversationArtifactStatus, ConversationAssistantIdentityResponse, ConversationListResponse, ConversationMcpStatus, ConversationMcpStatusKind, ConversationResponse, ConversationRuntimeStateKind, ConversationRuntimeSummary, CreateConversationRequest, - ListConversationsQuery, ListMessagesQuery, MessageListResponse, MessageResponse, MessageSearchItem, - MessageSearchResponse, SearchMessagesQuery, SendMessageRequest, SendMessageResponse, + EnsureConversationRuntimeResponse, ListConversationsQuery, ListMessagesQuery, MessageListResponse, MessageResponse, + MessageSearchItem, MessageSearchResponse, SearchMessagesQuery, SendMessageRequest, SendMessageResponse, UpdateConversationArtifactRequest, UpdateConversationRequest, }; pub use cron::{ diff --git a/crates/aionui-app/tests/acp_config_options_e2e.rs b/crates/aionui-app/tests/acp_config_options_e2e.rs index 50afa24e5..2e39b9351 100644 --- a/crates/aionui-app/tests/acp_config_options_e2e.rs +++ b/crates/aionui-app/tests/acp_config_options_e2e.rs @@ -15,13 +15,16 @@ fn create_body() -> serde_json::Value { }) } -async fn create_and_warmup_conversation(app: &mut axum::Router, token: &str, csrf: &str) -> String { +async fn create_conversation(app: &mut axum::Router, token: &str, csrf: &str) -> String { let req = json_with_token("POST", "/api/conversations", create_body(), token, csrf); let resp = app.clone().oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::CREATED); let json = body_json(resp).await; - let id = json["data"]["id"].as_str().unwrap().to_owned(); + json["data"]["id"].as_str().unwrap().to_owned() +} +async fn create_and_warmup_conversation(app: &mut axum::Router, token: &str, csrf: &str) -> String { + let id = create_conversation(app, token, csrf).await; let req = json_with_token( "POST", &format!("/api/conversations/{id}/warmup"), @@ -69,6 +72,94 @@ async fn config_options_returns_active_agent_snapshot() { assert_eq!(json["data"]["config_options"][0]["current_value"], "mock-model"); } +#[tokio::test] +async fn runtime_ensure_requires_auth() { + let (app, _services) = build_app_with_mock_agents().await; + let csrf = "csrf-test"; + + let resp = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/conversations/conv-1/runtime/ensure") + .header("x-csrf-token", csrf) + .header("cookie", format!("aionui-csrf-token={csrf}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_eq!(json["code"], "UNAUTHORIZED"); +} + +#[tokio::test] +async fn runtime_ensure_requires_csrf() { + let (mut app, services) = build_app_with_mock_agents().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let id = create_conversation(&mut app, &token, &csrf).await; + + let req = Request::builder() + .method("POST") + .uri(format!("/api/conversations/{id}/runtime/ensure")) + .header("authorization", format!("Bearer {token}")) + .body(Body::empty()) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + let json = body_json(resp).await; + assert_eq!(json["code"], "CSRF_INVALID"); +} + +#[tokio::test] +async fn runtime_ensure_recovers_missing_agent_and_returns_config_snapshot() { + let (mut app, services) = build_app_with_mock_agents().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let id = create_conversation(&mut app, &token, &csrf).await; + + let req = json_with_token( + "POST", + &format!("/api/conversations/{id}/runtime/ensure"), + json!(null), + &token, + &csrf, + ); + let resp = app.oneshot(req).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::OK); + let json = body_json(resp).await; + assert_eq!(json["success"], true); + assert_eq!(json["data"]["recovered"], true); + assert_eq!(json["data"]["runtime"]["has_task"], true); + assert_eq!(json["data"]["config_options"][0]["id"], "model"); + assert_eq!(json["data"]["config_options"][0]["current_value"], "mock-model"); +} + +#[tokio::test] +async fn runtime_ensure_uses_existing_agent_without_recovery() { + let (mut app, services) = build_app_with_mock_agents().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let id = create_and_warmup_conversation(&mut app, &token, &csrf).await; + + let req = json_with_token( + "POST", + &format!("/api/conversations/{id}/runtime/ensure"), + json!(null), + &token, + &csrf, + ); + let resp = app.oneshot(req).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::OK); + let json = body_json(resp).await; + assert_eq!(json["data"]["recovered"], false); + assert_eq!(json["data"]["runtime"]["has_task"], true); + assert_eq!(json["data"]["config_options"][0]["id"], "model"); +} + #[tokio::test] async fn set_config_option_requires_csrf() { let (mut app, services) = build_app_with_mock_agents().await; diff --git a/crates/aionui-conversation/src/routes.rs b/crates/aionui-conversation/src/routes.rs index 0f99388e5..f48da15da 100644 --- a/crates/aionui-conversation/src/routes.rs +++ b/crates/aionui-conversation/src/routes.rs @@ -10,9 +10,9 @@ use aionui_api_types::{ ActiveCountResponse, ApiResponse, ApprovalCheckQuery, ApprovalCheckResponse, CancelConversationRequest, CancelConversationResponse, CloneConversationRequest, ConfirmRequest, ConfirmationListResponse, ConversationArtifactListResponse, ConversationArtifactResponse, ConversationListResponse, ConversationResponse, - CreateConversationRequest, ListConversationsQuery, ListMessagesQuery, MessageListResponse, MessageResponse, - MessageSearchResponse, SearchMessagesQuery, SendMessageRequest, SendMessageResponse, - UpdateConversationArtifactRequest, UpdateConversationRequest, + CreateConversationRequest, EnsureConversationRuntimeResponse, ListConversationsQuery, ListMessagesQuery, + MessageListResponse, MessageResponse, MessageSearchResponse, SearchMessagesQuery, SendMessageRequest, + SendMessageResponse, UpdateConversationArtifactRequest, UpdateConversationRequest, }; use aionui_auth::CurrentUser; use aionui_common::ApiError; @@ -105,6 +105,7 @@ pub fn conversation_routes(state: ConversationRouterState) -> Router { .route("/api/conversations/{id}/artifacts/{artifactId}", patch(update_artifact)) .route("/api/conversations/{id}/cancel", post(cancel)) .route("/api/conversations/{id}/warmup", post(warmup)) + .route("/api/conversations/{id}/runtime/ensure", post(ensure_runtime)) .route("/api/conversations/{id}/active-lease", post(active_lease)) // Confirmation system .route("/api/conversations/{id}/confirmations", get(list_confirmations)) @@ -318,6 +319,19 @@ async fn warmup( Ok(Json(ApiResponse::success())) } +async fn ensure_runtime( + State(state): State, + Extension(user): Extension, + Path(id): Path, +) -> Result>, ApiError> { + let response = state + .service + .ensure_runtime(&user.id, &id, &state.task_manager) + .await + .map_err(ApiError::from)?; + Ok(Json(ApiResponse::ok(response))) +} + async fn active_lease( State(state): State, Extension(user): Extension, diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index df8df4a60..c35c39424 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -19,9 +19,10 @@ use aionui_api_types::{ ConfirmRequest, ConfirmationListResponse, ConversationArtifactKind, ConversationArtifactListResponse, ConversationArtifactResponse, ConversationArtifactStatus, ConversationListResponse, ConversationMcpStatus, ConversationMcpStatusKind, ConversationResponse, ConversationRuntimeSummary, CreateConversationRequest, - ListConversationsQuery, ListMessagesQuery, MessageListResponse, MessageResponse, MessageSearchResponse, - SearchMessagesQuery, SendMessageRequest, SendMessageResponse, SessionMcpServer, SessionMcpTransport, - TeamSessionBinding, UpdateConversationArtifactRequest, UpdateConversationRequest, WebSocketMessage, + EnsureConversationRuntimeResponse, ListConversationsQuery, ListMessagesQuery, MessageListResponse, MessageResponse, + MessageSearchResponse, SearchMessagesQuery, SendMessageRequest, SendMessageResponse, SessionMcpServer, + SessionMcpTransport, TeamSessionBinding, UpdateConversationArtifactRequest, UpdateConversationRequest, + WebSocketMessage, }; use aionui_common::{ AgentKillReason, AgentType, ConversationSource, ConversationStatus, ErrorChain, MessageType, OnConversationDelete, @@ -2914,6 +2915,43 @@ impl ConversationService { conversation_id: &str, task_manager: &Arc, ) -> Result<(), ConversationError> { + let _ = self + .ensure_runtime_agent(user_id, conversation_id, task_manager, "warmup") + .await?; + debug!("Agent warmed up"); + Ok(()) + } + + #[tracing::instrument(skip_all, fields(user_id = %user_id, conversation_id = %conversation_id))] + pub async fn ensure_runtime( + &self, + user_id: &str, + conversation_id: &str, + task_manager: &Arc, + ) -> Result { + let (agent, recovered) = self + .ensure_runtime_agent(user_id, conversation_id, task_manager, "runtime_ensure") + .await?; + let config_options = agent + .get_config_options() + .await + .map_err(ConversationError::from)? + .config_options; + + Ok(EnsureConversationRuntimeResponse { + recovered, + config_options, + runtime: self.runtime_summary_for(conversation_id).await, + }) + } + + async fn ensure_runtime_agent( + &self, + user_id: &str, + conversation_id: &str, + task_manager: &Arc, + phase: &'static str, + ) -> Result<(AgentInstance, bool), ConversationError> { let row = self .conversation_repo .get(conversation_id) @@ -2925,6 +2963,11 @@ impl ConversationService { reject_deprecated_runtime_row(&row)?; + if let Some(agent) = task_manager.get_task(conversation_id) { + debug!(conversation_id, phase, "Conversation runtime already active"); + return Ok((agent, false)); + } + let build_opts = self.build_task_options(&row).await?; self.ensure_workspace_skill_links(&row, &build_opts).await; let stored_workspace = build_opts.context.workspace.stored_path.clone(); @@ -2939,7 +2982,7 @@ impl ConversationService { backend = "openclaw", error_kind = "openclaw_gateway_unreachable", port = 18789_u16, - phase = "warmup", + phase, "OpenClaw Gateway unreachable during ACP startup" ); let detail = send_error @@ -2957,8 +3000,8 @@ impl ConversationService { self.maybe_persist_workspace(conversation_id, &stored_workspace, agent.workspace()) .await?; - debug!("Agent warmed up"); - Ok(()) + info!(conversation_id, phase, "Conversation runtime recovered"); + Ok((agent, true)) } } diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index 75f64ce64..e8cf5695c 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -3055,6 +3055,51 @@ async fn get_config_options_returns_active_agent_snapshot() { assert_eq!(result.config_options[0].current_value.as_deref(), Some("gpt-5.5")); } +#[tokio::test] +async fn ensure_runtime_recovers_missing_agent_and_returns_snapshot() { + let task_mgr = Arc::new(MockTaskManager::new()); + let (svc, _broadcaster, _repo) = make_service_with_mock_task_manager(task_mgr.clone()); + let conv = svc.create("user_1", make_create_req()).await.unwrap(); + + let result = svc + .ensure_runtime("user_1", &conv.id, &(task_mgr.clone() as Arc)) + .await + .unwrap(); + + assert!(result.recovered); + assert!(result.runtime.has_task); + assert!(task_mgr.get_task(&conv.id).is_some()); + assert!(result.config_options.is_empty()); +} + +#[tokio::test] +async fn ensure_runtime_uses_existing_agent_snapshot_without_recovery() { + let task_mgr = Arc::new(MockTaskManager::new()); + let (svc, _broadcaster, _repo) = make_service_with_mock_task_manager(task_mgr.clone()); + let conv = svc.create("user_1", make_create_req()).await.unwrap(); + let agent = MockAgent::new(&conv.id).with_config_options(vec![AcpConfigOptionDto { + id: "model".to_owned(), + name: Some("Model".to_owned()), + label: None, + description: None, + category: Some("model".to_owned()), + option_type: "select".to_owned(), + current_value: Some("gpt-5.5".to_owned()), + options: Vec::new(), + }]); + task_mgr.insert_agent(&conv.id, AgentInstance::Mock(Arc::new(agent))); + + let result = svc + .ensure_runtime("user_1", &conv.id, &(task_mgr.clone() as Arc)) + .await + .unwrap(); + + assert!(!result.recovered); + assert!(result.runtime.has_task); + assert_eq!(result.config_options[0].id, "model"); + assert_eq!(result.config_options[0].current_value.as_deref(), Some("gpt-5.5")); +} + #[tokio::test] async fn set_config_option_returns_observed_confirmation() { let task_mgr = Arc::new(MockTaskManager::new()); From 54f96c65243a8cba5164f5262b6198d4b48e5dda Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 2 Jul 2026 12:38:50 +0800 Subject: [PATCH 3/5] fix(team): wait for runtime rebuild kills --- crates/aionui-ai-agent/src/agent_task.rs | 5 +- crates/aionui-ai-agent/src/task_manager.rs | 75 +++++ crates/aionui-team/src/provisioning.rs | 318 ++++++++++++++++++++- crates/aionui-team/src/service.rs | 2 +- 4 files changed, 397 insertions(+), 3 deletions(-) diff --git a/crates/aionui-ai-agent/src/agent_task.rs b/crates/aionui-ai-agent/src/agent_task.rs index b059f0c40..15578e4ed 100644 --- a/crates/aionui-ai-agent/src/agent_task.rs +++ b/crates/aionui-ai-agent/src/agent_task.rs @@ -235,7 +235,10 @@ impl AgentInstance { Self::Acp(m) => m.kill_and_wait(reason), Self::Aionrs(m) => m.kill_and_wait(reason), #[cfg(any(test, feature = "test-support"))] - Self::Mock(_) => Box::pin(std::future::ready(())), + Self::Mock(m) => { + let _ = m.kill(reason); + Box::pin(std::future::ready(())) + } } } diff --git a/crates/aionui-ai-agent/src/task_manager.rs b/crates/aionui-ai-agent/src/task_manager.rs index 8063f1ef1..701f9e3ac 100644 --- a/crates/aionui-ai-agent/src/task_manager.rs +++ b/crates/aionui-ai-agent/src/task_manager.rs @@ -173,6 +173,22 @@ impl IWorkerTaskManager for WorkerTaskManagerImpl { if let Some(agent) = slot.get() { return agent.kill_and_wait(reason); } + return Box::pin(async move { + match slot + .get_or_try_init(|| async { Err(AgentError::internal("task slot removed before initialization")) }) + .await + { + Ok(agent) => agent.clone().kill_and_wait(reason).await, + Err(error) => { + debug!( + conversation_id = %id, + ?reason, + error = %ErrorChain(&error), + "Kill requested for task slot that did not finish initialization" + ); + } + } + }); } Box::pin(std::future::ready(())) } @@ -283,6 +299,7 @@ mod tests { workspace: String, status: Option, last_activity: AtomicI64, + killed: Arc, event_tx: broadcast::Sender, } @@ -295,10 +312,16 @@ mod tests { workspace: "/tmp/test".to_owned(), status, last_activity: AtomicI64::new(now_ms()), + killed: Arc::new(std::sync::atomic::AtomicUsize::new(0)), event_tx, } } + fn with_kill_counter(mut self, killed: Arc) -> Self { + self.killed = killed; + self + } + fn with_agent_type(mut self, t: AgentType) -> Self { self.agent_type = t; self @@ -340,6 +363,7 @@ mod tests { Ok(()) } fn kill(&self, _reason: Option) -> Result<(), AgentError> { + self.killed.fetch_add(1, Ordering::SeqCst); Ok(()) } } @@ -521,6 +545,57 @@ mod tests { assert_eq!(mgr.active_count(), 1); } + #[tokio::test] + async fn kill_and_wait_waits_for_in_flight_build_and_kills_result() { + let build_started = Arc::new(tokio::sync::Notify::new()); + let release_build = Arc::new(tokio::sync::Notify::new()); + let killed = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + + let factory: AgentFactory = Arc::new({ + let build_started = Arc::clone(&build_started); + let release_build = Arc::clone(&release_build); + let killed = Arc::clone(&killed); + move |opts: BuildTaskOptions| { + let build_started = Arc::clone(&build_started); + let release_build = Arc::clone(&release_build); + let killed = Arc::clone(&killed); + async move { + build_started.notify_one(); + release_build.notified().await; + Ok(mock_instance( + MockAgent::new(opts.conversation_id(), None).with_kill_counter(killed), + )) + } + .boxed() + } + }); + let mgr = Arc::new(WorkerTaskManagerImpl::new(factory)); + + let build = { + let mgr = Arc::clone(&mgr); + tokio::spawn(async move { mgr.get_or_build_task("conv-1", make_options("conv-1")).await }) + }; + build_started.notified().await; + + let wait = mgr.kill_and_wait("conv-1", Some(AgentKillReason::TeamMcpRebuild)); + let wait_task = tokio::spawn(async move { + wait.await; + }); + + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + assert!( + !wait_task.is_finished(), + "kill_and_wait must wait for the in-flight build before returning" + ); + + release_build.notify_one(); + build.await.unwrap().unwrap(); + wait_task.await.unwrap(); + + assert_eq!(killed.load(Ordering::SeqCst), 1); + assert_eq!(mgr.active_count(), 0); + } + #[tokio::test] async fn get_task_finds_existing() { let mgr = make_manager(); diff --git a/crates/aionui-team/src/provisioning.rs b/crates/aionui-team/src/provisioning.rs index 880463060..f893585ef 100644 --- a/crates/aionui-team/src/provisioning.rs +++ b/crates/aionui-team/src/provisioning.rs @@ -336,7 +336,9 @@ impl TeamAgentProvisioner { ) -> Result<(), TeamError> { let team_id = mcp_stdio_cfg.team_id.clone(); self.write_team_mcp_runtime_config(agent, mcp_stdio_cfg).await?; - let _ = task_manager.kill(&agent.conversation_id, Some(AgentKillReason::TeamMcpRebuild)); + task_manager + .kill_and_wait(&agent.conversation_id, Some(AgentKillReason::TeamMcpRebuild)) + .await; self.conversation_port .warmup_agent_process(user_id, &agent.conversation_id, task_manager) .await @@ -581,3 +583,317 @@ impl TeamAgentProvisioner { None } } + +#[cfg(test)] +mod tests { + use super::*; + use aionui_ai_agent::types::BuildTaskOptions; + use aionui_ai_agent::{AgentError, AgentInstance}; + use aionui_db::models::{ + AgentMetadataRow, AssistantDefinitionRow, AssistantOverlayRow, Provider, UpdateAgentAvailabilitySnapshotParams, + UpdateAgentHandshakeParams, UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, + UpsertAssistantOverlayParams, + }; + use aionui_db::{CreateProviderParams, DbError, UpdateProviderParams}; + use std::sync::Mutex; + use tokio::sync::Notify; + + struct RecordingProvisioningPort { + events: Arc>>, + } + + #[async_trait] + impl TeamConversationProvisioningPort for RecordingProvisioningPort { + async fn create_team_conversation( + &self, + _request: TeamConversationCreateRequest, + ) -> Result { + Err(TeamError::InvalidRequest("unused".into())) + } + + async fn conversation_workspace(&self, _conversation_id: &str) -> Result, TeamError> { + Ok(None) + } + + async fn conversation_assistant_id(&self, _conversation_id: &str) -> Result, TeamError> { + Ok(None) + } + + async fn create_team_temp_workspace(&self, _team_id: &str) -> Result { + Err(TeamError::InvalidRequest("unused".into())) + } + + async fn patch_runtime_config( + &self, + _conversation_id: &str, + _patch: serde_json::Value, + ) -> Result<(), TeamError> { + self.events.lock().unwrap().push("patch"); + Ok(()) + } + + async fn save_acp_runtime_mode(&self, _conversation_id: &str, _mode: &str) -> Result<(), TeamError> { + Ok(()) + } + + async fn warmup_agent_process( + &self, + _user_id: &str, + _conversation_id: &str, + _task_manager: &Arc, + ) -> Result<(), TeamError> { + self.events.lock().unwrap().push("warmup"); + Ok(()) + } + + async fn delete_team_conversation(&self, _user_id: &str, _conversation_id: &str) -> Result<(), TeamError> { + Ok(()) + } + } + + struct BlockingKillTaskManager { + events: Arc>>, + kill_started: Arc, + release_kill: Arc, + } + + #[async_trait] + impl IWorkerTaskManager for BlockingKillTaskManager { + fn get_task(&self, _conversation_id: &str) -> Option { + None + } + + async fn get_or_build_task( + &self, + _conversation_id: &str, + _options: BuildTaskOptions, + ) -> Result { + Err(AgentError::internal("unused")) + } + + fn kill(&self, _conversation_id: &str, _reason: Option) -> Result<(), AgentError> { + self.events.lock().unwrap().push("kill_sync"); + self.kill_started.notify_one(); + Ok(()) + } + + fn kill_and_wait( + &self, + _conversation_id: &str, + _reason: Option, + ) -> std::pin::Pin + Send>> { + let events = Arc::clone(&self.events); + let kill_started = Arc::clone(&self.kill_started); + let release_kill = Arc::clone(&self.release_kill); + Box::pin(async move { + events.lock().unwrap().push("kill_wait_start"); + kill_started.notify_one(); + release_kill.notified().await; + events.lock().unwrap().push("kill_wait_done"); + }) + } + + async fn clear(&self) {} + + fn active_count(&self) -> usize { + 0 + } + + fn collect_idle(&self, _idle_threshold_ms: aionui_common::TimestampMs) -> Vec { + Vec::new() + } + } + + struct UnusedAgentMetadataRepo; + + #[async_trait] + impl IAgentMetadataRepository for UnusedAgentMetadataRepo { + async fn list_all(&self) -> Result, DbError> { + Ok(Vec::new()) + } + async fn get(&self, _id: &str) -> Result, DbError> { + Ok(None) + } + async fn find_by_source_and_name( + &self, + _agent_source: &str, + _name: &str, + ) -> Result, DbError> { + Ok(None) + } + async fn find_builtin_by_backend(&self, _backend: &str) -> Result, DbError> { + Ok(None) + } + async fn upsert(&self, _params: &UpsertAgentMetadataParams<'_>) -> Result { + Err(DbError::Init("unused".into())) + } + async fn apply_handshake( + &self, + _id: &str, + _params: &UpdateAgentHandshakeParams<'_>, + ) -> Result, DbError> { + Ok(None) + } + async fn update_availability_snapshot( + &self, + _id: &str, + _params: &UpdateAgentAvailabilitySnapshotParams<'_>, + ) -> Result, DbError> { + Ok(None) + } + async fn update_agent_overrides( + &self, + _id: &str, + _command_override: Option<&str>, + _env_override: Option<&str>, + ) -> Result<(), DbError> { + Ok(()) + } + async fn set_enabled(&self, _id: &str, _enabled: bool) -> Result { + Ok(false) + } + async fn delete(&self, _id: &str) -> Result { + Ok(false) + } + } + + struct UnusedAssistantDefinitionRepo; + + #[async_trait] + impl IAssistantDefinitionRepository for UnusedAssistantDefinitionRepo { + async fn list(&self) -> Result, DbError> { + Ok(Vec::new()) + } + async fn get_by_assistant_id(&self, _assistant_id: &str) -> Result, DbError> { + Ok(None) + } + async fn get_by_id(&self, _id: &str) -> Result, DbError> { + Ok(None) + } + async fn get_by_source_ref( + &self, + _source: &str, + _source_ref: &str, + ) -> Result, DbError> { + Ok(None) + } + async fn upsert( + &self, + _params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + Err(DbError::Init("unused".into())) + } + async fn soft_delete(&self, _id: &str, _deleted_at: i64) -> Result { + Ok(false) + } + } + + struct UnusedAssistantOverlayRepo; + + #[async_trait] + impl IAssistantOverlayRepository for UnusedAssistantOverlayRepo { + async fn get(&self, _assistant_definition_id: &str) -> Result, DbError> { + Ok(None) + } + async fn list(&self) -> Result, DbError> { + Ok(Vec::new()) + } + async fn upsert(&self, _params: &UpsertAssistantOverlayParams<'_>) -> Result { + Err(DbError::Init("unused".into())) + } + async fn delete(&self, _assistant_definition_id: &str) -> Result { + Ok(false) + } + } + + struct EmptyProviderRepo; + + #[async_trait] + impl IProviderRepository for EmptyProviderRepo { + async fn list(&self) -> Result, DbError> { + Ok(Vec::new()) + } + async fn find_by_id(&self, _id: &str) -> Result, DbError> { + Ok(None) + } + async fn create(&self, _params: CreateProviderParams<'_>) -> Result { + Err(DbError::Init("unused".into())) + } + async fn update(&self, _id: &str, _params: UpdateProviderParams<'_>) -> Result { + Err(DbError::Init("unused".into())) + } + async fn delete(&self, _id: &str) -> Result<(), DbError> { + Ok(()) + } + } + + fn test_provisioner(events: Arc>>) -> TeamAgentProvisioner { + TeamAgentProvisioner::new( + Arc::new(crate::test_utils::MockTeamRepo::new()), + Arc::new(UnusedAgentMetadataRepo), + Arc::new(UnusedAssistantDefinitionRepo), + Arc::new(UnusedAssistantOverlayRepo), + Arc::new(EmptyProviderRepo), + Arc::new(RecordingProvisioningPort { events }), + ) + } + + fn test_agent() -> TeamAgent { + TeamAgent { + slot_id: "slot-1".into(), + name: "Agent".into(), + role: TeammateRole::Teammate, + conversation_id: "conv-1".into(), + backend: "claude".into(), + model: "sonnet".into(), + assistant_id: None, + status: None, + conversation_type: None, + cli_path: None, + } + } + + fn test_mcp_config() -> TeamMcpStdioConfig { + TeamMcpStdioConfig { + team_id: "team-1".into(), + port: 12345, + token: "token".into(), + slot_id: "slot-1".into(), + binary_path: "/tmp/aioncore".into(), + } + } + + #[tokio::test] + async fn attach_agent_process_waits_for_kill_before_warmup() { + let events = Arc::new(Mutex::new(Vec::new())); + let kill_started = Arc::new(Notify::new()); + let release_kill = Arc::new(Notify::new()); + let provisioner = test_provisioner(Arc::clone(&events)); + let task_manager: Arc = Arc::new(BlockingKillTaskManager { + events: Arc::clone(&events), + kill_started: Arc::clone(&kill_started), + release_kill: Arc::clone(&release_kill), + }); + + let attach = tokio::spawn(async move { + provisioner + .attach_agent_process("user-1", &test_agent(), test_mcp_config(), &task_manager) + .await + }); + kill_started.notified().await; + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + + assert!( + !events.lock().unwrap().contains(&"warmup"), + "agent warmup must wait until the previous task is fully killed" + ); + + release_kill.notify_one(); + attach.await.unwrap().unwrap(); + + assert_eq!( + events.lock().unwrap().as_slice(), + ["patch", "kill_wait_start", "kill_wait_done", "warmup"] + ); + } +} diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 7ba7b4878..269825994 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -474,7 +474,7 @@ impl TeamSessionService { /// Flow (mcp.md §4.3): /// 1. Start `TeamSession` (opens the MCP TCP server). /// 2. For each agent: persist `team_mcp_stdio_config` into - /// `conversation.extra` → `task_manager.kill(conv_id, TeamMcpRebuild)` + /// `conversation.extra` → `task_manager.kill_and_wait(conv_id, TeamMcpRebuild)` /// → `TeamConversationProvisioningPort::warmup_agent_process(...)` /// rebuilds the ACP process with /// the new extra. From 22044a66c867929d3befde5953dda0a140907ce0 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 2 Jul 2026 13:28:38 +0800 Subject: [PATCH 4/5] fix(runtime): remove warmup endpoint --- crates/aionui-ai-agent/src/factory/acp.rs | 2 +- .../aionui-ai-agent/src/manager/acp/agent.rs | 2 +- crates/aionui-ai-agent/src/task_manager.rs | 2 +- .../tests/acp_config_options_e2e.rs | 32 +++++++++++++++---- .../aionui-app/tests/agent_integration_e2e.rs | 6 ++-- crates/aionui-app/tests/common/mod.rs | 2 +- crates/aionui-app/tests/message_e2e.rs | 26 ++++----------- crates/aionui-conversation/src/routes.rs | 14 -------- crates/aionui-team/docs/frontend-guide.md | 2 +- crates/aionui-team/src/event_loop.rs | 2 +- 10 files changed, 40 insertions(+), 50 deletions(-) diff --git a/crates/aionui-ai-agent/src/factory/acp.rs b/crates/aionui-ai-agent/src/factory/acp.rs index 5639f3b33..4c37b4e4b 100644 --- a/crates/aionui-ai-agent/src/factory/acp.rs +++ b/crates/aionui-ai-agent/src/factory/acp.rs @@ -154,7 +154,7 @@ pub(super) async fn build( arc.set_session_id(sid).await; } - // Open the ACP session eagerly so `POST /warmup` returns only after + // Open the ACP session eagerly so runtime preparation returns only after // session/new (or claude-meta-resume / session/load) and the first // reconcile pass have completed. Matches aionrs factory behaviour: // the caller sees "warmed up" == "ready for PUT /mode | /model". diff --git a/crates/aionui-ai-agent/src/manager/acp/agent.rs b/crates/aionui-ai-agent/src/manager/acp/agent.rs index 83c7ccabb..6d346eeb1 100644 --- a/crates/aionui-ai-agent/src/manager/acp/agent.rs +++ b/crates/aionui-ai-agent/src/manager/acp/agent.rs @@ -977,7 +977,7 @@ impl AcpAgentManager { } /// Pre-open the ACP session without sending a prompt. Called by the - /// factory after `AcpAgentManager::build` so `POST /warmup` returns + /// factory after `AcpAgentManager::build` so runtime preparation returns /// only after the session is ready to accept `set_mode` / `set_model` /// / `prompt`. Idempotent — if already opened, returns immediately. #[tracing::instrument(skip_all, fields(conversation_id = %self.params.conversation_id))] diff --git a/crates/aionui-ai-agent/src/task_manager.rs b/crates/aionui-ai-agent/src/task_manager.rs index 701f9e3ac..9358ef236 100644 --- a/crates/aionui-ai-agent/src/task_manager.rs +++ b/crates/aionui-ai-agent/src/task_manager.rs @@ -37,7 +37,7 @@ pub trait IWorkerTaskManager: Send + Sync { /// Concurrent callers with the same `conversation_id` block on a shared /// [`OnceCell`] so the factory runs at most once per conversation — /// avoiding the race where two concurrent HTTP requests (e.g. - /// `/messages` + `/warmup`) would each spawn their own CLI process and + /// `/messages` + `/runtime/ensure`) would each spawn their own CLI process and /// ACP connection, with one of them leaking. async fn get_or_build_task( &self, diff --git a/crates/aionui-app/tests/acp_config_options_e2e.rs b/crates/aionui-app/tests/acp_config_options_e2e.rs index 2e39b9351..7efd00a95 100644 --- a/crates/aionui-app/tests/acp_config_options_e2e.rs +++ b/crates/aionui-app/tests/acp_config_options_e2e.rs @@ -23,11 +23,11 @@ async fn create_conversation(app: &mut axum::Router, token: &str, csrf: &str) -> json["data"]["id"].as_str().unwrap().to_owned() } -async fn create_and_warmup_conversation(app: &mut axum::Router, token: &str, csrf: &str) -> String { +async fn create_and_ensure_runtime_conversation(app: &mut axum::Router, token: &str, csrf: &str) -> String { let id = create_conversation(app, token, csrf).await; let req = json_with_token( "POST", - &format!("/api/conversations/{id}/warmup"), + &format!("/api/conversations/{id}/runtime/ensure"), json!(null), token, csrf, @@ -55,7 +55,7 @@ async fn config_options_requires_auth() { async fn config_options_returns_active_agent_snapshot() { let (mut app, services) = build_app_with_mock_agents().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let id = create_and_warmup_conversation(&mut app, &token, &csrf).await; + let id = create_and_ensure_runtime_conversation(&mut app, &token, &csrf).await; let resp = app .oneshot(get_with_token( @@ -114,6 +114,24 @@ async fn runtime_ensure_requires_csrf() { assert_eq!(json["code"], "CSRF_INVALID"); } +#[tokio::test] +async fn legacy_warmup_route_is_removed() { + let (mut app, services) = build_app_with_mock_agents().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let id = create_conversation(&mut app, &token, &csrf).await; + + let req = json_with_token( + "POST", + &format!("/api/conversations/{id}/warmup"), + json!(null), + &token, + &csrf, + ); + let resp = app.oneshot(req).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::NOT_FOUND); +} + #[tokio::test] async fn runtime_ensure_recovers_missing_agent_and_returns_config_snapshot() { let (mut app, services) = build_app_with_mock_agents().await; @@ -142,7 +160,7 @@ async fn runtime_ensure_recovers_missing_agent_and_returns_config_snapshot() { async fn runtime_ensure_uses_existing_agent_without_recovery() { let (mut app, services) = build_app_with_mock_agents().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let id = create_and_warmup_conversation(&mut app, &token, &csrf).await; + let id = create_and_ensure_runtime_conversation(&mut app, &token, &csrf).await; let req = json_with_token( "POST", @@ -164,7 +182,7 @@ async fn runtime_ensure_uses_existing_agent_without_recovery() { async fn set_config_option_requires_csrf() { let (mut app, services) = build_app_with_mock_agents().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let id = create_and_warmup_conversation(&mut app, &token, &csrf).await; + let id = create_and_ensure_runtime_conversation(&mut app, &token, &csrf).await; let req = Request::builder() .method("PUT") @@ -186,7 +204,7 @@ async fn set_config_option_requires_csrf() { async fn set_config_option_returns_observed_confirmation() { let (mut app, services) = build_app_with_mock_agents().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let id = create_and_warmup_conversation(&mut app, &token, &csrf).await; + let id = create_and_ensure_runtime_conversation(&mut app, &token, &csrf).await; let req = json_with_token( "PUT", @@ -207,7 +225,7 @@ async fn set_config_option_returns_observed_confirmation() { async fn old_mode_and_model_routes_are_removed() { let (mut app, services) = build_app_with_mock_agents().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let id = create_and_warmup_conversation(&mut app, &token, &csrf).await; + let id = create_and_ensure_runtime_conversation(&mut app, &token, &csrf).await; let resp = app .clone() diff --git a/crates/aionui-app/tests/agent_integration_e2e.rs b/crates/aionui-app/tests/agent_integration_e2e.rs index 19368a29c..7b5a83df2 100644 --- a/crates/aionui-app/tests/agent_integration_e2e.rs +++ b/crates/aionui-app/tests/agent_integration_e2e.rs @@ -478,14 +478,14 @@ async fn stop_stream_with_mock_agent() { } #[tokio::test] -async fn warmup_with_mock_agent() { +async fn runtime_ensure_with_mock_agent() { let (mut app, services, _mock_tm) = build_app_with_mock_tasks().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "Pass123!").await; - let conv_id = create_conversation(&mut app, &token, &csrf, "Warmup Test").await; + let conv_id = create_conversation(&mut app, &token, &csrf, "Runtime Ensure Test").await; let req = json_with_token( "POST", - &format!("/api/conversations/{conv_id}/warmup"), + &format!("/api/conversations/{conv_id}/runtime/ensure"), json!({}), &token, &csrf, diff --git a/crates/aionui-app/tests/common/mod.rs b/crates/aionui-app/tests/common/mod.rs index 57c044969..fa69e2ab4 100644 --- a/crates/aionui-app/tests/common/mod.rs +++ b/crates/aionui-app/tests/common/mod.rs @@ -120,7 +120,7 @@ pub async fn build_app_with_mock_version( /// Build app with a mock worker task manager that returns noop agents. /// -/// Use for tests that exercise session/warmup paths (team ensure_session, +/// Use for tests that exercise runtime preparation paths (team ensure_session, /// send_message) where spawning a real CLI process is not feasible. pub async fn build_app_with_mock_agents() -> (axum::Router, AppServices) { let db = aionui_db::init_database_memory().await.unwrap(); diff --git a/crates/aionui-app/tests/message_e2e.rs b/crates/aionui-app/tests/message_e2e.rs index 349d4ca30..9eb7e75e7 100644 --- a/crates/aionui-app/tests/message_e2e.rs +++ b/crates/aionui-app/tests/message_e2e.rs @@ -878,16 +878,16 @@ async fn t2_2_stop_stream_requires_auth() { assert_eq!(resp.status(), StatusCode::FORBIDDEN); } -// ── T2.3: Warmup ──────────────────────────────────────────────────── +// ── T2.3: Runtime ensure ──────────────────────────────────────────── #[tokio::test] -async fn t2_3_warmup_conversation_not_found() { +async fn t2_3_runtime_ensure_conversation_not_found() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; let req = common::json_with_token( "POST", - "/api/conversations/non-existent/warmup", + "/api/conversations/non-existent/runtime/ensure", json!({}), &token, &csrf, @@ -897,10 +897,10 @@ async fn t2_3_warmup_conversation_not_found() { } #[tokio::test] -async fn t2_3b_warmup_legacy_workspace_with_whitespace_succeeds() { +async fn t2_3b_runtime_ensure_legacy_workspace_with_whitespace_succeeds() { let (mut app, services) = build_app_with_mock_agents().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let conv_id = create_conversation(&mut app, &token, &csrf, "Legacy Warmup").await; + let conv_id = create_conversation(&mut app, &token, &csrf, "Runtime Ensure").await; let dir = tempfile::tempdir().unwrap(); let workspace = dir.path().join("my project"); std::fs::create_dir(&workspace).unwrap(); @@ -908,7 +908,7 @@ async fn t2_3b_warmup_legacy_workspace_with_whitespace_succeeds() { let req = common::json_with_token( "POST", - &format!("/api/conversations/{conv_id}/warmup"), + &format!("/api/conversations/{conv_id}/runtime/ensure"), json!({}), &token, &csrf, @@ -916,17 +916,3 @@ async fn t2_3b_warmup_legacy_workspace_with_whitespace_succeeds() { let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); } - -#[tokio::test] -async fn t2_3_warmup_requires_auth() { - let (app, _services) = build_app().await; - - let req = axum::http::Request::builder() - .method("POST") - .uri("/api/conversations/some-id/warmup") - .header("content-type", "application/json") - .body(Body::empty()) - .unwrap(); - let resp = app.oneshot(req).await.unwrap(); - assert_eq!(resp.status(), StatusCode::FORBIDDEN); -} diff --git a/crates/aionui-conversation/src/routes.rs b/crates/aionui-conversation/src/routes.rs index f48da15da..f99f6359c 100644 --- a/crates/aionui-conversation/src/routes.rs +++ b/crates/aionui-conversation/src/routes.rs @@ -104,7 +104,6 @@ pub fn conversation_routes(state: ConversationRouterState) -> Router { .route("/api/conversations/{id}/artifacts", get(list_artifacts)) .route("/api/conversations/{id}/artifacts/{artifactId}", patch(update_artifact)) .route("/api/conversations/{id}/cancel", post(cancel)) - .route("/api/conversations/{id}/warmup", post(warmup)) .route("/api/conversations/{id}/runtime/ensure", post(ensure_runtime)) .route("/api/conversations/{id}/active-lease", post(active_lease)) // Confirmation system @@ -306,19 +305,6 @@ async fn cancel( Ok(Json(ApiResponse::ok(response))) } -async fn warmup( - State(state): State, - Extension(user): Extension, - Path(id): Path, -) -> Result>, ApiError> { - state - .service - .warmup(&user.id, &id, &state.task_manager) - .await - .map_err(ApiError::from)?; - Ok(Json(ApiResponse::success())) -} - async fn ensure_runtime( State(state): State, Extension(user): Extension, diff --git a/crates/aionui-team/docs/frontend-guide.md b/crates/aionui-team/docs/frontend-guide.md index 304c0e0da..5fa0f6bd4 100644 --- a/crates/aionui-team/docs/frontend-guide.md +++ b/crates/aionui-team/docs/frontend-guide.md @@ -65,7 +65,7 @@ Team 创建是显式行为：用户通过 Team UI 或 `POST /api/teams` 创建 | D29a-2 caller role==Lead 校验 | ✅ | | | D29a-3 name normalize + 唯一性 | ✅ | | | D29a-4 backend 白名单校验 | ✅ | | -| **D29b spawn_agent 真实落地** | ✅ | **已合入** — conversation 创建 + slot 分配 + kill/warmup agent | +| **D29b spawn_agent 真实落地** | ✅ | **已合入** — conversation 创建 + slot 分配 + kill/rebuild agent | | D29d-1 `team.agentSpawned` WS 事件 | ✅ | spawn 成功后广播 | | D29e MCP dispatch 接通 session | ✅ | `exec_spawn_agent` 改成调 `TeamSession::spawn_agent` | | D30a-1 shutdown_approved/rejected 字符串拦截 | ✅ | | diff --git a/crates/aionui-team/src/event_loop.rs b/crates/aionui-team/src/event_loop.rs index 62558b307..3c127864d 100644 --- a/crates/aionui-team/src/event_loop.rs +++ b/crates/aionui-team/src/event_loop.rs @@ -186,7 +186,7 @@ async fn run_event_loop( match execute_turn(&ctx, &input).await { Some(turn) => finalize_turn(&ctx, turn, &input).await, - None => break, // Turn not started (guard/warmup); retry on next signal + None => break, // Turn not started (guard/runtime preparation); retry on next signal } } } From 3c76a16f21598d55f70f54b46c5afa30a0a9d579 Mon Sep 17 00:00:00 2001 From: zynx <> Date: Thu, 2 Jul 2026 14:42:34 +0800 Subject: [PATCH 5/5] fix(runtime): remove config options GET endpoint --- .../tests/acp_config_options_e2e.rs | 23 +++---------------- crates/aionui-conversation/src/routes_aux.rs | 15 ++---------- 2 files changed, 5 insertions(+), 33 deletions(-) diff --git a/crates/aionui-app/tests/acp_config_options_e2e.rs b/crates/aionui-app/tests/acp_config_options_e2e.rs index 7efd00a95..f37b9fa9c 100644 --- a/crates/aionui-app/tests/acp_config_options_e2e.rs +++ b/crates/aionui-app/tests/acp_config_options_e2e.rs @@ -5,7 +5,7 @@ use axum::http::{Request, StatusCode}; use serde_json::json; use tower::ServiceExt; -use common::{body_json, build_app_with_mock_agents, get_request, get_with_token, json_with_token, setup_and_login}; +use common::{body_json, build_app_with_mock_agents, get_with_token, json_with_token, setup_and_login}; fn create_body() -> serde_json::Value { json!({ @@ -39,20 +39,7 @@ async fn create_and_ensure_runtime_conversation(app: &mut axum::Router, token: & } #[tokio::test] -async fn config_options_requires_auth() { - let (app, _services) = build_app_with_mock_agents().await; - - let resp = app - .oneshot(get_request("/api/conversations/conv-1/config-options")) - .await - .unwrap(); - assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); - let json = body_json(resp).await; - assert_eq!(json["code"], "UNAUTHORIZED"); -} - -#[tokio::test] -async fn config_options_returns_active_agent_snapshot() { +async fn legacy_config_options_get_route_is_removed() { let (mut app, services) = build_app_with_mock_agents().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; let id = create_and_ensure_runtime_conversation(&mut app, &token, &csrf).await; @@ -65,11 +52,7 @@ async fn config_options_returns_active_agent_snapshot() { .await .unwrap(); - assert_eq!(resp.status(), StatusCode::OK); - let json = body_json(resp).await; - assert_eq!(json["success"], true); - assert_eq!(json["data"]["config_options"][0]["id"], "model"); - assert_eq!(json["data"]["config_options"][0]["current_value"], "mock-model"); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); } #[tokio::test] diff --git a/crates/aionui-conversation/src/routes_aux.rs b/crates/aionui-conversation/src/routes_aux.rs index 22405927b..5a56e27dd 100644 --- a/crates/aionui-conversation/src/routes_aux.rs +++ b/crates/aionui-conversation/src/routes_aux.rs @@ -2,8 +2,8 @@ use crate::state::ConversationRouterState; use aionui_api_types::{ - ApiResponse, GetConfigOptionsResponse, SetConfigOptionRequest, SetConfigOptionResponse, SideQuestionRequest, - SideQuestionResponse, SlashCommandItem, WorkspaceBrowseQuery, WorkspaceEntry, + ApiResponse, SetConfigOptionRequest, SetConfigOptionResponse, SideQuestionRequest, SideQuestionResponse, + SlashCommandItem, WorkspaceBrowseQuery, WorkspaceEntry, }; use aionui_auth::CurrentUser; use aionui_common::ApiError; @@ -19,7 +19,6 @@ pub fn conversation_ops_routes(state: ConversationRouterState) -> Router { .route("/api/conversations/{id}/side-question", post(side_question)) .route("/api/conversations/{id}/slash-commands", get(get_slash_commands)) .route("/api/conversations/{id}/usage", get(get_usage)) - .route("/api/conversations/{id}/config-options", get(get_config_options)) .route( "/api/conversations/{id}/config-options/{option_id}", put(set_config_option), @@ -30,16 +29,6 @@ pub fn conversation_ops_routes(state: ConversationRouterState) -> Router { // ── Route handlers ───────────────────────────────────────────────── -async fn get_config_options( - State(state): State, - Extension(_user): Extension, - Path(id): Path, -) -> Result>, ApiError> { - Ok(Json(ApiResponse::ok( - state.service.get_config_options(&id).await.map_err(ApiError::from)?, - ))) -} - async fn set_config_option( State(state): State, Extension(_user): Extension,