agentmux_srv/ambient/
call.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! The one path every ambient Haiku call takes: admit through the gateway,
5//! optionally wait for a concurrency permit, do the caller's own preparation,
6//! then run the CLI and judge the reply.
7//!
8//! ```ignore
9//! let slot = call::admit(&purpose::ACTIVITY_SUMMARY, block_id, generation).await?;
10//! // ... read the block, build the prompt (slot.is_superseded() is available) ...
11//! let reply = slot.run(&target, &prompt, |raw| reply::judge_line(raw, |t| validate::accept_line(t, &limits))).await;
12//! ```
13//!
14//! The slot holds the gateway guard and the permit, so every return path
15//! (including an early `?` while preparing) releases both, and a newer request
16//! for the same key still cancels this one's CLI. Every admitted call ends in
17//! exactly one recorded outcome, taken from the judge's [`Verdict`].
18
19use tokio::sync::{Semaphore, SemaphorePermit};
20use tokio_util::sync::CancellationToken;
21
22use super::limits::Class;
23use super::outcome::{Outcome, Timing};
24use super::purpose::Purpose;
25use super::reply::Verdict;
26use super::{gateway, Admission, AmbientCallGuard, AmbientCallKey};
27use crate::agents::TokenCounts;
28use crate::backend::obj::{self, MetaMapType};
29
30/// Which CLI to run, and the block meta carrying its auth env (`cmd:env`).
31pub struct CliTarget {
32    pub cli_path: String,
33    pub meta: MetaMapType,
34}
35
36impl CliTarget {
37    /// From a block's meta; `None` when it names no CLI (`cmd`).
38    pub fn from_meta(meta: &MetaMapType) -> Option<Self> {
39        let cli_path = obj::meta_get_string(meta, "cmd", "");
40        if cli_path.is_empty() {
41            return None;
42        }
43        Some(Self { cli_path, meta: meta.clone() })
44    }
45}
46
47/// What a call produced. `text` is already judged: empty means there is nothing
48/// to use (the model skipped, the reply was refused, or the call failed or was
49/// superseded). `tokens` is reported whenever the CLI reported usage, including
50/// when the text is withheld, because the spend happened either way.
51#[derive(Debug, Default)]
52pub struct Reply {
53    pub text: String,
54    pub tokens: Option<TokenCounts>,
55    /// Why the CLI call failed, if it did. Already logged at debug level.
56    pub error: Option<String>,
57}
58
59impl Reply {
60    /// `Some((text, tokens))` only when there is text to use.
61    pub fn into_option(self) -> Option<(String, Option<TokenCounts>)> {
62        if self.text.is_empty() {
63            None
64        } else {
65            Some((self.text, self.tokens))
66        }
67    }
68}
69
70/// An admitted call. Dropping it (any path) releases the gateway entry and the
71/// permit.
72pub struct Slot {
73    purpose: &'static Purpose,
74    entity_id: String,
75    cancel: CancellationToken,
76    _permit: Option<SemaphorePermit<'static>>,
77    _guard: AmbientCallGuard<'static>,
78    /// Set once this call's outcome is recorded, so `Drop` records `NotRun` only
79    /// for a call given up without one: every admitted call ends in exactly one
80    /// outcome.
81    recorded: bool,
82    /// Time spent waiting for a concurrency permit after admission.
83    queued_ms: u64,
84}
85
86impl Drop for Slot {
87    fn drop(&mut self) {
88        if !self.recorded {
89            super::outcome::record(self.purpose.tag, &self.entity_id, Outcome::NotRun, None);
90        }
91    }
92}
93
94/// Admit a call for `purpose` on `entity_id` at `generation`, in the purpose's
95/// own class (`limits`); `None` if it is stale on arrival, or was superseded while
96/// queued. It waits for a permit raced against cancellation, so a request
97/// superseded in the queue never spawns a CLI at all.
98pub async fn admit(purpose: &'static Purpose, entity_id: impl Into<String>, generation: u64) -> Option<Slot> {
99    let class = purpose.class;
100    admit_with_limit(purpose, entity_id, generation, Some(class.semaphore()), class.max_queue_wait()).await
101}
102
103/// [`admit`] in a class the caller chooses, for the naming calls: interactive
104/// for a subagent that just started or a row the user opened, background for the
105/// backlog pass over historical ones. With no queue limit: a name arrives as a
106/// `subagent:named` event, not as an RPC reply someone stops waiting for, and an
107/// eager name is claimed once (`naming_triggered`), so one that gave up in the
108/// queue would never come.
109pub async fn admit_as(
110    purpose: &'static Purpose,
111    entity_id: impl Into<String>,
112    generation: u64,
113    class: Class,
114) -> Option<Slot> {
115    admit_with_limit(purpose, entity_id, generation, Some(class.semaphore()), None).await
116}
117
118async fn admit_with_limit(
119    purpose: &'static Purpose,
120    entity_id: impl Into<String>,
121    generation: u64,
122    limit: Option<&'static Semaphore>,
123    max_wait: Option<std::time::Duration>,
124) -> Option<Slot> {
125    let entity_id = entity_id.into();
126    let guard = match gateway().admit(AmbientCallKey::new(entity_id.clone(), purpose.tag), generation) {
127        Admission::Proceed(guard) => guard,
128        Admission::StaleOnArrival => {
129            super::outcome::record(purpose.tag, &entity_id, Outcome::Superseded, None);
130            return None;
131        }
132    };
133    let cancel = guard.cancellation();
134    let admitted_at = std::time::Instant::now();
135    let permit = match limit {
136        None => None,
137        Some(sem) => {
138            let gave_up = async {
139                match max_wait {
140                    Some(wait) => tokio::time::sleep(wait).await,
141                    None => std::future::pending().await,
142                }
143            };
144            let permit = tokio::select! {
145                biased;
146                // Superseded while queued for a permit: never spawned.
147                _ = cancel.cancelled() => Err(Outcome::Superseded),
148                permit = sem.acquire() => permit.map_err(|_| Outcome::Superseded),
149                // Queued too long: whoever asked has stopped waiting (`limits`).
150                _ = gave_up => Err(Outcome::NotRun),
151            };
152            match permit {
153                Ok(permit) => Some(permit),
154                Err(outcome) => {
155                    super::outcome::record(purpose.tag, &entity_id, outcome, None);
156                    return None;
157                }
158            }
159        }
160    };
161    let queued_ms = admitted_at.elapsed().as_millis() as u64;
162    Some(Slot { purpose, entity_id, cancel, _permit: permit, _guard: guard, recorded: false, queued_ms })
163}
164
165impl Slot {
166    /// Whether a newer request for the same key has been admitted since this one.
167    pub fn is_superseded(&self) -> bool {
168        self.cancel.is_cancelled()
169    }
170
171    /// Give the call up without running it, recording why (for example
172    /// `EmptyDigest`), instead of the `NotRun` a plain drop records.
173    pub fn abandon(mut self, outcome: Outcome) {
174        super::outcome::record(self.purpose.tag, &self.entity_id, outcome, None);
175        self.recorded = true;
176    }
177
178    /// Run the CLI with `prompt`, within the purpose's time limit, and let `judge`
179    /// decide what the reply is worth (`reply::judge_line` for a one-line reply in
180    /// the `ANSWER:`/`SKIP` format). Never fails: a CLI error, a cancellation, a
181    /// skip and a refused reply all come back as a `Reply` with empty text.
182    pub async fn run(mut self, target: &CliTarget, prompt: &str, judge: impl Fn(&str) -> Verdict) -> Reply {
183        self.run_held(target, prompt, judge).await
184    }
185
186    /// [`Slot::run`], keeping the slot, and with it the gateway's hold on this
187    /// call's generation, for a caller that acts on the reply and must not if a
188    /// newer request arrived meanwhile ([`Slot::is_superseded`]).
189    pub async fn run_held(&mut self, target: &CliTarget, prompt: &str, judge: impl Fn(&str) -> Verdict) -> Reply {
190        let started = std::time::Instant::now();
191        let result = super::cli::invoke_haiku(
192            &target.cli_path,
193            prompt,
194            &target.meta,
195            self.cancel.clone(),
196            self.purpose.timeout,
197        )
198        .await;
199        let timing = Timing { queued_ms: self.queued_ms, run_ms: started.elapsed().as_millis() as u64 };
200        self.recorded = true;
201        match result {
202            Err(error) => {
203                let outcome = super::outcome::classify_error(&error, self.cancel.is_cancelled());
204                tracing::debug!(purpose = self.purpose.tag, entity = %self.entity_id, error = %error, "ambient call failed");
205                super::outcome::record_timed(self.purpose.tag, &self.entity_id, outcome, None, Some(timing));
206                Reply { text: String::new(), tokens: None, error: Some(error) }
207            }
208            Ok((raw, tokens)) => {
209                let (text, outcome) = match judge(&raw) {
210                    Verdict::Use(text) => (text, Outcome::Accepted),
211                    Verdict::Skip => (String::new(), Outcome::Skipped),
212                    Verdict::Reject(reason) => (String::new(), Outcome::Rejected(reason)),
213                };
214                let logged = self.purpose.logs_reply.then_some(raw.as_str());
215                super::outcome::record_timed(self.purpose.tag, &self.entity_id, outcome, logged, Some(timing));
216                // Spent whatever the verdict, so counted whatever the verdict.
217                if let Some(tokens) = &tokens {
218                    super::spend::report(self.purpose.tag, tokens);
219                }
220                Reply { text, tokens, error: None }
221            }
222        }
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use std::time::Duration;
230
231    /// Test purposes, each its own tag so parallel tests never share a gateway
232    /// key or an outcome counter.
233    macro_rules! test_purpose {
234        ($name:ident, $tag:literal) => {
235            static $name: Purpose =
236                Purpose { tag: $tag, timeout: Duration::from_secs(1), class: Class::Interactive, logs_reply: true };
237        };
238    }
239    test_purpose!(OUTCOMES, "test_purpose_slot_outcomes");
240    test_purpose!(STALE, "t_stale");
241    test_purpose!(NEWER, "t_newer");
242    test_purpose!(SIBLINGS, "t_sib");
243    test_purpose!(PERMIT, "t_permit");
244    test_purpose!(QUEUE, "t_queue");
245    test_purpose!(GIVE_UP, "t_give_up");
246    test_purpose!(REDO, "t_redo");
247
248    fn leaked_semaphore(permits: usize) -> &'static Semaphore {
249        Box::leak(Box::new(Semaphore::new(permits)))
250    }
251
252    /// An admitted call given up before running (no block, no CLI path) still ends
253    /// in one outcome, and an explicit `abandon` records its own outcome instead of
254    /// `not_run`, never both.
255    #[tokio::test]
256    async fn every_admitted_call_ends_in_exactly_one_outcome() {
257        let count = |label: &str| {
258            crate::ambient::outcome::snapshot()
259                .get(OUTCOMES.tag)
260                .and_then(|by| by.get(label).copied())
261                .unwrap_or(0)
262        };
263        let dropped = admit_with_limit(&OUTCOMES, "e1", 1, None, None).await.unwrap();
264        drop(dropped);
265        assert_eq!(count("not_run"), 1);
266
267        let abandoned = admit_with_limit(&OUTCOMES, "e2", 1, None, None).await.unwrap();
268        abandoned.abandon(Outcome::EmptyDigest);
269        assert_eq!(count("empty_digest"), 1);
270        assert_eq!(count("not_run"), 1, "abandon records once, not twice");
271    }
272
273    #[tokio::test]
274    async fn a_stale_request_is_not_admitted() {
275        let _newer = admit_with_limit(&STALE, "call-stale", 5, None, None).await.expect("first is admitted");
276        assert!(admit_with_limit(&STALE, "call-stale", 5, None, None).await.is_none(), "same generation");
277        assert!(admit_with_limit(&STALE, "call-stale", 4, None, None).await.is_none(), "older generation");
278    }
279
280    #[tokio::test]
281    async fn a_newer_request_cancels_the_older_ones_cli() {
282        let first = admit_with_limit(&NEWER, "call-newer", 1, None, None).await.unwrap();
283        assert!(!first.is_superseded());
284        let _second = admit_with_limit(&NEWER, "call-newer", 2, None, None).await.unwrap();
285        // A slot still held (`run_held`) can tell, and so not act on its reply.
286        assert!(first.is_superseded());
287    }
288
289    #[tokio::test]
290    async fn different_entities_do_not_cancel_each_other() {
291        let first = admit_with_limit(&SIBLINGS, "call-a", 1, None, None).await.unwrap();
292        let _second = admit_with_limit(&SIBLINGS, "call-b", 1, None, None).await.unwrap();
293        assert!(!first.is_superseded());
294    }
295
296    #[tokio::test]
297    async fn the_permit_is_held_until_the_slot_drops() {
298        let sem = leaked_semaphore(1);
299        let slot = admit_with_limit(&PERMIT, "call-permit", 1, Some(sem), None).await.unwrap();
300        assert_eq!(sem.available_permits(), 0);
301        drop(slot);
302        assert_eq!(sem.available_permits(), 1);
303    }
304
305    #[tokio::test]
306    async fn a_call_queued_past_its_limit_gives_up_without_running() {
307        let sem = leaked_semaphore(1);
308        let holder = admit_with_limit(&GIVE_UP, "call-hold", 1, Some(sem), None).await.unwrap();
309        let wait = std::time::Duration::from_millis(30);
310        let queued = admit_with_limit(&GIVE_UP, "call-wait", 1, Some(sem), Some(wait)).await;
311        assert!(queued.is_none(), "gave up in the queue");
312        drop(holder);
313        assert_eq!(sem.available_permits(), 1, "no permit leaked");
314        // With a permit free, the same limit doesn't get in the way.
315        assert!(admit_with_limit(&GIVE_UP, "call-wait", 2, Some(sem), Some(wait)).await.is_some());
316    }
317
318    #[tokio::test]
319    async fn a_request_superseded_while_queued_never_gets_a_permit() {
320        let sem = leaked_semaphore(1);
321        let holder = admit_with_limit(&QUEUE, "call-holder", 1, Some(sem), None).await.unwrap();
322
323        // Queues behind the holder (a different entity, so the holder is not cancelled).
324        let queued = tokio::spawn(admit_with_limit(&QUEUE, "call-queued", 1, Some(sem), None));
325        tokio::task::yield_now().await;
326        // A newer request for the queued key supersedes it while it waits.
327        let newer = admit_with_limit(&QUEUE, "call-queued", 2, None, None).await.unwrap();
328
329        assert!(queued.await.unwrap().is_none(), "superseded in the queue");
330        drop(holder);
331        drop(newer);
332        assert_eq!(sem.available_permits(), 1, "no permit leaked");
333    }
334
335    #[tokio::test]
336    async fn dropping_the_slot_lets_the_same_generation_be_admitted_again_only_if_newer() {
337        let slot = admit_with_limit(&REDO, "call-redo", 3, None, None).await.unwrap();
338        drop(slot);
339        // The gateway clears the in-flight entry on drop, so the next request admits.
340        assert!(admit_with_limit(&REDO, "call-redo", 3, None, None).await.is_some());
341    }
342
343    #[test]
344    fn a_reply_with_no_text_has_no_option() {
345        assert!(Reply::default().into_option().is_none());
346        let r = Reply { text: "x".into(), tokens: None, error: None };
347        assert_eq!(r.into_option().map(|(t, _)| t), Some("x".to_string()));
348    }
349
350    #[test]
351    fn a_target_needs_a_cmd() {
352        let mut meta = MetaMapType::new();
353        assert!(CliTarget::from_meta(&meta).is_none());
354        meta.insert("cmd".into(), serde_json::json!("/usr/bin/claude"));
355        assert_eq!(CliTarget::from_meta(&meta).unwrap().cli_path, "/usr/bin/claude");
356    }
357}
agentmux_srv/ambient/call.rs

agentmux_srv/ambient/
call.rs
