agentmux_srv/ambient/
mod.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Ambient Model Call (AMC) gateway — the single mandatory path for every
5//! non-user-driven ("ambient") model call in the backend: cheap, background
6//! LLM calls that augment the UX (e.g. summarizing what an agent is doing
7//! for the pane header / Swarm tree) rather than answering a user's own
8//! turn. See `docs/specs/SPEC_AMBIENT_MODEL_CALLS_FRAMEWORK_2026_07_03.md`.
9//!
10//! Every ambient call is keyed by `(entity_id, purpose)` and carries a
11//! caller-supplied generation (monotonic per entity — e.g. a turn counter).
12//! The gateway:
13//!   - cancels a lower-generation in-flight call for the same key the moment
14//!     a higher-generation one is admitted (real subprocess cancellation via
15//!     the returned `CancellationToken`, not just a discard-on-arrival check)
16//!   - rejects a request whose generation is not strictly greater than the
17//!     highest one already admitted for that key (stale-on-arrival — the
18//!     caller does no work at all, not even a discarded one)
19//!
20//! This does NOT track token usage itself — callers record usage through the
21//! normal per-call-site accounting path (tagged by purpose), same as any
22//! other RPC result. The gateway's job is coalescing/cancellation only.
23
24pub mod call;
25pub mod cli;
26pub mod digest;
27pub mod limits;
28pub mod outcome;
29pub mod prompt;
30pub mod purpose;
31pub mod reply;
32pub mod sanitize;
33pub mod spend;
34pub mod tasks;
35pub mod title;
36pub mod validate;
37
38use std::collections::HashMap;
39use std::sync::{Mutex, OnceLock};
40
41use tokio_util::sync::CancellationToken;
42
43/// Identifies one class of ambient call for a specific entity, e.g.
44/// `(block_id, "activity_summary")`. `purpose` is a short stable string —
45/// also suitable for tagging cost/usage dashboards by call category.
46#[derive(Debug, Clone, PartialEq, Eq, Hash)]
47pub struct AmbientCallKey {
48    pub entity_id: String,
49    pub purpose: &'static str,
50}
51
52impl AmbientCallKey {
53    pub fn new(entity_id: impl Into<String>, purpose: &'static str) -> Self {
54        Self { entity_id: entity_id.into(), purpose }
55    }
56}
57
58struct InFlight {
59    generation: u64,
60    cancel: CancellationToken,
61}
62
63#[derive(Default)]
64struct GatewayState {
65    inflight: HashMap<AmbientCallKey, InFlight>,
66}
67
68pub struct AmbientGateway {
69    state: Mutex<GatewayState>,
70}
71
72static GATEWAY: OnceLock<AmbientGateway> = OnceLock::new();
73
74/// The process-wide ambient-call gateway singleton.
75pub fn gateway() -> &'static AmbientGateway {
76    GATEWAY.get_or_init(|| AmbientGateway { state: Mutex::new(GatewayState::default()) })
77}
78
79/// Outcome of admitting a request for a key at a given generation.
80pub enum Admission<'a> {
81    /// This is the newest request seen for the key — proceed. Holding the
82    /// guard for the duration of the call gives access to its cancellation
83    /// token; dropping the guard (end of scope, any return path) clears the
84    /// in-flight entry unless a newer request has already replaced it.
85    Proceed(AmbientCallGuard<'a>),
86    /// A request at this generation-or-newer is already in flight, or has
87    /// already completed, for this key — do no work at all.
88    StaleOnArrival,
89}
90
91impl AmbientGateway {
92    /// Admit a request for `key` at `generation`. If an older in-flight
93    /// request exists for the same key, it is cancelled (its cancellation
94    /// token is triggered — the caller holding that token is responsible
95    /// for actually killing its subprocess). If `generation` is not
96    /// strictly greater than the generation already recorded for this key,
97    /// returns `StaleOnArrival` without mutating any state.
98    pub fn admit(&self, key: AmbientCallKey, generation: u64) -> Admission<'_> {
99        let mut state = self.state.lock().unwrap();
100        if let Some(existing) = state.inflight.get(&key) {
101            if generation <= existing.generation {
102                return Admission::StaleOnArrival;
103            }
104            existing.cancel.cancel();
105        }
106        let cancel = CancellationToken::new();
107        state.inflight.insert(
108            key.clone(),
109            InFlight { generation, cancel: cancel.clone() },
110        );
111        Admission::Proceed(AmbientCallGuard {
112            gateway: self,
113            key,
114            generation,
115            cancel,
116        })
117    }
118}
119
120/// Held for the duration of one admitted ambient call.
121pub struct AmbientCallGuard<'a> {
122    gateway: &'a AmbientGateway,
123    key: AmbientCallKey,
124    generation: u64,
125    cancel: CancellationToken,
126}
127
128impl AmbientCallGuard<'_> {
129    /// Cancellation token for this call. Race it against the subprocess
130    /// (`tokio::select!`) and kill the child if it fires — a newer request
131    /// for the same key has superseded this one.
132    pub fn cancellation(&self) -> CancellationToken {
133        self.cancel.clone()
134    }
135}
136
137impl Drop for AmbientCallGuard<'_> {
138    fn drop(&mut self) {
139        let mut state = self.gateway.state.lock().unwrap();
140        // Only clear if we're still the entry of record — a newer request
141        // may have already replaced it (and its own guard owns the clear).
142        if let Some(entry) = state.inflight.get(&self.key) {
143            if entry.generation == self.generation {
144                state.inflight.remove(&self.key);
145            }
146        }
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    fn key(id: &str) -> AmbientCallKey {
155        AmbientCallKey::new(id, "activity_summary")
156    }
157
158    #[test]
159    fn first_request_is_admitted() {
160        let gw = AmbientGateway { state: Mutex::new(GatewayState::default()) };
161        match gw.admit(key("block-1"), 1) {
162            Admission::Proceed(_) => {}
163            Admission::StaleOnArrival => panic!("first request should be admitted"),
164        };
165    }
166
167    #[test]
168    fn older_or_equal_generation_is_stale_on_arrival() {
169        let gw = AmbientGateway { state: Mutex::new(GatewayState::default()) };
170        let _guard = match gw.admit(key("block-1"), 5) {
171            Admission::Proceed(g) => g,
172            Admission::StaleOnArrival => panic!("gen 5 should be admitted"),
173        };
174        match gw.admit(key("block-1"), 5) {
175            Admission::StaleOnArrival => {}
176            Admission::Proceed(_) => panic!("equal generation must be rejected"),
177        }
178        match gw.admit(key("block-1"), 3) {
179            Admission::StaleOnArrival => {}
180            Admission::Proceed(_) => panic!("older generation must be rejected"),
181        };
182    }
183
184    #[test]
185    fn newer_generation_cancels_the_older_inflight_request() {
186        let gw = AmbientGateway { state: Mutex::new(GatewayState::default()) };
187        let guard1 = match gw.admit(key("block-1"), 1) {
188            Admission::Proceed(g) => g,
189            Admission::StaleOnArrival => panic!("gen 1 should be admitted"),
190        };
191        let cancel1 = guard1.cancellation();
192        assert!(!cancel1.is_cancelled());
193
194        let _guard2 = match gw.admit(key("block-1"), 2) {
195            Admission::Proceed(g) => g,
196            Admission::StaleOnArrival => panic!("gen 2 should be admitted"),
197        };
198        assert!(cancel1.is_cancelled(), "admitting gen 2 must cancel gen 1's token");
199    }
200
201    #[test]
202    fn different_keys_do_not_interfere() {
203        let gw = AmbientGateway { state: Mutex::new(GatewayState::default()) };
204        let guard1 = match gw.admit(key("block-1"), 1) {
205            Admission::Proceed(g) => g,
206            Admission::StaleOnArrival => panic!("block-1 gen 1 should be admitted"),
207        };
208        match gw.admit(key("block-2"), 1) {
209            Admission::Proceed(_) => {}
210            Admission::StaleOnArrival => panic!("block-2 gen 1 should be admitted independently"),
211        }
212        assert!(!guard1.cancellation().is_cancelled(), "unrelated key must not cancel block-1");
213    }
214
215    #[test]
216    fn dropping_the_guard_clears_the_slot_for_a_later_retry_at_the_same_generation() {
217        let gw = AmbientGateway { state: Mutex::new(GatewayState::default()) };
218        {
219            let _guard = match gw.admit(key("block-1"), 1) {
220                Admission::Proceed(g) => g,
221                Admission::StaleOnArrival => panic!("gen 1 should be admitted"),
222            };
223        } // guard dropped here — slot cleared
224        match gw.admit(key("block-1"), 1) {
225            Admission::Proceed(_) => {}
226            Admission::StaleOnArrival => {
227                panic!("gen 1 should be admittable again once the prior call finished")
228            }
229        };
230    }
231
232    #[test]
233    fn stale_guard_dropping_does_not_clear_a_newer_inflight_entry() {
234        let gw = AmbientGateway { state: Mutex::new(GatewayState::default()) };
235        let guard1 = match gw.admit(key("block-1"), 1) {
236            Admission::Proceed(g) => g,
237            Admission::StaleOnArrival => panic!("gen 1 should be admitted"),
238        };
239        let _guard2 = match gw.admit(key("block-1"), 2) {
240            Admission::Proceed(g) => g,
241            Admission::StaleOnArrival => panic!("gen 2 should be admitted"),
242        };
243        drop(guard1); // superseded guard drops — must NOT clear gen 2's slot
244        match gw.admit(key("block-1"), 2) {
245            Admission::StaleOnArrival => {}
246            Admission::Proceed(_) => {
247                panic!("gen 2 is still in flight; dropping the stale gen-1 guard must not clear it")
248            }
249        };
250    }
251}
agentmux_srv/ambient/mod.rs

agentmux_srv/ambient/
mod.rs
