agentmux_srv/muxbus/
pkce.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! PKCE Authorization Code flow for Cognito desktop login.
5//!
6//! Flow (cloud-relayed callback — no loopback listener; see
7//! docs/specs/SPEC_MUXBUS_CLOUD_RELAYED_LOGIN_CALLBACK_2026_08_15.md):
8//!   1. Generate code_verifier + code_challenge (S256).
9//!   2. Register the flow's state with the muxbus login relay.
10//!   3. Open browser to Cognito hosted UI; it redirects to the hosted
11//!      /desktop-callback page, which posts {state, code} to the relay.
12//!   4. Poll the relay for the code (single-read on the relay side).
13//!   5. Exchange code for tokens via /oauth2/token (PKCE verifier never
14//!      left this process).
15//!   6. Decode email from id_token claims (no re-verification needed —
16//!      we fetched the token directly from Cognito).
17
18use base64::Engine as _;
19use sha2::{Digest, Sha256};
20
21use crate::backend::storage::muxbus::MuxBusCredentials;
22use crate::util::open_browser;
23
24pub const LOGIN_TIMEOUT_SECS: u64 = 300;
25
26#[derive(Debug)]
27pub struct PkceResult {
28    pub credentials: MuxBusCredentials,
29}
30
31/// Base URL of the muxbus REST API hosting the login relay.
32///
33/// Single definition shared with the tier-4 outbound relay
34/// (`crate::muxbus::relay::rest_base_url`) — two copies of "where is the
35/// cloud" is precisely the drift the network DRYness pass set out to remove.
36use crate::muxbus::relay::rest_base_url as relay_base_url;
37
38// Only one login flow may exist per process. Two concurrent flows both bind
39// the fixed callback port, and (on Windows, where SO_REUSEADDR permits binding
40// a port another live socket already holds) the browser's redirect lands on
41// whichever listener the OS picks — a 400 state-mismatch coin flip. The newest
42// attempt is the one the user actually wants: abort the predecessor.
43//
44// Each flow gets a monotonic generation id (NEXT_LOGIN_ID), stored alongside
45// its abort handle. A prior design tracked "was this abort a user Cancel?"
46// in a single shared CANCELLED_BY_USER flag instead of per-attempt — that
47// was a real bug (Codex + reagent review, PR #2661, 2026-08-19): if flow A
48// is cancelled but a *different* flow B starts and is itself superseded by
49// C before A's own `task.await` resumes, nothing guarantees which of A/B
50// reads the shared flag first — the two flows' error messages could swap.
51// Keying the cancellation record on the generation id makes that
52// structurally impossible: each flow only ever looks up its OWN id.
53static NEXT_LOGIN_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
54static ACTIVE_LOGIN: std::sync::Mutex<Option<(u64, tokio::task::AbortHandle)>> =
55    std::sync::Mutex::new(None);
56static CANCELLED_IDS: std::sync::Mutex<Vec<u64>> = std::sync::Mutex::new(Vec::new());
57
58/// Abort the in-flight login flow, if any. Returns false if there was none
59/// (a harmless no-op — the UI can call this without checking state first).
60pub fn cancel_active_login() -> bool {
61    let mut guard = ACTIVE_LOGIN.lock().unwrap();
62    match guard.take() {
63        Some((id, handle)) => {
64            // ACTIVE_LOGIN is only ever overwritten by the NEXT login's
65            // critical section — nothing clears it when a flow resolves
66            // naturally (`Ok` or a real, non-cancelled `Err`). So this can
67            // find a stale entry whose task already finished on its own,
68            // possibly long ago; `AbortHandle::abort()` on a finished task
69            // is a documented no-op. Reporting `true` here for that case
70            // would tell the caller "cancelled" for a login that actually
71            // succeeded/failed by itself, and would push `id` into
72            // CANCELLED_IDS, which nothing on the natural-completion path
73            // in `run_pkce_login` ever pops — leaking one entry per
74            // occurrence for the life of the process (reagent P1, PR
75            // #2661, re-review 2026-08-19). Checking `is_finished()` first
76            // closes the common case (a stale entry sitting there for any
77            // length of time before this call). A narrower race remains —
78            // the task can still finish in the gap between this check and
79            // `abort()` below taking effect — closed instead by
80            // `run_pkce_login` cleaning up its own id on every resolution
81            // path, not just the cancelled one.
82            if handle.is_finished() {
83                return false;
84            }
85            CANCELLED_IDS.lock().unwrap().push(id);
86            handle.abort();
87            true
88        }
89        None => false,
90    }
91}
92
93pub async fn run_pkce_login(
94    cognito_domain: &str,
95    client_id: &str,
96    http_client: &reqwest::Client,
97) -> Result<PkceResult, String> {
98    let my_id = NEXT_LOGIN_ID.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
99
100    // Abort-predecessor and register-successor must be one critical section:
101    // with separate lock acquisitions, two concurrent muxbus.login RPCs can
102    // both take() None and neither aborts the other (reagent P1 on this PR).
103    // tokio::spawn is synchronous and non-blocking, so holding the std mutex
104    // across it is fine — the guard never spans an await.
105    let task = {
106        let mut guard = ACTIVE_LOGIN.lock().unwrap();
107        if let Some((_, prev)) = guard.take() {
108            prev.abort();
109        }
110        let task = tokio::spawn(run_pkce_login_inner(
111            cognito_domain.to_string(),
112            client_id.to_string(),
113            http_client.clone(),
114        ));
115        *guard = Some((my_id, task.abort_handle()));
116        task
117    };
118    match task.await {
119        Ok(r) => {
120            // The flow resolved naturally. `cancel_active_login()`'s
121            // `is_finished()` guard closes the common case, but a narrow
122            // race remains: the task can complete in the gap between that
123            // check and its `abort()` call taking effect (tokio: an abort
124            // request against an already-completing task is ignored, the
125            // task "will run to completion" per its own docs), which would
126            // have pushed `my_id` into CANCELLED_IDS just before this arm
127            // runs. Nothing else on this path ever pops it — clean it up
128            // here so it can't leak.
129            CANCELLED_IDS.lock().unwrap().retain(|&id| id != my_id);
130            r
131        }
132        Err(e) if e.is_cancelled() => {
133            let mut cancelled = CANCELLED_IDS.lock().unwrap();
134            if let Some(pos) = cancelled.iter().position(|&id| id == my_id) {
135                cancelled.remove(pos);
136                Err("sign-in cancelled".to_string())
137            } else {
138                Err("this login attempt was superseded by a newer one".to_string())
139            }
140        }
141        Err(e) => {
142            // Same residual-race cleanup as the Ok(r) arm above, for a
143            // flow that failed on its own (not via cancellation).
144            CANCELLED_IDS.lock().unwrap().retain(|&id| id != my_id);
145            Err(format!("login task failed: {e}"))
146        }
147    }
148}
149
150async fn run_pkce_login_inner(
151    cognito_domain: String,
152    client_id: String,
153    http_client: reqwest::Client,
154) -> Result<PkceResult, String> {
155    // 1. Generate code_verifier (43–128 chars of URL-safe chars)
156    let v1 = uuid::Uuid::new_v4();
157    let v2 = uuid::Uuid::new_v4();
158    let mut raw = [0u8; 32];
159    raw[..16].copy_from_slice(v1.as_bytes());
160    raw[16..].copy_from_slice(v2.as_bytes());
161    let code_verifier = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw);
162
163    // 2. code_challenge = BASE64URL(SHA-256(code_verifier))
164    let mut hasher = Sha256::new();
165    hasher.update(code_verifier.as_bytes());
166    let digest = hasher.finalize();
167    let code_challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest);
168
169    // 3. State (CSRF)
170    let state = uuid::Uuid::new_v4().to_string();
171
172    // 4. Register the flow with the cloud relay BEFORE opening the browser.
173    //    The hosted /desktop-callback page (agentmux-cloud's login-relay page)
174    //    posts {state, code} there; step 7 polls it back out. No loopback
175    //    listener: the browser never touches 127.0.0.1, and concurrent
176    //    instances can't collide — each flow polls its own state key.
177    //    See docs/specs/SPEC_MUXBUS_CLOUD_RELAYED_LOGIN_CALLBACK_2026_08_15.md.
178    let relay_base = relay_base_url();
179    let redirect_uri = format!("{relay_base}/desktop-callback");
180    let create_resp = http_client
181        .post(format!("{relay_base}/api/login-relay"))
182        .json(&serde_json::json!({ "state": state }))
183        .send()
184        .await
185        .map_err(|e| format!("could not reach the login relay: {e}"))?;
186    if !create_resp.status().is_success() {
187        let status = create_resp.status();
188        let body = create_resp.text().await.unwrap_or_default();
189        return Err(format!("login relay refused the flow ({status}): {body}"));
190    }
191
192    // 5. Build auth URL
193    let scopes = "openid+email+profile+https%3A%2F%2Fmuxbus.agentmux.ai%2Fread+https%3A%2F%2Fmuxbus.agentmux.ai%2Fwrite";
194    let auth_url = format!(
195        "{cognito_domain}/oauth2/authorize\
196         ?response_type=code\
197         &client_id={client_id}\
198         &redirect_uri={redirect_uri_enc}\
199         &scope={scopes}\
200         &code_challenge={code_challenge}\
201         &code_challenge_method=S256\
202         &state={state}",
203        redirect_uri_enc = percent_encode(&redirect_uri),
204    );
205
206    // 6. Open browser (suppressed under test — the supersede test drives
207    //    real flows that would otherwise launch real browser windows)
208    if !cfg!(test) {
209        open_browser(&auth_url);
210    }
211
212    tracing::info!(
213        cognito_domain = cognito_domain,
214        relay = relay_base,
215        "muxbus: PKCE login started, awaiting relayed browser callback"
216    );
217
218    // 7. Poll the relay until the code arrives (5-min overall timeout, matching
219    //    the relay record's TTL). Transient relay/network errors are tolerated
220    //    and re-polled; a 404 is fatal — the flow expired or was already
221    //    consumed (single-read on the relay side).
222    let code = tokio::time::timeout(
223        std::time::Duration::from_secs(LOGIN_TIMEOUT_SECS),
224        async {
225            loop {
226                tokio::time::sleep(std::time::Duration::from_secs(1)).await;
227                let resp = match http_client
228                    .get(format!("{relay_base}/api/login-relay/{state}"))
229                    .send()
230                    .await
231                {
232                    Ok(r) => r,
233                    Err(_) => continue, // transient network problem — keep polling
234                };
235                if resp.status() == reqwest::StatusCode::NOT_FOUND {
236                    return Err(
237                        "login attempt expired or was completed elsewhere — please try again"
238                            .to_string(),
239                    );
240                }
241                if !resp.status().is_success() {
242                    continue; // transient server error — keep polling
243                }
244                let body: serde_json::Value = match resp.json().await {
245                    Ok(v) => v,
246                    Err(_) => continue,
247                };
248                match body["status"].as_str() {
249                    Some("ready") => match body["code"].as_str() {
250                        Some(c) => return Ok(c.to_string()),
251                        None => return Err("relay returned ready without a code".to_string()),
252                    },
253                    // Cognito denial forwarded by the hosted page — fail fast
254                    // rather than polling out the rest of the 5 minutes.
255                    Some("failed") => {
256                        let err = body["error"].as_str().unwrap_or("unknown");
257                        return Err(format!("Cognito returned error: {err}"));
258                    }
259                    // "pending" or anything unrecognized — keep polling until
260                    // the overall timeout says otherwise.
261                    _ => continue,
262                }
263            }
264        },
265    )
266    .await
267    .map_err(|_| "login timed out (5 min) — please try again".to_string())??;
268
269    // 8. Exchange code for tokens
270    let token_url = format!("{cognito_domain}/oauth2/token");
271    let params = [
272        ("grant_type", "authorization_code"),
273        ("client_id", client_id.as_str()),
274        ("code", &code),
275        ("redirect_uri", &redirect_uri),
276        ("code_verifier", &code_verifier),
277    ];
278
279    let resp = http_client
280        .post(&token_url)
281        .form(&params)
282        .send()
283        .await
284        .map_err(|e| format!("token exchange request failed: {e}"))?;
285
286    if !resp.status().is_success() {
287        let body = resp.text().await.unwrap_or_default();
288        return Err(format!("token exchange failed: {body}"));
289    }
290
291    let token_json: serde_json::Value = resp
292        .json()
293        .await
294        .map_err(|e| format!("token response parse failed: {e}"))?;
295
296    let access_token = token_json["access_token"]
297        .as_str()
298        .unwrap_or("")
299        .to_string();
300    let refresh_token = token_json["refresh_token"]
301        .as_str()
302        .unwrap_or("")
303        .to_string();
304    let id_token = token_json["id_token"].as_str().unwrap_or("").to_string();
305    let expires_in = token_json["expires_in"].as_i64().unwrap_or(3600);
306    let expires_at = (agentmux_common::time::now_secs())
307        + expires_in;
308
309    // 9. Extract email + sub from id_token payload (no re-verification needed)
310    let (user_email, user_sub) = extract_jwt_claims(&id_token);
311
312    tracing::info!(
313        email = user_email,
314        "muxbus: PKCE login succeeded"
315    );
316
317    Ok(PkceResult {
318        credentials: MuxBusCredentials {
319            cognito_domain: cognito_domain.to_string(),
320            client_id: client_id.to_string(),
321            access_token,
322            refresh_token,
323            id_token,
324            expires_at,
325            user_email,
326            user_sub,
327        },
328    })
329}
330
331/// Typed so callers (the credential broker's registered `refresh` closure)
332/// can distinguish "worth retrying" from "this refresh_token is dead, only
333/// a fresh login fixes it" without string-sniffing an error message.
334/// `Rejected` specifically means the token endpoint responded but rejected
335/// the request (Cognito's `invalid_grant`-class 4xx for a revoked/expired
336/// refresh_token) — the credential itself is the problem. `Network`/
337/// `ParseFailed` mean the ATTEMPT failed, not the credential.
338#[derive(Debug)]
339pub enum RefreshTokenError {
340    NoRefreshToken,
341    Network(String),
342    Rejected { status: u16, body: String },
343    ParseFailed(String),
344}
345
346impl std::fmt::Display for RefreshTokenError {
347    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
348        match self {
349            RefreshTokenError::NoRefreshToken => write!(f, "no refresh token stored"),
350            RefreshTokenError::Network(e) => write!(f, "refresh request failed: {e}"),
351            RefreshTokenError::Rejected { status, body } => {
352                write!(f, "token refresh failed ({status}): {body}")
353            }
354            RefreshTokenError::ParseFailed(e) => write!(f, "refresh response parse failed: {e}"),
355        }
356    }
357}
358
359pub async fn refresh_token(
360    creds: &MuxBusCredentials,
361    http_client: &reqwest::Client,
362) -> Result<MuxBusCredentials, RefreshTokenError> {
363    if creds.refresh_token.is_empty() {
364        return Err(RefreshTokenError::NoRefreshToken);
365    }
366    let token_url = format!("{}/oauth2/token", creds.cognito_domain);
367    let params = [
368        ("grant_type", "refresh_token"),
369        ("client_id", creds.client_id.as_str()),
370        ("refresh_token", creds.refresh_token.as_str()),
371    ];
372    let resp = http_client
373        .post(&token_url)
374        .form(&params)
375        .send()
376        .await
377        .map_err(|e| RefreshTokenError::Network(e.to_string()))?;
378    if !resp.status().is_success() {
379        let status = resp.status().as_u16();
380        let body = resp.text().await.unwrap_or_default();
381        return Err(RefreshTokenError::Rejected { status, body });
382    }
383    let json: serde_json::Value = resp
384        .json()
385        .await
386        .map_err(|e| RefreshTokenError::ParseFailed(e.to_string()))?;
387
388    let access_token = json["access_token"].as_str().unwrap_or("").to_string();
389    let expires_in = json["expires_in"].as_i64().unwrap_or(3600);
390    let expires_at = (agentmux_common::time::now_secs())
391        + expires_in;
392    // Cognito doesn't rotate refresh_token on refresh — keep existing
393    let new_id_token = json["id_token"]
394        .as_str()
395        .map(|s| s.to_string())
396        .unwrap_or_else(|| creds.id_token.clone());
397
398    Ok(MuxBusCredentials {
399        cognito_domain: creds.cognito_domain.clone(),
400        client_id: creds.client_id.clone(),
401        access_token,
402        refresh_token: creds.refresh_token.clone(),
403        id_token: new_id_token,
404        expires_at,
405        user_email: creds.user_email.clone(),
406        user_sub: creds.user_sub.clone(),
407    })
408}
409
410fn percent_encode(s: &str) -> String {
411    let mut out = String::with_capacity(s.len() * 3);
412    for b in s.bytes() {
413        match b {
414            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9'
415            | b'-' | b'_' | b'.' | b'~' => out.push(b as char),
416            _ => {
417                out.push('%');
418                out.push_str(&format!("{b:02X}"));
419            }
420        }
421    }
422    out
423}
424
425/// The `sub` claim of a Cognito token (unverified decode: this names the
426/// account a token we already hold belongs to, it doesn't authenticate it).
427/// Empty if the token isn't a JWT or has no `sub`.
428pub(crate) fn token_sub(token: &str) -> String {
429    extract_jwt_claims(token).1
430}
431
432fn extract_jwt_claims(token: &str) -> (String, String) {
433    let payload = token.splitn(3, '.').nth(1).unwrap_or("");
434    let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD
435        .decode(payload)
436        .or_else(|_| base64::engine::general_purpose::STANDARD.decode(payload))
437        .unwrap_or_default();
438    let json: serde_json::Value =
439        serde_json::from_slice(&decoded).unwrap_or_default();
440    let email = json["email"].as_str().unwrap_or("").to_string();
441    let sub = json["sub"].as_str().unwrap_or("").to_string();
442    (email, sub)
443}
444
445
446/// Serializes every test that touches process-global state: the login statics
447/// above and the `AGENTMUX_MUXBUS_REST_URL` / `AGENTMUX_MUXBUS_WS_URL`
448/// environment overrides (`cloud_subscriber`'s session test shares it).
449#[cfg(test)]
450pub(crate) static TEST_SERIAL: std::sync::Mutex<()> = std::sync::Mutex::new(());
451
452#[cfg(test)]
453mod tests {
454    use super::*;
455    use tokio::io::{AsyncReadExt, AsyncWriteExt};
456
457    /// Minimal fake login relay: 201 on POST /api/login-relay, an eternally
458    /// "pending" JSON body on GET polls. Keeps real flows parked at the
459    /// poll stage with zero external network dependency.
460    async fn spawn_fake_relay() -> String {
461        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
462        let addr = listener.local_addr().unwrap();
463        tokio::spawn(async move {
464            loop {
465                let Ok((mut stream, _)) = listener.accept().await else { return };
466                tokio::spawn(async move {
467                    let mut buf = [0u8; 2048];
468                    let Ok(n) = stream.read(&mut buf).await else { return };
469                    let req = std::str::from_utf8(&buf[..n]).unwrap_or("");
470                    let body = if req.starts_with("POST /api/login-relay") {
471                        r#"{"status":"pending","ttl_seconds":300}"#
472                    } else {
473                        r#"{"status":"pending"}"#
474                    };
475                    let resp = format!(
476                        "HTTP/1.1 {} 
477Content-Type: application/json
478Content-Length: {}
479Connection: close
480
481{}",
482                        if req.starts_with("POST") { "201 Created" } else { "200 OK" },
483                        body.len(),
484                        body
485                    );
486                    let _ = stream.write_all(resp.as_bytes()).await;
487                });
488            }
489        });
490        format!("http://{addr}")
491    }
492
493    /// Like `spawn_fake_relay`, but GET polls report `"failed"` immediately
494    /// — lets a test drive a flow to a *natural* (non-cancelled) resolution
495    /// quickly, instead of waiting out the real 5-minute timeout.
496    async fn spawn_fake_failing_relay() -> String {
497        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
498        let addr = listener.local_addr().unwrap();
499        tokio::spawn(async move {
500            loop {
501                let Ok((mut stream, _)) = listener.accept().await else { return };
502                tokio::spawn(async move {
503                    let mut buf = [0u8; 2048];
504                    let Ok(n) = stream.read(&mut buf).await else { return };
505                    let req = std::str::from_utf8(&buf[..n]).unwrap_or("");
506                    let body = if req.starts_with("POST /api/login-relay") {
507                        r#"{"status":"pending","ttl_seconds":300}"#
508                    } else {
509                        r#"{"status":"failed","error":"access_denied"}"#
510                    };
511                    let resp = format!(
512                        "HTTP/1.1 {}
513Content-Type: application/json
514Content-Length: {}
515Connection: close
516
517{}",
518                        if req.starts_with("POST") { "201 Created" } else { "200 OK" },
519                        body.len(),
520                        body
521                    );
522                    let _ = stream.write_all(resp.as_bytes()).await;
523                });
524            }
525        });
526        format!("http://{addr}")
527    }
528
529    // ACTIVE_LOGIN / CANCELLED_IDS are process-global statics — by
530    // design, only one login flow may exist per process. Cargo runs tests
531    // in this file on parallel threads of the same process by default, so
532    // any two tests that drive a real flow through those statics race each
533    // other unless serialized here.
534    use super::TEST_SERIAL;
535
536    // The superseded-login guarantee: starting flow B while flow A is pending
537    // must abort A rather than leaving two flows both polling (and both able
538    // to open browser tabs). Uses the real run_pkce_login entry so the
539    // abort-registry path itself is exercised; the fake relay keeps both
540    // flows parked at the poll stage.
541    #[tokio::test(flavor = "multi_thread")]
542    async fn second_login_supersedes_first() {
543        let _serial = TEST_SERIAL.lock().unwrap_or_else(|e| e.into_inner());
544        let relay = spawn_fake_relay().await;
545        std::env::set_var("AGENTMUX_MUXBUS_REST_URL", &relay);
546
547        let http = reqwest::Client::new();
548        let flow_a = tokio::spawn({
549            let http = http.clone();
550            async move { run_pkce_login("http://127.0.0.1:1", "test-client", &http).await }
551        });
552        // Give A time to register + start polling before B supersedes it.
553        tokio::time::sleep(std::time::Duration::from_millis(400)).await;
554
555        let flow_b = tokio::spawn({
556            let http = http.clone();
557            async move { run_pkce_login("http://127.0.0.1:1", "test-client", &http).await }
558        });
559        tokio::time::sleep(std::time::Duration::from_millis(400)).await;
560
561        let a = tokio::time::timeout(std::time::Duration::from_secs(2), flow_a)
562            .await
563            .expect("flow A should resolve promptly after being superseded")
564            .expect("flow A task must not panic");
565        assert!(
566            a.as_ref().is_err_and(|e| e.contains("superseded")),
567            "flow A should report being superseded, got: {a:?}"
568        );
569        assert!(!flow_b.is_finished(), "flow B should still be polling the relay");
570
571        // Cleanup: abort B's inner flow via the registry.
572        if let Some((_, h)) = ACTIVE_LOGIN.lock().unwrap().take() {
573            h.abort();
574        }
575        flow_b.abort();
576        std::env::remove_var("AGENTMUX_MUXBUS_REST_URL");
577    }
578
579    // A user-initiated Cancel must resolve the pending flow immediately
580    // (not wait out LOGIN_TIMEOUT_SECS) and report a "cancelled" error
581    // distinct from the supersede case above — otherwise the UI has no way
582    // to tell "you cancelled this" from "a newer login attempt replaced
583    // this", which is what the Cancel buttons (`useMuxBusStatus` in
584    // frontend/app/view/accounts/AgentMuxConnectPanel.tsx, also used by
585    // frontend/app/statusbar/HostPopover.tsx) need to show a quiet,
586    // non-error reset instead of a scary banner.
587    #[tokio::test(flavor = "multi_thread")]
588    async fn manual_cancel_resolves_promptly_with_cancelled_error() {
589        let _serial = TEST_SERIAL.lock().unwrap_or_else(|e| e.into_inner());
590        let relay = spawn_fake_relay().await;
591        std::env::set_var("AGENTMUX_MUXBUS_REST_URL", &relay);
592
593        let http = reqwest::Client::new();
594        let flow = tokio::spawn({
595            let http = http.clone();
596            async move { run_pkce_login("http://127.0.0.1:1", "test-client", &http).await }
597        });
598        // Give the flow time to register + start polling before cancelling.
599        tokio::time::sleep(std::time::Duration::from_millis(400)).await;
600
601        assert!(cancel_active_login(), "expected an active login to cancel");
602
603        let result = tokio::time::timeout(std::time::Duration::from_secs(2), flow)
604            .await
605            .expect("cancel should resolve the flow promptly, not after the 5-min timeout")
606            .expect("flow task must not panic");
607        assert!(
608            result.as_ref().is_err_and(|e| e.contains("cancelled")),
609            "expected a cancellation error, got: {result:?}"
610        );
611
612        std::env::remove_var("AGENTMUX_MUXBUS_REST_URL");
613    }
614
615    // Cancelling with no active login is a harmless no-op the UI can call
616    // defensively without checking state first.
617    #[tokio::test]
618    async fn cancel_with_no_active_login_is_a_noop() {
619        let _serial = TEST_SERIAL.lock().unwrap_or_else(|e| e.into_inner());
620        assert!(!cancel_active_login());
621    }
622
623    // Regression test for a real bug caught in code review (Codex + reagent,
624    // 2026-08-19, PR #2661): the original fix tracked "was this abort a user
625    // Cancel?" in a single process-global flag. Flow A gets cancelled by the
626    // user (flag set), then — before A's own `task.await` resumes and reads
627    // the flag — flow B starts and is itself superseded by flow C. Nothing
628    // guarantees which task's continuation the executor resumes first, so
629    // whichever of A/B checks the flag first "steals" it: the two flows'
630    // error messages ("cancelled" vs "superseded") can end up swapped. Fixed
631    // by keying the cancellation record on a per-attempt generation id
632    // instead of a single shared flag, so each flow can only ever observe
633    // its OWN cancellation record — this test asserts that invariant
634    // directly rather than trying to force a specific race outcome (which
635    // isn't reliably controllable from a test).
636    #[tokio::test(flavor = "multi_thread")]
637    async fn cancelling_one_flow_never_leaks_into_a_concurrent_unrelated_flow() {
638        let _serial = TEST_SERIAL.lock().unwrap_or_else(|e| e.into_inner());
639        let relay = spawn_fake_relay().await;
640        std::env::set_var("AGENTMUX_MUXBUS_REST_URL", &relay);
641        let http = reqwest::Client::new();
642
643        // Flow A: the user-cancel path.
644        let flow_a = tokio::spawn({
645            let http = http.clone();
646            async move { run_pkce_login("http://127.0.0.1:1", "test-client", &http).await }
647        });
648        tokio::time::sleep(std::time::Duration::from_millis(300)).await;
649        assert!(cancel_active_login(), "expected flow A to be active and cancellable");
650
651        // Flow B, then C superseding B — the "superseded" path, unrelated to
652        // A's cancel. Started immediately after A's cancel so A's own
653        // task.await may not have resumed yet when B/C run.
654        let flow_b = tokio::spawn({
655            let http = http.clone();
656            async move { run_pkce_login("http://127.0.0.1:1", "test-client", &http).await }
657        });
658        tokio::time::sleep(std::time::Duration::from_millis(300)).await;
659        let flow_c = tokio::spawn({
660            let http = http.clone();
661            async move { run_pkce_login("http://127.0.0.1:1", "test-client", &http).await }
662        });
663        tokio::time::sleep(std::time::Duration::from_millis(300)).await;
664
665        let a = tokio::time::timeout(std::time::Duration::from_secs(2), flow_a)
666            .await
667            .expect("flow A should resolve promptly")
668            .expect("flow A task must not panic");
669        let b = tokio::time::timeout(std::time::Duration::from_secs(2), flow_b)
670            .await
671            .expect("flow B should resolve promptly after being superseded by C")
672            .expect("flow B task must not panic");
673
674        assert!(
675            a.as_ref().is_err_and(|e| e.contains("cancelled")),
676            "flow A was user-cancelled, got: {a:?}"
677        );
678        assert!(
679            b.as_ref().is_err_and(|e| e.contains("superseded")),
680            "flow B was superseded by C (never cancelled), got: {b:?}"
681        );
682
683        // Cleanup: abort C's inner flow via the registry.
684        if let Some((_, h)) = ACTIVE_LOGIN.lock().unwrap().take() {
685            h.abort();
686        }
687        flow_c.abort();
688        std::env::remove_var("AGENTMUX_MUXBUS_REST_URL");
689    }
690
691    // Regression test for reagent's P1 re-review (PR #2661, 2026-08-19):
692    // ACTIVE_LOGIN was never cleared when a flow resolved naturally (`Ok`
693    // or a real, non-cancelled `Err`), so a `muxbus.login.cancel` call any
694    // time afterward found a stale entry pointing at a finished task's
695    // (harmless, no-op) AbortHandle — reporting a misleading "cancelled"
696    // success for a login that had already succeeded/failed on its own,
697    // and leaking one u64 into CANCELLED_IDS forever (nothing on the
698    // natural-completion path ever popped it). Fixed by having
699    // cancel_active_login() check AbortHandle::is_finished() before
700    // treating a take() as a real cancellation, plus having run_pkce_login
701    // clean up its own id from CANCELLED_IDS on every resolution path, not
702    // just the cancelled one.
703    #[tokio::test(flavor = "multi_thread")]
704    async fn cancelling_after_natural_completion_is_a_noop_and_does_not_leak() {
705        let _serial = TEST_SERIAL.lock().unwrap_or_else(|e| e.into_inner());
706        let relay = spawn_fake_failing_relay().await;
707        std::env::set_var("AGENTMUX_MUXBUS_REST_URL", &relay);
708
709        let http = reqwest::Client::new();
710        let result = run_pkce_login("http://127.0.0.1:1", "test-client", &http).await;
711        assert!(
712            result.as_ref().is_err_and(|e| e.contains("Cognito returned error")),
713            "expected the flow to fail on its own via the relay's failed status, got: {result:?}"
714        );
715
716        // The flow already resolved naturally — a cancel call now must be a
717        // no-op, not a misleading "cancelled" success.
718        assert!(
719            !cancel_active_login(),
720            "cancelling an already-resolved login must report nothing to cancel"
721        );
722        assert!(
723            CANCELLED_IDS.lock().unwrap().is_empty(),
724            "a naturally-resolved flow must never leave a CANCELLED_IDS entry behind"
725        );
726
727        std::env::remove_var("AGENTMUX_MUXBUS_REST_URL");
728    }
729}
agentmux_srv/muxbus/pkce.rs

agentmux_srv/muxbus/
pkce.rs
