monoloop_loop/transaction/
owned_process_registry.rs1use super::tool_capacity::ToolPermit;
8use super::tool_handler::ToolKillHandle;
9use std::sync::Mutex;
10
11struct RegistryEntry {
12 kill: ToolKillHandle,
13 #[allow(dead_code)]
15 permit: Option<ToolPermit>,
16}
17
18#[derive(Default)]
20pub struct OwnedProcessRegistry {
21 entries: Mutex<Vec<RegistryEntry>>,
22}
23
24impl OwnedProcessRegistry {
25 pub fn new() -> Self {
27 Self::default()
28 }
29
30 pub fn park(&self, kill: ToolKillHandle, permit: Option<ToolPermit>) {
36 debug_assert!(kill.is_process_isolated());
37 let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
38 entries.push(RegistryEntry { kill, permit });
39 }
40
41 pub fn shutdown_progress(&self) -> usize {
45 let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
46 for e in entries.iter() {
47 e.kill.kill();
48 }
49 entries.retain(|e| e.kill.has_join());
50 entries.len()
51 }
52
53 pub fn live_count(&self) -> usize {
55 let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
56 entries.retain(|e| e.kill.has_join());
58 entries.len()
59 }
60
61 pub fn is_empty(&self) -> bool {
63 self.live_count() == 0
64 }
65}
66
67#[cfg(test)]
68mod tests {
69 use super::*;
70 use crate::transaction::process_tool::ProcessIsolatedToolHandler;
71 use crate::transaction::tool_handler::ToolHandler;
72 use monoloop_contracts::{
73 ChannelId, SessionId, SessionKey, ToolCall, ToolCallContext, ToolId, ToolName,
74 TransactionId,
75 };
76 use std::time::{Duration, Instant};
77
78 fn call_ctx() -> (ToolCall, ToolCallContext) {
79 let call = ToolCall {
80 tool_name: ToolName::try_new("sleep").unwrap(),
81 tool_id: ToolId::try_new("sleep").unwrap(),
82 provider_tool_call_id: "p".into(),
83 arguments: serde_json::json!({}),
84 request_ordinal: 0,
85 };
86 let ctx = ToolCallContext {
87 transaction_id: TransactionId::generate(),
88 session_key: SessionKey::new(
89 ChannelId::try_new("llm").unwrap(),
90 SessionId::try_new("s").unwrap(),
91 ),
92 exchange_id: None,
93 tool_action_id: monoloop_contracts::ToolActionId::new("a"),
94 tool_id: ToolId::try_new("sleep").unwrap(),
95 deadline: Instant::now() + Duration::from_secs(5),
96 };
97 (call, ctx)
98 }
99
100 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
101 async fn registry_retains_until_reap_then_empties() {
102 let reg = OwnedProcessRegistry::new();
103 let handler = ProcessIsolatedToolHandler::sleep_until_killed(3600);
104 let (call, ctx) = call_ctx();
105 let handle = handler.start(call, ctx).expect("start");
106 let kill = handle.kill.expect("kill");
107 assert!(kill.has_join());
108 reg.park(kill.clone(), None);
109 assert_eq!(reg.live_count(), 1);
110 assert!(!reg.is_empty());
111 let _ = reg.shutdown_progress();
113 kill.join_timeout(Duration::from_secs(2))
115 .await
116 .expect("reaped");
117 assert_eq!(reg.shutdown_progress(), 0);
118 assert!(reg.is_empty());
119 }
120}