Skip to main content

everruns_engine/
native_async.rs

1//! Opt-in native-call coordinator for streaming hosts. Normal Reason/Act hosts
2//! keep synchronous execution unless they install this coordinator explicitly.
3
4use async_trait::async_trait;
5use everruns_provider::{
6    LlmResponseStream, LlmStreamEvent,
7    error::{AgentLoopError, Result},
8    native_async::{Delivery, NativeAsyncCheckpoint, NativeToolCall, PendingCallState},
9};
10use futures::{StreamExt, stream::FuturesUnordered};
11use std::{collections::HashMap, sync::Arc};
12use tokio::sync::{Mutex, Semaphore};
13
14/// An exclusive, durable conversation journal. Keep the ownership fence for the
15/// lifetime of the coordinator, including provider requests between pump calls.
16#[async_trait]
17pub trait NativeAsyncJournal: Send + Sync {
18    async fn load(&self) -> Result<NativeAsyncCheckpoint>;
19    /// Atomically persist and flush before returning success.
20    async fn save(&self, checkpoint: &NativeAsyncCheckpoint) -> Result<()>;
21    /// Renew shared ownership while streams/jobs are quiet. Failure stops jobs.
22    async fn heartbeat(&self) -> Result<()> {
23        Ok(())
24    }
25    async fn release(&self) -> Result<()> {
26        Ok(())
27    }
28}
29
30#[derive(Debug, Clone)]
31pub struct NativeCallPolicy {
32    pub allow_async: bool,
33    pub replay_safe: bool,
34    pub concurrency_class: Option<String>,
35}
36
37/// Hosts must use their ordinary tool authorization, argument validation,
38/// pre/post execution hooks, resource scoping and outbound limits here. Native
39/// metadata from the provider never grants permission to execute a tool.
40#[async_trait]
41pub trait NativeAsyncExecutor: Send + Sync + 'static {
42    async fn authorize(&self, call: &NativeToolCall) -> Result<NativeCallPolicy>;
43    async fn execute(&self, call: NativeToolCall) -> Result<String>;
44}
45
46// Tool RPCs must progress while journal writes await the same transport lock.
47// Ownership remains local: dropping the coordinator aborts every running task.
48struct OwnedJob {
49    id: String,
50    handle: tokio::task::JoinHandle<(String, Result<String>)>,
51}
52impl Drop for OwnedJob {
53    fn drop(&mut self) {
54        self.handle.abort();
55    }
56}
57impl std::future::Future for OwnedJob {
58    type Output = (String, Result<String>);
59    fn poll(
60        mut self: std::pin::Pin<&mut Self>,
61        context: &mut std::task::Context<'_>,
62    ) -> std::task::Poll<Self::Output> {
63        match std::pin::Pin::new(&mut self.handle).poll(context) {
64            std::task::Poll::Ready(Ok(result)) => std::task::Poll::Ready(result),
65            std::task::Poll::Ready(Err(_)) => {
66                std::task::Poll::Ready((self.id.clone(), Err(AgentLoopError::Cancelled)))
67            }
68            std::task::Poll::Pending => std::task::Poll::Pending,
69        }
70    }
71}
72
73pub struct NativeAsyncCoordinator {
74    journal: Box<dyn NativeAsyncJournal>,
75    executor: Arc<dyn NativeAsyncExecutor>,
76    checkpoint: NativeAsyncCheckpoint,
77    jobs: FuturesUnordered<OwnedJob>,
78    permits: Arc<Semaphore>,
79    classes: HashMap<String, Arc<Mutex<()>>>,
80    serialize_all: bool,
81    poisoned: bool,
82    last_heartbeat: tokio::time::Instant,
83}
84
85impl NativeAsyncCoordinator {
86    /// Recover only after acquiring the journal's exclusive ownership fence.
87    /// Safe calls are reauthorized and restarted; unsafe calls become explicit
88    /// interrupted outputs. Ambiguous HTTP delivery requires receipt recovery.
89    pub async fn open(
90        journal: Box<dyn NativeAsyncJournal>,
91        executor: Arc<dyn NativeAsyncExecutor>,
92        max_concurrency: usize,
93        parallel_tool_calls: bool,
94    ) -> Result<Self> {
95        let mut checkpoint = journal.load().await?;
96        checkpoint.recover()?;
97        journal.save(&checkpoint).await?;
98        let mut this = Self {
99            journal,
100            executor,
101            checkpoint,
102            jobs: FuturesUnordered::new(),
103            permits: Arc::new(Semaphore::new(max_concurrency.max(1))),
104            classes: HashMap::new(),
105            serialize_all: !parallel_tool_calls,
106            poisoned: false,
107            last_heartbeat: tokio::time::Instant::now(),
108        };
109        let queued: Vec<_> = this
110            .checkpoint
111            .order
112            .iter()
113            .filter_map(|id| this.checkpoint.calls.get(id))
114            .filter(|pending| pending.state == PendingCallState::Queued)
115            .map(|pending| pending.call.clone())
116            .collect();
117        for call in queued {
118            let policy = this.executor.authorize(&call).await?;
119            this.launch(call, policy).await?;
120        }
121        Ok(this)
122    }
123
124    pub fn checkpoint(&self) -> &NativeAsyncCheckpoint {
125        &self.checkpoint
126    }
127
128    fn healthy(&self) -> Result<()> {
129        if self.poisoned {
130            Err(AgentLoopError::store(
131                "native coordinator lost durable state; reopen under a fresh ownership fence",
132            ))
133        } else {
134            Ok(())
135        }
136    }
137
138    async fn save(&mut self) -> Result<()> {
139        match self.journal.save(&self.checkpoint).await {
140            Ok(()) => {
141                self.last_heartbeat = tokio::time::Instant::now();
142                Ok(())
143            }
144            Err(error) => {
145                self.poisoned = true;
146                self.jobs.clear();
147                Err(error)
148            }
149        }
150    }
151
152    async fn heartbeat(&mut self) -> Result<()> {
153        if let Err(error) = self.journal.heartbeat().await {
154            self.poisoned = true;
155            self.jobs.clear();
156            return Err(error);
157        }
158        self.last_heartbeat = tokio::time::Instant::now();
159        Ok(())
160    }
161
162    pub async fn persist_host_outcome(&mut self, outcome: serde_json::Value) -> Result<()> {
163        self.healthy()?;
164        if !self.checkpoint.can_complete() {
165            return Err(AgentLoopError::store("native outputs remain pending"));
166        }
167        self.checkpoint.host_outcome = Some(outcome);
168        self.save().await
169    }
170
171    pub async fn release(&mut self) -> Result<()> {
172        self.healthy()?;
173        if !self.checkpoint.can_complete() {
174            return Err(AgentLoopError::store(
175                "cannot release native conversation with pending outputs",
176            ));
177        }
178        self.journal.release().await?;
179        self.poisoned = true;
180        Ok(())
181    }
182
183    async fn launch(&mut self, call: NativeToolCall, policy: NativeCallPolicy) -> Result<()> {
184        if call.is_async() && !policy.allow_async {
185            return Err(AgentLoopError::config(
186                "tool is not authorized for native asynchronous execution",
187            ));
188        }
189        if self
190            .checkpoint
191            .calls
192            .get(call.id())
193            .is_some_and(|pending| pending.replay_safe != policy.replay_safe)
194        {
195            return Err(AgentLoopError::config(
196                "native replay policy changed; explicit reconciliation required",
197            ));
198        }
199        self.checkpoint.start(call.id())?;
200        self.save().await?;
201        let class = if self.serialize_all {
202            Some(String::new())
203        } else {
204            policy.concurrency_class
205        };
206        let lock = class.map(|class| {
207            self.classes
208                .entry(class)
209                .or_insert_with(|| Arc::new(Mutex::new(())))
210                .clone()
211        });
212        let permits = self.permits.clone();
213        let executor = self.executor.clone();
214        let call_id = call.id().to_owned();
215        let handle = tokio::spawn(async move {
216            let _class_guard = match lock {
217                Some(lock) => Some(lock.lock_owned().await),
218                None => None,
219            };
220            // Unreachable: the coordinator owns these permits and never
221            // closes them. Kept as a panic so a closed semaphore cannot
222            // silently lift the native-tool concurrency cap.
223            #[expect(
224                clippy::expect_used,
225                reason = "fail closed rather than lose the concurrency bound"
226            )]
227            let _permit = permits
228                .acquire()
229                .await
230                .expect("coordinator never closes permits");
231            let id = call.id().to_owned();
232            (id, executor.execute(call).await)
233        });
234        self.jobs.push(OwnedJob {
235            id: call_id,
236            handle,
237        });
238        Ok(())
239    }
240
241    pub async fn register(&mut self, call: NativeToolCall) -> Result<()> {
242        self.healthy()?;
243        let policy = self.executor.authorize(&call).await?;
244        if call.is_async() && !policy.allow_async {
245            return Err(AgentLoopError::config(
246                "tool is not authorized for native asynchronous execution",
247            ));
248        }
249        if !self.checkpoint.register(call.clone(), policy.replay_safe)? {
250            return Ok(());
251        }
252        self.save().await?;
253        // Only explicit async calls may run before the response is accepted.
254        if self.checkpoint.response_in_flight && !call.is_async() {
255            return Ok(());
256        }
257        self.launch(call, policy).await
258    }
259
260    async fn launch_synchronous_calls(&mut self) -> Result<()> {
261        let queued: Vec<_> = self
262            .checkpoint
263            .order
264            .iter()
265            .filter_map(|id| self.checkpoint.calls.get(id))
266            .filter(|pending| !pending.call.is_async() && pending.state == PendingCallState::Queued)
267            .map(|pending| pending.call.clone())
268            .collect();
269        for call in queued {
270            let policy = self.executor.authorize(&call).await?;
271            self.launch(call, policy).await?;
272        }
273        Ok(())
274    }
275
276    async fn settle(&mut self, id: String, result: Result<String>) -> Result<()> {
277        // Executors own disclosure: errors must already be safe for the model.
278        let output = result
279            .unwrap_or_else(|error| serde_json::json!({"error":error.to_string()}).to_string());
280        self.checkpoint.settle(&id, output)?;
281        self.save().await
282    }
283
284    pub async fn begin_transcript_response(&mut self, message_id: String) -> Result<()> {
285        self.healthy()?;
286        if self.checkpoint.transcript_message_id.is_some() {
287            return Err(AgentLoopError::store("native transcript is not committed"));
288        }
289        self.checkpoint.transcript_message_id = Some(message_id);
290        self.begin_response().await
291    }
292
293    pub async fn stage_transcript_result(&mut self, result: serde_json::Value) -> Result<()> {
294        self.healthy()?;
295        if self.checkpoint.response_in_flight || self.checkpoint.transcript_message_id.is_none() {
296            return Err(AgentLoopError::store(
297                "native response is not ready for transcript commit",
298            ));
299        }
300        if self.checkpoint.host_responses.len() + 1 != self.checkpoint.completed_responses as usize
301        {
302            return Err(AgentLoopError::store(
303                "native response summary count mismatch",
304            ));
305        }
306        self.checkpoint.host_responses.push(result);
307        self.save().await
308    }
309
310    pub async fn transcript_committed(&mut self, message_id: &str) -> Result<()> {
311        self.healthy()?;
312        if self.checkpoint.response_in_flight
313            || self.checkpoint.transcript_message_id.as_deref() != Some(message_id)
314        {
315            return Err(AgentLoopError::store("native transcript boundary mismatch"));
316        }
317        self.checkpoint.transcript_message_id = None;
318        self.save().await
319    }
320
321    /// Record request intent before a host opens its provider HTTP stream.
322    pub async fn begin_response(&mut self) -> Result<()> {
323        self.healthy()?;
324        if self.checkpoint.response_in_flight {
325            return Err(AgentLoopError::store(
326                "prior native response requires reconciliation",
327            ));
328        }
329        self.checkpoint.response_in_flight = true;
330        self.save().await
331    }
332
333    /// Drive jobs alongside one stream event, retaining call events for the
334    /// host's normal transcript pipeline. A completed response is not a turn end.
335    pub async fn next_response_event(
336        &mut self,
337        stream: &mut LlmResponseStream,
338    ) -> Result<LlmStreamEvent> {
339        self.healthy()?;
340        let result = self.next_response_event_inner(stream).await;
341        if result.is_err() {
342            self.poisoned = true;
343            self.jobs.clear();
344        }
345        result
346    }
347
348    async fn next_response_event_inner(
349        &mut self,
350        stream: &mut LlmResponseStream,
351    ) -> Result<LlmStreamEvent> {
352        if !self.checkpoint.response_in_flight {
353            return Err(AgentLoopError::store(
354                "native response has no persisted request intent",
355            ));
356        }
357        loop {
358            tokio::select! {
359                _ = tokio::time::sleep_until(self.last_heartbeat + std::time::Duration::from_secs(10)) => self.heartbeat().await?,
360                completed = self.jobs.next(), if !self.jobs.is_empty() => {
361                    if let Some((id, result)) = completed {
362                        self.settle(id, result).await?;
363                    }
364                }
365                event = stream.next() => {
366                    let event = event.ok_or_else(|| AgentLoopError::llm("native response stream ended before completion"))??;
367                    match &event {
368                        LlmStreamEvent::NativeToolCall(call) => self.register(call.clone()).await?,
369                        LlmStreamEvent::ToolCalls(calls) => for call in calls {
370                            self.register(NativeToolCall::Function { call_id: call.id.clone(), name: call.name.clone(), arguments: serde_json::to_string(&call.arguments).map_err(|error| AgentLoopError::config(error.to_string()))?, asynchronous: false }).await?;
371                        },
372                        LlmStreamEvent::Done(metadata) => {
373                            let id = metadata.response_id.clone().ok_or_else(|| AgentLoopError::llm("native response omitted response ID"))?;
374                            if metadata.finish_reason.as_deref().is_some_and(|reason| !matches!(reason, "stop" | "tool_calls" | "end_turn")) { return Err(AgentLoopError::llm("native response did not finish successfully")); }
375                            if self.checkpoint.delivery.is_some() { self.checkpoint.acknowledge_delivery(id)?; } else { self.checkpoint.response_completed(id)?; }
376                            self.checkpoint.response_in_flight = false;
377                            self.save().await?;
378                            self.launch_synchronous_calls().await?;
379                        }
380                        LlmStreamEvent::Error(error) => return Err(AgentLoopError::llm(error.to_string())),
381                        _ => {}
382                    }
383                    return Ok(event);
384                }
385            }
386        }
387    }
388
389    /// Consume a response while jobs execute. Returns once the provider response
390    /// is complete, even if async work remains. Independent follow-up responses
391    /// may be pumped before waiting for jobs. `observe` receives prose/reasoning
392    /// unchanged; only executable call events are consumed by the coordinator.
393    pub async fn pump(
394        &mut self,
395        stream: LlmResponseStream,
396        mut observe: impl FnMut(LlmStreamEvent),
397    ) -> Result<()> {
398        self.healthy()?;
399        if self.checkpoint.response_in_flight {
400            return Err(AgentLoopError::store(
401                "prior native response requires reconciliation",
402            ));
403        }
404        let result = self.pump_inner(stream, &mut observe).await;
405        if result.is_err() {
406            self.poisoned = true;
407            self.jobs.clear();
408        }
409        result
410    }
411
412    async fn pump_inner(
413        &mut self,
414        mut stream: LlmResponseStream,
415        mut observe: impl FnMut(LlmStreamEvent),
416    ) -> Result<()> {
417        self.begin_response().await?;
418        loop {
419            let event = self.next_response_event(&mut stream).await?;
420            match event {
421                LlmStreamEvent::NativeToolCall(_) | LlmStreamEvent::ToolCalls(_) => {}
422                LlmStreamEvent::Done(_) => {
423                    observe(event);
424                    return Ok(());
425                }
426                other => observe(other),
427            }
428        }
429    }
430
431    /// Wait for one completion; outputs may be delivered out of launch order.
432    pub async fn wait_next(&mut self) -> Result<bool> {
433        self.healthy()?;
434        let mut lease_tick = tokio::time::interval(std::time::Duration::from_secs(10));
435        loop {
436            tokio::select! {
437                _ = lease_tick.tick() => self.heartbeat().await?,
438                completed = self.jobs.next() => {
439                    if let Some((id, result)) = completed {
440                        self.settle(id, result).await?;
441                        return Ok(true);
442                    }
443                    return Ok(false);
444                }
445            }
446        }
447    }
448
449    /// Persist the delivery intent before the caller submits its HTTP request.
450    /// Synchronous calls must all finish before the provider can continue.
451    pub async fn prepare_delivery(&mut self) -> Result<Option<Delivery>> {
452        self.healthy()?;
453        while self.checkpoint.calls.values().any(|pending| {
454            !pending.call.is_async()
455                && matches!(
456                    pending.state,
457                    PendingCallState::Queued | PendingCallState::Running
458                )
459        }) {
460            if !self.wait_next().await? {
461                return Err(AgentLoopError::store("synchronous call has no running job"));
462            }
463        }
464        let delivery = self.checkpoint.prepare_delivery()?.cloned();
465        self.save().await?;
466        Ok(delivery)
467    }
468
469    /// Dropping owned futures cancels local execution. Persist cancellation
470    /// outputs; the conversation remains incomplete until they are delivered.
471    pub async fn cancel(&mut self) -> Result<()> {
472        self.healthy()?;
473        for job in self.jobs.iter() {
474            job.handle.abort();
475        }
476        while let Some((id, result)) = self.jobs.next().await {
477            if !matches!(result, Err(AgentLoopError::Cancelled)) {
478                let output = result.unwrap_or_else(|error| {
479                    serde_json::json!({"error":error.to_string()}).to_string()
480                });
481                self.checkpoint.settle(&id, output)?;
482            }
483        }
484        self.checkpoint.cancel();
485        self.save().await
486    }
487}
488
489impl NativeAsyncCoordinator {
490    /// Drive HTTP response continuations through completion, returning only when
491    /// every accepted call's output has a provider receipt. The caller supplies
492    /// request construction; no transport or model defaults are chosen here.
493    pub async fn run<Request, RequestFuture>(
494        &mut self,
495        max_responses: usize,
496        mut request: Request,
497        mut observe: impl FnMut(LlmStreamEvent),
498    ) -> Result<()>
499    where
500        Request: FnMut(Option<Delivery>, Option<String>) -> RequestFuture,
501        RequestFuture: std::future::Future<Output = Result<LlmResponseStream>>,
502    {
503        self.healthy()?;
504        for _ in 0..max_responses {
505            let mut delivery = self.prepare_delivery().await?;
506            // On a restart with a completed response, resume its pending jobs
507            // before making a continuation request with no new input.
508            if delivery.is_none()
509                && self.checkpoint.latest_response_id.is_some()
510                && !self.jobs.is_empty()
511            {
512                self.wait_next().await?;
513                delivery = self.prepare_delivery().await?;
514            }
515            let response = request(delivery, self.checkpoint.latest_response_id.clone());
516            tokio::pin!(response);
517            let mut lease_tick = tokio::time::interval(std::time::Duration::from_secs(10));
518            let stream = loop {
519                tokio::select! {
520                    result = &mut response => break result?,
521                    _ = lease_tick.tick() => self.heartbeat().await?,
522                }
523            };
524            self.pump(stream, &mut observe).await?;
525            if self.checkpoint.can_complete() {
526                return Ok(());
527            }
528            if !self
529                .checkpoint
530                .calls
531                .values()
532                .any(|pending| matches!(pending.state, PendingCallState::Ready { .. }))
533            {
534                self.wait_next().await?;
535            }
536        }
537        Err(AgentLoopError::config(
538            "native async response limit reached; pending work remains checkpointed",
539        ))
540    }
541}
542
543#[cfg(test)]
544mod tests {
545    use super::*;
546    use everruns_provider::LlmCompletionMetadata;
547    use std::{
548        sync::{
549            Mutex as StdMutex,
550            atomic::{AtomicUsize, Ordering},
551        },
552        time::Duration,
553    };
554
555    #[derive(Clone, Default)]
556    struct MemoryJournal(Arc<StdMutex<NativeAsyncCheckpoint>>);
557    #[async_trait]
558    impl NativeAsyncJournal for MemoryJournal {
559        async fn load(&self) -> Result<NativeAsyncCheckpoint> {
560            Ok(self.0.lock().unwrap().clone())
561        }
562        async fn save(&self, checkpoint: &NativeAsyncCheckpoint) -> Result<()> {
563            *self.0.lock().unwrap() = checkpoint.clone();
564            Ok(())
565        }
566    }
567    #[derive(Default)]
568    struct Executor {
569        started: AtomicUsize,
570        active: AtomicUsize,
571        maximum: AtomicUsize,
572    }
573    #[async_trait]
574    impl NativeAsyncExecutor for Executor {
575        async fn authorize(&self, call: &NativeToolCall) -> Result<NativeCallPolicy> {
576            if call.name() == "forbidden" {
577                return Err(AgentLoopError::tool("not authorized"));
578            }
579            Ok(NativeCallPolicy {
580                allow_async: true,
581                replay_safe: true,
582                concurrency_class: None,
583            })
584        }
585        async fn execute(&self, call: NativeToolCall) -> Result<String> {
586            self.started.fetch_add(1, Ordering::SeqCst);
587            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
588            self.maximum.fetch_max(active, Ordering::SeqCst);
589            struct Active<'a>(&'a AtomicUsize);
590            impl Drop for Active<'_> {
591                fn drop(&mut self) {
592                    self.0.fetch_sub(1, Ordering::SeqCst);
593                }
594            }
595            let _active = Active(&self.active);
596            tokio::time::sleep(Duration::from_millis(if call.id() == "slow" {
597                100
598            } else {
599                1
600            }))
601            .await;
602            Ok(call.id().into())
603        }
604    }
605    fn call(id: &str, asynchronous: bool) -> NativeToolCall {
606        NativeToolCall::Function {
607            call_id: id.into(),
608            name: "lookup".into(),
609            arguments: "{}".into(),
610            asynchronous,
611        }
612    }
613    fn done(id: &str) -> LlmStreamEvent {
614        LlmStreamEvent::Done(Box::new({
615            let mut metadata = LlmCompletionMetadata::default();
616            metadata.response_id = Some(id.into());
617            metadata
618        }))
619    }
620    fn stream(events: Vec<LlmStreamEvent>) -> LlmResponseStream {
621        Box::pin(futures::stream::iter(events.into_iter().map(Ok)))
622    }
623
624    #[tokio::test]
625    async fn synchronous_calls_wait_for_successful_response_completion() {
626        for rejected in [false, true] {
627            let executor = Arc::new(Executor::default());
628            let mut coordinator = NativeAsyncCoordinator::open(
629                Box::new(MemoryJournal::default()),
630                executor.clone(),
631                2,
632                true,
633            )
634            .await
635            .unwrap();
636            coordinator.begin_response().await.unwrap();
637            let mut terminal = done("response");
638            if rejected && let LlmStreamEvent::Done(metadata) = &mut terminal {
639                metadata.finish_reason = Some("length".into());
640            }
641            let mut response = stream(vec![
642                LlmStreamEvent::NativeToolCall(call("sync", false)),
643                terminal,
644            ]);
645            coordinator
646                .next_response_event(&mut response)
647                .await
648                .unwrap();
649            assert_eq!(
650                coordinator.checkpoint().calls["sync"].state,
651                PendingCallState::Queued,
652                "synchronous calls must not gain early execution from native opt-in"
653            );
654            let result = coordinator.next_response_event(&mut response).await;
655            if rejected {
656                assert!(result.is_err());
657                assert!(
658                    coordinator.checkpoint().clone().recover().is_err(),
659                    "recovery must not execute rejected calls"
660                );
661                assert_eq!(executor.started.load(Ordering::SeqCst), 0);
662            } else {
663                result.unwrap();
664                assert!(coordinator.wait_next().await.unwrap());
665                assert_eq!(executor.started.load(Ordering::SeqCst), 1);
666            }
667        }
668    }
669
670    #[tokio::test]
671    async fn early_dispatch_mixed_calls_out_of_order_and_independent_work() {
672        let journal = MemoryJournal::default();
673        let executor = Arc::new(Executor::default());
674        let mut coordinator =
675            NativeAsyncCoordinator::open(Box::new(journal.clone()), executor.clone(), 2, true)
676                .await
677                .unwrap();
678        let (sender, receiver) = futures::channel::mpsc::unbounded();
679        sender
680            .unbounded_send(Ok(LlmStreamEvent::NativeToolCall(call("slow", true))))
681            .unwrap();
682        let observed = executor.clone();
683        let producer = async move {
684            // The stream is not complete yet; the tool must already be running.
685            while observed.started.load(Ordering::SeqCst) == 0 {
686                tokio::task::yield_now().await;
687            }
688            sender
689                .unbounded_send(Ok(LlmStreamEvent::NativeToolCall(call("fast", true))))
690                .unwrap();
691            sender
692                .unbounded_send(Ok(LlmStreamEvent::TextDelta("independent answer".into())))
693                .unwrap();
694            sender.unbounded_send(Ok(done("launch"))).unwrap();
695        };
696        let mut text = String::new();
697        let (result, _) = tokio::join!(
698            coordinator.pump(Box::pin(receiver), |event| {
699                if let LlmStreamEvent::TextDelta(delta) = event {
700                    text.push_str(&delta);
701                }
702            }),
703            producer
704        );
705        result.unwrap();
706        assert_eq!(text, "independent answer");
707        assert!(!coordinator.checkpoint().can_complete());
708        coordinator.wait_next().await.unwrap();
709        coordinator
710            .pump(stream(vec![done("independent_followup")]), |_| {})
711            .await
712            .unwrap();
713        let delivery = coordinator.prepare_delivery().await.unwrap().unwrap();
714        assert_eq!(delivery.previous_response_id, "independent_followup");
715        assert_eq!(delivery.call_ids, vec!["fast"]);
716        coordinator
717            .pump(
718                stream(vec![
719                    LlmStreamEvent::NativeToolCall(call("sync", false)),
720                    done("fast_receipt"),
721                ]),
722                |_| {},
723            )
724            .await
725            .unwrap();
726        let next = coordinator.prepare_delivery().await.unwrap().unwrap();
727        assert!(next.call_ids.contains(&"sync".to_string()));
728        coordinator
729            .pump(stream(vec![done("sync_receipt")]), |_| {})
730            .await
731            .unwrap();
732        while coordinator.wait_next().await.unwrap() {}
733        coordinator.prepare_delivery().await.unwrap();
734        coordinator
735            .pump(stream(vec![done("final")]), |_| {})
736            .await
737            .unwrap();
738        assert!(coordinator.checkpoint().can_complete());
739        assert!(executor.maximum.load(Ordering::SeqCst) <= 2);
740    }
741
742    #[tokio::test]
743    async fn restart_cancellation_duplicate_calls_and_serial_limit() {
744        let journal = MemoryJournal::default();
745        let executor = Arc::new(Executor::default());
746        let mut coordinator =
747            NativeAsyncCoordinator::open(Box::new(journal.clone()), executor.clone(), 8, false)
748                .await
749                .unwrap();
750        coordinator.register(call("slow", true)).await.unwrap();
751        coordinator.register(call("slow", true)).await.unwrap();
752        coordinator.register(call("fast", true)).await.unwrap();
753        coordinator
754            .pump(stream(vec![done("before_restart")]), |_| {})
755            .await
756            .unwrap();
757        drop(coordinator);
758        let mut recovered =
759            NativeAsyncCoordinator::open(Box::new(journal), executor.clone(), 8, false)
760                .await
761                .unwrap();
762        while recovered.wait_next().await.unwrap() {}
763        assert!(executor.maximum.load(Ordering::SeqCst) <= 1);
764        let delivery = recovered.prepare_delivery().await.unwrap().unwrap();
765        assert_eq!(delivery.call_ids.len(), 2);
766        recovered
767            .pump(stream(vec![done("receipt")]), |_| {})
768            .await
769            .unwrap();
770        recovered.register(call("cancel", true)).await.unwrap();
771        recovered.cancel().await.unwrap();
772        assert!(!recovered.wait_next().await.unwrap());
773        assert!(!recovered.checkpoint().can_complete());
774        assert!(
775            recovered.prepare_delivery().await.unwrap().unwrap().input[0]["output"]
776                .as_str()
777                .unwrap()
778                .contains("cancelled")
779        );
780    }
781
782    #[tokio::test]
783    async fn failed_journal_write_prevents_dispatch() {
784        struct FailsAfterOpen(AtomicUsize);
785        #[async_trait]
786        impl NativeAsyncJournal for FailsAfterOpen {
787            async fn load(&self) -> Result<NativeAsyncCheckpoint> {
788                Ok(NativeAsyncCheckpoint::default())
789            }
790            async fn save(&self, _: &NativeAsyncCheckpoint) -> Result<()> {
791                if self.0.fetch_add(1, Ordering::SeqCst) == 0 {
792                    Ok(())
793                } else {
794                    Err(AgentLoopError::store("disk unavailable"))
795                }
796            }
797        }
798        let executor = Arc::new(Executor::default());
799        let mut coordinator = NativeAsyncCoordinator::open(
800            Box::new(FailsAfterOpen(AtomicUsize::new(0))),
801            executor.clone(),
802            1,
803            true,
804        )
805        .await
806        .unwrap();
807        assert!(
808            coordinator
809                .register(call("must_not_run", true))
810                .await
811                .is_err()
812        );
813        assert_eq!(executor.started.load(Ordering::SeqCst), 0);
814        assert!(coordinator.wait_next().await.is_err());
815        assert!(
816            coordinator
817                .register(call("must_not_run", true))
818                .await
819                .is_err()
820        );
821    }
822
823    #[tokio::test]
824    async fn authorization_and_incomplete_stream_do_not_silently_finish() {
825        let journal = MemoryJournal::default();
826        let executor = Arc::new(Executor::default());
827        let mut coordinator =
828            NativeAsyncCoordinator::open(Box::new(journal.clone()), executor.clone(), 1, true)
829                .await
830                .unwrap();
831        let forbidden = NativeToolCall::Function {
832            call_id: "bad".into(),
833            name: "forbidden".into(),
834            arguments: "{}".into(),
835            asynchronous: true,
836        };
837        assert!(coordinator.register(forbidden).await.is_err());
838        assert!(coordinator.checkpoint().calls.is_empty());
839        assert!(
840            coordinator
841                .pump(
842                    stream(vec![LlmStreamEvent::NativeToolCall(call("pending", true))]),
843                    |_| {}
844                )
845                .await
846                .is_err()
847        );
848        drop(coordinator);
849        assert!(
850            NativeAsyncCoordinator::open(Box::new(journal), executor, 1, true)
851                .await
852                .is_err()
853        );
854    }
855    #[tokio::test]
856    async fn native_journal_and_jobs_share_transport_without_deadlocking() {
857        struct Journal {
858            memory: MemoryJournal,
859            transport: Arc<Mutex<()>>,
860        }
861        #[async_trait]
862        impl NativeAsyncJournal for Journal {
863            async fn load(&self) -> Result<NativeAsyncCheckpoint> {
864                self.memory.load().await
865            }
866            async fn save(&self, state: &NativeAsyncCheckpoint) -> Result<()> {
867                let _transport = self.transport.lock().await;
868                self.memory.save(state).await
869            }
870        }
871        struct Tool {
872            transport: Arc<Mutex<()>>,
873            started: Arc<tokio::sync::Notify>,
874        }
875        #[async_trait]
876        impl NativeAsyncExecutor for Tool {
877            async fn authorize(&self, _: &NativeToolCall) -> Result<NativeCallPolicy> {
878                Ok(NativeCallPolicy {
879                    allow_async: true,
880                    replay_safe: true,
881                    concurrency_class: None,
882                })
883            }
884            async fn execute(&self, call: NativeToolCall) -> Result<String> {
885                let _transport = self.transport.lock().await;
886                self.started.notify_one();
887                tokio::time::sleep(Duration::from_millis(20)).await;
888                Ok(call.id().into())
889            }
890        }
891        let transport = Arc::new(Mutex::new(()));
892        let started = Arc::new(tokio::sync::Notify::new());
893        let journal = Journal {
894            memory: MemoryJournal::default(),
895            transport: transport.clone(),
896        };
897        let tool = Arc::new(Tool {
898            transport,
899            started: started.clone(),
900        });
901        let mut coordinator = NativeAsyncCoordinator::open(Box::new(journal), tool, 2, true)
902            .await
903            .unwrap();
904        let (sender, receiver) = futures::channel::mpsc::unbounded();
905        sender
906            .unbounded_send(Ok(LlmStreamEvent::NativeToolCall(call("first", true))))
907            .unwrap();
908        let sender_task = tokio::spawn(async move {
909            started.notified().await;
910            sender
911                .unbounded_send(Ok(LlmStreamEvent::NativeToolCall(call("second", true))))
912                .unwrap();
913            sender.unbounded_send(Ok(done("response"))).unwrap();
914        });
915        tokio::time::timeout(
916            Duration::from_secs(1),
917            coordinator.pump(Box::pin(receiver), |_| {}),
918        )
919        .await
920        .expect("journal writes must not stop the job that owns their transport")
921        .unwrap();
922        sender_task.await.unwrap();
923        while coordinator.wait_next().await.unwrap() {}
924        assert!(
925            coordinator
926                .checkpoint()
927                .calls
928                .values()
929                .all(|pending| matches!(pending.state, PendingCallState::Ready { .. }))
930        );
931    }
932}