harn-vm 0.10.63

//! Per-tool-call cancellation registry.
//!
//! When the host or a Harn script wants to abort a *specific* in-flight tool
//! call (e.g. user clicks "stop" on a runaway `git push --force`), they reach
//! through here. The registry keys cancellation handles by
//! `(session_id, call_id)` so a caller can target one call without
//! tearing down the whole session.
//!
//! ## Lifecycle
//!
//! 1. `host_agent_dispatch_tool_call` calls [`register`] right before it
//!    runs the underlying tool, holding the returned [`Handle`].
//! 2. The dispatch wraps the tool execution in a `select!` against
//!    [`Handle::cancelled`], so a triggered cancellation drops the
//!    pending tool future immediately.
//! 3. When the dispatch returns (either normally or via cancellation),
//!    the [`Guard`] returned by [`register`] auto-unregisters on drop.
//!
//! ## Triggering cancellation
//!
//! - **Harn:** the public `cancel_in_flight_tool_call` builtin registered
//!   in [`crate::stdlib::agent_sessions`].
//! - **Bridge stdin:** the `session/cancel_tool_call` notification
//!   handled in [`crate::bridge`].
//! - **ACP:** the `session/cancel_tool_call` method exposed by
//!   `crates/harn-serve/src/adapters/acp`.
//!
//! All three paths funnel into [`cancel`].
//!
//! Storage is thread-local because the VM runs on a current-thread tokio
//! worker — bridge stdin and ACP adapter share that worker via
//! `tokio::task::LocalSet`, so they can read the same registry without
//! cross-thread sync.

use std::cell::RefCell;
use std::collections::HashMap;
#[cfg(debug_assertions)]
use std::collections::HashSet;
use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(debug_assertions)]
use std::sync::OnceLock;
use std::sync::{Arc, Mutex};

use tokio::sync::Notify;

/// Reason + flags captured when a tool call is cancelled.
#[derive(Clone, Debug)]
pub struct CancellationDetails {
    pub reason: String,
    pub inject_reminder: bool,
}

#[derive(Debug)]
struct Inner {
    cancelled: AtomicBool,
    completed: AtomicBool,
    details: Mutex<Option<CancellationDetails>>,
    notify: Notify,
    completion_notify: Notify,
}

/// Per-tool-call cancellation handle. Cheap to clone (`Arc` inside).
#[derive(Clone, Debug)]
pub struct Handle {
    pub session_id: String,
    pub call_id: String,
    pub tool_name: String,
    inner: Arc<Inner>,
}

impl Handle {
    fn new(session_id: String, call_id: String, tool_name: String) -> Self {
        Self {
            session_id,
            call_id,
            tool_name,
            inner: Arc::new(Inner {
                cancelled: AtomicBool::new(false),
                completed: AtomicBool::new(false),
                details: Mutex::new(None),
                notify: Notify::new(),
                completion_notify: Notify::new(),
            }),
        }
    }

    pub fn is_cancelled(&self) -> bool {
        self.inner.cancelled.load(Ordering::SeqCst)
    }

    pub fn is_completed(&self) -> bool {
        self.inner.completed.load(Ordering::SeqCst)
    }

    /// Mark dispatch finished — called by the [`Guard`] on drop. Wakes any
    /// caller awaiting [`Handle::completed`].
    fn mark_completed(&self) {
        if !self.inner.completed.swap(true, Ordering::SeqCst) {
            self.inner.completion_notify.notify_waiters();
        }
    }

    /// Mark this call cancelled. Returns `true` if this call was the first
    /// to trigger cancellation (the caller can then push a reminder, etc.).
    pub fn cancel(&self, reason: impl Into<String>, inject_reminder: bool) -> bool {
        if self.inner.cancelled.swap(true, Ordering::SeqCst) {
            return false;
        }
        let reason = reason.into();
        let mut details = self
            .inner
            .details
            .lock()
            .unwrap_or_else(|err| err.into_inner());
        *details = Some(CancellationDetails {
            reason,
            inject_reminder,
        });
        drop(details);
        self.inner.notify.notify_waiters();
        true
    }

    pub fn details(&self) -> Option<CancellationDetails> {
        self.inner
            .details
            .lock()
            .unwrap_or_else(|err| err.into_inner())
            .clone()
    }

