Skip to main content

talon_http/
lanes.rs

1//! Lanes and per-lane SSE writer bookkeeping.
2//!
3//! SSE writers are tracked per lane. A single global FIFO cannot correlate a
4//! response with the request that produced it: concurrent callers, or an
5//! investigation firing while an analyst has the chat open, would have their
6//! output written to whichever writer happened to be current.
7//!
8//! Life of a writer (one per `POST /message`):
9//! 1. queued on its lane (`pending`) before the message reaches the host;
10//! 2. promoted to the lane's `active` writer when the host signals a turn
11//!    start (`set_typing(jid, true)`);
12//! 3. receives every `send_message(jid, ..)` chunk while active;
13//! 4. closed with a `done` event on `on_turn_complete` (or, as a safety net,
14//!    `set_typing(jid, false)` when the container exits);
15//! 5. removed whenever the client disconnects.
16//!
17//! A turn start that arrives while the lane still has an active writer (a
18//! follow-up piped into a running container) never replaces it mid-answer;
19//! it is remembered and the next queued writer is promoted when the current
20//! turn completes.
21
22use std::collections::{HashMap, VecDeque};
23use std::fmt;
24use std::sync::atomic::{AtomicU64, Ordering};
25use std::sync::{Mutex, MutexGuard};
26
27use serde_json::Value;
28use sha2::{Digest, Sha256};
29use tokio::sync::mpsc;
30
31/// Session/IPC key prefix of a per-user lane
32/// (`copilot-u-<id prefix>-<hash>`, see [`UserId::session_key`]).
33pub const USER_SESSION_KEY_PREFIX: &str = "copilot-u-";
34/// Session/IPC key of the investigation lane. Its state directory
35/// (`<data_dir>/sessions/copilot-investigate`, `/home/agent/.talon` inside
36/// the container) is private to that lane and holds the eval records
37/// served by `GET /evals/{alert_id}`.
38pub const INVESTIGATE_SESSION_KEY: &str = "copilot-investigate";
39/// Max characters of a CoPilot user id (`^[A-Za-z0-9_-]{1,48}$`).
40pub const MAX_USER_ID_CHARS: usize = 48;
41/// Characters of the lowercased user id kept in a per-user session/IPC key
42/// (a readable hint only; uniqueness comes from the hash).
43pub const USER_KEY_ID_CHARS: usize = 20;
44/// Hex digits of `sha256(exact user id)` appended to a per-user key:
45/// 32 hex = 128 bits, so finding an id that shares another user's key is a
46/// second-preimage attack on (truncated) SHA-256, not a 2^32 search (N-8).
47pub const USER_KEY_HASH_HEX: usize = 32;
48/// Max length of a group folder / lane key (`^[A-Za-z0-9][A-Za-z0-9_-]{0,63}$`).
49const MAX_LANE_KEY_CHARS: usize = 64;
50
51// A per-user key must always be a valid group folder.
52const _: () = assert!(
53    USER_SESSION_KEY_PREFIX.len() + USER_KEY_ID_CHARS + 1 + USER_KEY_HASH_HEX <= MAX_LANE_KEY_CHARS
54);
55const _: () = assert!(USER_KEY_HASH_HEX % 2 == 0 && USER_KEY_HASH_HEX <= 64);
56/// Frames buffered per SSE writer before chunks are dropped.
57pub(crate) const WRITER_BUFFER: usize = 256;
58
59/// The JIDs this channel hands out, derived from the base JID.
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub struct LaneIds {
62    /// Shared chat lane (`http:copilot`); also the anonymous lane.
63    pub base: String,
64    /// Investigation lane (`http:copilot:investigate`).
65    pub investigate: String,
66    /// Prefix of per-user lanes (`http:copilot:u:`).
67    pub user_prefix: String,
68}
69
70impl LaneIds {
71    /// Derive the lane JIDs from the base JID.
72    pub fn new(base: &str) -> Self {
73        Self {
74            base: base.to_string(),
75            investigate: format!("{base}:investigate"),
76            user_prefix: format!("{base}:u:"),
77        }
78    }
79
80    /// JID of the lane owned by a CoPilot user id (the exact id).
81    pub fn user_lane(&self, user_id: &str) -> String {
82        format!("{}{user_id}", self.user_prefix)
83    }
84
85    /// Whether `jid` is a per-user lane.
86    pub fn is_user_lane(&self, jid: &str) -> bool {
87        jid.len() > self.user_prefix.len() && jid.starts_with(&self.user_prefix)
88    }
89
90    /// The base lane and every `<base>:` sub-lane.
91    pub fn owns(&self, jid: &str) -> bool {
92        jid == self.base
93            || (jid.len() > self.base.len()
94                && jid.starts_with(&self.base)
95                && jid.as_bytes()[self.base.len()] == b':')
96    }
97}
98
99/// JID of the lane owned by `user_id` under the default base JID
100/// (`http:copilot:u:<id>`).
101pub fn lane_jid_for_user(user_id: &str) -> String {
102    LaneIds::new(crate::config::DEFAULT_BASE_JID).user_lane(user_id)
103}
104
105/// Whether `id` is an acceptable CoPilot user id: 1-48 characters from
106/// `[A-Za-z0-9_-]`.
107pub fn is_valid_user_id(id: &str) -> bool {
108    !id.is_empty()
109        && id.len() <= MAX_USER_ID_CHARS
110        && id
111            .bytes()
112            .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-')
113}
114
115/// A request named a `user_id` that is not acceptable. Ids are never
116/// "cleaned up": stripping characters would merge distinct users
117/// (`bob.smith` and `bobsmith`) into one lane (MYTH-12).
118#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
119#[error("user_id must be a string or integer matching ^[A-Za-z0-9_-]{{1,48}}$")]
120pub struct InvalidUserId;
121
122/// A validated CoPilot user id, kept exactly as sent (case included).
123///
124/// The lane JID embeds the exact id (`http:copilot:u:<id>`); the session and
125/// IPC key ([`Self::session_key`]) is derived so that ids differing only in
126/// case cannot share a directory on a case-insensitive filesystem.
127#[derive(Debug, Clone, PartialEq, Eq, Hash)]
128pub struct UserId(String);
129
130impl UserId {
131    /// Validate a user id string (no trimming, no stripping).
132    pub fn parse(raw: &str) -> Result<Self, InvalidUserId> {
133        if is_valid_user_id(raw) {
134            Ok(Self(raw.to_string()))
135        } else {
136            Err(InvalidUserId)
137        }
138    }
139
140    /// The `user_id` field of a request body: missing, `null` or `""` mean
141    /// "no user" (the shared anonymous lane); a valid string or integer is
142    /// accepted; anything else is rejected.
143    pub fn from_json(raw: Option<&Value>) -> Result<Option<Self>, InvalidUserId> {
144        match raw {
145            None | Some(Value::Null) => Ok(None),
146            Some(Value::String(s)) if s.is_empty() => Ok(None),
147            Some(Value::String(s)) => Self::parse(s).map(Some),
148            Some(Value::Number(n)) if n.is_i64() || n.is_u64() => {
149                Self::parse(&n.to_string()).map(Some)
150            }
151            Some(_) => Err(InvalidUserId),
152        }
153    }
154
155    /// The exact id.
156    pub fn as_str(&self) -> &str {
157        &self.0
158    }
159
160    /// Session/IPC key of this user's lane:
161    /// `copilot-u-<lowercased id, first 20 chars>-<first 32 hex of sha256(exact id)>`,
162    /// e.g. `copilot-u-bob-81b637d8fcd2c6da6359e6963113a117`.
163    ///
164    /// At most 63 characters of `[a-z0-9_-]`, so always a valid group
165    /// folder. The 128-bit hash of the *exact* id keeps `Bob` and `bob` (and
166    /// ids sharing a 20-char prefix) apart even where the filesystem folds
167    /// case, and makes it infeasible to pick an id whose key collides with
168    /// another user's (and so share their session, IPC namespace and lane).
169    pub fn session_key(&self) -> String {
170        let prefix: String = self
171            .0
172            .chars()
173            .take(USER_KEY_ID_CHARS)
174            .map(|c| c.to_ascii_lowercase())
175            .collect();
176        let digest = Sha256::digest(self.0.as_bytes());
177        let hash: String = digest
178            .iter()
179            .take(USER_KEY_HASH_HEX / 2)
180            .map(|b| format!("{b:02x}"))
181            .collect();
182        format!("{USER_SESSION_KEY_PREFIX}{prefix}-{hash}")
183    }
184}
185
186impl fmt::Display for UserId {
187    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188        f.write_str(&self.0)
189    }
190}
191
192/// One frame on an SSE stream.
193#[derive(Debug, Clone, PartialEq, Eq)]
194pub(crate) enum SseFrame {
195    /// `{"type":"text","content":...}`
196    Text(String),
197    /// `{"type":"done"}` — the turn finished; the stream closes after it.
198    Done,
199}
200
201impl SseFrame {
202    /// JSON payload of the `data:` line. Keys are always ordered `type`
203    /// first, regardless of serde_json's map-ordering feature.
204    pub(crate) fn to_json(&self) -> String {
205        match self {
206            SseFrame::Text(content) => format!(
207                r#"{{"type":"text","content":{}}}"#,
208                Value::String(content.clone())
209            ),
210            SseFrame::Done => r#"{"type":"done"}"#.to_string(),
211        }
212    }
213}
214
215struct Writer {
216    id: u64,
217    tx: mpsc::Sender<SseFrame>,
218}
219
220#[derive(Default)]
221struct LaneStreams {
222    pending: VecDeque<Writer>,
223    active: Option<Writer>,
224    /// Turn starts signalled while a writer was still active; each one
225    /// promotes the next pending writer when the current turn completes.
226    owed: usize,
227    /// A turn is in progress on this lane.
228    in_turn: bool,
229}
230
231impl LaneStreams {
232    fn writer_count(&self) -> usize {
233        self.pending.len() + usize::from(self.active.is_some())
234    }
235
236    fn is_idle(&self) -> bool {
237        self.writer_count() == 0 && !self.in_turn && self.owed == 0
238    }
239
240    fn clamp_owed(&mut self) {
241        self.owed = self.owed.min(self.pending.len());
242    }
243}
244
245#[derive(Default)]
246struct RegistryInner {
247    lanes: HashMap<String, LaneStreams>,
248    total_writers: usize,
249}
250
251/// Why a stream could not be opened.
252#[derive(Debug, Clone, Copy, PartialEq, Eq)]
253pub(crate) enum OpenError {
254    /// Too many requests already queued on this lane.
255    LaneBusy,
256    /// Too many streams open across all lanes.
257    TooManyStreams,
258}
259
260/// What happened to a text chunk.
261#[derive(Debug, Clone, Copy, PartialEq, Eq)]
262pub(crate) enum SendOutcome {
263    /// Queued on the active writer.
264    Delivered,
265    /// The lane has no active writer; the chunk was dropped.
266    NoWriter,
267    /// The active writer's buffer is full; the chunk was dropped.
268    Overflow,
269    /// The client went away; the writer was removed.
270    ClientGone,
271}
272
273/// Counts for `/status` and `/metrics`.
274#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
275pub(crate) struct StreamCounts {
276    pub pending: usize,
277    pub active: usize,
278    pub lanes_in_turn: usize,
279}
280
281/// Per-lane SSE writer registry. All operations are synchronous and hold the
282/// lock briefly; no `.await` happens under it.
283#[derive(Default)]
284pub(crate) struct StreamRegistry {
285    inner: Mutex<RegistryInner>,
286    next_id: AtomicU64,
287}
288
289impl StreamRegistry {
290    fn lock(&self) -> MutexGuard<'_, RegistryInner> {
291        self.inner.lock().unwrap_or_else(|e| e.into_inner())
292    }
293
294    /// Queue a new writer on `jid`.
295    pub(crate) fn open(
296        &self,
297        jid: &str,
298        max_per_lane: usize,
299        max_total: usize,
300    ) -> Result<(u64, mpsc::Receiver<SseFrame>), OpenError> {
301        let mut inner = self.lock();
302        if inner.total_writers >= max_total {
303            return Err(OpenError::TooManyStreams);
304        }
305        let lane = inner.lanes.entry(jid.to_string()).or_default();
306        if lane.writer_count() >= max_per_lane {
307            return Err(OpenError::LaneBusy);
308        }
309        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
310        let (tx, rx) = mpsc::channel(WRITER_BUFFER);
311        lane.pending.push_back(Writer { id, tx });
312        inner.total_writers += 1;
313        Ok((id, rx))
314    }
315
316    /// Forget writer `id` (client disconnected or stream finished).
317    pub(crate) fn remove(&self, jid: &str, id: u64) {
318        let mut inner = self.lock();
319        let mut removed = 0;
320        if let Some(lane) = inner.lanes.get_mut(jid) {
321            let before = lane.pending.len();
322            lane.pending.retain(|w| w.id != id);
323            removed += before - lane.pending.len();
324            if lane.active.as_ref().is_some_and(|w| w.id == id) {
325                lane.active = None;
326                removed += 1;
327            }
328            lane.clamp_owed();
329            if lane.is_idle() {
330                inner.lanes.remove(jid);
331            }
332        }
333        inner.total_writers = inner.total_writers.saturating_sub(removed);
334    }
335
336    /// `set_typing(jid, true)`: a turn starts (or a follow-up was piped into
337    /// the running container). Returns true if a writer was promoted.
338    pub(crate) fn turn_started(&self, jid: &str) -> bool {
339        let mut inner = self.lock();
340        let lane = inner.lanes.entry(jid.to_string()).or_default();
341        lane.in_turn = true;
342        if lane.active.is_some() {
343            if lane.owed < lane.pending.len() {
344                lane.owed += 1;
345            }
346            return false;
347        }
348        match lane.pending.pop_front() {
349            Some(w) => {
350                lane.active = Some(w);
351                lane.clamp_owed();
352                true
353            }
354            None => false,
355        }
356    }
357
358    /// `on_turn_complete(jid)`: close the active writer; promote the next
359    /// one if a follow-up turn was already signalled. Returns true if a
360    /// writer was closed.
361    pub(crate) fn turn_complete(&self, jid: &str) -> bool {
362        let mut inner = self.lock();
363        let mut closed = false;
364        if let Some(lane) = inner.lanes.get_mut(jid) {
365            if let Some(w) = lane.active.take() {
366                let _ = w.tx.try_send(SseFrame::Done);
367                closed = true;
368            }
369            if lane.owed > 0 && !lane.pending.is_empty() {
370                lane.owed -= 1;
371                lane.active = lane.pending.pop_front();
372            } else {
373                lane.owed = 0;
374                lane.in_turn = false;
375            }
376            if lane.is_idle() {
377                inner.lanes.remove(jid);
378            }
379        }
380        if closed {
381            inner.total_writers = inner.total_writers.saturating_sub(1);
382        }
383        closed
384    }
385
386    /// `set_typing(jid, false)`: the container exited. Safety net that
387    /// closes a still-active writer. Returns true if one was closed.
388    pub(crate) fn turn_ended(&self, jid: &str) -> bool {
389        let mut inner = self.lock();
390        let mut closed = false;
391        if let Some(lane) = inner.lanes.get_mut(jid) {
392            if let Some(w) = lane.active.take() {
393                let _ = w.tx.try_send(SseFrame::Done);
394                closed = true;
395            }
396            lane.owed = 0;
397            lane.in_turn = false;
398            if lane.is_idle() {
399                inner.lanes.remove(jid);
400            }
401        }
402        if closed {
403            inner.total_writers = inner.total_writers.saturating_sub(1);
404        }
405        closed
406    }
407
408    /// Write a text chunk to the lane's active writer. Never falls back to
409    /// another lane's writer.
410    pub(crate) fn send_text(&self, jid: &str, text: &str) -> SendOutcome {
411        let mut inner = self.lock();
412        let Some(lane) = inner.lanes.get_mut(jid) else {
413            return SendOutcome::NoWriter;
414        };
415        let Some(w) = lane.active.as_ref() else {
416            return SendOutcome::NoWriter;
417        };
418        match w.tx.try_send(SseFrame::Text(text.to_string())) {
419            Ok(()) => SendOutcome::Delivered,
420            Err(mpsc::error::TrySendError::Full(_)) => SendOutcome::Overflow,
421            Err(mpsc::error::TrySendError::Closed(_)) => {
422                lane.active = None;
423                inner.total_writers = inner.total_writers.saturating_sub(1);
424                SendOutcome::ClientGone
425            }
426        }
427    }
428
429    /// Whether the lane has an active or queued writer (an open stream).
430    pub(crate) fn has_streams(&self, jid: &str) -> bool {
431        self.lock()
432            .lanes
433            .get(jid)
434            .is_some_and(|l| l.writer_count() > 0)
435    }
436
437    /// Writers queued (not yet active) on `jid`.
438    pub(crate) fn pending_len(&self, jid: &str) -> usize {
439        self.lock().lanes.get(jid).map_or(0, |l| l.pending.len())
440    }
441
442    /// Whether `jid` has an active writer.
443    #[cfg(test)]
444    pub(crate) fn has_active(&self, jid: &str) -> bool {
445        self.lock()
446            .lanes
447            .get(jid)
448            .is_some_and(|l| l.active.is_some())
449    }
450
451    /// Aggregate counts.
452    pub(crate) fn counts(&self) -> StreamCounts {
453        let inner = self.lock();
454        let mut c = StreamCounts::default();
455        for lane in inner.lanes.values() {
456            c.pending += lane.pending.len();
457            c.active += usize::from(lane.active.is_some());
458            c.lanes_in_turn += usize::from(lane.in_turn);
459        }
460        c
461    }
462
463    /// Drop a lane's bookkeeping (idle sweep). Open writers are dropped, which
464    /// ends their streams.
465    pub(crate) fn forget_lane(&self, jid: &str) {
466        let mut inner = self.lock();
467        if let Some(lane) = inner.lanes.remove(jid) {
468            let n = lane.writer_count();
469            inner.total_writers = inner.total_writers.saturating_sub(n);
470        }
471    }
472
473    /// Close every stream (shutdown). Active writers get a `done` frame.
474    pub(crate) fn close_all(&self) {
475        let mut inner = self.lock();
476        for (_, lane) in inner.lanes.drain() {
477            if let Some(w) = lane.active {
478                let _ = w.tx.try_send(SseFrame::Done);
479            }
480        }
481        inner.total_writers = 0;
482    }
483}
484
485#[cfg(test)]
486mod tests {
487    use super::*;
488    use serde_json::json;
489
490    const COPILOT_JID: &str = "http:copilot";
491    const INVESTIGATE_JID: &str = "http:copilot:investigate";
492
493    fn ids() -> LaneIds {
494        LaneIds::new(COPILOT_JID)
495    }
496
497    struct Enqueued {
498        id: u64,
499        rx: mpsc::Receiver<SseFrame>,
500    }
501
502    impl Enqueued {
503        /// Text payloads received so far.
504        fn texts(&mut self) -> Vec<String> {
505            self.drain()
506                .into_iter()
507                .filter_map(|f| match f {
508                    SseFrame::Text(t) => Some(t),
509                    SseFrame::Done => None,
510                })
511                .collect()
512        }
513
514        fn drain(&mut self) -> Vec<SseFrame> {
515            let mut out = Vec::new();
516            while let Ok(f) = self.rx.try_recv() {
517                out.push(f);
518            }
519            out
520        }
521
522        /// Whether the stream was ended (sender dropped, buffer drained).
523        fn ended(&mut self) -> bool {
524            matches!(
525                self.rx.try_recv(),
526                Err(mpsc::error::TryRecvError::Disconnected)
527            )
528        }
529    }
530
531    fn enqueue(reg: &StreamRegistry, jid: &str) -> Enqueued {
532        let (id, rx) = reg.open(jid, 16, 1024).unwrap();
533        Enqueued { id, rx }
534    }
535
536    // ---- writer isolation ----------------------------------------------------
537
538    #[test]
539    fn writes_a_lane_response_only_to_that_lanes_writer() {
540        let reg = StreamRegistry::default();
541        let mut chat = enqueue(&reg, COPILOT_JID);
542        let mut investigate = enqueue(&reg, INVESTIGATE_JID);
543        reg.turn_started(COPILOT_JID);
544        reg.turn_started(INVESTIGATE_JID);
545        reg.send_text(COPILOT_JID, "answer for the analyst");
546        reg.send_text(INVESTIGATE_JID, "investigation report");
547        assert_eq!(chat.texts(), vec!["answer for the analyst"]);
548        assert_eq!(investigate.texts(), vec!["investigation report"]);
549    }
550
551    #[test]
552    fn drops_output_for_a_lane_with_no_writer_instead_of_borrowing_another() {
553        let reg = StreamRegistry::default();
554        let mut chat = enqueue(&reg, COPILOT_JID);
555        reg.turn_started(COPILOT_JID);
556        assert_eq!(
557            reg.send_text(INVESTIGATE_JID, "investigation report"),
558            SendOutcome::NoWriter
559        );
560        assert!(chat.texts().is_empty());
561        assert!(!chat.ended());
562    }
563
564    #[test]
565    fn one_lanes_turn_start_does_not_steal_another_lanes_queued_writer() {
566        let reg = StreamRegistry::default();
567        let mut chat = enqueue(&reg, COPILOT_JID);
568        reg.turn_started(INVESTIGATE_JID);
569        reg.send_text(COPILOT_JID, "still waiting");
570        assert!(chat.texts().is_empty());
571        reg.turn_started(COPILOT_JID);
572        reg.send_text(COPILOT_JID, "answer for the analyst");
573        assert_eq!(chat.texts(), vec!["answer for the analyst"]);
574    }
575
576    #[test]
577    fn closes_only_the_completing_lanes_stream() {
578        let reg = StreamRegistry::default();
579        let mut chat = enqueue(&reg, COPILOT_JID);
580        let mut investigate = enqueue(&reg, INVESTIGATE_JID);
581        reg.turn_started(COPILOT_JID);
582        reg.turn_started(INVESTIGATE_JID);
583        assert!(reg.turn_complete(INVESTIGATE_JID));
584        assert_eq!(investigate.drain(), vec![SseFrame::Done]);
585        assert!(investigate.ended());
586        assert!(chat.drain().is_empty());
587        assert!(!chat.ended());
588    }
589
590    #[test]
591    fn serialises_two_callers_on_the_same_lane() {
592        let reg = StreamRegistry::default();
593        let mut first = enqueue(&reg, COPILOT_JID);
594        let mut second = enqueue(&reg, COPILOT_JID);
595        reg.turn_started(COPILOT_JID);
596        reg.send_text(COPILOT_JID, "first answer");
597        reg.turn_complete(COPILOT_JID);
598        reg.turn_started(COPILOT_JID);
599        reg.send_text(COPILOT_JID, "second answer");
600        assert_eq!(first.texts(), vec!["first answer"]);
601        assert_eq!(second.texts(), vec!["second answer"]);
602    }
603
604    // ---- new behaviour -------------------------------------------------------
605
606    #[test]
607    fn follow_up_turn_start_does_not_hijack_the_active_writer() {
608        // Replacing the active writer when a piped follow-up signals a new
609        // turn would land the rest of answer #1 in caller #2's stream.
610        let reg = StreamRegistry::default();
611        let mut first = enqueue(&reg, COPILOT_JID);
612        reg.turn_started(COPILOT_JID);
613        let mut second = enqueue(&reg, COPILOT_JID);
614        reg.turn_started(COPILOT_JID); // follow-up piped into the container
615        reg.send_text(COPILOT_JID, "rest of first answer");
616        reg.turn_complete(COPILOT_JID);
617        reg.send_text(COPILOT_JID, "second answer");
618        reg.turn_complete(COPILOT_JID);
619        assert_eq!(first.texts(), vec!["rest of first answer"]);
620        assert!(first.ended());
621        assert_eq!(second.texts(), vec!["second answer"]);
622        assert!(second.ended());
623        assert_eq!(reg.counts(), StreamCounts::default());
624    }
625
626    #[test]
627    fn container_exit_closes_the_active_writer_as_a_safety_net() {
628        let reg = StreamRegistry::default();
629        let mut w = enqueue(&reg, COPILOT_JID);
630        reg.turn_started(COPILOT_JID);
631        assert!(reg.turn_ended(COPILOT_JID));
632        assert_eq!(w.drain(), vec![SseFrame::Done]);
633        assert!(w.ended());
634        assert!(!reg.turn_ended(COPILOT_JID));
635    }
636
637    #[test]
638    fn disconnect_removes_pending_and_active_writers() {
639        let reg = StreamRegistry::default();
640        let a = enqueue(&reg, COPILOT_JID);
641        let b = enqueue(&reg, COPILOT_JID);
642        reg.turn_started(COPILOT_JID);
643        assert!(reg.has_active(COPILOT_JID));
644        reg.remove(COPILOT_JID, a.id);
645        assert!(!reg.has_active(COPILOT_JID));
646        assert_eq!(reg.pending_len(COPILOT_JID), 1);
647        reg.remove(COPILOT_JID, b.id);
648        assert!(!reg.has_streams(COPILOT_JID));
649        // A dropped receiver is detected on the next write.
650        let c = enqueue(&reg, COPILOT_JID);
651        reg.turn_started(COPILOT_JID);
652        drop(c);
653        assert_eq!(reg.send_text(COPILOT_JID, "x"), SendOutcome::ClientGone);
654        reg.turn_ended(COPILOT_JID);
655        assert_eq!(reg.counts(), StreamCounts::default());
656    }
657
658    #[test]
659    fn caps_are_enforced() {
660        let reg = StreamRegistry::default();
661        assert!(reg.open(COPILOT_JID, 1, 10).is_ok());
662        assert_eq!(
663            reg.open(COPILOT_JID, 1, 10).err(),
664            Some(OpenError::LaneBusy)
665        );
666        assert!(reg.open("http:copilot:u:1", 1, 2).is_ok());
667        assert_eq!(
668            reg.open("http:copilot:u:2", 1, 2).err(),
669            Some(OpenError::TooManyStreams)
670        );
671    }
672
673    #[test]
674    fn full_buffer_drops_chunks_without_blocking() {
675        let reg = StreamRegistry::default();
676        let mut w = enqueue(&reg, COPILOT_JID);
677        reg.turn_started(COPILOT_JID);
678        for _ in 0..WRITER_BUFFER {
679            assert_eq!(reg.send_text(COPILOT_JID, "x"), SendOutcome::Delivered);
680        }
681        assert_eq!(reg.send_text(COPILOT_JID, "x"), SendOutcome::Overflow);
682        assert_eq!(w.texts().len(), WRITER_BUFFER);
683    }
684
685    #[test]
686    fn close_all_ends_every_stream() {
687        let reg = StreamRegistry::default();
688        let mut a = enqueue(&reg, COPILOT_JID);
689        let mut b = enqueue(&reg, INVESTIGATE_JID);
690        reg.turn_started(COPILOT_JID);
691        reg.close_all();
692        assert_eq!(a.drain(), vec![SseFrame::Done]);
693        assert!(a.ended());
694        assert!(b.ended());
695        assert_eq!(reg.counts(), StreamCounts::default());
696    }
697
698    #[test]
699    fn frames_serialize_like_the_ts_channel() {
700        assert_eq!(
701            SseFrame::Text("hi \"x\"\n".into()).to_json(),
702            r#"{"type":"text","content":"hi \"x\"\n"}"#
703        );
704        assert_eq!(SseFrame::Done.to_json(), r#"{"type":"done"}"#);
705    }
706
707    // ---- lane ownership + user id validation -------------------------------
708
709    #[test]
710    fn owns_the_chat_lane_and_every_sub_lane() {
711        let l = ids();
712        assert!(l.owns(COPILOT_JID));
713        assert!(l.owns(INVESTIGATE_JID));
714        assert!(l.owns("http:copilot:u:42"));
715    }
716
717    #[test]
718    fn does_not_own_another_channels_jid() {
719        let l = ids();
720        assert!(!l.owns("webhook:copilot"));
721        assert!(!l.owns("telegram:123"));
722        assert!(!l.owns("http:copilotx"));
723        assert!(!l.owns("http:copilo"));
724    }
725
726    fn parse(v: Value) -> Result<Option<String>, InvalidUserId> {
727        UserId::from_json(Some(&v)).map(|o| o.map(|u| u.as_str().to_string()))
728    }
729
730    #[test]
731    fn accepts_ids_copilot_actually_issues() {
732        assert_eq!(parse(json!(42)).unwrap().as_deref(), Some("42"));
733        assert_eq!(
734            parse(json!("analyst-7")).unwrap().as_deref(),
735            Some("analyst-7")
736        );
737        assert_eq!(parse(json!("a_b-C9")).unwrap().as_deref(), Some("a_b-C9"));
738        let max = "x".repeat(MAX_USER_ID_CHARS);
739        assert_eq!(parse(json!(max.clone())).unwrap(), Some(max));
740    }
741
742    #[test]
743    fn rejects_ids_instead_of_stripping_them_into_a_neighbour() {
744        // Stripping these instead would let `bob.smith` and `bobsmith` share
745        // a lane.
746        for raw in [
747            "bob.smith",
748            "../../etc/passwd",
749            "a:b",
750            " bob",
751            "bob ",
752            "ünïcödé-1",
753            "a/b",
754            "///",
755        ] {
756            assert_eq!(parse(json!(raw)), Err(InvalidUserId), "{raw:?}");
757        }
758        assert_eq!(
759            parse(json!("x".repeat(MAX_USER_ID_CHARS + 1))),
760            Err(InvalidUserId)
761        );
762        for v in [json!({ "id": 1 }), json!(true), json!(["1"]), json!(1.5)] {
763            assert_eq!(parse(v.clone()), Err(InvalidUserId), "{v}");
764        }
765        assert!(UserId::parse("bob.smith").is_err());
766        assert!(is_valid_user_id("bobsmith"));
767    }
768
769    #[test]
770    fn absent_ids_map_to_the_anon_lane() {
771        assert_eq!(UserId::from_json(None), Ok(None));
772        assert_eq!(parse(Value::Null), Ok(None));
773        assert_eq!(parse(json!("")), Ok(None));
774    }
775
776    #[test]
777    fn session_key_has_the_documented_format() {
778        let bob = UserId::parse("bob").unwrap();
779        // sha256("bob") = 81b637d8fcd2c6da6359e6963113a117...
780        assert_eq!(
781            bob.session_key(),
782            "copilot-u-bob-81b637d8fcd2c6da6359e6963113a117"
783        );
784        let upper = UserId::parse("Bob").unwrap();
785        // sha256("Bob") = cd9fb1e148ccd8442e5aa74904cc73bf...
786        assert_eq!(
787            upper.session_key(),
788            "copilot-u-bob-cd9fb1e148ccd8442e5aa74904cc73bf"
789        );
790        assert_eq!(
791            UserId::parse("42").unwrap().session_key(),
792            "copilot-u-42-73475cb40a568e8da8a045ced110137e"
793        );
794    }
795
796    #[test]
797    fn session_key_carries_a_128_bit_hash_of_the_exact_id() {
798        // N-8: a 32-bit suffix let a caller grind an id (~2^32 tries) whose
799        // key equals a victim's whenever both shared the truncated prefix.
800        for raw in ["bob", "Bob", "42", &"Q".repeat(MAX_USER_ID_CHARS)] {
801            let id = UserId::parse(raw).unwrap();
802            let key = id.session_key();
803            let (head, hash) = key.rsplit_once('-').unwrap();
804            assert_eq!(hash.len(), USER_KEY_HASH_HEX, "{key}");
805            assert!(hash.len() * 4 >= 128, "{key}");
806            assert!(
807                hash.bytes()
808                    .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
809            );
810            let digest = Sha256::digest(raw.as_bytes());
811            let full: String = digest.iter().map(|b| format!("{b:02x}")).collect();
812            assert_eq!(hash, &full[..USER_KEY_HASH_HEX]);
813            let prefix = head.strip_prefix(USER_SESSION_KEY_PREFIX).unwrap();
814            assert!(prefix.len() <= USER_KEY_ID_CHARS, "{key}");
815            assert!(raw.to_ascii_lowercase().starts_with(prefix), "{key}");
816        }
817    }
818
819    #[test]
820    fn ids_sharing_a_long_prefix_get_distinct_full_hash_keys() {
821        // Ids longer than the readable prefix are told apart by the hash
822        // alone (the old format kept 40 characters + 32 bits).
823        let base = "victim-account-with-a-long-identifier-00";
824        let victim = UserId::parse(&format!("{base}zz")).unwrap().session_key();
825        assert_eq!(
826            victim.len(),
827            USER_SESSION_KEY_PREFIX.len() + USER_KEY_ID_CHARS + 1 + USER_KEY_HASH_HEX
828        );
829        let mut seen = std::collections::HashSet::new();
830        seen.insert(victim.clone());
831        for i in 0..2_000u32 {
832            let k = UserId::parse(&format!("{base}{i:x}"))
833                .unwrap()
834                .session_key();
835            assert_eq!(k.len(), victim.len());
836            assert!(seen.insert(k), "duplicate key for suffix {i:x}");
837        }
838    }
839
840    #[test]
841    fn session_keys_never_collide_on_case_insensitive_filesystems() {
842        let ids = [
843            "bob",
844            "Bob",
845            "BOB",
846            "bOb",
847            "bobsmith",
848            "bob_smith",
849            "bob-smith",
850        ];
851        let mut keys: Vec<String> = ids
852            .iter()
853            .map(|i| UserId::parse(i).unwrap().session_key().to_ascii_lowercase())
854            .collect();
855        keys.sort();
856        keys.dedup();
857        assert_eq!(keys.len(), ids.len(), "{keys:?}");
858        // Ids sharing the readable prefix still get distinct keys.
859        let base = "a".repeat(USER_KEY_ID_CHARS);
860        let k1 = UserId::parse(&format!("{base}1")).unwrap().session_key();
861        let k2 = UserId::parse(&format!("{base}2")).unwrap().session_key();
862        assert_ne!(k1, k2);
863    }
864
865    #[test]
866    fn every_accepted_id_yields_a_valid_lane_key() {
867        for raw in [
868            "42".to_string(),
869            "analyst-7".to_string(),
870            "-leading-dash".to_string(),
871            "Z".repeat(MAX_USER_ID_CHARS),
872            "_".repeat(MAX_USER_ID_CHARS),
873            "-".repeat(MAX_USER_ID_CHARS),
874            "9".repeat(USER_KEY_ID_CHARS + 1),
875        ] {
876            let key = UserId::parse(&raw).unwrap().session_key();
877            assert!(key.len() <= MAX_LANE_KEY_CHARS, "{key}");
878            assert_eq!(
879                key.len(),
880                USER_SESSION_KEY_PREFIX.len()
881                    + raw.len().min(USER_KEY_ID_CHARS)
882                    + 1
883                    + USER_KEY_HASH_HEX
884            );
885            assert!(key.starts_with(USER_SESSION_KEY_PREFIX));
886            assert!(talon_core::paths::is_valid_group_folder(&key), "{key}");
887            assert_eq!(key, key.to_ascii_lowercase());
888        }
889    }
890
891    #[test]
892    fn lane_jid_helpers() {
893        assert_eq!(lane_jid_for_user("42"), "http:copilot:u:42");
894        let l = ids();
895        assert!(l.is_user_lane("http:copilot:u:42"));
896        assert!(!l.is_user_lane("http:copilot:u:"));
897        assert!(!l.is_user_lane(INVESTIGATE_JID));
898        assert_eq!(l.investigate, INVESTIGATE_JID);
899    }
900}