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