car_engine/

tool_handles.rs

//! Per-runtime registry of detached tool invocations (C2).
//!
//! A `ToolCall` action with a detached [`car_ir::tool_stream::ToolInvocationMode`]
//! (`streaming` / `long_running`) is *started*, not awaited: `dispatch`
//! registers the invocation here, hands the caller a [`ToolHandle`] as
//! the action's output, and the DAG proceeds. Chunks the executor
//! yields are buffered per handle (drained by [`ToolHandleRegistry::poll`])
//! and fanned out on a broadcast channel so a server can forward them to
//! subscribed clients as `tools.stream.event` notifications.
//!
//! Lifecycle: `register` → zero or more `push_chunk` → a terminal chunk
//! (`done` / `error`) or `cancel` seals the entry. A sealed entry stays
//! queryable until a `poll` observes it *after* draining its last chunks;
//! that poll removes it (the "read your final answer once" contract —
//! mirrors how idempotency results are consumed).

use std::collections::HashMap;
use std::sync::Arc;

use car_ir::{ToolHandle, ToolStatus, ToolStreamChunk, ToolStreamEvent};
use serde::{Deserialize, Serialize};
use tokio::sync::{broadcast, Mutex};
use tokio_util::sync::CancellationToken;

/// Broadcast capacity for [`ToolStreamEvent`] fanout. A lagging
/// subscriber loses oldest events (chunks remain available via `poll`,
/// which reads the buffer, not the broadcast).
const EVENT_CHANNEL_CAP: usize = 256;

/// Cap on chunks buffered per handle awaiting a `poll` (linus review D2).
/// A subscriber-only client (the documented low-latency pattern) never
/// polls, so an uncapped buffer is an OOM with a long-running tool as
/// the trigger. Oldest chunks are dropped first; the drop count is
/// surfaced on the next poll so a late poller knows the gap exists.
const MAX_BUFFERED_CHUNKS: usize = 1024;

/// How long a terminal (done/failed/cancelled) entry stays queryable when
/// nobody polls it (linus review D2). Reaped opportunistically on
/// register/poll/cancel — a never-polled cancel must not live for the
/// session's lifetime.
const TERMINAL_TTL: std::time::Duration = std::time::Duration::from_secs(15 * 60);

/// What `poll` returns: the drained chunks since the last poll plus the
/// invocation's current status and, once terminal, its final result or
/// error extracted from the terminal chunk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolPollResult {
    pub handle: String,
    pub tool: String,
    pub action_id: String,
    pub status: ToolStatus,
    /// Chunks produced since the previous poll (oldest first). Terminal
    /// chunk included when the stream just finished.
    pub chunks: Vec<ToolStreamChunk>,
    /// Chunks dropped from the buffer because it hit its cap before this
    /// poll drained it (0 for callers that poll promptly). The event
    /// broadcast may still have delivered them live.
    #[serde(default, skip_serializing_if = "is_zero")]
    pub dropped_chunks: u64,
    /// Set once status is `succeeded` (from the `done` chunk's payload).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub result: Option<serde_json::Value>,
    /// Set once status is `failed` (from the `error` chunk's message).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

fn is_zero(n: &u64) -> bool {
    *n == 0
}

struct HandleEntry {
    tool: String,
    action_id: String,
    status: ToolStatus,
    buffered: Vec<ToolStreamChunk>,
    /// Chunks evicted from `buffered` at the cap since the last poll.
    dropped_chunks: u64,
    result: Option<serde_json::Value>,
    error: Option<String>,
    cancel: CancellationToken,
    /// True once a poll has observed the entry in a terminal state —
    /// the next terminal observation removes it (so a caller always
    /// sees the terminal status at least once).
    drained_terminal: bool,
    /// When the entry reached a terminal state; drives the TTL reap.
    sealed_at: Option<std::time::Instant>,
}

/// Registry of live (and recently-terminal) detached tool invocations.
/// One per [`crate::executor::Runtime`]; cheap to share via `Arc`.
pub struct ToolHandleRegistry {
    entries: Mutex<HashMap<String, HandleEntry>>,
    events: broadcast::Sender<ToolStreamEvent>,
}

impl Default for ToolHandleRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl ToolHandleRegistry {
    pub fn new() -> Self {
        let (events, _) = broadcast::channel(EVENT_CHANNEL_CAP);
        Self {
            entries: Mutex::new(HashMap::new()),
            events,
        }
    }

