1use crate::registry::{EmptyToolRegistry, ResolveToolRequest, ToolRegistry, ToolResolution};
4use crate::subscription::{CanonicalEventSubscription, SubscriptionStatus};
5use crate::tools::{NoToolRuntime, ToolRuntime};
6use monoloop_contracts::{
7 CanonicalUnit, CanonicalUnitEvent, InterpretationEnd, InterpreterOutputEvent, LoopEnd,
8 LoopEndKind, LoopError, LoopId, LoopLimits, LoopOutputEvent, LoopScope, MonoloopRunId,
9 OutboundToolOutcome, OutboundToolResult, ToolActionId, ToolExecutionId, ToolRequestState,
10 UnitId,
11};
12use std::collections::HashMap;
13use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
14use std::sync::Arc;
15use tokio::sync::{mpsc, oneshot, Mutex};
16
17pub struct StartLoop {
19 pub monoloop_run_id: MonoloopRunId,
21 pub loop_id: LoopId,
23 pub scope: LoopScope,
25 pub subscription: CanonicalEventSubscription,
27 pub tool_registry: Arc<dyn ToolRegistry>,
29 pub tool_runtime: Arc<dyn ToolRuntime>,
31 pub output_capacity: usize,
33 pub limits: LoopLimits,
35}
36
37pub struct LoopHandle {
39 pub loop_id: LoopId,
41 pub control: LoopControl,
43 pub health: LoopHealth,
45 pub completion: LoopCompletion,
47 pub output: Mutex<Option<mpsc::Receiver<LoopOutputEvent>>>,
52}
53
54impl LoopHandle {
55 pub async fn take_output(&self) -> mpsc::Receiver<LoopOutputEvent> {
57 let mut guard = self.output.lock().await;
58 guard.take().expect("LoopHandle output taken twice")
59 }
60}
61
62#[derive(Clone)]
64pub struct LoopControl {
65 cancelled: Arc<AtomicBool>,
66 notify: Arc<tokio::sync::Notify>,
67}
68
69impl LoopControl {
70 fn new() -> Self {
71 Self {
72 cancelled: Arc::new(AtomicBool::new(false)),
73 notify: Arc::new(tokio::sync::Notify::new()),
74 }
75 }
76
77 pub fn cancel(&self) {
79 self.cancelled.store(true, Ordering::SeqCst);
80 self.notify.notify_waiters();
81 }
82
83 fn is_cancelled(&self) -> bool {
84 self.cancelled.load(Ordering::SeqCst)
85 }
86}
87
88#[derive(Clone, Debug, Default)]
90pub struct LoopHealth {
91 pub events_received: Arc<AtomicU64>,
93 pub tools_unavailable: Arc<AtomicU64>,
95}
96
97pub struct LoopCompletion {
99 rx: Mutex<Option<oneshot::Receiver<LoopEnd>>>,
100}
101
102impl LoopCompletion {
103 pub async fn wait(self) -> LoopEnd {
105 let rx = self.take_receiver().await;
106 Self::recv_end(rx).await
107 }
108
109 pub async fn take_receiver(&self) -> oneshot::Receiver<LoopEnd> {
111 let mut guard = self.rx.lock().await;
112 guard.take().expect("LoopCompletion polled twice")
113 }
114
115 pub async fn wait_ref(&self) -> LoopEnd {
117 let rx = self.take_receiver().await;
118 Self::recv_end(rx).await
119 }
120
121 async fn recv_end(rx: oneshot::Receiver<LoopEnd>) -> LoopEnd {
122 rx.await.unwrap_or(LoopEnd {
123 monoloop_run_id: MonoloopRunId::new("unknown"),
124 loop_id: LoopId::new("unknown"),
125 kind: LoopEndKind::InvariantFailed,
126 delivery_events_received: 0,
127 duplicate_events: 0,
128 tools_unavailable: 0,
129 outbound_results_emitted: 0,
130 safe_diagnostics: vec!["completion dropped".into()],
131 })
132 }
133}
134
135pub type LoopRunFuture = std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>;
137
138#[derive(Clone, Debug, Default)]
140pub struct DefaultLoopRuntime;
141
142impl DefaultLoopRuntime {
143 pub fn new() -> Self {
145 Self
146 }
147
148 pub fn prepare(&self, request: StartLoop) -> Result<(LoopHandle, LoopRunFuture), LoopError> {
153 prepare_loop(request)
154 }
155
156 pub fn prepare_empty(
158 &self,
159 monoloop_run_id: MonoloopRunId,
160 loop_id: LoopId,
161 scope: LoopScope,
162 subscription: CanonicalEventSubscription,
163 limits: LoopLimits,
164 ) -> Result<(LoopHandle, LoopRunFuture), LoopError> {
165 self.prepare(StartLoop {
166 monoloop_run_id,
167 loop_id,
168 scope,
169 subscription,
170 tool_registry: Arc::new(EmptyToolRegistry::new()),
171 tool_runtime: Arc::new(NoToolRuntime::new()),
172 output_capacity: limits.max_output_queue,
173 limits,
174 })
175 }
176}
177
178fn prepare_loop(request: StartLoop) -> Result<(LoopHandle, LoopRunFuture), LoopError> {
179 let control = LoopControl::new();
180 let health = LoopHealth::default();
181 let (out_tx, out_rx) = mpsc::channel(request.output_capacity.max(1));
182 let (end_tx, end_rx) = oneshot::channel();
183
184 let loop_id = request.loop_id.clone();
185 let control_task = control.clone();
186 let health_task = health.clone();
187
188 let fut = Box::pin(async move {
189 let mut owner = LoopOwner {
190 monoloop_run_id: request.monoloop_run_id,
191 loop_id: request.loop_id,
192 scope: request.scope,
193 registry: request.tool_registry,
194 runtime: request.tool_runtime,
195 limits: request.limits,
196 out_tx,
197 control: control_task,
198 health: health_task,
199 last_seq: 0,
200 actions: HashMap::new(),
201 dedup: HashMap::new(),
202 delivery_events: 0,
203 duplicates: 0,
204 tools_unavailable: 0,
205 outbound_results: 0,
206 diagnostics: Vec::new(),
207 ended: false,
208 };
209 owner.run(request.subscription, end_tx).await;
210 });
211
212 Ok((
213 LoopHandle {
214 loop_id,
215 control,
216 health,
217 completion: LoopCompletion {
218 rx: Mutex::new(Some(end_rx)),
219 },
220 output: Mutex::new(Some(out_rx)),
221 },
222 fut,
223 ))
224}
225
226#[derive(Clone, Copy, Debug, PartialEq, Eq)]
227enum ActionState {
228 ObservedWaiting,
229 RequestReady,
230 Unavailable,
231 Incomplete,
232}
233
234struct ActionRecord {
235 #[allow(dead_code)]
236 tool_action_id: ToolActionId,
237 #[allow(dead_code)]
238 unit_id: UnitId,
239 last_generation: u64,
240 state: ActionState,
241 dispatched: bool,
242}
243
244struct LoopOwner {
245 monoloop_run_id: MonoloopRunId,
246 loop_id: LoopId,
247 scope: LoopScope,
248 registry: Arc<dyn ToolRegistry>,
249 runtime: Arc<dyn ToolRuntime>,
250 limits: LoopLimits,
251 out_tx: mpsc::Sender<LoopOutputEvent>,
252 control: LoopControl,
253 health: LoopHealth,
254 last_seq: u64,
255 actions: HashMap<String, ActionRecord>,
256 dedup: HashMap<String, u64>,
257 delivery_events: u64,
258 duplicates: u64,
259 tools_unavailable: u64,
260 outbound_results: u64,
261 diagnostics: Vec<String>,
262 ended: bool,
263}
264
265impl LoopOwner {
266 async fn run(
267 &mut self,
268 mut subscription: CanonicalEventSubscription,
269 end_tx: oneshot::Sender<LoopEnd>,
270 ) {
271 loop {
272 if self.control.is_cancelled() {
273 self.finish(LoopEndKind::Cancelled, end_tx).await;
274 return;
275 }
276
277 tokio::select! {
278 biased;
279 _ = self.control.notify.notified() => {
280 if self.control.is_cancelled() {
281 self.finish(LoopEndKind::Cancelled, end_tx).await;
282 return;
283 }
284 }
285 msg = subscription.recv() => {
286 match msg {
287 None => {
288 self.finish(LoopEndKind::Drained, end_tx).await;
289 return;
290 }
291 Some(Err(SubscriptionStatus::Gap(_)))
292 | Some(Err(SubscriptionStatus::Lost)) => {
293 self.finish(LoopEndKind::SubscriptionLost, end_tx).await;
294 return;
295 }
296 Some(Err(SubscriptionStatus::Opened | SubscriptionStatus::Closing)) => {}
297 Some(Ok(delivered)) => {
298 if let Err(kind) = self.on_delivered(delivered).await {
299 self.finish(kind, end_tx).await;
300 return;
301 }
302 }
303 }
304 }
305 }
306 }
307 }
308
309 async fn on_delivered(
310 &mut self,
311 delivered: crate::subscription::DeliveredEvent,
312 ) -> Result<(), LoopEndKind> {
313 if self.last_seq > 0 && delivered.delivery_sequence != self.last_seq + 1 {
315 self.diag(format!(
316 "delivery gap: expected {}, got {}",
317 self.last_seq + 1,
318 delivered.delivery_sequence
319 ));
320 return Err(LoopEndKind::SubscriptionLost);
321 }
322 self.last_seq = delivered.delivery_sequence;
323 self.delivery_events += 1;
324 self.health.events_received.fetch_add(1, Ordering::Relaxed);
325
326 match delivered.event {
327 InterpreterOutputEvent::Unit(ev) => self.on_unit(*ev).await,
328 InterpreterOutputEvent::Ended(end) => {
329 self.on_interpretation_end(end).await?;
330 Ok(())
332 }
333 }
334 }
335
336 async fn on_unit(&mut self, event: CanonicalUnitEvent) -> Result<(), LoopEndKind> {
337 let snap = event.snapshot();
338 if !self.in_scope(snap) {
339 return Ok(());
341 }
342
343 let CanonicalUnit::Tool(tool) = &snap.unit else {
345 return Ok(());
346 };
347 let key = format!(
348 "{}:{}",
349 snap.interpretation_id.as_str(),
350 tool.tool_action_id.as_str()
351 );
352 let dig_key = format!("{}:{}", key, snap.unit_generation);
353 if let Some(prev) = self.dedup.get(&dig_key) {
354 if *prev == snap.unit_generation {
355 self.duplicates += 1;
356 return Ok(());
357 }
358 }
359 if self.dedup.len() >= self.limits.max_dedup_entries {
360 return Err(LoopEndKind::InvariantFailed);
361 }
362 self.dedup.insert(dig_key, snap.unit_generation);
363
364 match tool.request_state {
365 ToolRequestState::Assembling => {
366 self.track_waiting(
367 &key,
368 tool.tool_action_id.clone(),
369 snap.unit_id.clone(),
370 snap.unit_generation,
371 );
372 }
373 ToolRequestState::Ready => {
374 self.on_request_ready(
375 &key,
376 tool.tool_action_id.clone(),
377 snap.unit_id.clone(),
378 snap.unit_generation,
379 tool.tool_name.clone(),
380 tool.request_payload.clone(),
381 snap,
382 )
383 .await?;
384 }
385 ToolRequestState::Incomplete | ToolRequestState::Malformed => {
386 let rec = self.actions.entry(key).or_insert_with(|| ActionRecord {
387 tool_action_id: tool.tool_action_id.clone(),
388 unit_id: snap.unit_id.clone(),
389 last_generation: 0,
390 state: ActionState::Incomplete,
391 dispatched: false,
392 });
393 if snap.unit_generation >= rec.last_generation {
394 rec.last_generation = snap.unit_generation;
395 rec.state = ActionState::Incomplete;
396 }
397 }
398 }
399 Ok(())
400 }
401
402 fn track_waiting(
403 &mut self,
404 key: &str,
405 tool_action_id: ToolActionId,
406 unit_id: UnitId,
407 generation: u64,
408 ) {
409 let rec = self
410 .actions
411 .entry(key.to_string())
412 .or_insert_with(|| ActionRecord {
413 tool_action_id,
414 unit_id,
415 last_generation: 0,
416 state: ActionState::ObservedWaiting,
417 dispatched: false,
418 });
419 if generation < rec.last_generation {
420 return; }
422 rec.last_generation = generation;
423 if rec.state != ActionState::Unavailable && !rec.dispatched {
424 rec.state = ActionState::ObservedWaiting;
425 }
426 }
427
428 #[allow(clippy::too_many_arguments)]
429 async fn on_request_ready(
430 &mut self,
431 key: &str,
432 tool_action_id: ToolActionId,
433 unit_id: UnitId,
434 generation: u64,
435 tool_name: Option<String>,
436 request_payload: Option<String>,
437 snap: &monoloop_contracts::CanonicalUnitSnapshot,
438 ) -> Result<(), LoopEndKind> {
439 if self.actions.len() >= self.limits.max_tool_actions && !self.actions.contains_key(key) {
440 return Err(LoopEndKind::InvariantFailed);
441 }
442
443 let rec = self
444 .actions
445 .entry(key.to_string())
446 .or_insert_with(|| ActionRecord {
447 tool_action_id: tool_action_id.clone(),
448 unit_id: unit_id.clone(),
449 last_generation: 0,
450 state: ActionState::RequestReady,
451 dispatched: false,
452 });
453
454 if generation < rec.last_generation {
455 return Ok(()); }
457 rec.last_generation = generation;
458
459 if rec.dispatched {
461 self.duplicates += 1;
462 return Ok(());
463 }
464
465 let Some(name) = tool_name else {
466 self.diag("ToolRequestReady missing tool name".into());
467 return Ok(());
468 };
469 let Some(payload) = request_payload else {
470 self.diag("ToolRequestReady missing payload".into());
471 return Ok(());
472 };
473
474 rec.state = ActionState::RequestReady;
475 rec.dispatched = true;
476
477 if self
478 .out_tx
479 .send(LoopOutputEvent::ToolDispatchRequested {
480 tool_action_id: tool_action_id.clone(),
481 request_generation: generation,
482 })
483 .await
484 .is_err()
485 {
486 return Err(LoopEndKind::OutputFailed);
487 }
488
489 let resolution = self
490 .registry
491 .resolve(ResolveToolRequest {
492 tool_action_id: tool_action_id.clone(),
493 tool_name: name.clone(),
494 request_payload: payload.clone(),
495 })
496 .await
497 .map_err(|_| LoopEndKind::InvariantFailed)?;
498
499 match resolution {
500 ToolResolution::Unavailable(reason) => {
501 rec.state = ActionState::Unavailable;
502 self.tools_unavailable += 1;
503 self.health
504 .tools_unavailable
505 .fetch_add(1, Ordering::Relaxed);
506
507 if self
508 .out_tx
509 .send(LoopOutputEvent::ToolUnavailable {
510 tool_action_id: tool_action_id.clone(),
511 reason,
512 })
513 .await
514 .is_err()
515 {
516 return Err(LoopEndKind::OutputFailed);
517 }
518
519 let result = OutboundToolResult {
520 outbound_result_id: uuid::Uuid::new_v4().to_string(),
521 monoloop_run_id: self.monoloop_run_id.clone(),
522 loop_id: self.loop_id.clone(),
523 source_interpretation_id: snap.interpretation_id.clone(),
524 source_connection_id: snap.connection_id.clone(),
525 external_session_id: snap.external_session_id.clone(),
526 tool_action_id,
527 request_generation: generation,
528 tool_execution_id: None,
529 outcome: OutboundToolOutcome::ToolUnavailable,
530 payload: format!("{reason:?}"),
531 source_unit_id: unit_id,
532 };
533 self.outbound_results += 1;
534 if self
535 .out_tx
536 .send(LoopOutputEvent::OutboundToolResult(result))
537 .await
538 .is_err()
539 {
540 return Err(LoopEndKind::OutputFailed);
541 }
542 }
544 ToolResolution::Available(_) => {
545 let execution_id = ToolExecutionId::generate();
546 let handle = match self.runtime.start(crate::tools::StartToolExecution {
547 execution_id: execution_id.clone(),
548 tool_action_id: tool_action_id.as_str().to_string(),
549 tool_name: name,
550 request_payload: payload,
551 request_generation: generation,
552 }) {
553 Ok(h) => h,
554 Err(e) => {
555 self.diag(format!("tool runtime start failed: {e}"));
556 if self
557 .out_tx
558 .send(LoopOutputEvent::OutboundToolResult(OutboundToolResult {
559 outbound_result_id: uuid::Uuid::new_v4().to_string(),
560 monoloop_run_id: self.monoloop_run_id.clone(),
561 loop_id: self.loop_id.clone(),
562 source_interpretation_id: snap.interpretation_id.clone(),
563 source_connection_id: snap.connection_id.clone(),
564 external_session_id: snap.external_session_id.clone(),
565 tool_action_id,
566 request_generation: generation,
567 tool_execution_id: None,
568 outcome: OutboundToolOutcome::ExecutionFailed,
569 payload: e.0,
570 source_unit_id: unit_id,
571 }))
572 .await
573 .is_err()
574 {
575 return Err(LoopEndKind::OutputFailed);
576 }
577 self.outbound_results += 1;
578 return Ok(());
579 }
580 };
581
582 let terminal = if let Some(rx) = handle.completion {
583 rx.await.unwrap_or(crate::tools::ToolRuntimeTerminal {
584 outcome: OutboundToolOutcome::ExecutionLost,
585 payload: "completion lost".into(),
586 })
587 } else {
588 crate::tools::ToolRuntimeTerminal {
589 outcome: OutboundToolOutcome::ExecutionFailed,
590 payload: "runtime returned no completion".into(),
591 }
592 };
593
594 self.outbound_results += 1;
595 if self
596 .out_tx
597 .send(LoopOutputEvent::OutboundToolResult(OutboundToolResult {
598 outbound_result_id: uuid::Uuid::new_v4().to_string(),
599 monoloop_run_id: self.monoloop_run_id.clone(),
600 loop_id: self.loop_id.clone(),
601 source_interpretation_id: snap.interpretation_id.clone(),
602 source_connection_id: snap.connection_id.clone(),
603 external_session_id: snap.external_session_id.clone(),
604 tool_action_id,
605 request_generation: generation,
606 tool_execution_id: Some(execution_id),
607 outcome: terminal.outcome,
608 payload: terminal.payload,
609 source_unit_id: unit_id,
610 }))
611 .await
612 .is_err()
613 {
614 return Err(LoopEndKind::OutputFailed);
615 }
616 }
617 }
618 Ok(())
619 }
620
621 async fn on_interpretation_end(&mut self, _end: InterpretationEnd) -> Result<(), LoopEndKind> {
622 Ok(())
624 }
625
626 fn in_scope(&self, snap: &monoloop_contracts::CanonicalUnitSnapshot) -> bool {
627 if self.scope.accept_all_in_run {
628 return true;
629 }
630 if !self.scope.accepted_interpretation_ids.is_empty()
631 && !self
632 .scope
633 .accepted_interpretation_ids
634 .iter()
635 .any(|id| id.as_str() == snap.interpretation_id.as_str())
636 {
637 return false;
638 }
639 if !self.scope.accepted_connection_ids.is_empty()
640 && !self
641 .scope
642 .accepted_connection_ids
643 .iter()
644 .any(|id| id.as_str() == snap.connection_id.as_str())
645 {
646 return false;
647 }
648 if !self.scope.accepted_external_session_ids.is_empty() {
649 match &snap.external_session_id {
650 Some(ext) => {
651 if !self
652 .scope
653 .accepted_external_session_ids
654 .iter()
655 .any(|id| id.as_str() == ext.as_str())
656 {
657 return false;
658 }
659 }
660 None => return false,
661 }
662 }
663 true
664 }
665
666 fn diag(&mut self, msg: String) {
667 if self.diagnostics.len() < 32 {
668 self.diagnostics.push(msg);
669 }
670 }
671
672 async fn finish(&mut self, kind: LoopEndKind, end_tx: oneshot::Sender<LoopEnd>) {
673 if self.ended {
674 return;
675 }
676 self.ended = true;
677 let end = LoopEnd {
678 monoloop_run_id: self.monoloop_run_id.clone(),
679 loop_id: self.loop_id.clone(),
680 kind,
681 delivery_events_received: self.delivery_events,
682 duplicate_events: self.duplicates,
683 tools_unavailable: self.tools_unavailable,
684 outbound_results_emitted: self.outbound_results,
685 safe_diagnostics: self.diagnostics.clone(),
686 };
687 let _ = self
688 .out_tx
689 .send(LoopOutputEvent::LoopEnded(end.clone()))
690 .await;
691 let _ = end_tx.send(end);
692 }
693}