agentmux_srv/identity/
auth_session.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! In-memory manager for pre-launch OAuth sessions.
5//!
6//! See `docs/specs/SPEC_PRE_LAUNCH_OAUTH_FLOW_2026_05_14.md` §7.
7//!
8//! Each session represents one user-initiated "Connect with OAuth"
9//! attempt. The frontend creates a session via `start_session`,
10//! polls via `poll_session`, optionally pastes a callback URL (the
11//! `auth.submitcallback` handler forwards it via `send_to_stdin`),
12//! and cancels via `cancel_session`.
13//!
14//! PR A scope: the session-state machine + per-line stdout
15//! interpretation + lifecycle (timeout, cancel, cleanup). The
16//! actual CLI spawn lives in the handler (`server/identity_auth_spawn.rs`,
17//! so it can use AppState's CLI resolver) but emits frames into this
18//! module via `record_line` / `finish_success` / `finish_failure`.
19//! That keeps this module pure and testable.
20
21use std::collections::HashMap;
22use std::sync::{Arc, Mutex};
23use std::time::{Duration, Instant};
24
25use serde::{Deserialize, Serialize};
26
27use super::auth_patterns::{match_line, AuthPatternMatch};
28
29/// Wall-clock cap on a single auth session. Past this, the session
30/// transitions to `Failed` and its CLI is killed, whether or not anyone is
31/// polling it (`sweep`).
32const SESSION_TIMEOUT_SECS: u64 = 600;
33
34/// How long a finished session stays pollable before `sweep` drops it. The
35/// UI stops polling once it sees the terminal state, so this only has to
36/// outlast one poll interval; it's generous.
37const FINISHED_RETENTION_SECS: u64 = 300;
38
39/// How often the sweeper runs (`spawn_sweeper`).
40const SWEEP_INTERVAL_SECS: u64 = 30;
41
42#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, ts_rs::TS)]
43#[ts(export, export_to = "../../../frontend/types/rpc/")]
44#[serde(
45    tag = "status",
46    rename_all = "kebab-case",
47    rename_all_fields = "camelCase"
48)]
49pub enum AuthSessionStatus {
50    /// CLI is spawned, we're waiting for it to emit either a URL,
51    /// a device code, or a success line.
52    Pending,
53    /// CLI emitted an OAuth URL. Frontend surfaces this to the user
54    /// for paste-into-browser if auto-open failed.
55    UrlAvailable { auth_url: String },
56    /// CLI emitted a device code. Frontend renders it prominently.
57    CodeEmitted {
58        device_code: String,
59        verification_url: String,
60    },
61    /// CLI authenticated successfully. The handler captured the
62    /// credentials and created the bundle.
63    Success {
64        bundle_id: String,
65        /// Best-effort — extracted from the CLI's "logged in as..."
66        /// line. Used by the frontend to name the bundle.
67        email: Option<String>,
68        /// Set only for a direct-account session (issue #1624 PR-C
69        /// Part B) — the newly-minted or reused `IdentityAccount` id,
70        /// with no bundle involved. `bundle_id` is an empty string in
71        /// that mode; existing bundle-mode callers never populate
72        /// this field, so it's skipped on the wire when absent.
73        #[serde(skip_serializing_if = "Option::is_none")]
74        #[ts(optional)]
75        account_id: Option<String>,
76    },
77    /// Auth attempt failed. `error` is a short human-readable phrase
78    /// suitable to render inline.
79    Failed { error: String },
80}
81
82impl AuthSessionStatus {
83    pub fn is_terminal(&self) -> bool {
84        matches!(self, Self::Success { .. } | Self::Failed { .. })
85    }
86}
87
88#[derive(Debug)]
89struct Session {
90    provider_id: String,
91    status: AuthSessionStatus,
92    /// Last URL or code we surfaced — kept across polls so the
93    /// frontend can repaint without re-receiving on every tick.
94    captured_url: Option<String>,
95    captured_device_code: Option<(String, String)>,
96    captured_email: Option<String>,
97    started_at: Instant,
98    /// When the session reached a terminal state; `sweep` drops it
99    /// `FINISHED_RETENTION_SECS` later.
100    finished_at: Option<Instant>,
101    /// The account directory this login writes to. Only one live login per
102    /// directory: starting another cancels this one (`start_session`).
103    exclusive_key: Option<String>,
104    /// All stdout/stderr lines we've seen, in order. Used for the
105    /// diagnostic "show me what the CLI said" panel and the
106    /// integration tests.
107    transcript: Vec<String>,
108}
109
110impl Session {
111    fn new(provider_id: String, exclusive_key: Option<String>) -> Self {
112        Self {
113            provider_id,
114            status: AuthSessionStatus::Pending,
115            captured_url: None,
116            captured_device_code: None,
117            captured_email: None,
118            started_at: Instant::now(),
119            finished_at: None,
120            exclusive_key,
121            transcript: Vec::new(),
122        }
123    }
124
125    /// Move to a terminal state, unless already in one. True if it moved.
126    fn finish(&mut self, status: AuthSessionStatus) -> bool {
127        if self.status.is_terminal() {
128            return false;
129        }
130        self.status = status;
131        self.finished_at = Some(Instant::now());
132        true
133    }
134
135    fn timed_out(&self) -> bool {
136        self.started_at.elapsed() > Duration::from_secs(SESSION_TIMEOUT_SECS)
137    }
138}
139
140/// Public-facing handle returned by `start_session`. The handler
141/// owns the spawned CLI's child handle (so it can stdin-inject the
142/// pasted callback URL and kill on cancel); this manager owns the
143/// session state.
144#[derive(Debug, Clone, Serialize, ts_rs::TS)]
145#[ts(export, export_to = "../../../frontend/types/rpc/")]
146#[serde(rename_all = "camelCase")]
147pub struct StartSessionResult {
148    pub session_id: String,
149    /// If the CLI emitted the URL synchronously (before this call
150    /// returns) — rare but happens for fast providers. Usually
151    /// `None`; the frontend polls and picks it up on the first tick.
152    pub auth_url: Option<String>,
153}
154
155/// Snapshot returned by `poll_session`. Mirrors `AuthSessionStatus`
156/// plus the provider id for renderer dispatch.
157#[derive(Debug, Clone, Serialize, ts_rs::TS)]
158#[ts(export, export_to = "../../../frontend/types/rpc/")]
159#[serde(rename_all = "camelCase")]
160pub struct PollSessionResult {
161    pub provider_id: String,
162    #[serde(flatten)]
163    pub status: AuthSessionStatus,
164}
165
166/// Per-session "process refs" — the join handle for the drain task
167/// and the stdin sender for the callback-paste-back path. Held
168/// outside the `sessions` Mutex so cancellation can `.abort()` /
169/// drop without holding the state lock across an await.
170#[derive(Default)]
171struct ProcessRefs {
172    drain_tasks: HashMap<String, tokio::task::JoinHandle<()>>,
173    stdin_senders: HashMap<String, tokio::sync::mpsc::Sender<String>>,
174    /// PID of a PTY-backed auth subprocess. Aborting `drain_tasks`
175    /// can't reach into a `spawn_blocking` wait task, so cancel_session
176    /// kills the child by PID for PTY transports.
177    pty_pids: HashMap<String, u32>,
178}
179
180#[derive(Default)]
181pub struct AuthSessionManager {
182    sessions: Arc<Mutex<HashMap<String, Session>>>,
183    process_refs: Arc<Mutex<ProcessRefs>>,
184}
185
186impl AuthSessionManager {
187    pub fn new() -> Self {
188        Self::default()
189    }
190
191    /// Allocate a new session id and store the initial Pending state.
192    /// The caller (handler) is responsible for spawning the CLI and
193    /// feeding stdout lines into `record_line`.
194    ///
195    /// `exclusive_key` is the account directory the login writes to. A live
196    /// session for the same directory is cancelled first: two CLIs logging
197    /// in to one directory at once race on its credential files (the
198    /// desktop host's single login slot has the same rule).
199    pub fn start_session(&self, provider_id: String, exclusive_key: Option<String>) -> StartSessionResult {
200        let session_id = format!("auth-{}", uuid::Uuid::new_v4());
201        // Replacing the live session and inserting this one happen under one
202        // lock, so two concurrent starts for the same directory can't both
203        // find nothing to replace. The replaced CLIs are stopped after it.
204        let superseded: Vec<String> = {
205            let mut sessions = self.sessions.lock().unwrap();
206            // Every session for the directory: one already finished (say a
207            // success line was seen) can still have its CLI running.
208            let superseded: Vec<String> = match exclusive_key.as_deref() {
209                Some(key) => sessions
210                    .iter_mut()
211                    .filter(|(_, s)| s.exclusive_key.as_deref() == Some(key))
212                    .map(|(id, s)| {
213                        s.finish(AuthSessionStatus::Failed {
214                            error: "replaced by a newer sign-in for the same account".to_string(),
215                        });
216                        id.clone()
217                    })
218                    .collect(),
219                None => Vec::new(),
220            };
221            sessions.insert(session_id.clone(), Session::new(provider_id, exclusive_key));
222            superseded
223        };
224        for id in &superseded {
225            self.stop_process(id);
226        }
227        StartSessionResult {
228            session_id,
229            auth_url: None,
230        }
231    }
232
233    /// Feed a single line of CLI stdout/stderr into the session.
234    /// Returns the pattern match (if any) so the handler can decide
235    /// to forward the URL elsewhere (e.g. broker event).
236    pub fn record_line(&self, session_id: &str, line: &str) -> Option<AuthPatternMatch> {
237        let mut sessions = self.sessions.lock().unwrap();
238        let session = sessions.get_mut(session_id)?;
239        session.transcript.push(line.to_string());
240        let m = match_line(&session.provider_id, line)?;
241        match &m {
242            AuthPatternMatch::OAuthUrl(url) => {
243                if session.captured_url.is_none() {
244                    session.captured_url = Some(url.clone());
245                    if matches!(session.status, AuthSessionStatus::Pending) {
246                        session.status = AuthSessionStatus::UrlAvailable {
247                            auth_url: url.clone(),
248                        };
249                    }
250                }
251            }
252            AuthPatternMatch::DeviceCode {
253                code,
254                verification_url,
255            } => {
256                if session.captured_device_code.is_none() {
257                    session.captured_device_code =
258                        Some((code.clone(), verification_url.clone()));
259                    if matches!(
260                        session.status,
261                        AuthSessionStatus::Pending | AuthSessionStatus::UrlAvailable { .. }
262                    ) {
263                        session.status = AuthSessionStatus::CodeEmitted {
264                            device_code: code.clone(),
265                            verification_url: verification_url.clone(),
266                        };
267                    }
268                }
269            }
270            AuthPatternMatch::LoginSuccess { email } => {
271                // Don't transition to Success here — the handler is
272                // responsible for confirming with authCheckCommand
273                // AND persisting the bundle row before declaring
274                // success. We just record the email for later.
275                if session.captured_email.is_none() {
276                    session.captured_email = email.clone();
277                }
278            }
279        }
280        Some(m)
281    }
282
283    /// Handler-side completion hook. Called when the CLI exits or
284    /// auth check confirms success. Transitions the session to a
285    /// terminal state. `account_id` is `None` for the legacy
286    /// bundle-mode path; `Some(id)` for a direct-account session
287    /// (issue #1624 PR-C Part B) — see `AuthSessionStatus::Success`.
288    pub fn finish_success(
289        &self,
290        session_id: &str,
291        bundle_id: String,
292        account_id: Option<String>,
293    ) -> bool {
294        let mut sessions = self.sessions.lock().unwrap();
295        let Some(session) = sessions.get_mut(session_id) else {
296            return false;
297        };
298        let email = session.captured_email.clone();
299        session.finish(AuthSessionStatus::Success {
300            bundle_id,
301            email,
302            account_id,
303        })
304    }
305
306    pub fn finish_failure(&self, session_id: &str, error: String) -> bool {
307        let mut sessions = self.sessions.lock().unwrap();
308        let Some(session) = sessions.get_mut(session_id) else {
309            return false;
310        };
311        session.finish(AuthSessionStatus::Failed { error })
312    }
313
314    /// Poll a session's current status. A session past
315    /// SESSION_TIMEOUT_SECS is ended first (Failed, CLI killed).
316    pub fn poll_session(&self, session_id: &str) -> Option<PollSessionResult> {
317        let timed_out = {
318            let sessions = self.sessions.lock().unwrap();
319            let session = sessions.get(session_id)?;
320            !session.status.is_terminal() && session.timed_out()
321        };
322        if timed_out {
323            self.end_session(session_id, timeout_error());
324        }
325        let sessions = self.sessions.lock().unwrap();
326        let session = sessions.get(session_id)?;
327        Some(PollSessionResult {
328            provider_id: session.provider_id.clone(),
329            status: session.status.clone(),
330        })
331    }
332
333    /// End sessions past their timeout (Failed, CLI killed) and drop
334    /// finished ones nobody has polled for FINISHED_RETENTION_SECS. Runs
335    /// every SWEEP_INTERVAL_SECS (`spawn_sweeper`), so a login whose UI went
336    /// away doesn't leave its CLI running or its session behind.
337    pub fn sweep(&self) {
338        let (expired, stale): (Vec<String>, Vec<String>) = {
339            let sessions = self.sessions.lock().unwrap();
340            let expired = sessions
341                .iter()
342                .filter(|(_, s)| !s.status.is_terminal() && s.timed_out())
343                .map(|(id, _)| id.clone())
344                .collect();
345            let stale = sessions
346                .iter()
347                .filter(|(_, s)| {
348                    s.finished_at
349                        .is_some_and(|t| t.elapsed() > Duration::from_secs(FINISHED_RETENTION_SECS))
350                })
351                .map(|(id, _)| id.clone())
352                .collect();
353            (expired, stale)
354        };
355        for id in &expired {
356            tracing::info!(session_id = %id, "auth session timed out: ending it");
357            self.end_session(id, timeout_error());
358        }
359        if !stale.is_empty() {
360            // A finished session's CLI can still be running (a success line
361            // seen, no exit yet): stop it before its session, and with it the
362            // directory key a later sign-in would find it by, goes.
363            for id in &stale {
364                self.stop_process(id);
365            }
366            let mut sessions = self.sessions.lock().unwrap();
367            for id in &stale {
368                sessions.remove(id);
369            }
370        }
371    }
372
373    /// Run `sweep` every SWEEP_INTERVAL_SECS for as long as srv runs.
374    pub fn spawn_sweeper(self: &Arc<Self>) {
375        let mgr = Arc::clone(self);
376        tokio::spawn(async move {
377            let mut tick = tokio::time::interval(Duration::from_secs(SWEEP_INTERVAL_SECS));
378            loop {
379                tick.tick().await;
380                mgr.sweep();
381            }
382        });
383    }
384
385    /// Cancel a session: transition state to Failed + abort the
386    /// spawned drain task (which kills the child via kill_on_drop).
387    /// Returns true if the state transition happened (false on
388    /// unknown / already-terminal session — the process refs are
389    /// torn down either way so a re-fired cancel is a no-op).
390    pub fn cancel_session(&self, session_id: &str) -> bool {
391        self.end_session(session_id, "cancelled by user".to_string())
392    }
393
394    /// Fail a session with `error` and stop its CLI: abort the drain task
395    /// (the pipe path's child goes with it, `kill_on_drop`), drop its stdin,
396    /// and kill a PTY child by PID. True if the state changed; the process
397    /// refs go either way, so ending a session twice is harmless.
398    fn end_session(&self, session_id: &str, error: String) -> bool {
399        let transitioned = self.finish_failure(session_id, error);
400        self.stop_process(session_id);
401        transitioned
402    }
403
404    /// Stop a session's CLI and drop its process refs: abort the drain task
405    /// (the pipe path's child goes with it, `kill_on_drop`), drop its stdin,
406    /// and kill a PTY child by PID.
407    fn stop_process(&self, session_id: &str) {
408        let mut refs = self.process_refs.lock().unwrap();
409        if let Some(handle) = refs.drain_tasks.remove(session_id) {
410            handle.abort();
411        }
412        refs.stdin_senders.remove(session_id);
413        // PTY transport: abort() doesn't reach the spawn_blocking
414        // wait task, so kill the child by PID. Pipes path doesn't
415        // need this — its tokio Child uses kill_on_drop.
416        if let Some(pid) = refs.pty_pids.remove(session_id) {
417            if let Err(e) = kill_pid(pid) {
418                tracing::warn!(pid, session_id, error = %e, "end_session: kill_pid failed");
419            } else {
420                tracing::info!(pid, session_id, "end_session: PTY child killed");
421            }
422        }
423    }
424
425    /// Has this session ended (or gone)? A login confirmed after its session
426    /// was replaced, cancelled or timed out must not be saved: a newer
427    /// sign-in owns the directory now.
428    pub fn ended(&self, session_id: &str) -> bool {
429        self.is_terminal(session_id)
430    }
431
432    fn is_terminal(&self, session_id: &str) -> bool {
433        self.sessions
434            .lock()
435            .unwrap()
436            .get(session_id)
437            .is_none_or(|s| s.status.is_terminal())
438    }
439
440    /// Register the PID of a PTY-backed auth subprocess so
441    /// `cancel_session` can terminate it. Called by `auth.start`
442    /// after spawning a PTY login (in addition to `attach_process`).
443    pub fn attach_pty_pid(&self, session_id: &str, pid: u32) {
444        self.process_refs.lock().unwrap().pty_pids.insert(session_id.to_string(), pid);
445        // Registered first, checked second: a session ended (replaced,
446        // cancelled, timed out) before this, whose end found nothing to stop,
447        // has its CLI stopped here; one ended after finds it registered.
448        if self.is_terminal(session_id) {
449            self.stop_process(session_id);
450        }
451    }
452
453    /// Register the drain task + stdin sender for a session. Called
454    /// by the handler immediately after spawning the CLI.
455    pub fn attach_process(
456        &self,
457        session_id: &str,
458        drain_task: tokio::task::JoinHandle<()>,
459        stdin_sender: tokio::sync::mpsc::Sender<String>,
460    ) {
461        {
462            let mut refs = self.process_refs.lock().unwrap();
463            refs.drain_tasks.insert(session_id.to_string(), drain_task);
464            refs.stdin_senders.insert(session_id.to_string(), stdin_sender);
465        }
466        // As in `attach_pty_pid`: a session that already ended stops now.
467        if self.is_terminal(session_id) {
468            self.stop_process(session_id);
469        }
470    }
471
472    /// Forward a pasted callback URL to the spawned CLI's stdin.
473    /// Returns true if the session has an attached stdin sender;
474    /// false if the session never spawned a CLI or the sender's
475    /// receiver was dropped.
476    pub async fn send_to_stdin(&self, session_id: &str, line: String) -> bool {
477        let sender = {
478            let refs = self.process_refs.lock().unwrap();
479            refs.stdin_senders.get(session_id).cloned()
480        };
481        match sender {
482            Some(s) => s.send(line).await.is_ok(),
483            None => false,
484        }
485    }
486
487    /// Drop a session's process refs without aborting them. Called
488    /// by the drain task itself when it exits normally so terminal
489    /// state lookups still work but the resources are reclaimed.
490    pub fn detach_process(&self, session_id: &str) {
491        let mut refs = self.process_refs.lock().unwrap();
492        refs.drain_tasks.remove(session_id);
493        refs.stdin_senders.remove(session_id);
494        refs.pty_pids.remove(session_id);
495    }
496
497    /// The account email scraped from this session's login transcript, if the
498    /// provider printed one.
499    ///
500    /// Read *before* `finish_success`, which consumes the session — the
501    /// account is persisted first (`identity_auth_spawn`) and needs the email
502    /// at that point to record it on the account
503    /// (`SPEC_ACCOUNT_EMAIL_IN_ARMORY_2026_09_23.md`). `None` when the
504    /// provider reported no email, which is not an error: providers whose CLI
505    /// does not surface one simply have no email to show.
506    pub fn captured_email(&self, session_id: &str) -> Option<String> {
507        self.sessions
508            .lock()
509            .unwrap()
510            .get(session_id)
511            .and_then(|s| s.captured_email.clone())
512    }
513
514    /// Read the full transcript of captured stdout/stderr lines.
515    /// Used by integration tests; exposed for completeness even
516    /// though no production caller currently reads it.
517    #[allow(dead_code)]
518    pub fn transcript(&self, session_id: &str) -> Option<Vec<String>> {
519        self.sessions
520            .lock()
521            .unwrap()
522            .get(session_id)
523            .map(|s| s.transcript.clone())
524    }
525
526    /// Inject a force-elapsed `started_at` for the given session —
527    /// test helper so we can exercise the timeout transition without
528    /// actually waiting SESSION_TIMEOUT_SECS.
529    #[cfg(test)]
530    fn force_age(&self, session_id: &str) {
531        if let Some(s) = self.sessions.lock().unwrap().get_mut(session_id) {
532            s.started_at = Instant::now() - Duration::from_secs(SESSION_TIMEOUT_SECS + 1);
533        }
534    }
535
536    /// Test helper: make a finished session look FINISHED_RETENTION_SECS old.
537    #[cfg(test)]
538    fn force_finished_long_ago(&self, session_id: &str) {
539        if let Some(s) = self.sessions.lock().unwrap().get_mut(session_id) {
540            s.finished_at = Some(Instant::now() - Duration::from_secs(FINISHED_RETENTION_SECS + 1));
541        }
542    }
543}
544
545fn timeout_error() -> String {
546    format!("auth session timed out after {SESSION_TIMEOUT_SECS}s")
547}
548
549/// Best-effort kill of a child process by PID — `SIGTERM` on Unix,
550/// `taskkill /F /T /PID` on Windows. This used to be a private copy that
551/// described itself as a "mirror of the cef-side helper"; both now share
552/// `agentmux_common::process::kill_pid`.
553use agentmux_common::process::kill_pid;
554
555#[cfg(test)]
556mod tests {
557    use super::*;
558
559    fn mgr() -> AuthSessionManager {
560        AuthSessionManager::new()
561    }
562
563    #[test]
564    fn start_creates_pending_session() {
565        let m = mgr();
566        let r = m.start_session("claude".to_string(), None);
567        assert!(!r.session_id.is_empty());
568        assert!(r.auth_url.is_none());
569        let p = m.poll_session(&r.session_id).expect("session exists");
570        assert_eq!(p.provider_id, "claude");
571        assert!(matches!(p.status, AuthSessionStatus::Pending));
572    }
573
574    #[test]
575    fn url_line_transitions_to_url_available() {
576        let m = mgr();
577        let r = m.start_session("claude".to_string(), None);
578        let _ = m.record_line(
579            &r.session_id,
580            "Open https://console.anthropic.com/oauth/authorize?state=xyz",
581        );
582        let p = m.poll_session(&r.session_id).unwrap();
583        match p.status {
584            AuthSessionStatus::UrlAvailable { auth_url } => {
585                assert!(auth_url.contains("anthropic.com/oauth"));
586            }
587            other => panic!("expected UrlAvailable, got {other:?}"),
588        }
589    }
590
591    #[test]
592    fn device_code_line_transitions_to_code_emitted() {
593        let m = mgr();
594        let r = m.start_session("copilot".to_string(), None);
595        let _ = m.record_line(&r.session_id, "! Copy your one-time code: ABCD-1234");
596        let p = m.poll_session(&r.session_id).unwrap();
597        match p.status {
598            AuthSessionStatus::CodeEmitted {
599                device_code,
600                verification_url,
601            } => {
602                assert_eq!(device_code, "ABCD-1234");
603                assert_eq!(verification_url, "https://github.com/login/device");
604            }
605            other => panic!("expected CodeEmitted, got {other:?}"),
606        }
607    }
608
609    #[test]
610    fn login_success_line_does_not_transition_state_alone() {
611        // The CLI saying "logged in" isn't enough — the handler
612        // confirms via authCheck before declaring Success.
613        let m = mgr();
614        let r = m.start_session("claude".to_string(), None);
615        let _ = m.record_line(&r.session_id, "Successfully logged in as asaf@example.com");
616        let p = m.poll_session(&r.session_id).unwrap();
617        // Still pending — handler hasn't called finish_success yet.
618        assert!(matches!(p.status, AuthSessionStatus::Pending));
619    }
620
621    #[test]
622    fn finish_success_carries_email_from_transcript() {
623        let m = mgr();
624        let r = m.start_session("claude".to_string(), None);
625        let _ = m.record_line(&r.session_id, "Successfully logged in as asaf@example.com");
626        assert!(m.finish_success(&r.session_id, "bundle-1".to_string(), None));
627        let p = m.poll_session(&r.session_id).unwrap();
628        match p.status {
629            AuthSessionStatus::Success { bundle_id, email, account_id } => {
630                assert_eq!(bundle_id, "bundle-1");
631                assert_eq!(email.as_deref(), Some("asaf@example.com"));
632                assert_eq!(account_id, None);
633            }
634            other => panic!("expected Success, got {other:?}"),
635        }
636    }
637
638    #[test]
639    fn cancel_transitions_to_failed() {
640        let m = mgr();
641        let r = m.start_session("claude".to_string(), None);
642        assert!(m.cancel_session(&r.session_id));
643        let p = m.poll_session(&r.session_id).unwrap();
644        match p.status {
645            AuthSessionStatus::Failed { error } => assert!(error.contains("cancelled")),
646            other => panic!("expected Failed, got {other:?}"),
647        }
648    }
649
650    #[test]
651    // FLAKY on low-uptime machines (CI runners): force_age does
652    // `Instant::now() - Duration::from_secs(SESSION_TIMEOUT_SECS + 1)`, which
653    // PANICS when monotonic uptime is less than the timeout (can't construct an
654    // Instant that far in the past). Passes locally (high uptime), fails on a
655    // freshly-booted CI runner. Production is unaffected — it never subtracts from
656    // Instant. Ignored to unblock CI; fix = make the timeout mockable (deadline
657    // model / injectable clock) so force_age doesn't underflow.
658    // SPEC_CI_TEST_RUNNER_2026_06_22.md §6.4.
659    #[ignore = "force_age underflows Instant on low-uptime CI runners; make the timeout mockable, then un-ignore"]
660    fn timeout_transitions_pending_to_failed_on_poll() {
661        let m = mgr();
662        let r = m.start_session("claude".to_string(), None);
663        m.force_age(&r.session_id);
664        let p = m.poll_session(&r.session_id).unwrap();
665        match p.status {
666            AuthSessionStatus::Failed { error } => assert!(error.contains("timed out")),
667            other => panic!("expected timeout Failed, got {other:?}"),
668        }
669    }
670
671    #[test]
672    fn terminal_states_cannot_be_re_transitioned() {
673        let m = mgr();
674        let r = m.start_session("claude".to_string(), None);
675        assert!(m.finish_success(&r.session_id, "bundle-1".to_string(), None));
676        // Second finish_failure is a no-op.
677        assert!(!m.finish_failure(&r.session_id, "should be ignored".to_string()));
678        let p = m.poll_session(&r.session_id).unwrap();
679        assert!(matches!(p.status, AuthSessionStatus::Success { .. }));
680    }
681
682    #[test]
683    fn multiple_url_lines_keep_the_first_url() {
684        // If the CLI emits the URL again later (some do), we don't
685        // overwrite — the first URL is what the user saw + pasted.
686        let m = mgr();
687        let r = m.start_session("claude".to_string(), None);
688        let _ = m.record_line(
689            &r.session_id,
690            "Open https://console.anthropic.com/oauth/authorize?state=first",
691        );
692        let _ = m.record_line(
693            &r.session_id,
694            "Open https://console.anthropic.com/oauth/authorize?state=second",
695        );
696        let p = m.poll_session(&r.session_id).unwrap();
697        match p.status {
698            AuthSessionStatus::UrlAvailable { auth_url } => {
699                assert!(auth_url.contains("state=first"));
700            }
701            _ => panic!("expected UrlAvailable"),
702        }
703    }
704
705    #[test]
706    fn status_serializes_with_camelcase_field_names() {
707        // Codex P2 on PR #840: the per-variant fields used to
708        // serialize as auth_url / device_code / etc. (snake_case)
709        // while the rest of the wire is camelCase. Now uniformly
710        // camelCase via rename_all_fields.
711        let s = AuthSessionStatus::UrlAvailable {
712            auth_url: "https://example.com/oauth".to_string(),
713        };
714        let v = serde_json::to_value(&s).unwrap();
715        assert_eq!(
716            v,
717            serde_json::json!({
718                "status": "url-available",
719                "authUrl": "https://example.com/oauth"
720            })
721        );
722
723        let s = AuthSessionStatus::CodeEmitted {
724            device_code: "ABCD-1234".to_string(),
725            verification_url: "https://github.com/login/device".to_string(),
726        };
727        let v = serde_json::to_value(&s).unwrap();
728        assert_eq!(
729            v,
730            serde_json::json!({
731                "status": "code-emitted",
732                "deviceCode": "ABCD-1234",
733                "verificationUrl": "https://github.com/login/device"
734            })
735        );
736
737        let s = AuthSessionStatus::Success {
738            bundle_id: "bundle-1".to_string(),
739            email: Some("asaf@example.com".to_string()),
740            account_id: None,
741        };
742        let v = serde_json::to_value(&s).unwrap();
743        assert_eq!(
744            v,
745            serde_json::json!({
746                "status": "success",
747                "bundleId": "bundle-1",
748                "email": "asaf@example.com"
749            })
750        );
751
752        // Direct-account mode (issue #1624 PR-C Part B) — accountId
753        // appears on the wire; bundleId is present but empty.
754        let s = AuthSessionStatus::Success {
755            bundle_id: String::new(),
756            email: Some("asaf@example.com".to_string()),
757            account_id: Some("acc-1".to_string()),
758        };
759        let v = serde_json::to_value(&s).unwrap();
760        assert_eq!(
761            v,
762            serde_json::json!({
763                "status": "success",
764                "bundleId": "",
765                "email": "asaf@example.com",
766                "accountId": "acc-1"
767            })
768        );
769    }
770
771    #[test]
772    fn unknown_session_polls_to_none() {
773        let m = mgr();
774        assert!(m.poll_session("does-not-exist").is_none());
775    }
776
777    #[test]
778    fn transcript_records_all_lines_including_non_matching() {
779        let m = mgr();
780        let r = m.start_session("claude".to_string(), None);
781        let _ = m.record_line(&r.session_id, "Starting auth flow...");
782        let _ = m.record_line(&r.session_id, "Open https://console.anthropic.com/oauth/authorize");
783        let _ = m.record_line(&r.session_id, "Waiting for callback...");
784        let t = m.transcript(&r.session_id).unwrap();
785        assert_eq!(t.len(), 3);
786        assert_eq!(t[0], "Starting auth flow...");
787        assert!(t[1].contains("anthropic.com/oauth"));
788        assert_eq!(t[2], "Waiting for callback...");
789    }
790
791    #[tokio::test]
792    async fn sweep_ends_a_timed_out_session_nobody_polls_and_drops_its_stdin() {
793        let m = mgr();
794        let r = m.start_session("claude".to_string(), None);
795        let (tx, _rx) = tokio::sync::mpsc::channel::<String>(1);
796        m.attach_process(&r.session_id, tokio::spawn(async {}), tx);
797        m.force_age(&r.session_id);
798
799        m.sweep();
800
801        let status = m.sessions.lock().unwrap().get(&r.session_id).unwrap().status.clone();
802        assert!(matches!(status, AuthSessionStatus::Failed { ref error } if error.contains("timed out")), "{status:?}");
803        assert!(!m.send_to_stdin(&r.session_id, "code".to_string()).await, "stdin should be gone");
804    }
805
806    #[test]
807    fn sweep_drops_a_session_finished_long_ago_and_keeps_a_recent_one() {
808        let m = mgr();
809        let old = m.start_session("claude".to_string(), None);
810        let recent = m.start_session("claude".to_string(), None);
811        m.finish_failure(&old.session_id, "x".to_string());
812        m.finish_failure(&recent.session_id, "x".to_string());
813        m.force_finished_long_ago(&old.session_id);
814
815        m.sweep();
816
817        assert!(m.poll_session(&old.session_id).is_none());
818        assert!(m.poll_session(&recent.session_id).is_some());
819    }
820
821    #[test]
822    fn a_new_login_for_the_same_account_dir_replaces_the_live_one() {
823        let m = mgr();
824        let dir = Some("/home/u/.agentmux/accounts/claude-1".to_string());
825        let first = m.start_session("claude".to_string(), dir.clone());
826        let other = m.start_session("claude".to_string(), Some("/elsewhere".to_string()));
827        let second = m.start_session("claude".to_string(), dir);
828
829        let first_status = m.poll_session(&first.session_id).unwrap().status;
830        assert!(
831            matches!(first_status, AuthSessionStatus::Failed { ref error } if error.contains("replaced")),
832            "{first_status:?}"
833        );
834        assert_eq!(m.poll_session(&second.session_id).unwrap().status, AuthSessionStatus::Pending);
835        assert_eq!(m.poll_session(&other.session_id).unwrap().status, AuthSessionStatus::Pending);
836    }
837
838    #[tokio::test]
839    async fn a_cli_attached_after_its_session_was_replaced_is_stopped_at_once() {
840        let m = mgr();
841        let dir = Some("/home/u/.agentmux/accounts/claude-1".to_string());
842        let older = m.start_session("claude".to_string(), dir.clone());
843        // The newer start replaces the older session before its CLI exists.
844        let _newer = m.start_session("claude".to_string(), dir);
845        // The older request's CLI arrives afterwards.
846        let (tx, _rx) = tokio::sync::mpsc::channel::<String>(1);
847        m.attach_process(&older.session_id, tokio::spawn(std::future::pending::<()>()), tx);
848
849        assert!(!m.send_to_stdin(&older.session_id, "code".to_string()).await, "the replaced CLI keeps no stdin");
850        assert!(m.process_refs.lock().unwrap().drain_tasks.get(&older.session_id).is_none());
851    }
852
853    #[tokio::test]
854    async fn a_new_login_stops_the_cli_of_a_finished_session_for_the_same_dir() {
855        let m = mgr();
856        let dir = Some("/home/u/.agentmux/accounts/claude-1".to_string());
857        let done = m.start_session("claude".to_string(), dir.clone());
858        let (tx, _rx) = tokio::sync::mpsc::channel::<String>(1);
859        m.attach_process(&done.session_id, tokio::spawn(std::future::pending::<()>()), tx);
860        // Success was recorded, but its CLI hasn't exited yet.
861        m.finish_success(&done.session_id, String::new(), Some("acct".to_string()));
862
863        let _newer = m.start_session("claude".to_string(), dir);
864
865        assert!(m.process_refs.lock().unwrap().drain_tasks.get(&done.session_id).is_none(), "its CLI was stopped");
866        assert!(matches!(m.poll_session(&done.session_id).unwrap().status, AuthSessionStatus::Success { .. }), "a finished session keeps its result");
867    }
868
869    #[tokio::test]
870    async fn sweep_stops_a_still_running_cli_before_dropping_its_finished_session() {
871        let m = mgr();
872        let r = m.start_session("claude".to_string(), Some("/d".to_string()));
873        let (tx, _rx) = tokio::sync::mpsc::channel::<String>(1);
874        m.attach_process(&r.session_id, tokio::spawn(std::future::pending::<()>()), tx);
875        m.finish_success(&r.session_id, String::new(), None);
876        m.force_finished_long_ago(&r.session_id);
877
878        m.sweep();
879
880        assert!(m.poll_session(&r.session_id).is_none());
881        assert!(m.process_refs.lock().unwrap().drain_tasks.get(&r.session_id).is_none(), "its CLI was stopped");
882    }
883
884    #[test]
885    fn ended_is_true_once_replaced_and_for_an_unknown_session() {
886        let m = mgr();
887        let first = m.start_session("claude".to_string(), Some("/d".to_string()));
888        assert!(!m.ended(&first.session_id));
889        let _second = m.start_session("claude".to_string(), Some("/d".to_string()));
890        assert!(m.ended(&first.session_id));
891        assert!(m.ended("auth-nope"));
892    }
893}
agentmux_srv/identity/auth_session.rs

agentmux_srv/identity/
auth_session.rs