    /// Subscribe to the live event fanout (every chunk from every handle,
    /// wrapped in a [`ToolStreamEvent`]). Used by the WS layer to forward
    /// `tools.stream.event` notifications.
    pub fn subscribe(&self) -> broadcast::Receiver<ToolStreamEvent> {
        self.events.subscribe()
    }

    /// Register a new detached invocation. Returns the opaque handle and
    /// the cancel token the drain task must select on.
    pub async fn register(&self, tool: &str, action_id: &str) -> (ToolHandle, CancellationToken) {
        // Full 128-bit id (linus review Q3): the previous 12-hex-char
        // truncation made a collision a silent cross-wiring of two live
        // invocations (plain insert overwrites). Full width is free.
        let id = uuid::Uuid::new_v4().simple().to_string();
        let cancel = CancellationToken::new();
        let entry = HandleEntry {
            tool: tool.to_string(),
            action_id: action_id.to_string(),
            status: ToolStatus::Running,
            buffered: Vec::new(),
            dropped_chunks: 0,
            result: None,
            error: None,
            cancel: cancel.clone(),
            drained_terminal: false,
            sealed_at: None,
        };
        let mut entries = self.entries.lock().await;
        Self::reap_expired(&mut entries);
        entries.insert(id.clone(), entry);
        drop(entries);
        (ToolHandle::new(id), cancel)
    }

    /// Drop terminal entries older than [`TERMINAL_TTL`] (linus review
    /// D2): a never-polled terminal handle must not live for the
    /// session's lifetime. Called opportunistically under the lock from
    /// register/poll/cancel — no background task needed.
    fn reap_expired(entries: &mut HashMap<String, HandleEntry>) {
        entries.retain(|_, e| match e.sealed_at {
            Some(at) => at.elapsed() < TERMINAL_TTL,
            None => true,
        });
    }

    /// Append a chunk from the executor's stream. A terminal chunk seals
    /// the status (`done` → `succeeded` + result, `error` → `failed` +
    /// message). Chunks arriving after a terminal state (e.g. a racing
    /// executor after cancel) are dropped. Every accepted chunk is also
    /// broadcast as a [`ToolStreamEvent`].
    pub async fn push_chunk(&self, handle_id: &str, chunk: ToolStreamChunk) {
        let mut entries = self.entries.lock().await;
        let Some(entry) = entries.get_mut(handle_id) else {
            return;
        };
        if entry.status.is_terminal() {
            return;
        }
        match &chunk {
            ToolStreamChunk::Done { result } => {
                entry.status = ToolStatus::Succeeded;
                entry.result = result.clone();
                entry.sealed_at = Some(std::time::Instant::now());
            }
            ToolStreamChunk::Error { message } => {
                entry.status = ToolStatus::Failed;
                entry.error = Some(message.clone());
                entry.sealed_at = Some(std::time::Instant::now());
            }
            _ => {}
        }
        // Bounded buffer (linus review D2): drop-oldest at the cap and
        // count the drops. A terminal chunk always fits (it just landed
        // in status/result above regardless).
        if entry.buffered.len() >= MAX_BUFFERED_CHUNKS {
            entry.buffered.remove(0);
            entry.dropped_chunks += 1;
        }
        entry.buffered.push(chunk.clone());
        // Broadcast after the buffer write so a poll racing the event
        // never observes the event before the chunk is drainable.
        let _ = self.events.send(ToolStreamEvent {
            handle: ToolHandle::new(handle_id.to_string()),
            chunk,
        });
    }

    /// The executor's stream ended without a terminal chunk — a dropped
    /// sender (panicked task, closed pipe). Treated as failure, unless the
    /// entry was already sealed (normal after a `done`/`error`/cancel).
    pub async fn mark_stream_closed(&self, handle_id: &str) {
        let mut entries = self.entries.lock().await;
        if let Some(entry) = entries.get_mut(handle_id) {
            if !entry.status.is_terminal() {
                entry.status = ToolStatus::Failed;
                entry.error = Some("tool stream closed without a terminal chunk".to_string());
                entry.sealed_at = Some(std::time::Instant::now());
            }
        }
    }

