nerve_ipc_core/
request_table.rs

1//! Per-connection request lifecycle and cancellation tracking.
2//!
3//! Each accepted connection (UDS or WebSocket) owns a [`RequestTable`] that
4//! tracks in-flight requests for that connection only. Request IDs are scoped
5//! to the connection: the same ID on two different connections refers to two
6//! independent requests, and cancelling one cannot affect the other.
7//!
8//! # Lifecycle
9//!
10//! The intended sequence for every request is:
11//!
12//! 1. **Insert** — call [`RequestTable::insert`] when the request begins
13//!    (on `SearchQuery` or `AgentTaskStart` reception).
14//! 2. **Check** — call [`RequestTable::is_cancelled`] before each unit of
15//!    work (before streaming a result token, before an HTTP call, etc.).
16//! 3. **Cancel** — call [`RequestTable::cancel`] when a `Cancel` frame
17//!    arrives. The AI daemon will notice on the next cancellation check.
18//! 4. **Remove** — call [`RequestTable::remove`] when the request is
19//!    complete, regardless of whether it was cancelled.
20//!
21//! Entries are not removed automatically; callers are responsible for cleanup.
22
23use std::collections::HashMap;
24
25use nerve_protocol::types::RequestId;
26
27/// State of a request tracked by the [`RequestTable`].
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum RequestState {
30    /// The request is active and has not been cancelled.
31    Active,
32    /// The request has been cancelled. The AI daemon should stop work for this
33    /// request at the next opportunity.
34    Cancelled,
35}
36
37/// Tracks in-flight requests for a single connection.
38///
39/// Each accepted connection (UDS or WebSocket) has its own `RequestTable`.
40/// Request IDs are scoped to the connection: the same ID on two different
41/// connections refers to two independent requests.
42///
43/// # Examples
44///
45/// ```
46/// use nerve_ipc_core::RequestTable;
47/// use nerve_protocol::types::RequestId;
48///
49/// let mut table = RequestTable::new();
50/// let id = RequestId(42);
51///
52/// // 1. Register when the request begins.
53/// assert!(table.insert(id));
54///
55/// // 2. Check before producing results.
56/// assert!(!table.is_cancelled(id));
57///
58/// // 3. Mark cancelled when a Cancel message arrives.
59/// table.cancel(id);
60/// assert!(table.is_cancelled(id));
61/// assert!(!table.is_active(id));
62///
63/// // 4. Remove on completion.
64/// table.remove(id);
65/// assert!(table.is_empty());
66/// ```
67#[derive(Default)]
68pub struct RequestTable {
69    requests: HashMap<RequestId, RequestState>,
70}
71
72impl RequestTable {
73    /// Create a new empty request table.
74    #[inline]
75    pub fn new() -> Self {
76        Self {
77            requests: HashMap::new(),
78        }
79    }
80
81    /// Register a new request as [`RequestState::Active`].
82    ///
83    /// Returns `false` if `request_id` is already present; the existing entry
84    /// is left unchanged in that case.
85    pub fn insert(&mut self, request_id: RequestId) -> bool {
86        match self.requests.get(&request_id) {
87            Some(_) => false,
88            None => {
89                self.requests.insert(request_id, RequestState::Active);
90                true
91            }
92        }
93    }
94
95    /// Mark an existing request as [`RequestState::Cancelled`].
96    ///
97    /// Returns `false` if `request_id` is not present in the table.
98    pub fn cancel(&mut self, request_id: RequestId) -> bool {
99        match self.requests.get_mut(&request_id) {
100            Some(state) => {
101                *state = RequestState::Cancelled;
102                true
103            }
104            None => false,
105        }
106    }
107
108    /// Returns `true` if the request has been cancelled.
109    pub fn is_cancelled(&self, request_id: RequestId) -> bool {
110        matches!(
111            self.requests.get(&request_id),
112            Some(RequestState::Cancelled)
113        )
114    }
115
116    /// Remove the request from the table, regardless of its state.
117    pub fn remove(&mut self, request_id: RequestId) {
118        self.requests.remove(&request_id);
119    }
120
121    /// Number of requests currently tracked (active and cancelled combined).
122    pub fn len(&self) -> usize {
123        self.requests.len()
124    }
125
126    /// Returns `true` if no requests are currently tracked.
127    pub fn is_empty(&self) -> bool {
128        self.requests.is_empty()
129    }
130
131    /// Returns `true` if `request_id` is present in the table, regardless of
132    /// its state.
133    pub fn contains(&self, request_id: RequestId) -> bool {
134        self.requests.contains_key(&request_id)
135    }
136
137    /// Returns `true` if the request is present and in the
138    /// [`RequestState::Active`] state.
139    pub fn is_active(&self, request_id: RequestId) -> bool {
140        matches!(self.requests.get(&request_id), Some(RequestState::Active))
141    }
142}
nerve_ipc_core/request_table.rs

nerve_ipc_core/
request_table.rs
