agentmux_srv/backend/pending_background_pids.rs
1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! In-memory holding pen for a declared-background task's OS pid when it
5//! arrives before `db_background_tasks` has a row to attach it to.
6//!
7//! `agentmux-bashwrap` publishes its own pid over MPS essentially at process
8//! start (before the wrapped command even begins running) — see
9//! `bash_wrap.rs`. The `db_background_tasks` row, by contrast, is only
10//! created once the FRONTEND observes the tool call as an accepted
11//! background launch (requires the literal "Command running in background
12//! with ID:" tool-result prefix) and pushes `docknodestatus`, which is a
13//! full extra round-trip through the CLI's own stream-json processing. The
14//! pid publish routinely wins that race.
15//!
16//! `background_task_set_pid` (`background_tasks.rs`) is a bare
17//! `UPDATE ... WHERE id = ?` — it silently no-ops (`Ok(false)`) when the
18//! row doesn't exist yet, and the pid would be lost forever with no retry
19//! that could ever succeed (retrying the same too-early write doesn't
20//! help — the row still won't exist).
21//!
22//! [`set_or_stash`](PendingBackgroundPids::set_or_stash) and
23//! [`observe_and_apply`](PendingBackgroundPids::observe_and_apply) are the
24//! two sides of the fix, called from `websocket.rs`'s
25//! `COMMAND_BACKGROUND_TASK_PID` and `COMMAND_DOCK_NODE_STATUS` handlers
26//! respectively. Both run their DB call (via the `try_set`/`apply` closure)
27//! WHILE HOLDING this module's own lock, not as a separate step after it —
28//! an earlier version of this module exposed plain `stash`/`take` methods
29//! instead, with each handler doing its own DB call first and the
30//! stash/take second, unlocked in between. That let the two handlers'
31//! steps interleave (pid handler's DB call sees no row → before it stashes,
32//! the observe handler's own take() already ran and found nothing →
33//! observe handler's row now exists but nothing will ever re-check it →
34//! pid handler's stash lands after the one take() that would have
35//! consumed it, and sits unclaimed until the TTL drops it). Making the DB
36//! call part of the SAME critical section as the buffer mutation closes
37//! that window: whichever handler's critical section runs second always
38//! observes the first one's effect, DB row or stash, correctly. See
39//! docs/specs/SPEC_BACKGROUND_TASK_PID_CAPTURE_2026_08_20.md and the
40//! Codex/reagentx findings on PR #2681 this was written to address.
41
42use std::collections::HashMap;
43
44use parking_lot::Mutex;
45
46/// How long a stashed pid waits for its row before being dropped. Generous
47/// relative to the observed race window (a network round-trip plus one
48/// SQLite write, normally well under a second) but still bounded — a
49/// declared-background call whose launch is never actually accepted as
50/// background client-side (e.g. the command finished fast enough that
51/// Claude's harness returned synchronously despite `run_in_background:
52/// true` — see issue #2519) would otherwise leak an entry here forever,
53/// since `background_task_observe` would then never fire for that id.
54const PENDING_PID_TTL_MS: i64 = 5 * 60 * 1000;
55
56struct PendingEntry {
57 pid: i64,
58 stashed_at_ms: i64,
59}
60
61#[derive(Default)]
62pub struct PendingBackgroundPids {
63 inner: Mutex<HashMap<String, PendingEntry>>,
64}
65
66impl PendingBackgroundPids {
67 pub fn new() -> Self {
68 Self::default()
69 }
70
71 fn evict_expired_locked(guard: &mut HashMap<String, PendingEntry>, now_ms: i64) {
72 guard.retain(|_, e| now_ms.saturating_sub(e.stashed_at_ms) < PENDING_PID_TTL_MS);
73 }
74
75 /// Called from the pid-arrival side (`COMMAND_BACKGROUND_TASK_PID`).
76 /// Runs `try_set(pid)` (the `background_task_set_pid` DB call) while
77 /// holding this module's lock; if it reports the row doesn't exist yet
78 /// (`Ok(false)`), stashes the pid instead of losing it. Atomic with
79 /// `observe_and_apply` below for the same `id` — see the module doc
80 /// comment for why that matters.
81 pub fn set_or_stash<E>(
82 &self,
83 id: &str,
84 pid: i64,
85 now_ms: i64,
86 try_set: impl FnOnce(i64) -> Result<bool, E>,
87 ) -> Result<(), E> {
88 let mut guard = self.inner.lock();
89 Self::evict_expired_locked(&mut guard, now_ms);
90 match try_set(pid)? {
91 true => {
92 guard.remove(id);
93 }
94 false => {
95 guard.insert(id.to_string(), PendingEntry { pid, stashed_at_ms: now_ms });
96 }
97 }
98 Ok(())
99 }
100
101 /// Called from the row-creation side (`COMMAND_DOCK_NODE_STATUS`),
102 /// after its own `background_task_observe(...)` call. Checks for a
103 /// pid stashed ahead of the row's existence and, if present and not
104 /// expired, applies it via `apply` (another `background_task_set_pid`
105 /// call) while holding this module's lock. Atomic with `set_or_stash`
106 /// above for the same `id`.
107 pub fn observe_and_apply<E>(
108 &self,
109 id: &str,
110 now_ms: i64,
111 apply: impl FnOnce(i64) -> Result<bool, E>,
112 ) -> Result<(), E> {
113 let mut guard = self.inner.lock();
114 Self::evict_expired_locked(&mut guard, now_ms);
115 if let Some(entry) = guard.remove(id) {
116 apply(entry.pid)?;
117 }
118 Ok(())
119 }
120}
121
122#[cfg(test)]
123mod tests {
124 use super::*;
125
126 #[derive(Debug, PartialEq)]
127 struct TestErr;
128
129 #[test]
130 fn set_or_stash_applies_directly_when_the_row_already_exists() {
131 let p = PendingBackgroundPids::new();
132 let mut set_calls = Vec::new();
133 let r: Result<(), TestErr> = p.set_or_stash("t1", 4242, 1000, |pid| {
134 set_calls.push(pid);
135 Ok(true) // row exists, update succeeded
136 });
137 assert!(r.is_ok());
138 assert_eq!(set_calls, vec![4242]);
139 // Nothing stashed — a later observe_and_apply finds nothing to apply.
140 let mut applied = Vec::new();
141 let _: Result<(), TestErr> = p.observe_and_apply("t1", 1001, |pid| {
142 applied.push(pid);
143 Ok(true)
144 });
145 assert!(applied.is_empty());
146 }
147
148 #[test]
149 fn set_or_stash_then_observe_and_apply_delivers_the_pid() {
150 // The exact race this module exists to close: pid arrives before
151 // the row does.
152 let p = PendingBackgroundPids::new();
153 let r: Result<(), TestErr> = p.set_or_stash("t1", 4242, 1000, |_pid| Ok(false)); // no row yet
154 assert!(r.is_ok());
155
156 let mut applied = Vec::new();
157 let r2: Result<(), TestErr> = p.observe_and_apply("t1", 1050, |pid| {
158 applied.push(pid);
159 Ok(true)
160 });
161 assert!(r2.is_ok());
162 assert_eq!(applied, vec![4242]);
163
164 // Consumed — a second observe_and_apply finds nothing.
165 let mut applied2 = Vec::new();
166 let _: Result<(), TestErr> = p.observe_and_apply("t1", 1051, |pid| {
167 applied2.push(pid);
168 Ok(true)
169 });
170 assert!(applied2.is_empty());
171 }
172
173 #[test]
174 fn observe_and_apply_before_any_pid_arrives_is_a_no_op() {
175 let p = PendingBackgroundPids::new();
176 let mut applied = Vec::new();
177 let _: Result<(), TestErr> = p.observe_and_apply("t1", 1000, |pid| {
178 applied.push(pid);
179 Ok(true)
180 });
181 assert!(applied.is_empty());
182 }
183
184 #[test]
185 fn set_or_stash_propagates_the_db_error_without_stashing() {
186 let p = PendingBackgroundPids::new();
187 let r: Result<(), TestErr> = p.set_or_stash("t1", 4242, 1000, |_pid| Err(TestErr));
188 assert_eq!(r, Err(TestErr));
189 // Nothing stashed despite the error.
190 let mut applied = Vec::new();
191 let _: Result<(), TestErr> = p.observe_and_apply("t1", 1001, |pid| {
192 applied.push(pid);
193 Ok(true)
194 });
195 assert!(applied.is_empty());
196 }
197
198 #[test]
199 fn observe_and_apply_propagates_the_db_error_but_still_consumes_the_stash() {
200 // Matches the module's stated semantics: this is a best-effort
201 // relay, not something worth retrying indefinitely — a stashed
202 // pid is consumed on the attempt regardless of outcome.
203 let p = PendingBackgroundPids::new();
204 let _: Result<(), TestErr> = p.set_or_stash("t1", 4242, 1000, |_pid| Ok(false));
205 let r: Result<(), TestErr> = p.observe_and_apply("t1", 1001, |_pid| Err(TestErr));
206 assert_eq!(r, Err(TestErr));
207 let mut applied = Vec::new();
208 let _: Result<(), TestErr> = p.observe_and_apply("t1", 1002, |pid| {
209 applied.push(pid);
210 Ok(true)
211 });
212 assert!(applied.is_empty());
213 }
214
215 #[test]
216 fn a_stash_past_ttl_is_not_applied() {
217 let p = PendingBackgroundPids::new();
218 let _: Result<(), TestErr> = p.set_or_stash("t1", 4242, 1000, |_pid| Ok(false));
219 let mut applied = Vec::new();
220 let _: Result<(), TestErr> = p.observe_and_apply("t1", 1000 + PENDING_PID_TTL_MS, |pid| {
221 applied.push(pid);
222 Ok(true)
223 });
224 assert!(applied.is_empty());
225 }
226
227 #[test]
228 fn set_or_stash_overwrites_a_previous_pending_pid_for_the_same_id() {
229 let p = PendingBackgroundPids::new();
230 let _: Result<(), TestErr> = p.set_or_stash("t1", 111, 1000, |_pid| Ok(false));
231 let _: Result<(), TestErr> = p.set_or_stash("t1", 222, 1001, |_pid| Ok(false));
232 let mut applied = Vec::new();
233 let _: Result<(), TestErr> = p.observe_and_apply("t1", 1002, |pid| {
234 applied.push(pid);
235 Ok(true)
236 });
237 assert_eq!(applied, vec![222]);
238 }
239
240 #[test]
241 fn different_ids_coexist_independently() {
242 let p = PendingBackgroundPids::new();
243 let _: Result<(), TestErr> = p.set_or_stash("t1", 1, 1000, |_pid| Ok(false));
244 let _: Result<(), TestErr> = p.set_or_stash("t2", 2, 1000, |_pid| Ok(false));
245 let mut applied = Vec::new();
246 let _: Result<(), TestErr> = p.observe_and_apply("t1", 1000, |pid| {
247 applied.push(pid);
248 Ok(true)
249 });
250 let _: Result<(), TestErr> = p.observe_and_apply("t2", 1000, |pid| {
251 applied.push(pid);
252 Ok(true)
253 });
254 assert_eq!(applied, vec![1, 2]);
255 }
256}