Skip to content

Commit 1fa7a54

Browse files
authored
fix(runtime): protect active ACP tasks from idle cleanup (#561)
## Summary - Add a process-local active lease registry and authenticated, CSRF-protected conversation/team active lease endpoints. - Expand ACP idle cleanup to reclaim stale warmup-only tasks while skipping conversations with active foreground leases. - Add conversation runtime ensure, remove legacy warmup/config-options HTTP routes, and wait for team runtime rebuild kills before warmup. ## Test Plan - cargo fmt --all -- --check - cargo clippy -p aionui-ai-agent -p aionui-conversation -p aionui-team -p aionui-app -- -D warnings - cargo test -p aionui-ai-agent -p aionui-conversation -p aionui-team -p aionui-app - cargo test --workspace - just push -u origin aionissue/task-20260701-002-004 ## Reviewer Notes - The active lease endpoints are state-changing routes with empty request bodies and must remain behind auth and CSRF protection. - Lease state is process-local and has no database migration or persistence behavior. - This backend change is intended to ship with the matching AionUi runtime ensure and foreground lease renewal changes. --------- Co-authored-by: zynx <>
1 parent 3840b77 commit 1fa7a54

28 files changed

Lines changed: 1470 additions & 106 deletions
Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
use aionui_common::{TimestampMs, now_ms};
2+
use dashmap::DashMap;
3+
4+
pub const ACTIVE_LEASE_TTL_MS: TimestampMs = 90_000;
5+
6+
#[derive(Debug)]
7+
pub struct ActiveLeaseRegistry {
8+
leases: DashMap<String, TimestampMs>,
9+
ttl_ms: TimestampMs,
10+
}
11+
12+
impl ActiveLeaseRegistry {
13+
pub fn new() -> Self {
14+
Self::with_ttl_ms(ACTIVE_LEASE_TTL_MS)
15+
}
16+
17+
pub fn with_ttl_ms(ttl_ms: TimestampMs) -> Self {
18+
Self {
19+
leases: DashMap::new(),
20+
ttl_ms,
21+
}
22+
}
23+
24+
pub fn renew(&self, conversation_id: &str) -> TimestampMs {
25+
let expires_at = now_ms().saturating_add(self.ttl_ms);
26+
self.leases.insert(conversation_id.to_owned(), expires_at);
27+
expires_at
28+
}
29+
30+
pub fn renew_many<'a>(&self, conversation_ids: impl IntoIterator<Item = &'a str>) -> (usize, TimestampMs) {
31+
let expires_at = now_ms().saturating_add(self.ttl_ms);
32+
let mut count = 0;
33+
for conversation_id in conversation_ids {
34+
self.leases.insert(conversation_id.to_owned(), expires_at);
35+
count += 1;
36+
}
37+
(count, expires_at)
38+
}
39+
40+
pub fn active_until(&self, conversation_id: &str) -> Option<TimestampMs> {
41+
let expires_at = *self.leases.get(conversation_id)?;
42+
if expires_at > now_ms() {
43+
Some(expires_at)
44+
} else {
45+
self.leases.remove(conversation_id);
46+
None
47+
}
48+
}
49+
50+
pub fn is_active(&self, conversation_id: &str) -> bool {
51+
self.active_until(conversation_id).is_some()
52+
}
53+
}
54+
55+
impl Default for ActiveLeaseRegistry {
56+
fn default() -> Self {
57+
Self::new()
58+
}
59+
}
60+
61+
#[cfg(test)]
62+
mod tests {
63+
use super::*;
64+
use std::time::Duration;
65+
66+
#[test]
67+
fn renew_sets_active_lease() {
68+
let registry = ActiveLeaseRegistry::new();
69+
70+
let expires_at = registry.renew("conv-1");
71+
72+
assert_eq!(registry.active_until("conv-1"), Some(expires_at));
73+
assert!(registry.is_active("conv-1"));
74+
}
75+
76+
#[test]
77+
fn renew_many_sets_all_leases_and_returns_count() {
78+
let registry = ActiveLeaseRegistry::new();
79+
80+
let (count, expires_at) = registry.renew_many(["conv-1", "conv-2"]);
81+
82+
assert_eq!(count, 2);
83+
assert_eq!(registry.active_until("conv-1"), Some(expires_at));
84+
assert_eq!(registry.active_until("conv-2"), Some(expires_at));
85+
}
86+
87+
#[test]
88+
fn active_lookup_returns_none_for_missing_lease() {
89+
let registry = ActiveLeaseRegistry::new();
90+
91+
assert_eq!(registry.active_until("missing"), None);
92+
assert!(!registry.is_active("missing"));
93+
}
94+
95+
#[test]
96+
fn active_lookup_lazily_removes_expired_lease() {
97+
let registry = ActiveLeaseRegistry::with_ttl_ms(1);
98+
registry.renew("conv-1");
99+
100+
std::thread::sleep(Duration::from_millis(5));
101+
102+
assert_eq!(registry.active_until("conv-1"), None);
103+
assert!(!registry.leases.contains_key("conv-1"));
104+
}
105+
}

