nerve-ipc-core 0.1.1

//! Per-connection request lifecycle and cancellation tracking.
//!
//! Each accepted connection (UDS or WebSocket) owns a [`RequestTable`] that
//! tracks in-flight requests for that connection only. Request IDs are scoped
//! to the connection: the same ID on two different connections refers to two
//! independent requests, and cancelling one cannot affect the other.
//!
//! # Lifecycle
//!
//! The intended sequence for every request is:
//!
//! 1. **Insert** — call [`RequestTable::insert`] when the request begins
//!    (on `SearchQuery` or `AgentTaskStart` reception).
//! 2. **Check** — call [`RequestTable::is_cancelled`] before each unit of
//!    work (before streaming a result token, before an HTTP call, etc.).
//! 3. **Cancel** — call [`RequestTable::cancel`] when a `Cancel` frame
//!    arrives. The AI daemon will notice on the next cancellation check.
//! 4. **Remove** — call [`RequestTable::remove`] when the request is
//!    complete, regardless of whether it was cancelled.
//!
//! Entries are not removed automatically; callers are responsible for cleanup.

use std::collections::HashMap;

use nerve_protocol::types::RequestId;

/// State of a request tracked by the [`RequestTable`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RequestState {
    /// The request is active and has not been cancelled.
    Active,
    /// The request has been cancelled. The AI daemon should stop work for this
    /// request at the next opportunity.
    Cancelled,
}

/// Tracks in-flight requests for a single connection.
///
/// Each accepted connection (UDS or WebSocket) has its own `RequestTable`.
/// Request IDs are scoped to the connection: the same ID on two different
/// connections refers to two independent requests.
///
/// # Examples
///
/// ```
/// use nerve_ipc_core::RequestTable;
/// use nerve_protocol::types::RequestId;
///
/// let mut table = RequestTable::new();
/// let id = RequestId(42);
///
/// // 1. Register when the request begins.
/// assert!(table.insert(id));
///
/// // 2. Check before producing results.
/// assert!(!table.is_cancelled(id));
///
/// // 3. Mark cancelled when a Cancel message arrives.
/// table.cancel(id);
/// assert!(table.is_cancelled(id));
/// assert!(!table.is_active(id));
///
/// // 4. Remove on completion.
/// table.remove(id);
/// assert!(table.is_empty());
/// ```
#[derive(Default)]
pub struct RequestTable {
    requests: HashMap<RequestId, RequestState>,
}

impl RequestTable {
    /// Create a new empty request table.
    #[inline]
    pub fn new() -> Self {
        Self {
            requests: HashMap::new(),
        }
    }

    /// Register a new request as [`RequestState::Active`].
    ///
    /// Returns `false` if `request_id` is already present; the existing entry
    /// is left unchanged in that case.
    pub fn insert(&mut self, request_id: RequestId) -> bool {
        match self.requests.get(&request_id) {
            Some(_) => false,
            None => {
                self.requests.insert(request_id, RequestState::Active);
                true
            }
        }
    }

    /// Mark an existing request as [`RequestState::Cancelled`].
    ///
    /// Returns `false` if `request_id` is not present in the table.
    pub fn cancel(&mut self, request_id: RequestId) -> bool {
        match self.requests.get_mut(&request_id) {
            Some(state) => {
                *state = RequestState::Cancelled;
                true
            }
            None => false,
        }
    }

    /// Returns `true` if the request has been cancelled.
    pub fn is_cancelled(&self, request_id: RequestId) -> bool {
        matches!(
            self.requests.get(&request_id),
            Some(RequestState::Cancelled)
        )
    }

    /// Remove the request from the table, regardless of its state.
    pub fn remove(&mut self, request_id: RequestId) {
        self.requests.remove(&request_id);
    }

    /// Number of requests currently tracked (active and cancelled combined).
    pub fn len(&self) -> usize {
        self.requests.len()
    }

    /// Returns `true` if no requests are currently tracked.
    pub fn is_empty(&self) -> bool {
        self.requests.is_empty()
    }

    /// Returns `true` if `request_id` is present in the table, regardless of
    /// its state.
    pub fn contains(&self, request_id: RequestId) -> bool {
        self.requests.contains_key(&request_id)
    }

    /// Returns `true` if the request is present and in the
    /// [`RequestState::Active`] state.
    pub fn is_active(&self, request_id: RequestId) -> bool {
        matches!(self.requests.get(&request_id), Some(RequestState::Active))
    }
}