    /// Request cancellation: fires the cancel token (the drain task and a
    /// cooperative executor observe it) and seals the entry as
    /// `cancelled` unless already terminal. Returns false for an unknown
    /// handle.
    pub async fn cancel(&self, handle_id: &str) -> bool {
        let mut entries = self.entries.lock().await;
        Self::reap_expired(&mut entries);
        let Some(entry) = entries.get_mut(handle_id) else {
            return false;
        };
        entry.cancel.cancel();
        if !entry.status.is_terminal() {
            entry.status = ToolStatus::Cancelled;
            entry.sealed_at = Some(std::time::Instant::now());
        }
        true
    }

    /// Cancel every live invocation (linus review D3): wired into the
    /// daemon's session teardown so a dropped WS connection doesn't
    /// orphan running detached tools — each session's runtime (and thus
    /// registry) is unreachable after disconnect, so anything still
    /// running would be uncancellable until daemon restart. Returns how
    /// many live invocations were cancelled.
    pub async fn cancel_all(&self) -> usize {
        let mut entries = self.entries.lock().await;
        let mut n = 0;
        for e in entries.values_mut() {
            if !e.status.is_terminal() {
                e.cancel.cancel();
                e.status = ToolStatus::Cancelled;
                e.sealed_at = Some(std::time::Instant::now());
                n += 1;
            }
        }
        n
    }

    /// Drain buffered chunks + report status. `None` for an unknown (or
    /// already fully-consumed) handle. The FIRST poll that observes the
    /// entry in a terminal state marks it consumed (it also drains any
    /// remaining chunks); the next poll removes it — so a caller always
    /// gets to see the terminal status at least once.
    pub async fn poll(&self, handle_id: &str) -> Option<ToolPollResult> {
        let mut entries = self.entries.lock().await;
        Self::reap_expired(&mut entries);
        let entry = entries.get_mut(handle_id)?;
        let chunks = std::mem::take(&mut entry.buffered);
        let dropped = std::mem::take(&mut entry.dropped_chunks);
        let res = ToolPollResult {
            handle: handle_id.to_string(),
            tool: entry.tool.clone(),
            action_id: entry.action_id.clone(),
            status: entry.status,
            chunks,
            dropped_chunks: dropped,
            result: entry.result.clone(),
            error: entry.error.clone(),
        };
        if entry.status.is_terminal() {
            if entry.drained_terminal {
                entries.remove(handle_id);
            } else {
                entry.drained_terminal = true;
            }
        }
        Some(res)
    }

    /// Current status without draining (used by tests and liveness checks).
    pub async fn status(&self, handle_id: &str) -> Option<ToolStatus> {
        self.entries.lock().await.get(handle_id).map(|e| e.status)
    }

    /// Number of live (non-terminal) invocations.
    pub async fn live_count(&self) -> usize {
        self.entries
            .lock()
            .await
            .values()
            .filter(|e| !e.status.is_terminal())
            .count()
    }
}

/// Spawn the drain task tying an executor's chunk stream to the registry:
/// forwards chunks until a terminal chunk, stream close, or cancellation.
/// Returns immediately; the task owns the receiver.
pub fn spawn_drain(
    registry: Arc<ToolHandleRegistry>,
    handle_id: String,
    rx: tokio::sync::mpsc::Receiver<ToolStreamChunk>,
    cancel: CancellationToken,
) {
    drop(spawn_drain_task(registry, handle_id, rx, cancel));
}

