1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
//! Out-of-band session close handle.
//!
//! `SessionCloseHandle` is an `Arc`-shareable substruct that owns just the
//! fields needed to terminate an `AgentSession` from outside (typically from
//! the parent [`Agent`](super::Agent)'s session registry).
//!
//! `AgentSession` carries one of these via `Arc<SessionCloseHandle>`; the
//! parent `Agent` stores a `Weak<SessionCloseHandle>` in its registry. When
//! the user drops the session, the handle drops too and the registry's
//! `Weak` becomes dangling — pruned on the next `list_sessions()` /
//! `close_session()` call.
//!
//! Sharing the close mechanics through a single `close()` method on this
//! struct guarantees `AgentSession::close()` and `Agent::close_session(id)`
//! perform exactly the same cleanup.
use crate::hitl::ConfirmationProvider;
use crate::hooks::HookExecutor;
use crate::run::InMemoryRunStore;
use crate::subagent_task_tracker::{InMemorySubagentTaskTracker, SubagentStatus};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
const QUEUE_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
/// Bundle of `Arc`-shared session state needed to perform a graceful close
/// from anywhere holding (a clone of) the handle.
pub(crate) struct SessionCloseHandle {
pub(crate) session_id: String,
/// Tripped on first `close()` call; subsequent calls become no-ops and
/// `AgentSession::send`/`stream` fast-fail.
pub(crate) closed: Arc<AtomicBool>,
/// Session-level parent token. All in-flight run/subagent tokens are
/// `child_token()` of this.
pub(crate) session_cancel: CancellationToken,
/// Per-run cancel-token slot (currently active run's token, if any).
/// Populated by the run lifecycle.
pub(crate) cancel_token: Arc<Mutex<Option<CancellationToken>>>,
/// Current run id (matches `cancel_token` when set).
pub(crate) current_run_id: Arc<Mutex<Option<String>>>,
pub(crate) run_store: Arc<InMemoryRunStore>,
pub(crate) subagent_tasks: Arc<InMemorySubagentTaskTracker>,
pub(crate) confirmation_manager: Option<Arc<dyn ConfirmationProvider>>,
pub(crate) hook_executor: Option<Arc<dyn HookExecutor>>,
/// Optional session-owned lane queue. Close first shuts it down to reject
/// new work, then drains commands admitted before the close boundary.
pub(crate) command_queue: Option<Arc<crate::session_lane_queue::SessionLaneQueue>>,
/// Session-owned MCP source. Inherited/global managers are deliberately
/// absent so closing one session cannot tear down shared connections.
pub(crate) mcp_manager: Arc<crate::mcp::manager::McpManager>,
/// Session tool executor used to unwind exact local MCP registrations at
/// close while restoring their captured shadows.
pub(crate) tool_executor: Arc<crate::tools::ToolExecutor>,
/// Serializes session-local extension mutation. `closed` is the admission
/// bit; close takes this mutex before its final MCP cleanup so a mutation
/// admitted before close either rolls back or is subsequently cleaned up.
pub(crate) extension_mutation: Mutex<()>,
/// Linearizes immediate, non-async capability changes (dynamic tools,
/// workers, hooks, and commands) with the close admission boundary.
pub(crate) immediate_extension_mutation: std::sync::Mutex<()>,
/// Exact MCP wrapper identities installed by each session-local server.
pub(crate) mcp_tool_ownership:
std::sync::Mutex<super::session_extensions::SessionMcpToolOwnership>,
}
impl SessionCloseHandle {
/// Return whether `close()` has already been called.
pub(crate) fn is_closed(&self) -> bool {
self.closed.load(Ordering::Acquire)
}
/// Perform the full session close sequence. Idempotent: subsequent calls
/// are no-ops and are guaranteed not to panic.
///
/// Sequence (see [`AgentSession::close`](super::AgentSession::close)
/// for the public-facing contract):
/// 1. Flip the `closed` flag so further `send`/`stream` fast-fail;
/// 2. Stop the lane queue from accepting new commands;
/// 3. Fire the session-level cancellation token so every derived run
/// and subagent task token fires;
/// 4. Mark the active run as `Cancelled` in the run store and notify the
/// configured hook executor;
/// 5. Mark every still-running delegated subagent task as `Cancelled`
/// in the tracker;
/// 6. Cancel pending HITL tool confirmations so blocked tool callers
/// receive a rejection instead of hanging.
/// 7. Drain commands admitted before queue shutdown;
/// 8. Disconnect MCP servers owned by this session only.
pub(crate) async fn close(&self) {
{
let _mutation = self
.immediate_extension_mutation
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if self.closed.swap(true, Ordering::AcqRel) {
return;
}
}
// 1. Establish the queue admission boundary before the first
// cancellation await so close cannot race with a new queued command.
if let Some(queue) = &self.command_queue {
queue.shutdown().await;
}
// 2. Fire the session-level token so children cascade.
self.session_cancel.cancel();
// 3. Mark the active run cancelled and notify the hook executor. The
// per-run token has already fired via step 1.
let had_active_token = self.cancel_token.lock().await.is_some();
if had_active_token {
if let Some(run_id) = self.current_run_id.lock().await.clone() {
let _ = self.run_store.mark_cancelled(&run_id).await;
if let Some(hook) = &self.hook_executor {
hook.record_run_cancelled(&run_id, &self.session_id, Some("cancelled by host"))
.await;
}
}
}
// 4. Mark every still-running subagent task cancelled.
let pending: Vec<String> = self
.subagent_tasks
.list_for_parent(&self.session_id)
.await
.into_iter()
.filter(|task| task.status == SubagentStatus::Running)
.map(|task| task.task_id)
.collect();
for task_id in pending {
let _ = self.subagent_tasks.cancel(&task_id).await;
}
// 5. Cancel pending HITL confirmations.
if let Some(manager) = &self.confirmation_manager {
let _ = manager.cancel_all().await;
}
// 6. Let work admitted before shutdown finish after cancellation has
// propagated. A misbehaving command cannot block close indefinitely.
if let Some(queue) = &self.command_queue {
if let Err(error) = queue.drain(QUEUE_DRAIN_TIMEOUT).await {
tracing::warn!(
session_id = %self.session_id,
error = %error,
"Session lane queue did not drain before close timeout"
);
}
}
// 7. Wait for any extension mutation admitted before `closed` was
// flipped, then release only session-owned MCP resources. Removing all
// configurations (not only connected clients) also cleans up a failed
// add transaction. Agent-global and host-owned inherited managers have
// independent lifetimes.
let _extension_mutation = self.extension_mutation.lock().await;
let mut server_names: std::collections::HashSet<String> = self
.mcp_manager
.all_configs()
.await
.into_iter()
.map(|config| config.name)
.collect();
server_names.extend(self.mcp_manager.list_connected().await);
server_names.extend(
self.mcp_tool_ownership
.lock()
.unwrap_or_else(|poison| poison.into_inner())
.server_names(),
);
for name in server_names {
if let Err(error) = self.mcp_manager.remove_server(&name).await {
tracing::warn!(
session_id = %self.session_id,
server = %name,
error = %error,
"Failed to clean up session MCP transport during close"
);
}
self.mcp_tool_ownership
.lock()
.unwrap_or_else(|poison| poison.into_inner())
.remove(&name, &self.tool_executor);
}
tracing::info!(session_id = %self.session_id, "AgentSession closed");
}
pub(crate) fn mutate_immediate<T>(
&self,
mutate: impl FnOnce() -> T,
) -> crate::error::Result<T> {
let _mutation = self
.immediate_extension_mutation
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if self.closed.load(Ordering::Acquire) {
return Err(crate::error::CodeError::SessionClosed {
session_id: self.session_id.clone(),
});
}
Ok(mutate())
}
}