1use 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#[async_trait]
17pub trait NativeAsyncJournal: Send + Sync {
18 async fn load(&self) -> Result<NativeAsyncCheckpoint>;
19 async fn save(&self, checkpoint: &NativeAsyncCheckpoint) -> Result<()>;
21 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#[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
46struct 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 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 #[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 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 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 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 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 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 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 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 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 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 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 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}