crates/aionui-ai-agent/src/agent_task.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -235,7 +235,10 @@ impl AgentInstance {
235235
Self::Acp(m) => m.kill_and_wait(reason),
236236
Self::Aionrs(m) => m.kill_and_wait(reason),
237237
#[cfg(any(test, feature = "test-support"))]
238-
Self::Mock(_) => Box::pin(std::future::ready(())),
238+
Self::Mock(m) => {
239+
let _ = m.kill(reason);
240+
Box::pin(std::future::ready(()))
241+
}
239242
}
240243
}
241244

crates/aionui-ai-agent/src/factory/acp.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ pub(super) async fn build(
151151
arc.set_session_id(sid).await;
152152
}
153153

154-
// Open the ACP session eagerly so `POST /warmup` returns only after
154+
// Open the ACP session eagerly so runtime preparation returns only after
155155
// session/new (or claude-meta-resume / session/load) and the first
156156
// reconcile pass have completed. Matches aionrs factory behaviour:
157157
// the caller sees "warmed up" == "ready for PUT /mode | /model".

crates/aionui-ai-agent/src/idle_scanner.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ const SCAN_INTERVAL_SECS: u64 = 60;
1515
/// Start the background idle agent scanner.
1616
///
1717
/// Periodically scans active tasks and kills ACP agents that have been
18-
/// idle (finished + no activity) beyond the configured threshold.
18+
/// idle (finished or warmup-only + no activity) beyond the configured threshold.
1919
///
2020
/// The scanner runs until the provided `shutdown` signal resolves.
2121
pub fn start_idle_scanner(
@@ -59,7 +59,7 @@ fn scan_and_cleanup(manager: &Arc<dyn IWorkerTaskManager>, threshold_ms: i64) {
5959
let idle_ids = manager.collect_idle(threshold_ms);
6060

6161
if idle_ids.is_empty() {
62-
debug!(active = manager.active_count(), "Idle scan: no idle agents found");
62+
debug!(active_count = manager.active_count(), "Idle scan: no idle agents found");
6363
return;
6464
}
6565

crates/aionui-ai-agent/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#![warn(clippy::disallowed_types)]
22

33
//! AI agent lifecycle, worker task dispatch, and skill management.
4+
pub mod active_lease;
45
pub(crate) mod agent_runtime;
56
pub mod agent_task;
67
pub mod capability;
@@ -21,6 +22,7 @@ pub mod shared_kernel;
2122
pub mod task_manager;
2223
pub mod types;
2324

25+
pub use active_lease::{ACTIVE_LEASE_TTL_MS, ActiveLeaseRegistry};
2426
pub use agent_runtime::AgentRuntime;
2527
#[cfg(any(test, feature = "test-support"))]
2628
pub use agent_task::IMockAgent;

crates/aionui-ai-agent/src/manager/acp/agent.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -970,7 +970,7 @@ impl AcpAgentManager {
970970
}
971971

972972
/// Pre-open the ACP session without sending a prompt. Called by the
973-
/// factory after `AcpAgentManager::build` so `POST /warmup` returns
973+
/// factory after `AcpAgentManager::build` so runtime preparation returns
974974
/// only after the session is ready to accept `set_mode` / `set_model`
975975
/// / `prompt`. Idempotent — if already opened, returns immediately.
976976
#[tracing::instrument(skip_all, fields(conversation_id = %self.params.conversation_id))]

0 commit comments

Comments
 (0)