    pub fn reason(&self) -> Option<String> {
        self.details().map(|d| d.reason)
    }

    /// Future that resolves when this call is cancelled. Safe to await
    /// even if cancellation has already fired (returns immediately).
    pub async fn cancelled(&self) {
        if self.is_cancelled() {
            return;
        }
        let notified = self.inner.notify.notified();
        if self.is_cancelled() {
            return;
        }
        notified.await;
    }

    /// Future that resolves when dispatch finishes and the [`Guard`]
    /// drops. Lets a caller observe whether a cancellation actually
    /// landed (used by `cancel_in_flight_tool_call`'s `timeout_ms`).
    pub async fn completed(&self) {
        if self.is_completed() {
            return;
        }
        let notified = self.inner.completion_notify.notified();
        if self.is_completed() {
            return;
        }
        notified.await;
    }
}

thread_local! {
    static REGISTRY: RefCell<HashMap<(String, String), Handle>> =
        RefCell::new(HashMap::new());
}

// The production registry intentionally stays thread-local: every VM and its
// cancellation ingress run on one LocalSet thread, while separate embedded VMs
// may reuse the same session/call ids. In debug builds, mirror only registration
// ownership process-wide so a future runtime refactor cannot silently turn a
// cross-thread cancellation into `NotFound`.
#[cfg(debug_assertions)]
type RegistrationKey = (String, String);

#[cfg(debug_assertions)]
fn affinity_owners() -> &'static Mutex<HashMap<RegistrationKey, HashSet<std::thread::ThreadId>>> {
    static OWNERS: OnceLock<Mutex<HashMap<RegistrationKey, HashSet<std::thread::ThreadId>>>> =
        OnceLock::new();
    OWNERS.get_or_init(|| Mutex::new(HashMap::new()))
}

#[cfg(debug_assertions)]
fn track_affinity(session_id: &str, call_id: &str, owner: std::thread::ThreadId) {
    affinity_owners()
        .lock()
        .unwrap_or_else(|error| error.into_inner())
        .entry((session_id.to_string(), call_id.to_string()))
        .or_default()
        .insert(owner);
}

#[cfg(debug_assertions)]
fn untrack_affinity(session_id: &str, call_id: &str, owner: std::thread::ThreadId) {
    let key = (session_id.to_string(), call_id.to_string());
    let mut owners_by_key = affinity_owners()
        .lock()
        .unwrap_or_else(|error| error.into_inner());
    if let Some(owners) = owners_by_key.get_mut(&key) {
        owners.remove(&owner);
        if owners.is_empty() {
            owners_by_key.remove(&key);
        }
    }
}

#[cfg(debug_assertions)]
fn assert_no_foreign_registration(session_id: &str, call_id: &str) {
    let current = std::thread::current().id();
    let key = (session_id.to_string(), call_id.to_string());
    let owners_by_key = affinity_owners()
        .lock()
        .unwrap_or_else(|error| error.into_inner());
    let foreign_owner = owners_by_key
        .get(&key)
        .and_then(|owners| owners.iter().find(|owner| **owner != current));
    debug_assert!(
        foreign_owner.is_none(),
        "tool-call cancellation registry affinity violated: \
         ({session_id}, {call_id}) is registered on {foreign_owner:?}, \
         but lookup ran on {current:?}"
    );
}

/// RAII guard that unregisters the (session_id, call_id) entry on drop
/// and signals the matching handle's `completed` future.
pub struct Guard {
    session_id: String,
    call_id: String,
    handle: Handle,
    #[cfg(debug_assertions)]
    owner_thread: std::thread::ThreadId,
}

impl Drop for Guard {
    fn drop(&mut self) {
        if self.call_id.is_empty() {
            return;
        }
        #[cfg(debug_assertions)]
        {
            let current = std::thread::current().id();
            debug_assert_eq!(
                self.owner_thread, current,
                "tool-call cancellation guard dropped off its registration thread"
            );
            untrack_affinity(&self.session_id, &self.call_id, self.owner_thread);
        }
        self.handle.mark_completed();
        REGISTRY.with(|registry| {
            registry
                .borrow_mut()
                .remove(&(self.session_id.clone(), self.call_id.clone()));
        });
    }
}

