oxicode_agent/mcp/lifecycle.rs
1//! Channel-based MCP server lifecycle management.
2//!
3//! The previous design attempted to put lifecycle timers (idle disconnect,
4//! keep-alive health checks) inside the same `McpManagerInner` that is guarded
5//! by a `tokio::sync::Mutex`. This caused a deadlock risk: an idle timer
6//! callback firing on a background task would need to re-acquire the mutex,
7//! while the agent thread might also be waiting for the callback.
8//!
9//! To break this, lifecycle events are sent over an `mpsc` channel to a
10//! dedicated background task that owns the timer handles. The background
11//! task only holds a `Weak<McpManager>` (no reference cycles) and re-enters
12//! the manager through a public async method that acquires the mutex on its
13//! own — which can never deadlock with the agent thread because the agent
14//! thread is not waiting for the lifecycle task while holding the lock.
15//!
16//! ```text
17//! [Agent task] [Lifecycle task] [McpManager]
18//! │ │ │
19//! │ call_tool() │ │
20//! ├─ lock(inner) ──────────────┼──────────────────────────► │
21//! │ do work │ │
22//! ├─ unlock │ │
23//! │ │ │
24//! │ │ (idle timer fires) │
25//! │ ├─ Weak::upgrade() ───────► │
26//! │ │ disconnect_server() │
27//! │ │ ├─ lock(inner) ───────► │
28//! │ │ ├─ unlock │
29//! │ │ │
30//! │ ← no contention: agent is not holding the lock when │
31//! │ the lifecycle task tries to acquire it. │
32//! ```
33
34use std::collections::HashMap;
35use std::sync::Weak;
36use std::time::Duration;
37use tokio::sync::mpsc;
38use tokio::task::JoinHandle;
39
40use crate::mcp::McpManager;
41
42/// Lifecycle events sent from [`McpManager`] to the background task.
43#[derive(Debug, Clone)]
44pub enum LifecycleEvent {
45 /// Start (or reset) an idle-disconnect timer for a server.
46 StartIdleTimer {
47 /// Name of the server to schedule the timer for.
48 server: String,
49 /// Idle duration before the server is disconnected.
50 timeout: Duration,
51 },
52 /// Cancel any pending idle-disconnect timer for a server.
53 CancelIdleTimer {
54 /// Name of the server whose timer should be cancelled.
55 server: String,
56 },
57 /// Start periodic health checks for a keep-alive server.
58 StartHealthCheck {
59 /// Name of the server to health-check.
60 server: String,
61 },
62 /// Stop all timers for a server (it was disconnected).
63 ServerStopped {
64 /// Name of the server that was stopped.
65 server: String,
66 },
67 /// Stop everything (process shutdown).
68 Shutdown,
69}
70
71/// Sender side of the lifecycle channel.
72pub type LifecycleTx = mpsc::UnboundedSender<LifecycleEvent>;
73/// Receiver side of the lifecycle channel.
74pub type LifecycleRx = mpsc::UnboundedReceiver<LifecycleEvent>;
75
76/// Create a new lifecycle event channel.
77pub fn channel() -> (LifecycleTx, LifecycleRx) {
78 mpsc::unbounded_channel()
79}
80
81/// The background task that owns all timer handles.
82///
83/// Created by [`McpManager::spawn()`]. Exits when the `LifecycleRx` returns
84/// `None` (the manager was dropped) or when a `Shutdown` event is received.
85pub async fn lifecycle_event_loop(mut rx: LifecycleRx, manager: Weak<McpManager>) {
86 let mut idle_timers: HashMap<String, JoinHandle<()>> = HashMap::new();
87 let mut health_handles: HashMap<String, JoinHandle<()>> = HashMap::new();
88
89 while let Some(event) = rx.recv().await {
90 match event {
91 LifecycleEvent::StartIdleTimer { server, timeout } => {
92 if let Some(h) = idle_timers.remove(&server) {
93 h.abort();
94 }
95 let mgr = manager.clone();
96 let srv = server.clone();
97 idle_timers.insert(
98 server,
99 tokio::spawn(async move {
100 tokio::time::sleep(timeout).await;
101 if let Some(m) = mgr.upgrade() {
102 // Best-effort: log but never panic on disconnect
103 // failure (the agent might be using the server).
104 if let Err(e) = m.disconnect_server(&srv).await {
105 tracing::debug!("MCP: idle-disconnect for '{}' failed: {}", srv, e);
106 }
107 }
108 }),
109 );
110 }
111 LifecycleEvent::CancelIdleTimer { server } => {
112 if let Some(h) = idle_timers.remove(&server) {
113 h.abort();
114 }
115 }
116 LifecycleEvent::StartHealthCheck { server } => {
117 if let Some(h) = health_handles.remove(&server) {
118 h.abort();
119 }
120 let mgr = manager.clone();
121 let srv = server.clone();
122 health_handles.insert(
123 server,
124 tokio::spawn(async move {
125 // Health-check every 30s; on failure, attempt one
126 // reconnect and stop on success.
127 let interval = Duration::from_secs(30);
128 loop {
129 tokio::time::sleep(interval).await;
130 let Some(m) = mgr.upgrade() else {
131 break;
132 };
133 match m.health_check_and_reconnect(&srv).await {
134 Ok(()) => continue,
135 Err(e) => {
136 tracing::warn!(
137 "MCP: health check for keep-alive server '{}' failed: {}",
138 srv,
139 e
140 );
141 // Give up after a single reconnect attempt
142 // (we don't want to spam the system log).
143 // The next tool call will trigger lazy
144 // reconnection.
145 break;
146 }
147 }
148 }
149 }),
150 );
151 }
152 LifecycleEvent::ServerStopped { server } => {
153 if let Some(h) = idle_timers.remove(&server) {
154 h.abort();
155 }
156 if let Some(h) = health_handles.remove(&server) {
157 h.abort();
158 }
159 }
160 LifecycleEvent::Shutdown => break,
161 }
162 }
163
164 // Cleanup on exit
165 for (_, h) in idle_timers {
166 h.abort();
167 }
168 for (_, h) in health_handles {
169 h.abort();
170 }
171}
172
173#[cfg(test)]
174mod tests {
175 use super::*;
176
177 #[tokio::test]
178 async fn channel_round_trip() {
179 let (tx, mut rx) = channel();
180 tx.send(LifecycleEvent::CancelIdleTimer {
181 server: "test".into(),
182 })
183 .unwrap();
184 let event = rx.recv().await.unwrap();
185 match event {
186 LifecycleEvent::CancelIdleTimer { server } => assert_eq!(server, "test"),
187 _ => panic!("wrong event"),
188 }
189 }
190
191 #[tokio::test]
192 async fn lifecycle_event_loop_runs_to_completion() {
193 // Without a real McpManager we just verify that a spawned task runs
194 // to completion. The timer callback will see `Weak::upgrade() == None`.
195 let (tx, rx) = channel();
196 let manager: Weak<McpManager> = Weak::new();
197 let task = tokio::spawn(lifecycle_event_loop(rx, manager));
198
199 tx.send(LifecycleEvent::CancelIdleTimer { server: "x".into() })
200 .unwrap();
201
202 drop(tx);
203 task.await.unwrap();
204 }
205}