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 let _permit = permits
221 .acquire()
222 .await
223 .expect("coordinator never closes permits");
224 let id = call.id().to_owned();
225 (id, executor.execute(call).await)
226 });
227 self.jobs.push(OwnedJob {
228 id: call_id,
229 handle,
230 });
231 Ok(())
232 }
233
234 pub async fn register(&mut self, call: NativeToolCall) -> Result<()> {
235 self.healthy()?;
236 let policy = self.executor.authorize(&call).await?;
237 if call.is_async() && !policy.allow_async {
238 return Err(AgentLoopError::config(
239 "tool is not authorized for native asynchronous execution",
240 ));
241 }
242 if !self.checkpoint.register(call.clone(), policy.replay_safe)? {
243 return Ok(());
244 }
245 self.save().await?;
246 if self.checkpoint.response_in_flight && !call.is_async() {
248 return Ok(());
249 }
250 self.launch(call, policy).await
251 }
252
253 async fn launch_synchronous_calls(&mut self) -> Result<()> {
254 let queued: Vec<_> = self
255 .checkpoint
256 .order
257 .iter()
258 .filter_map(|id| self.checkpoint.calls.get(id))
259 .filter(|pending| !pending.call.is_async() && pending.state == PendingCallState::Queued)
260 .map(|pending| pending.call.clone())
261 .collect();
262 for call in queued {
263 let policy = self.executor.authorize(&call).await?;
264 self.launch(call, policy).await?;
265 }
266 Ok(())
267 }
268
269 async fn settle(&mut self, id: String, result: Result<String>) -> Result<()> {
270 let output = result
272 .unwrap_or_else(|error| serde_json::json!({"error":error.to_string()}).to_string());
273 self.checkpoint.settle(&id, output)?;
274 self.save().await
275 }
276
277 pub async fn begin_transcript_response(&mut self, message_id: String) -> Result<()> {
278 self.healthy()?;
279 if self.checkpoint.transcript_message_id.is_some() {
280 return Err(AgentLoopError::store("native transcript is not committed"));
281 }
282 self.checkpoint.transcript_message_id = Some(message_id);
283 self.begin_response().await
284 }
285
286 pub async fn stage_transcript_result(&mut self, result: serde_json::Value) -> Result<()> {
287 self.healthy()?;
288 if self.checkpoint.response_in_flight || self.checkpoint.transcript_message_id.is_none() {
289 return Err(AgentLoopError::store(
290 "native response is not ready for transcript commit",
291 ));
292 }
293 if self.checkpoint.host_responses.len() + 1 != self.checkpoint.completed_responses as usize
294 {
295 return Err(AgentLoopError::store(
296 "native response summary count mismatch",
297 ));
298 }
299 self.checkpoint.host_responses.push(result);
300 self.save().await
301 }
302
303 pub async fn transcript_committed(&mut self, message_id: &str) -> Result<()> {
304 self.healthy()?;
305 if self.checkpoint.response_in_flight
306 || self.checkpoint.transcript_message_id.as_deref() != Some(message_id)
307 {
308 return Err(AgentLoopError::store("native transcript boundary mismatch"));
309 }
310 self.checkpoint.transcript_message_id = None;
311 self.save().await
312 }
313
314 pub async fn begin_response(&mut self) -> Result<()> {
316 self.healthy()?;
317 if self.checkpoint.response_in_flight {
318 return Err(AgentLoopError::store(
319 "prior native response requires reconciliation",
320 ));
321 }
322 self.checkpoint.response_in_flight = true;
323 self.save().await
324 }
325
326 pub async fn next_response_event(
329 &mut self,
330 stream: &mut LlmResponseStream,
331 ) -> Result<LlmStreamEvent> {
332 self.healthy()?;
333 let result = self.next_response_event_inner(stream).await;
334 if result.is_err() {
335 self.poisoned = true;
336 self.jobs.clear();
337 }
338 result
339 }
340
341 async fn next_response_event_inner(
342 &mut self,
343 stream: &mut LlmResponseStream,
344 ) -> Result<LlmStreamEvent> {
345 if !self.checkpoint.response_in_flight {
346 return Err(AgentLoopError::store(
347 "native response has no persisted request intent",
348 ));
349 }
350 loop {
351 tokio::select! {
352 _ = tokio::time::sleep_until(self.last_heartbeat + std::time::Duration::from_secs(10)) => self.heartbeat().await?,
353 completed = self.jobs.next(), if !self.jobs.is_empty() => {
354 let (id, result) = completed.expect("nonempty jobs");
355 self.settle(id, result).await?;
356 }
357 event = stream.next() => {
358 let event = event.ok_or_else(|| AgentLoopError::llm("native response stream ended before completion"))??;
359 match &event {
360 LlmStreamEvent::NativeToolCall(call) => self.register(call.clone()).await?,
361 LlmStreamEvent::ToolCalls(calls) => for call in calls {
362 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?;
363 },
364 LlmStreamEvent::Done(metadata) => {
365 let id = metadata.response_id.clone().ok_or_else(|| AgentLoopError::llm("native response omitted response ID"))?;
366 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")); }
367 if self.checkpoint.delivery.is_some() { self.checkpoint.acknowledge_delivery(id)?; } else { self.checkpoint.response_completed(id)?; }
368 self.checkpoint.response_in_flight = false;
369 self.save().await?;
370 self.launch_synchronous_calls().await?;
371 }
372 LlmStreamEvent::Error(error) => return Err(AgentLoopError::llm(error.to_string())),
373 _ => {}
374 }
375 return Ok(event);
376 }
377 }
378 }
379 }
380
381 pub async fn pump(
386 &mut self,
387 stream: LlmResponseStream,
388 mut observe: impl FnMut(LlmStreamEvent),
389 ) -> Result<()> {
390 self.healthy()?;
391 if self.checkpoint.response_in_flight {
392 return Err(AgentLoopError::store(
393 "prior native response requires reconciliation",
394 ));
395 }
396 let result = self.pump_inner(stream, &mut observe).await;
397 if result.is_err() {
398 self.poisoned = true;
399 self.jobs.clear();
400 }
401 result
402 }
403
404 async fn pump_inner(
405 &mut self,
406 mut stream: LlmResponseStream,
407 mut observe: impl FnMut(LlmStreamEvent),
408 ) -> Result<()> {
409 self.begin_response().await?;
410 loop {
411 let event = self.next_response_event(&mut stream).await?;
412 match event {
413 LlmStreamEvent::NativeToolCall(_) | LlmStreamEvent::ToolCalls(_) => {}
414 LlmStreamEvent::Done(_) => {
415 observe(event);
416 return Ok(());
417 }
418 other => observe(other),
419 }
420 }
421 }
422
423 pub async fn wait_next(&mut self) -> Result<bool> {
425 self.healthy()?;
426 let mut lease_tick = tokio::time::interval(std::time::Duration::from_secs(10));
427 loop {
428 tokio::select! {
429 _ = lease_tick.tick() => self.heartbeat().await?,
430 completed = self.jobs.next() => {
431 if let Some((id, result)) = completed {
432 self.settle(id, result).await?;
433 return Ok(true);
434 }
435 return Ok(false);
436 }
437 }
438 }
439 }
440
441 pub async fn prepare_delivery(&mut self) -> Result<Option<Delivery>> {
444 self.healthy()?;
445 while self.checkpoint.calls.values().any(|pending| {
446 !pending.call.is_async()
447 && matches!(
448 pending.state,
449 PendingCallState::Queued | PendingCallState::Running
450 )
451 }) {
452 if !self.wait_next().await? {
453 return Err(AgentLoopError::store("synchronous call has no running job"));
454 }
455 }
456 let delivery = self.checkpoint.prepare_delivery()?.cloned();
457 self.save().await?;
458 Ok(delivery)
459 }
460
461 pub async fn cancel(&mut self) -> Result<()> {
464 self.healthy()?;
465 for job in self.jobs.iter() {
466 job.handle.abort();
467 }
468 while let Some((id, result)) = self.jobs.next().await {
469 if !matches!(result, Err(AgentLoopError::Cancelled)) {
470 let output = result.unwrap_or_else(|error| {
471 serde_json::json!({"error":error.to_string()}).to_string()
472 });
473 self.checkpoint.settle(&id, output)?;
474 }
475 }
476 self.checkpoint.cancel();
477 self.save().await
478 }
479}
480
481impl NativeAsyncCoordinator {
482 pub async fn run<Request, RequestFuture>(
486 &mut self,
487 max_responses: usize,
488 mut request: Request,
489 mut observe: impl FnMut(LlmStreamEvent),
490 ) -> Result<()>
491 where
492 Request: FnMut(Option<Delivery>, Option<String>) -> RequestFuture,
493 RequestFuture: std::future::Future<Output = Result<LlmResponseStream>>,
494 {
495 self.healthy()?;
496 for _ in 0..max_responses {
497 let mut delivery = self.prepare_delivery().await?;
498 if delivery.is_none()
501 && self.checkpoint.latest_response_id.is_some()
502 && !self.jobs.is_empty()
503 {
504 self.wait_next().await?;
505 delivery = self.prepare_delivery().await?;
506 }
507 let response = request(delivery, self.checkpoint.latest_response_id.clone());
508 tokio::pin!(response);
509 let mut lease_tick = tokio::time::interval(std::time::Duration::from_secs(10));
510 let stream = loop {
511 tokio::select! {
512 result = &mut response => break result?,
513 _ = lease_tick.tick() => self.heartbeat().await?,
514 }
515 };
516 self.pump(stream, &mut observe).await?;
517 if self.checkpoint.can_complete() {
518 return Ok(());
519 }
520 if !self
521 .checkpoint
522 .calls
523 .values()
524 .any(|pending| matches!(pending.state, PendingCallState::Ready { .. }))
525 {
526 self.wait_next().await?;
527 }
528 }
529 Err(AgentLoopError::config(
530 "native async response limit reached; pending work remains checkpointed",
531 ))
532 }
533}
534
535#[cfg(test)]
536mod tests {
537 use super::*;
538 use everruns_provider::LlmCompletionMetadata;
539 use std::{
540 sync::{
541 Mutex as StdMutex,
542 atomic::{AtomicUsize, Ordering},
543 },
544 time::Duration,
545 };
546
547 #[derive(Clone, Default)]
548 struct MemoryJournal(Arc<StdMutex<NativeAsyncCheckpoint>>);
549 #[async_trait]
550 impl NativeAsyncJournal for MemoryJournal {
551 async fn load(&self) -> Result<NativeAsyncCheckpoint> {
552 Ok(self.0.lock().unwrap().clone())
553 }
554 async fn save(&self, checkpoint: &NativeAsyncCheckpoint) -> Result<()> {
555 *self.0.lock().unwrap() = checkpoint.clone();
556 Ok(())
557 }
558 }
559 #[derive(Default)]
560 struct Executor {
561 started: AtomicUsize,
562 active: AtomicUsize,
563 maximum: AtomicUsize,
564 }
565 #[async_trait]
566 impl NativeAsyncExecutor for Executor {
567 async fn authorize(&self, call: &NativeToolCall) -> Result<NativeCallPolicy> {
568 if call.name() == "forbidden" {
569 return Err(AgentLoopError::tool("not authorized"));
570 }
571 Ok(NativeCallPolicy {
572 allow_async: true,
573 replay_safe: true,
574 concurrency_class: None,
575 })
576 }
577 async fn execute(&self, call: NativeToolCall) -> Result<String> {
578 self.started.fetch_add(1, Ordering::SeqCst);
579 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
580 self.maximum.fetch_max(active, Ordering::SeqCst);
581 struct Active<'a>(&'a AtomicUsize);
582 impl Drop for Active<'_> {
583 fn drop(&mut self) {
584 self.0.fetch_sub(1, Ordering::SeqCst);
585 }
586 }
587 let _active = Active(&self.active);
588 tokio::time::sleep(Duration::from_millis(if call.id() == "slow" {
589 100
590 } else {
591 1
592 }))
593 .await;
594 Ok(call.id().into())
595 }
596 }
597 fn call(id: &str, asynchronous: bool) -> NativeToolCall {
598 NativeToolCall::Function {
599 call_id: id.into(),
600 name: "lookup".into(),
601 arguments: "{}".into(),
602 asynchronous,
603 }
604 }
605 fn done(id: &str) -> LlmStreamEvent {
606 LlmStreamEvent::Done(Box::new(LlmCompletionMetadata {
607 response_id: Some(id.into()),
608 ..Default::default()
609 }))
610 }
611 fn stream(events: Vec<LlmStreamEvent>) -> LlmResponseStream {
612 Box::pin(futures::stream::iter(events.into_iter().map(Ok)))
613 }
614
615 #[tokio::test]
616 async fn synchronous_calls_wait_for_successful_response_completion() {
617 for rejected in [false, true] {
618 let executor = Arc::new(Executor::default());
619 let mut coordinator = NativeAsyncCoordinator::open(
620 Box::new(MemoryJournal::default()),
621 executor.clone(),
622 2,
623 true,
624 )
625 .await
626 .unwrap();
627 coordinator.begin_response().await.unwrap();
628 let mut terminal = done("response");
629 if rejected && let LlmStreamEvent::Done(metadata) = &mut terminal {
630 metadata.finish_reason = Some("length".into());
631 }
632 let mut response = stream(vec![
633 LlmStreamEvent::NativeToolCall(call("sync", false)),
634 terminal,
635 ]);
636 coordinator
637 .next_response_event(&mut response)
638 .await
639 .unwrap();
640 assert_eq!(
641 coordinator.checkpoint().calls["sync"].state,
642 PendingCallState::Queued,
643 "synchronous calls must not gain early execution from native opt-in"
644 );
645 let result = coordinator.next_response_event(&mut response).await;
646 if rejected {
647 assert!(result.is_err());
648 assert!(
649 coordinator.checkpoint().clone().recover().is_err(),
650 "recovery must not execute rejected calls"
651 );
652 assert_eq!(executor.started.load(Ordering::SeqCst), 0);
653 } else {
654 result.unwrap();
655 assert!(coordinator.wait_next().await.unwrap());
656 assert_eq!(executor.started.load(Ordering::SeqCst), 1);
657 }
658 }
659 }
660
661 #[tokio::test]
662 async fn early_dispatch_mixed_calls_out_of_order_and_independent_work() {
663 let journal = MemoryJournal::default();
664 let executor = Arc::new(Executor::default());
665 let mut coordinator =
666 NativeAsyncCoordinator::open(Box::new(journal.clone()), executor.clone(), 2, true)
667 .await
668 .unwrap();
669 let (sender, receiver) = futures::channel::mpsc::unbounded();
670 sender
671 .unbounded_send(Ok(LlmStreamEvent::NativeToolCall(call("slow", true))))
672 .unwrap();
673 let observed = executor.clone();
674 let producer = async move {
675 while observed.started.load(Ordering::SeqCst) == 0 {
677 tokio::task::yield_now().await;
678 }
679 sender
680 .unbounded_send(Ok(LlmStreamEvent::NativeToolCall(call("fast", true))))
681 .unwrap();
682 sender
683 .unbounded_send(Ok(LlmStreamEvent::TextDelta("independent answer".into())))
684 .unwrap();
685 sender.unbounded_send(Ok(done("launch"))).unwrap();
686 };
687 let mut text = String::new();
688 let (result, _) = tokio::join!(
689 coordinator.pump(Box::pin(receiver), |event| {
690 if let LlmStreamEvent::TextDelta(delta) = event {
691 text.push_str(&delta);
692 }
693 }),
694 producer
695 );
696 result.unwrap();
697 assert_eq!(text, "independent answer");
698 assert!(!coordinator.checkpoint().can_complete());
699 coordinator.wait_next().await.unwrap();
700 coordinator
701 .pump(stream(vec![done("independent_followup")]), |_| {})
702 .await
703 .unwrap();
704 let delivery = coordinator.prepare_delivery().await.unwrap().unwrap();
705 assert_eq!(delivery.previous_response_id, "independent_followup");
706 assert_eq!(delivery.call_ids, vec!["fast"]);
707 coordinator
708 .pump(
709 stream(vec![
710 LlmStreamEvent::NativeToolCall(call("sync", false)),
711 done("fast_receipt"),
712 ]),
713 |_| {},
714 )
715 .await
716 .unwrap();
717 let next = coordinator.prepare_delivery().await.unwrap().unwrap();
718 assert!(next.call_ids.contains(&"sync".to_string()));
719 coordinator
720 .pump(stream(vec![done("sync_receipt")]), |_| {})
721 .await
722 .unwrap();
723 while coordinator.wait_next().await.unwrap() {}
724 coordinator.prepare_delivery().await.unwrap();
725 coordinator
726 .pump(stream(vec![done("final")]), |_| {})
727 .await
728 .unwrap();
729 assert!(coordinator.checkpoint().can_complete());
730 assert!(executor.maximum.load(Ordering::SeqCst) <= 2);
731 }
732
733 #[tokio::test]
734 async fn restart_cancellation_duplicate_calls_and_serial_limit() {
735 let journal = MemoryJournal::default();
736 let executor = Arc::new(Executor::default());
737 let mut coordinator =
738 NativeAsyncCoordinator::open(Box::new(journal.clone()), executor.clone(), 8, false)
739 .await
740 .unwrap();
741 coordinator.register(call("slow", true)).await.unwrap();
742 coordinator.register(call("slow", true)).await.unwrap();
743 coordinator.register(call("fast", true)).await.unwrap();
744 coordinator
745 .pump(stream(vec![done("before_restart")]), |_| {})
746 .await
747 .unwrap();
748 drop(coordinator);
749 let mut recovered =
750 NativeAsyncCoordinator::open(Box::new(journal), executor.clone(), 8, false)
751 .await
752 .unwrap();
753 while recovered.wait_next().await.unwrap() {}
754 assert!(executor.maximum.load(Ordering::SeqCst) <= 1);
755 let delivery = recovered.prepare_delivery().await.unwrap().unwrap();
756 assert_eq!(delivery.call_ids.len(), 2);
757 recovered
758 .pump(stream(vec![done("receipt")]), |_| {})
759 .await
760 .unwrap();
761 recovered.register(call("cancel", true)).await.unwrap();
762 recovered.cancel().await.unwrap();
763 assert!(!recovered.wait_next().await.unwrap());
764 assert!(!recovered.checkpoint().can_complete());
765 assert!(
766 recovered.prepare_delivery().await.unwrap().unwrap().input[0]["output"]
767 .as_str()
768 .unwrap()
769 .contains("cancelled")
770 );
771 }
772
773 #[tokio::test]
774 async fn failed_journal_write_prevents_dispatch() {
775 struct FailsAfterOpen(AtomicUsize);
776 #[async_trait]
777 impl NativeAsyncJournal for FailsAfterOpen {
778 async fn load(&self) -> Result<NativeAsyncCheckpoint> {
779 Ok(NativeAsyncCheckpoint::default())
780 }
781 async fn save(&self, _: &NativeAsyncCheckpoint) -> Result<()> {
782 if self.0.fetch_add(1, Ordering::SeqCst) == 0 {
783 Ok(())
784 } else {
785 Err(AgentLoopError::store("disk unavailable"))
786 }
787 }
788 }
789 let executor = Arc::new(Executor::default());
790 let mut coordinator = NativeAsyncCoordinator::open(
791 Box::new(FailsAfterOpen(AtomicUsize::new(0))),
792 executor.clone(),
793 1,
794 true,
795 )
796 .await
797 .unwrap();
798 assert!(
799 coordinator
800 .register(call("must_not_run", true))
801 .await
802 .is_err()
803 );
804 assert_eq!(executor.started.load(Ordering::SeqCst), 0);
805 assert!(coordinator.wait_next().await.is_err());
806 assert!(
807 coordinator
808 .register(call("must_not_run", true))
809 .await
810 .is_err()
811 );
812 }
813
814 #[tokio::test]
815 async fn authorization_and_incomplete_stream_do_not_silently_finish() {
816 let journal = MemoryJournal::default();
817 let executor = Arc::new(Executor::default());
818 let mut coordinator =
819 NativeAsyncCoordinator::open(Box::new(journal.clone()), executor.clone(), 1, true)
820 .await
821 .unwrap();
822 let forbidden = NativeToolCall::Function {
823 call_id: "bad".into(),
824 name: "forbidden".into(),
825 arguments: "{}".into(),
826 asynchronous: true,
827 };
828 assert!(coordinator.register(forbidden).await.is_err());
829 assert!(coordinator.checkpoint().calls.is_empty());
830 assert!(
831 coordinator
832 .pump(
833 stream(vec![LlmStreamEvent::NativeToolCall(call("pending", true))]),
834 |_| {}
835 )
836 .await
837 .is_err()
838 );
839 drop(coordinator);
840 assert!(
841 NativeAsyncCoordinator::open(Box::new(journal), executor, 1, true)
842 .await
843 .is_err()
844 );
845 }
846 #[tokio::test]
847 async fn native_journal_and_jobs_share_transport_without_deadlocking() {
848 struct Journal {
849 memory: MemoryJournal,
850 transport: Arc<Mutex<()>>,
851 }
852 #[async_trait]
853 impl NativeAsyncJournal for Journal {
854 async fn load(&self) -> Result<NativeAsyncCheckpoint> {
855 self.memory.load().await
856 }
857 async fn save(&self, state: &NativeAsyncCheckpoint) -> Result<()> {
858 let _transport = self.transport.lock().await;
859 self.memory.save(state).await
860 }
861 }
862 struct Tool {
863 transport: Arc<Mutex<()>>,
864 started: Arc<tokio::sync::Notify>,
865 }
866 #[async_trait]
867 impl NativeAsyncExecutor for Tool {
868 async fn authorize(&self, _: &NativeToolCall) -> Result<NativeCallPolicy> {
869 Ok(NativeCallPolicy {
870 allow_async: true,
871 replay_safe: true,
872 concurrency_class: None,
873 })
874 }
875 async fn execute(&self, call: NativeToolCall) -> Result<String> {
876 let _transport = self.transport.lock().await;
877 self.started.notify_one();
878 tokio::time::sleep(Duration::from_millis(20)).await;
879 Ok(call.id().into())
880 }
881 }
882 let transport = Arc::new(Mutex::new(()));
883 let started = Arc::new(tokio::sync::Notify::new());
884 let journal = Journal {
885 memory: MemoryJournal::default(),
886 transport: transport.clone(),
887 };
888 let tool = Arc::new(Tool {
889 transport,
890 started: started.clone(),
891 });
892 let mut coordinator = NativeAsyncCoordinator::open(Box::new(journal), tool, 2, true)
893 .await
894 .unwrap();
895 let (sender, receiver) = futures::channel::mpsc::unbounded();
896 sender
897 .unbounded_send(Ok(LlmStreamEvent::NativeToolCall(call("first", true))))
898 .unwrap();
899 let sender_task = tokio::spawn(async move {
900 started.notified().await;
901 sender
902 .unbounded_send(Ok(LlmStreamEvent::NativeToolCall(call("second", true))))
903 .unwrap();
904 sender.unbounded_send(Ok(done("response"))).unwrap();
905 });
906 tokio::time::timeout(
907 Duration::from_secs(1),
908 coordinator.pump(Box::pin(receiver), |_| {}),
909 )
910 .await
911 .expect("journal writes must not stop the job that owns their transport")
912 .unwrap();
913 sender_task.await.unwrap();
914 while coordinator.wait_next().await.unwrap() {}
915 assert!(
916 coordinator
917 .checkpoint()
918 .calls
919 .values()
920 .all(|pending| matches!(pending.state, PendingCallState::Ready { .. }))
921 );
922 }
923}