/// Register a fresh handle for an in-flight tool call.
///
/// Returns `None` when the call has no id (e.g. some legacy parse paths
/// omit it). The caller cannot be targeted by `cancel_in_flight_tool_call`
/// in that case, so cancellation is a no-op; that matches today's behavior
/// for the few code paths that produce id-less calls.
pub fn register(
    session_id: impl Into<String>,
    call_id: impl Into<String>,
    tool_name: impl Into<String>,
) -> Option<(Handle, Guard)> {
    let session_id = session_id.into();
    let call_id = call_id.into();
    let tool_name = tool_name.into();
    if call_id.is_empty() {
        return None;
    }
    let handle = Handle::new(session_id.clone(), call_id.clone(), tool_name);
    let guard = Guard {
        session_id: session_id.clone(),
        call_id: call_id.clone(),
        handle: handle.clone(),
        #[cfg(debug_assertions)]
        owner_thread: std::thread::current().id(),
    };
    REGISTRY.with(|registry| {
        registry
            .borrow_mut()
            .insert((session_id, call_id), handle.clone());
    });
    #[cfg(debug_assertions)]
    track_affinity(&guard.session_id, &guard.call_id, guard.owner_thread);
    Some((handle, guard))
}

pub fn lookup(session_id: &str, call_id: &str) -> Option<Handle> {
    let handle = REGISTRY.with(|registry| {
        registry
            .borrow()
            .get(&(session_id.to_string(), call_id.to_string()))
            .cloned()
    });
    #[cfg(debug_assertions)]
    if handle.is_none() {
        assert_no_foreign_registration(session_id, call_id);
    }
    handle
}

/// List all in-flight handles for a session — used by introspection and
/// tests; not by the hot path.
pub fn list_for_session(session_id: &str) -> Vec<Handle> {
    REGISTRY.with(|registry| {
        registry
            .borrow()
            .values()
            .filter(|handle| handle.session_id == session_id)
            .cloned()
            .collect()
    })
}

/// Outcome returned to callers of [`cancel`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CancelStatus {
    /// Cancellation flag was set on this call. The in-flight future will
    /// observe it and unwind.
    Cancelled,
    /// The call was already cancelled by an earlier request.
    AlreadyCancelled,
    /// No in-flight call matches (session_id, call_id) right now —
    /// either it never started, or it already completed.
    NotFound,
}

impl CancelStatus {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Cancelled => "cancelled",
            Self::AlreadyCancelled => "already_cancelled",
            Self::NotFound => "not_found",
        }
    }
}

/// Result of a cancellation request: the status plus, when a handle was
/// found, the tool's name and the live handle (so callers can wait for
/// the dispatch to actually unwind via [`Handle::completed`]).
#[derive(Clone, Debug)]
pub struct CancelOutcome {
    pub status: CancelStatus,
    pub tool_name: Option<String>,
    pub handle: Option<Handle>,
}

/// Trigger cancellation for one in-flight tool call.
pub fn cancel(
    session_id: &str,
    call_id: &str,
    reason: impl Into<String>,
    inject_reminder: bool,
) -> CancelOutcome {
    let Some(handle) = lookup(session_id, call_id) else {
        return CancelOutcome {
            status: CancelStatus::NotFound,
            tool_name: None,
            handle: None,
        };
    };
    let tool_name = Some(handle.tool_name.clone());
    let status = if handle.cancel(reason, inject_reminder) {
        CancelStatus::Cancelled
    } else {
        CancelStatus::AlreadyCancelled
    };
    CancelOutcome {
        status,
        tool_name,
        handle: Some(handle),
    }
}

/// Drop every in-flight cancellation handle on this thread. A handle is
/// normally unregistered by its [`Guard`] on drop, but an abandoned
/// dispatch (test timeout, panic during teardown) can leave a
/// `(session_id, call_id)` entry behind. The per-test reset path calls
/// this so a reused worker thread does not accumulate stale handles
/// across the suite.
pub fn reset_registry() {
    #[cfg(debug_assertions)]
    {
        let owner = std::thread::current().id();
        let removed_keys: Vec<_> =
            REGISTRY.with(|registry| registry.borrow_mut().drain().map(|entry| entry.0).collect());
        for (session_id, call_id) in removed_keys {
            untrack_affinity(&session_id, &call_id, owner);
        }
    }
    #[cfg(not(debug_assertions))]
    REGISTRY.with(|registry| registry.borrow_mut().clear());
}

