harn_vm/bridge/
control.rs1use std::sync::{
2 atomic::{AtomicBool, Ordering},
3 Arc,
4};
5
6use tokio::sync::Notify;
7
8use super::{HostBridge, HostBridgeInjectionState};
9use crate::tool_call_cancellations::{fresh_registry, CancellationRegistry};
10
11pub(super) struct DaemonIdleState {
13 idle: AtomicBool,
14 changed: Arc<Notify>,
15}
16
17impl Default for DaemonIdleState {
18 fn default() -> Self {
19 Self {
20 idle: AtomicBool::new(false),
21 changed: Arc::new(Notify::new()),
22 }
23 }
24}
25
26impl HostBridge {
27 pub fn set_daemon_idle(&self, idle: bool) {
28 if self.daemon_idle.idle.swap(idle, Ordering::SeqCst) != idle {
29 self.daemon_idle.changed.notify_waiters();
30 }
31 }
32
33 pub fn is_daemon_idle(&self) -> bool {
34 self.daemon_idle.idle.load(Ordering::SeqCst)
35 }
36
37 pub(crate) fn daemon_idle_notifier(&self) -> Arc<Notify> {
38 self.daemon_idle.changed.clone()
39 }
40}
41
42#[derive(Clone)]
48pub struct HostBridgeControlState {
49 pub(super) cancelled: Arc<AtomicBool>,
50 pub(super) cancel_notify: Arc<Notify>,
51 pub(super) queued_transcript_injections: HostBridgeInjectionState,
52 pub(super) tool_call_cancellations: Arc<CancellationRegistry>,
53}
54
55impl HostBridgeControlState {
56 pub fn new(
57 cancelled: Arc<AtomicBool>,
58 cancel_notify: Arc<Notify>,
59 queued_transcript_injections: HostBridgeInjectionState,
60 tool_call_cancellations: Arc<CancellationRegistry>,
61 ) -> Self {
62 Self {
63 cancelled,
64 cancel_notify,
65 queued_transcript_injections,
66 tool_call_cancellations,
67 }
68 }
69
70 pub(super) fn isolated(cancelled: Arc<AtomicBool>) -> Self {
71 Self::new(
72 cancelled,
73 Arc::new(Notify::new()),
74 HostBridgeInjectionState::default(),
75 fresh_registry(),
76 )
77 }
78}
79
80pub(super) fn handle_cancel_tool_call_notification(
82 registry: &CancellationRegistry,
83 params: &serde_json::Value,
84) {
85 let session_id = params
86 .get("sessionId")
87 .or_else(|| params.get("session_id"))
88 .and_then(serde_json::Value::as_str)
89 .unwrap_or_default();
90 let call_id = params
91 .get("toolCallId")
92 .or_else(|| params.get("tool_call_id"))
93 .or_else(|| params.get("callId"))
94 .or_else(|| params.get("call_id"))
95 .and_then(serde_json::Value::as_str)
96 .unwrap_or_default();
97 if call_id.is_empty() {
98 return;
99 }
100 let reason = params
101 .get("reason")
102 .and_then(serde_json::Value::as_str)
103 .unwrap_or("host cancelled in-flight tool call")
104 .to_string();
105 let inject_reminder = params
106 .get("injectReminder")
107 .or_else(|| params.get("inject_reminder"))
108 .and_then(serde_json::Value::as_bool)
109 .unwrap_or(true);
110 registry.cancel(session_id, call_id, reason, inject_reminder);
111}
112
113#[cfg(test)]
114mod tests {
115 use super::*;
116
117 #[test]
118 fn targeted_cancel_notification_uses_the_supplied_registry() {
119 let registry = Arc::new(CancellationRegistry::default());
120 let (handle, _guard) = registry.register("session", "call", "shell");
121
122 handle_cancel_tool_call_notification(
123 ®istry,
124 &serde_json::json!({
125 "sessionId": "session",
126 "toolCallId": "call",
127 "reason": "host stop",
128 "injectReminder": false,
129 }),
130 );
131
132 assert!(handle.is_cancelled());
133 assert_eq!(handle.reason().as_deref(), Some("host stop"));
134 }
135}