fn spawn_drain_task(
    registry: Arc<ToolHandleRegistry>,
    handle_id: String,
    mut rx: tokio::sync::mpsc::Receiver<ToolStreamChunk>,
    cancel: CancellationToken,
) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        loop {
            tokio::select! {
                _ = cancel.cancelled() => {
                    // Status was already sealed by `cancel()`; dropping rx
                    // is the signal a cooperative executor sees.
                    break;
                }
                chunk = rx.recv() => match chunk {
                    Some(c) => {
                        let terminal = c.is_terminal();
                        registry.push_chunk(&handle_id, c).await;
                        if terminal {
                            break;
                        }
                    }
                    None => {
                        registry.mark_stream_closed(&handle_id).await;
                        break;
                    }
                }
            }
        }
    })
}

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

    #[tokio::test]
    async fn chunks_buffer_and_drain_in_order() {
        let reg = ToolHandleRegistry::new();
        let (h, _tok) = reg.register("tail_log", "a1").await;
        reg.push_chunk(&h.id, ToolStreamChunk::Text { text: "one".into() })
            .await;
        reg.push_chunk(&h.id, ToolStreamChunk::Text { text: "two".into() })
            .await;
        let poll = reg.poll(&h.id).await.expect("known handle");
        assert_eq!(poll.status, ToolStatus::Running);
        assert_eq!(poll.chunks.len(), 2);
        // Drained — a second poll returns no chunks but stays queryable.
        let poll2 = reg.poll(&h.id).await.expect("still live");
        assert!(poll2.chunks.is_empty());
    }

    #[tokio::test]
    async fn done_chunk_seals_success_with_result() {
        let reg = ToolHandleRegistry::new();
        let (h, _tok) = reg.register("build", "a2").await;
        reg.push_chunk(
            &h.id,
            ToolStreamChunk::Done {
                result: Some(serde_json::json!({"exit": 0})),
            },
        )
        .await;
        let poll = reg.poll(&h.id).await.unwrap();
        assert_eq!(poll.status, ToolStatus::Succeeded);
        assert_eq!(poll.result, Some(serde_json::json!({"exit": 0})));
        // Late chunks after terminal are dropped.
        reg.push_chunk(
            &h.id,
            ToolStreamChunk::Text {
                text: "late".into(),
            },
        )
        .await;
        let poll2 = reg.poll(&h.id).await.unwrap();
        assert!(poll2.chunks.is_empty());
        // Terminal observed twice with empty buffer → entry consumed.
        assert!(reg.poll(&h.id).await.is_none());
    }

    #[tokio::test]
    async fn cancel_seals_cancelled_and_fires_token() {
        let reg = ToolHandleRegistry::new();
        let (h, tok) = reg.register("watch", "a3").await;
        assert!(reg.cancel(&h.id).await);
        assert!(tok.is_cancelled());
        assert_eq!(reg.status(&h.id).await, Some(ToolStatus::Cancelled));
        // Executor racing past the cancel can't overwrite the status.
        reg.push_chunk(&h.id, ToolStreamChunk::Done { result: None })
            .await;
        assert_eq!(reg.status(&h.id).await, Some(ToolStatus::Cancelled));
        assert!(!reg.cancel("nope").await);
    }

    #[tokio::test]
    async fn drain_task_forwards_until_terminal() {
        let reg = Arc::new(ToolHandleRegistry::new());
        let (h, tok) = reg.register("gen", "a4").await;
        let (tx, rx) = tokio::sync::mpsc::channel(8);
        let drain = spawn_drain_task(reg.clone(), h.id.clone(), rx, tok);
        tx.send(ToolStreamChunk::Progress {
            fraction: 0.5,
            message: Some("half".into()),
        })
        .await
        .unwrap();
        tx.send(ToolStreamChunk::Done { result: None })
            .await
            .unwrap();
        drain.await.expect("drain task must observe terminal chunk");
        let poll = reg.poll(&h.id).await.unwrap();
        assert_eq!(poll.status, ToolStatus::Succeeded);
        assert_eq!(poll.chunks.len(), 2);
    }

    #[tokio::test]
    async fn dropped_stream_without_terminal_is_failure() {
        let reg = Arc::new(ToolHandleRegistry::new());
        let (h, tok) = reg.register("flaky", "a5").await;
        let (tx, rx) = tokio::sync::mpsc::channel(8);
        let drain = spawn_drain_task(reg.clone(), h.id.clone(), rx, tok);
        tx.send(ToolStreamChunk::Text {
            text: "partial".into(),
        })
        .await
        .unwrap();
        drop(tx);
        drain
            .await
            .expect("drain task must observe stream closure without terminal chunk");
        let poll = reg.poll(&h.id).await.unwrap();
        assert_eq!(poll.status, ToolStatus::Failed);
        assert!(poll.error.unwrap().contains("without a terminal chunk"));
    }

    #[tokio::test]
    async fn events_broadcast_to_subscribers() {
        let reg = ToolHandleRegistry::new();
        let mut sub = reg.subscribe();
        let (h, _tok) = reg.register("emit", "a6").await;
        reg.push_chunk(&h.id, ToolStreamChunk::Text { text: "x".into() })
            .await;
        let ev = sub.recv().await.expect("event delivered");
        assert_eq!(ev.handle.id, h.id);
        assert!(matches!(ev.chunk, ToolStreamChunk::Text { .. }));
    }
}