/// Number of in-flight cancellation handles on this thread. Test-only.
#[cfg(test)]
pub fn registry_len() -> usize {
    REGISTRY.with(|registry| registry.borrow().len())
}

#[cfg(test)]
pub fn clear_registry_for_test() {
    reset_registry();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn handle_resolves_cancelled_future() {
        clear_registry_for_test();
        let (handle, _guard) = register("sess_1", "call_1", "bash").expect("registered");
        assert!(!handle.is_cancelled());
        let cancelled = handle.cancel("user requested stop", false);
        assert!(cancelled);
        // Awaiting after the fact should resolve immediately.
        handle.cancelled().await;
        assert!(handle.is_cancelled());
        assert_eq!(handle.reason().as_deref(), Some("user requested stop"));
    }

    #[tokio::test]
    async fn cancel_returns_not_found_when_missing() {
        clear_registry_for_test();
        let outcome = cancel("sess_unknown", "call_unknown", "irrelevant", false);
        assert_eq!(outcome.status, CancelStatus::NotFound);
        assert_eq!(outcome.tool_name, None);
    }

    #[tokio::test]
    async fn cancel_is_idempotent() {
        clear_registry_for_test();
        let (_handle, _guard) = register("sess", "call_2", "shell").expect("registered");
        let first = cancel("sess", "call_2", "first", false);
        let second = cancel("sess", "call_2", "second", false);
        assert_eq!(first.status, CancelStatus::Cancelled);
        assert_eq!(second.status, CancelStatus::AlreadyCancelled);
        assert_eq!(first.tool_name.as_deref(), Some("shell"));
        assert_eq!(second.tool_name.as_deref(), Some("shell"));
    }

    #[tokio::test]
    async fn guard_unregisters_on_drop() {
        clear_registry_for_test();
        {
            let _registration = register("sess", "call_g", "tool").expect("registered");
            assert!(lookup("sess", "call_g").is_some());
        }
        assert!(lookup("sess", "call_g").is_none());
    }

    #[test]
    fn cancelled_wakes_pending_waiter() {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("rt");
        let local = tokio::task::LocalSet::new();
        local.block_on(&rt, async {
            clear_registry_for_test();
            let (handle, _guard) = register("sess", "call_w", "tool").expect("registered");
            let waiter_handle = handle.clone();
            let task = tokio::task::spawn_local(async move { waiter_handle.cancelled().await });
            tokio::task::yield_now().await;
            handle.cancel("stopping", false);
            tokio::time::timeout(std::time::Duration::from_secs(1), task)
                .await
                .expect("task should resolve quickly")
                .expect("task did not panic");
        });
    }

    #[tokio::test]
    async fn list_for_session_filters_by_session() {
        clear_registry_for_test();
        let _r1 = register("sess_a", "call_x", "tool").expect("registered");
        let _r2 = register("sess_a", "call_y", "tool").expect("registered");
        let _r3 = register("sess_b", "call_x", "tool").expect("registered");
        let mut ids: Vec<String> = list_for_session("sess_a")
            .into_iter()
            .map(|h| h.call_id)
            .collect();
        ids.sort();
        assert_eq!(ids, vec!["call_x".to_string(), "call_y".to_string()]);
    }

    #[cfg(debug_assertions)]
    #[test]
    fn cross_thread_cancel_trips_the_affinity_invariant() {
        clear_registry_for_test();
        let (registered_tx, registered_rx) = std::sync::mpsc::sync_channel(0);
        let (release_tx, release_rx) = std::sync::mpsc::sync_channel(0);

        let result = std::thread::scope(|scope| {
            scope.spawn(move || {
                let _registration =
                    register("sess_affinity", "call_affinity", "tool").expect("registered");
                registered_tx.send(()).expect("announce registration");
                release_rx.recv().expect("wait for cancellation");
            });

            registered_rx.recv().expect("registration should be live");
            let result = std::panic::catch_unwind(|| {
                cancel("sess_affinity", "call_affinity", "cancelled", false)
            });
            release_tx.send(()).expect("release registration thread");
            result
        });

        assert!(
            result.is_err(),
            "cross-thread cancellation of a live registration must trip the debug affinity invariant"
        );
    }
}