Skip to main content

harn_vm/bridge/
control.rs

1use 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
11/// Observable idle boundary for one daemon-mode host bridge.
12pub(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/// Shared control state for one host-bridge execution domain.
43///
44/// Keeping cancellation, notification, queued injection, and targeted
45/// tool-call routing in one value prevents protocol adapters from assembling
46/// only part of the state a VM and its out-of-band control task must share.
47#[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
80/// Apply a host notification to the bridge's explicit cancellation address space.
81pub(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            &registry,
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}