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
//! A re-armable, per-turn interrupt handle for an [`AgentRunner`](super::AgentRunner).
//!
//! A TUI (or any caller) can ask the agent to abandon the **current turn** —
//! aborting an in-flight LLM generation — and return to awaiting the next input,
//! WITHOUT tearing down the session (conversation history is preserved). Because
//! a [`CancellationToken`] is one-shot, this handle swaps in a fresh token after
//! each interrupt so the next turn starts armed again.
//!
//! Semantics (interactive path): on interrupt the runner abandons whatever the
//! turn is doing and ends it cleanly, then waits for the next message via
//! `on_input`. An interrupt during an LLM generation aborts the stream; an
//! interrupt during a tool batch abandons it — synthesizing a result for every
//! in-flight `tool_use` (so no call is left unanswered) and killing any
//! subprocess via `kill_on_drop` — then the next LLM-call race ends the turn.
use std::sync::{Arc, Mutex};
use tokio_util::sync::CancellationToken;
/// A cloneable handle to interrupt the in-flight turn of an agent run.
#[derive(Clone)]
pub struct InterruptHandle {
inner: Arc<Mutex<CancellationToken>>,
}
impl Default for InterruptHandle {
fn default() -> Self {
Self {
inner: Arc::new(Mutex::new(CancellationToken::new())),
}
}
}
impl InterruptHandle {
/// Create a fresh, un-triggered handle.
pub fn new() -> Self {
Self::default()
}
/// Request that the agent abandon its current turn. Idempotent until the
/// runner [`rearm`](Self::rearm)s for the next turn.
pub fn interrupt(&self) {
self.inner.lock().expect("interrupt lock poisoned").cancel();
}
/// Whether the current turn's token has been triggered.
pub fn is_interrupted(&self) -> bool {
self.inner
.lock()
.expect("interrupt lock poisoned")
.is_cancelled()
}
/// A clone of the current turn's token, to race an LLM call against.
pub(crate) fn token(&self) -> CancellationToken {
self.inner.lock().expect("interrupt lock poisoned").clone()
}
/// Swap in a fresh token, arming the handle for the next turn. Called by the
/// runner right after it has handled an interrupt.
pub(crate) fn rearm(&self) {
*self.inner.lock().expect("interrupt lock poisoned") = CancellationToken::new();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn starts_un_triggered() {
assert!(!InterruptHandle::new().is_interrupted());
}
// Regression test for the live mid-tool-interrupt topology, minimized: a
// `select!` racing the per-turn token against a real SUBPROCESS, spawned on a
// JoinSet, on a SEPARATE thread with its OWN current_thread runtime (the
// hardest flavor — one scheduler thread), cancelled CROSS-THREAD, with a
// blocking sync-approval gate before the race. Proves the cancel arm wins
// (returns "cancelled") — if it ever returns "completed", cross-runtime
// interrupt over a subprocess is broken. Pairs with the end-to-end
// `runner::tests::interrupt_during_tool_batch_abandons_it_and_ends_turn`.
#[tokio::test(flavor = "multi_thread")]
async fn cross_runtime_interrupt_aborts_subprocess_tool_race() {
use std::time::Duration;
// Force the MAIN (this) runtime to register tokio::process / SIGCHLD first,
// exactly like a process that already has a primary runtime (the TUI).
let _ = tokio::process::Command::new("true").status().await;
let handle = InterruptHandle::new();
let h2 = handle.clone();
let (tx, rx) = std::sync::mpsc::channel();
let (appr_tx, appr_rx) = std::sync::mpsc::channel::<()>();
// Agent on its OWN current_thread runtime, on a separate thread (like the TUI).
std::thread::spawn(move || {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let outcome = rt.block_on(async move {
appr_rx.recv().unwrap(); // mimic the synchronous on_approval block
let token = h2.token();
tokio::select! {
biased;
_ = token.cancelled() => "cancelled",
_ = async {
let mut js = tokio::task::JoinSet::new();
js.spawn(async {
// Mirror BashTool EXACTLY: piped stdout/stderr +
// wait_with_output wrapped in tokio::time::timeout.
let child = tokio::process::Command::new("sh")
.arg("-c").arg("sleep 5")
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.spawn();
if let Ok(child) = child {
let _ = tokio::time::timeout(
Duration::from_secs(30),
child.wait_with_output(),
)
.await;
}
});
js.join_next().await;
} => "completed",
}
});
let _ = tx.send(outcome);
});
// Drive timing + the interrupt from a TASK on the MAIN runtime (like the UI).
tokio::time::sleep(Duration::from_millis(150)).await;
appr_tx.send(()).unwrap();
tokio::time::sleep(Duration::from_millis(250)).await;
handle.interrupt();
let outcome = tokio::task::spawn_blocking(move || rx.recv_timeout(Duration::from_secs(8)))
.await
.unwrap()
.expect("thread finished");
assert_eq!(
outcome, "cancelled",
"cross-thread interrupt must abort the subprocess race"
);
}
#[test]
fn interrupt_triggers_then_rearm_clears() {
let h = InterruptHandle::new();
let before = h.token();
h.interrupt();
assert!(h.is_interrupted());
assert!(
before.is_cancelled(),
"the live token reflects the interrupt"
);
h.rearm();
assert!(!h.is_interrupted(), "rearm arms a fresh token");
// The previously-handed-out token stays cancelled (it's the old turn's).
assert!(before.is_cancelled());
}
#[test]
fn clones_share_state() {
let h = InterruptHandle::new();
let clone = h.clone();
clone.interrupt();
assert!(h.is_interrupted(), "a clone interrupts the same run");
}
}