Add generated token auth to app-server WebSockets

codex-rs/app-server-transport/src/transport/auth.rs

@@ -1,7 +1,10 @@
 use anyhow::Context;
 use axum::http::HeaderMap;
 use axum::http::StatusCode;
+use axum::http::Uri;
 use axum::http::header::AUTHORIZATION;
+use base64::Engine;
+use base64::engine::general_purpose::URL_SAFE_NO_PAD;
 use clap::Args;
 use clap::ValueEnum;
 use codex_utils_absolute_path::AbsolutePathBuf;
@@ -10,17 +13,20 @@ use jsonwebtoken::Algorithm;
 use jsonwebtoken::DecodingKey;
 use jsonwebtoken::Validation;
 use jsonwebtoken::decode;
+use rand::TryRngCore;
+use rand::rngs::OsRng;
 use serde::Deserialize;
 use sha2::Digest;
 use sha2::Sha256;
 use std::io;
 use std::io::ErrorKind;
-use std::net::SocketAddr;
 use std::path::Path;
 use std::path::PathBuf;
 use time::OffsetDateTime;
+use tracing::warn;
 
 const DEFAULT_MAX_CLOCK_SKEW_SECONDS: u64 = 30;
+const GENERATED_QUERY_TOKEN_BYTES: usize = 32;
 const MIN_SIGNED_BEARER_SECRET_BYTES: usize = 32;
 const INVALID_AUTHORIZATION_HEADER_MESSAGE: &str = "invalid authorization header";
 
@@ -53,6 +59,13 @@ pub struct AppServerWebsocketAuthArgs {
     /// Maximum clock skew when validating signed JWT bearer tokens.
     #[arg(long = "ws-max-clock-skew-seconds", value_name = "SECONDS")]
     pub ws_max_clock_skew_seconds: Option<u64>,
+
+    /// Disable enforcement of the generated websocket query token.
+    ///
+    /// A token is still generated and printed. Missing or incorrect tokens are
+    /// accepted.
+    #[arg(long = "no-token-check")]
+    pub no_token_check: bool,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
@@ -61,13 +74,17 @@ pub enum WebsocketAuthCliMode {
     SignedBearerToken,
 }
 
-#[derive(Debug, Clone, Default, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct AppServerWebsocketAuthSettings {
-    pub config: Option<AppServerWebsocketAuthConfig>,
+    pub config: AppServerWebsocketAuthConfig,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
+#[allow(clippy::enum_variant_names)]
 pub enum AppServerWebsocketAuthConfig {
+    QueryToken {
+        enforce: bool,
+    },
     CapabilityToken {
         source: AppServerWebsocketCapabilityTokenSource,
     },
@@ -85,13 +102,26 @@ pub enum AppServerWebsocketCapabilityTokenSource {
     TokenSha256 { token_sha256: [u8; 32] },
 }
 
-#[derive(Clone, Debug, Default)]
-pub struct WebsocketAuthPolicy {
-    pub(crate) mode: Option<WebsocketAuthMode>,
+impl Default for AppServerWebsocketAuthSettings {
+    fn default() -> Self {
+        Self {
+            config: AppServerWebsocketAuthConfig::QueryToken { enforce: true },
+        }
+    }
 }
 
 #[derive(Clone, Debug)]
+pub struct WebsocketAuthPolicy {
+    pub(crate) mode: WebsocketAuthMode,
+}
+
+#[derive(Clone, Debug)]
+#[allow(clippy::enum_variant_names)]
 pub(crate) enum WebsocketAuthMode {
+    QueryToken {
+        token: String,
+        enforce: bool,
+    },
     CapabilityToken {
         token_sha256: [u8; 32],
     },
@@ -176,7 +206,7 @@ impl AppServerWebsocketAuthArgs {
                         );
                     }
                 };
-                Some(AppServerWebsocketAuthConfig::CapabilityToken { source })
+                AppServerWebsocketAuthConfig::CapabilityToken { source }
             }
             Some(WebsocketAuthCliMode::SignedBearerToken) => {
                 if self.ws_token_file.is_some() || self.ws_token_sha256.is_some() {
@@ -187,7 +217,7 @@ impl AppServerWebsocketAuthArgs {
                 let shared_secret_file = self.ws_shared_secret_file.context(
                     "`--ws-shared-secret-file` is required when `--ws-auth signed-bearer-token` is set",
                 )?;
-                Some(AppServerWebsocketAuthConfig::SignedBearerToken {
+                AppServerWebsocketAuthConfig::SignedBearerToken {
                     shared_secret_file: absolute_path_arg(
                         "--ws-shared-secret-file",
                         shared_secret_file,
@@ -197,7 +227,7 @@ impl AppServerWebsocketAuthArgs {
                     max_clock_skew_seconds: self
                         .ws_max_clock_skew_seconds
                         .unwrap_or(DEFAULT_MAX_CLOCK_SKEW_SECONDS),
-                })
+                }
             }
             None => {
                 if self.ws_token_file.is_some()
@@ -211,7 +241,9 @@ impl AppServerWebsocketAuthArgs {
                         "websocket auth flags require `--ws-auth capability-token` or `--ws-auth signed-bearer-token`"
                     );
                 }
-                None
+                AppServerWebsocketAuthConfig::QueryToken {
+                    enforce: !self.no_token_check,
+                }
             }
         };
 
@@ -222,26 +254,30 @@ impl AppServerWebsocketAuthArgs {
 pub fn policy_from_settings(
     settings: &AppServerWebsocketAuthSettings,
 ) -> io::Result<WebsocketAuthPolicy> {
-    let mode = match settings.config.as_ref() {
-        Some(AppServerWebsocketAuthConfig::CapabilityToken { source }) => match source {
+    let mode = match &settings.config {
+        AppServerWebsocketAuthConfig::QueryToken { enforce } => WebsocketAuthMode::QueryToken {
+            token: generate_query_token()?,
+            enforce: *enforce,
+        },
+        AppServerWebsocketAuthConfig::CapabilityToken { source } => match source {
             AppServerWebsocketCapabilityTokenSource::TokenFile { token_file } => {
                 let token = read_trimmed_secret(token_file.as_ref())?;
-                Some(WebsocketAuthMode::CapabilityToken {
+                WebsocketAuthMode::CapabilityToken {
                     token_sha256: sha256_digest(token.as_bytes()),
-                })
+                }
             }
             AppServerWebsocketCapabilityTokenSource::TokenSha256 { token_sha256 } => {
-                Some(WebsocketAuthMode::CapabilityToken {
+                WebsocketAuthMode::CapabilityToken {
                     token_sha256: *token_sha256,
-                })
+                }
             }
         },
-        Some(AppServerWebsocketAuthConfig::SignedBearerToken {
+        AppServerWebsocketAuthConfig::SignedBearerToken {
             shared_secret_file,
             issuer,
             audience,
             max_clock_skew_seconds,
-        }) => {
+        } => {
             let shared_secret = read_trimmed_secret(shared_secret_file.as_ref())?.into_bytes();
             validate_signed_bearer_secret(shared_secret_file.as_ref(), &shared_secret)?;
             let max_clock_skew_seconds = i64::try_from(*max_clock_skew_seconds).map_err(|_| {
@@ -250,37 +286,57 @@ pub fn policy_from_settings(
                     "websocket auth clock skew must fit in a signed 64-bit integer",
                 )
             })?;
-            Some(WebsocketAuthMode::SignedBearerToken {
+            WebsocketAuthMode::SignedBearerToken {
                 shared_secret,
                 issuer: issuer.clone(),
                 audience: audience.clone(),
                 max_clock_skew_seconds,
-            })
+            }
         }
-        None => None,
     };
 
     Ok(WebsocketAuthPolicy { mode })
 }
 
-pub(crate) fn is_unauthenticated_non_loopback_listener(
-    bind_address: SocketAddr,
-    policy: &WebsocketAuthPolicy,
-) -> bool {
-    !bind_address.ip().is_loopback() && policy.mode.is_none()
-}
-
 pub(crate) fn authorize_upgrade(
+    uri: &Uri,
     headers: &HeaderMap,
     policy: &WebsocketAuthPolicy,
 ) -> Result<(), WebsocketAuthError> {
-    let Some(mode) = policy.mode.as_ref() else {
-        return Ok(());
-    };
-
-    let token = bearer_token_from_headers(headers)?;
-    match mode {
+    match &policy.mode {
+        WebsocketAuthMode::QueryToken { token, enforce } => {
+            let query_token = match query_token_from_uri(uri) {
+                Ok(query_token) => query_token,
+                Err(reason) if !enforce => {
+                    warn!(
+                        %reason,
+                        "allowing websocket client because query token enforcement is disabled"
+                    );
+                    return Ok(());
+                }
+                Err(reason) => return Err(unauthorized(reason)),
+            };
+            let token_matches = query_token.as_deref() == Some(token.as_str());
+            if token_matches {
+                return Ok(());
+            }
+            let reason = if query_token.is_some() {
+                "invalid generated websocket query token"
+            } else {
+                "missing generated websocket query token"
+            };
+            if *enforce {
+                Err(unauthorized(reason))
+            } else {
+                warn!(
+                    %reason,
+                    "allowing websocket client because query token enforcement is disabled"
+                );
+                Ok(())
+            }
+        }
         WebsocketAuthMode::CapabilityToken { token_sha256 } => {
+            let token = bearer_token_from_headers(headers)?;
             let actual_sha256 = sha256_digest(token.as_bytes());
             if constant_time_eq_32(token_sha256, &actual_sha256) {
                 Ok(())
@@ -293,16 +349,54 @@ pub(crate) fn authorize_upgrade(
             issuer,
             audience,
             max_clock_skew_seconds,
-        } => verify_signed_bearer_token(
-            token,
-            shared_secret,
-            issuer.as_deref(),
-            audience.as_deref(),
-            *max_clock_skew_seconds,
-        ),
+        } => {
+            let token = bearer_token_from_headers(headers)?;
+            verify_signed_bearer_token(
+                token,
+                shared_secret,
+                issuer.as_deref(),
+                audience.as_deref(),
+                *max_clock_skew_seconds,
+            )
+        }
     }
 }
 
+impl WebsocketAuthPolicy {
+    pub(crate) fn generated_query_token(&self) -> Option<&str> {
+        match &self.mode {
+            WebsocketAuthMode::QueryToken { token, .. } => Some(token),
+            WebsocketAuthMode::CapabilityToken { .. }
+            | WebsocketAuthMode::SignedBearerToken { .. } => None,
+        }
+    }
+}
+
+fn generate_query_token() -> io::Result<String> {
+    let mut bytes = [0_u8; GENERATED_QUERY_TOKEN_BYTES];
+    let mut rng = OsRng;
+    rng.try_fill_bytes(&mut bytes)
+        .map_err(|err| io::Error::other(format!("failed to generate websocket token: {err}")))?;
+    Ok(URL_SAFE_NO_PAD.encode(bytes))
+}
+
+fn query_token_from_uri(uri: &Uri) -> Result<Option<String>, &'static str> {
+    let Some(query) = uri.query() else {
+        return Ok(None);
+    };
+    let mut token = None;
+    for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
+        if key != "token" {
+            continue;
+        }
+        if token.is_some() {
+            return Err("multiple websocket query tokens");
+        }
+        token = Some(value.into_owned());
+    }
+    Ok(token)
+}
+
 fn verify_signed_bearer_token(
     token: &str,
     shared_secret: &[u8],
@@ -481,27 +575,6 @@ mod tests {
         format!("{payload}.{signature}")
     }
 
-    #[test]
-    fn detects_unauthenticated_non_loopback_listener() {
-        let policy = WebsocketAuthPolicy::default();
-        assert!(is_unauthenticated_non_loopback_listener(
-            "0.0.0.0:8765".parse().unwrap(),
-            &policy,
-        ));
-        assert!(!is_unauthenticated_non_loopback_listener(
-            "127.0.0.1:8765".parse().unwrap(),
-            &policy,
-        ));
-        assert!(!is_unauthenticated_non_loopback_listener(
-            "0.0.0.0:8765".parse().unwrap(),
-            &WebsocketAuthPolicy {
-                mode: Some(WebsocketAuthMode::CapabilityToken {
-                    token_sha256: [0u8; 32],
-                }),
-            },
-        ));
-    }
-
     #[test]
     fn capability_token_args_require_token_file_or_hash() {
         let err = AppServerWebsocketAuthArgs {
@@ -530,11 +603,11 @@ mod tests {
         assert_eq!(
             settings,
             AppServerWebsocketAuthSettings {
-                config: Some(AppServerWebsocketAuthConfig::CapabilityToken {
+                config: AppServerWebsocketAuthConfig::CapabilityToken {
                     source: AppServerWebsocketCapabilityTokenSource::TokenSha256 {
                         token_sha256: [0xab; 32],
                     },
-                }),
+                },
             }
         );
     }
@@ -573,28 +646,79 @@ mod tests {
     #[test]
     fn capability_token_hash_policy_authorizes_matching_bearer_token() {
         let settings = AppServerWebsocketAuthSettings {
-            config: Some(AppServerWebsocketAuthConfig::CapabilityToken {
+            config: AppServerWebsocketAuthConfig::CapabilityToken {
                 source: AppServerWebsocketCapabilityTokenSource::TokenSha256 {
                     token_sha256: sha256_digest(b"super-secret-token"),
                 },
-            }),
+            },
         };
         let policy = policy_from_settings(&settings).expect("hash policy should build");
+        assert_eq!(policy.generated_query_token(), None);
+        let uri = Uri::from_static("/");
         let mut headers = HeaderMap::new();
         headers.insert(
             AUTHORIZATION,
             HeaderValue::from_static("Bearer super-secret-token"),
         );
-        authorize_upgrade(&headers, &policy).expect("matching token should authorize");
+        authorize_upgrade(&uri, &headers, &policy).expect("matching token should authorize");
 
         headers.insert(
             AUTHORIZATION,
             HeaderValue::from_static("Bearer wrong-token"),
         );
-        let err = authorize_upgrade(&headers, &policy).expect_err("wrong token should fail");
+        let err = authorize_upgrade(&uri, &headers, &policy).expect_err("wrong token should fail");
         assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
     }
 
+    #[test]
+    fn generated_query_token_is_required_by_default() {
+        let policy = policy_from_settings(&AppServerWebsocketAuthSettings::default())
+            .expect("generated token policy should build");
+        let token = policy
+            .generated_query_token()
+            .expect("query-token policy should expose its token");
+
+        let valid_uri: Uri = format!("/?token={token}").parse().expect("valid token URI");
+        authorize_upgrade(&valid_uri, &HeaderMap::new(), &policy)
+            .expect("matching query token should authorize");
+
+        for uri in [
+            Uri::from_static("/"),
+            Uri::from_static("/?token=wrong"),
+            Uri::from_static("/?token=one&token=two"),
+        ] {
+            let err = authorize_upgrade(&uri, &HeaderMap::new(), &policy)
+                .expect_err("missing or invalid query token should fail");
+            assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
+        }
+    }
+
+    #[test]
+    fn no_token_check_allows_invalid_tokens() {
+        let settings = AppServerWebsocketAuthArgs {
+            no_token_check: true,
+            ..Default::default()
+        }
+        .try_into_settings()
+        .expect("no-token-check args should parse");
+        assert_eq!(
+            settings,
+            AppServerWebsocketAuthSettings {
+                config: AppServerWebsocketAuthConfig::QueryToken { enforce: false },
+            }
+        );
+        let policy = policy_from_settings(&settings).expect("generated token policy should build");
+
+        for uri in [
+            Uri::from_static("/"),
+            Uri::from_static("/?token=wrong"),
+            Uri::from_static("/?token=one&token=two"),
+        ] {
+            authorize_upgrade(&uri, &HeaderMap::new(), &policy)
+                .expect("missing or invalid token should be allowed");
+        }
+    }
+
     #[test]
     fn signed_bearer_args_require_mode_when_mode_specific_flags_are_set() {
         let err = AppServerWebsocketAuthArgs {
@@ -624,13 +748,13 @@ mod tests {
         assert_eq!(
             settings,
             AppServerWebsocketAuthSettings {
-                config: Some(AppServerWebsocketAuthConfig::SignedBearerToken {
+                config: AppServerWebsocketAuthConfig::SignedBearerToken {
                     shared_secret_file: AbsolutePathBuf::from_absolute_path("/tmp/secret")
                         .expect("absolute path"),
                     issuer: Some("issuer".to_string()),
                     audience: None,
                     max_clock_skew_seconds: DEFAULT_MAX_CLOCK_SKEW_SECONDS,
-                }),
+                },
             }
         );
     }

codex-rs/app-server-transport/src/transport/websocket.rs

@@ -3,7 +3,6 @@ use super::ConnectionOrigin;
 use super::TransportEvent;
 use super::auth::WebsocketAuthPolicy;
 use super::auth::authorize_upgrade;
-use super::auth::is_unauthenticated_non_loopback_listener;
 use super::forward_incoming_message;
 use super::next_connection_id;
 use super::serialize_outgoing_message;
@@ -19,6 +18,7 @@ use axum::extract::ws::WebSocketUpgrade;
 use axum::http::HeaderMap;
 use axum::http::Request;
 use axum::http::StatusCode;
+use axum::http::Uri;
 use axum::http::header::ORIGIN;
 use axum::middleware;
 use axum::middleware::Next;
@@ -54,10 +54,14 @@ fn colorize(text: &str, style: Style) -> String {
 }
 
 #[allow(clippy::print_stderr)]
-fn print_websocket_startup_banner(addr: SocketAddr) {
+fn print_websocket_startup_banner(addr: SocketAddr, auth_policy: &WebsocketAuthPolicy) {
     let title = colorize("codex app-server (WebSockets)", Style::new().bold().cyan());
     let listening_label = colorize("listening on:", Style::new().dimmed());
-    let listen_url = colorize(&format!("ws://{addr}"), Style::new().green());
+    let raw_listen_url = match auth_policy.generated_query_token() {
+        Some(token) => format!("ws://{addr}/?token={token}"),
+        None => format!("ws://{addr}"),
+    };
+    let listen_url = colorize(&raw_listen_url, Style::new().green());
     let ready_label = colorize("readyz:", Style::new().dimmed());
     let ready_url = colorize(&format!("http://{addr}/readyz"), Style::new().green());
     let health_label = colorize("healthz:", Style::new().dimmed());
@@ -106,9 +110,10 @@ async fn websocket_upgrade_handler(
     websocket: WebSocketUpgrade,
     ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
     State(state): State<WebSocketListenerState>,
+    uri: Uri,
     headers: HeaderMap,
 ) -> impl IntoResponse {
-    if let Err(err) = authorize_upgrade(&headers, state.auth_policy.as_ref()) {
+    if let Err(err) = authorize_upgrade(&uri, &headers, state.auth_policy.as_ref()) {
         warn!(
             %peer_addr,
             message = err.message(),
@@ -132,17 +137,9 @@ pub async fn start_websocket_acceptor(
     shutdown_token: CancellationToken,
     auth_policy: WebsocketAuthPolicy,
 ) -> IoResult<JoinHandle<()>> {
-    if is_unauthenticated_non_loopback_listener(bind_address, &auth_policy) {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidInput,
-            format!(
-                "refusing to start non-loopback websocket listener {bind_address} without auth; configure `--ws-auth capability-token` or `--ws-auth signed-bearer-token`"
-            ),
-        ));
-    }
     let listener = TcpListener::bind(bind_address).await?;
     let local_addr = listener.local_addr()?;
-    print_websocket_startup_banner(local_addr);
+    print_websocket_startup_banner(local_addr, &auth_policy);
     info!("app-server websocket listening on ws://{local_addr}");
 
     let router = Router::new()

codex-rs/app-server/README.md

@@ -34,6 +34,20 @@ When running with `--listen ws://IP:PORT`, the same listener also serves basic H
 - `GET /healthz` returns `200 OK` when no `Origin` header is present.
 - Any request carrying an `Origin` header is rejected with `403 Forbidden`.
 
+By default, app-server uses query-token authentication: it generates a 256-bit connection token
+and prints the complete connection URL to stderr, for example:
+
+```text
+ws://127.0.0.1:4500/?token=<generated-token>
+```
+
+By default, clients must provide that exact `token` query parameter during the websocket upgrade.
+Treat the printed URL as a secret. Pass `--no-token-check` to accept clients that omit the token or
+provide an incorrect token; app-server still prints the tokenized URL.
+
+When `--ws-auth capability-token` or `--ws-auth signed-bearer-token` is configured, app-server
+prints the token-free listen URL and requires `Authorization: Bearer <token>` during upgrade.
+
 Websocket transport is currently experimental and unsupported. Do not rely on it for production workloads.
 
 The unix socket transport is intended for local app-server control-plane clients. `codex app-server proxy`

codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs

@@ -160,6 +160,68 @@ async fn websocket_transport_rejects_browser_origin_without_auth() -> Result<()>
     Ok(())
 }
 
+#[tokio::test]
+async fn websocket_transport_requires_generated_query_token() -> Result<()> {
+    let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await;
+    let codex_home = TempDir::new()?;
+    create_config_toml(codex_home.path(), &server.uri(), "never")?;
+
+    let (mut process, bind_addr, token) =
+        spawn_websocket_server_with_generated_token(codex_home.path(), "ws://127.0.0.1:0", &[])
+            .await?;
+
+    assert_websocket_query_connect_rejected(bind_addr, /*query_token*/ None).await?;
+    assert_websocket_query_connect_rejected(bind_addr, Some("wrong-token")).await?;
+
+    let mut ws = connect_websocket_with_query_token(bind_addr, &token).await?;
+    send_initialize_request(&mut ws, /*id*/ 1, "ws_generated_token_client").await?;
+    let init = read_response_for_id(&mut ws, /*id*/ 1).await?;
+    assert_eq!(init.id, RequestId::Integer(1));
+
+    process
+        .kill()
+        .await
+        .context("failed to stop websocket app-server process")?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn websocket_transport_no_token_check_accepts_invalid_tokens() -> Result<()> {
+    let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await;
+    let codex_home = TempDir::new()?;
+    create_config_toml(codex_home.path(), &server.uri(), "never")?;
+
+    let (mut process, bind_addr, token) = spawn_websocket_server_with_generated_token(
+        codex_home.path(),
+        "ws://127.0.0.1:0",
+        &["--no-token-check".to_string()],
+    )
+    .await?;
+
+    let mut without_token = connect_websocket(bind_addr).await?;
+    send_initialize_request(&mut without_token, /*id*/ 1, "ws_missing_token_client").await?;
+    read_response_for_id(&mut without_token, /*id*/ 1).await?;
+
+    let mut wrong_token = connect_websocket_with_query_token(bind_addr, "wrong-token").await?;
+    send_initialize_request(&mut wrong_token, /*id*/ 2, "ws_wrong_token_client").await?;
+    read_response_for_id(&mut wrong_token, /*id*/ 2).await?;
+
+    let mut matching_token = connect_websocket_with_query_token(bind_addr, &token).await?;
+    send_initialize_request(
+        &mut matching_token,
+        /*id*/ 3,
+        "ws_matching_token_client",
+    )
+    .await?;
+    read_response_for_id(&mut matching_token, /*id*/ 3).await?;
+
+    process
+        .kill()
+        .await
+        .context("failed to stop websocket app-server process")?;
+    Ok(())
+}
+
 #[tokio::test]
 async fn websocket_transport_rejects_missing_and_invalid_capability_tokens() -> Result<()> {
     let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await;
@@ -322,24 +384,24 @@ async fn websocket_transport_rejects_short_signed_bearer_secret_configuration()
 }
 
 #[tokio::test]
-async fn websocket_transport_rejects_unauthenticated_non_loopback_startup() -> Result<()> {
+async fn websocket_transport_authenticates_non_loopback_by_default() -> Result<()> {
     let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await;
     let codex_home = TempDir::new()?;
     create_config_toml(codex_home.path(), &server.uri(), "never")?;
 
-    let output =
-        run_websocket_server_to_completion_with_args(codex_home.path(), "ws://0.0.0.0:0", &[])
+    let (mut process, bind_addr, token) =
+        spawn_websocket_server_with_generated_token(codex_home.path(), "ws://0.0.0.0:0", &[])
             .await?;
-    assert!(
-        !output.status.success(),
-        "unauthenticated non-loopback listener should fail websocket server startup"
-    );
-    let stderr = String::from_utf8(output.stderr).context("stderr should be valid utf-8")?;
-    assert!(
-        stderr.contains("refusing to start non-loopback websocket listener"),
-        "unexpected stderr: {stderr}"
-    );
 
+    assert_websocket_query_connect_rejected(bind_addr, /*query_token*/ None).await?;
+    let mut ws = connect_websocket_with_query_token(bind_addr, &token).await?;
+    send_initialize_request(&mut ws, /*id*/ 1, "ws_non_loopback_token_client").await?;
+    read_response_for_id(&mut ws, /*id*/ 1).await?;
+
+    process
+        .kill()
+        .await
+        .context("failed to stop websocket app-server process")?;
     Ok(())
 }
 
@@ -376,7 +438,12 @@ async fn websocket_disconnect_keeps_last_subscribed_thread_loaded_until_idle_tim
 }
 
 pub(super) async fn spawn_websocket_server(codex_home: &Path) -> Result<(Child, SocketAddr)> {
-    spawn_websocket_server_with_args(codex_home, "ws://127.0.0.1:0", &[]).await
+    spawn_websocket_server_with_args(
+        codex_home,
+        "ws://127.0.0.1:0",
+        &["--no-token-check".to_string()],
+    )
+    .await
 }
 
 pub(super) async fn spawn_websocket_server_with_args(
@@ -384,6 +451,27 @@ pub(super) async fn spawn_websocket_server_with_args(
     listen_url: &str,
     extra_args: &[String],
 ) -> Result<(Child, SocketAddr)> {
+    let (process, bind_addr, _token) =
+        spawn_websocket_server_and_read_generated_token(codex_home, listen_url, extra_args).await?;
+    Ok((process, bind_addr))
+}
+
+async fn spawn_websocket_server_with_generated_token(
+    codex_home: &Path,
+    listen_url: &str,
+    extra_args: &[String],
+) -> Result<(Child, SocketAddr, String)> {
+    let (process, bind_addr, token) =
+        spawn_websocket_server_and_read_generated_token(codex_home, listen_url, extra_args).await?;
+    let token = token.context("websocket app-server did not print a generated query token")?;
+    Ok((process, bind_addr, token))
+}
+
+async fn spawn_websocket_server_and_read_generated_token(
+    codex_home: &Path,
+    listen_url: &str,
+    extra_args: &[String],
+) -> Result<(Child, SocketAddr, Option<String>)> {
     let program = codex_utils_cargo_bin::cargo_bin("codex-app-server")
         .context("should find app-server binary")?;
     let mut cmd = Command::new(program);
@@ -407,7 +495,7 @@ pub(super) async fn spawn_websocket_server_with_args(
         .context("failed to capture websocket app-server stderr")?;
     let mut stderr_reader = BufReader::new(stderr).lines();
     let deadline = Instant::now() + DEFAULT_READ_TIMEOUT;
-    let bind_addr = loop {
+    let (bind_addr, generated_token) = loop {
         let line = timeout(
             deadline.saturating_duration_since(Instant::now()),
             stderr_reader.next_line(),
@@ -436,13 +524,25 @@ pub(super) async fn spawn_websocket_server_with_args(
             stripped
         };
 
-        if let Some(bind_addr) = stripped_line
+        let Some(raw_url) = stripped_line
             .split_whitespace()
             .find_map(|token| token.strip_prefix("ws://"))
-            .and_then(|addr| addr.parse::<SocketAddr>().ok())
-        {
-            break bind_addr;
-        }
+        else {
+            continue;
+        };
+        let Ok(url) = url::Url::parse(&format!("ws://{raw_url}")) else {
+            continue;
+        };
+        let Ok(socket_addrs) = url.socket_addrs(|| None) else {
+            continue;
+        };
+        let Some(bind_addr) = socket_addrs.first().copied() else {
+            continue;
+        };
+        let generated_token = url
+            .query_pairs()
+            .find_map(|(key, value)| (key == "token").then(|| value.into_owned()));
+        break (bind_addr, generated_token);
     };
 
     tokio::spawn(async move {
@@ -451,7 +551,7 @@ pub(super) async fn spawn_websocket_server_with_args(
         }
     });
 
-    Ok((process, bind_addr))
+    Ok((process, bind_addr, generated_token))
 }
 
 pub(super) async fn connect_websocket(bind_addr: SocketAddr) -> Result<WsClient> {
@@ -478,6 +578,63 @@ pub(super) async fn connect_websocket_with_bearer(
     }
 }
 
+async fn connect_websocket_with_query_token(
+    bind_addr: SocketAddr,
+    query_token: &str,
+) -> Result<WsClient> {
+    let url = format!(
+        "ws://{}/?token={query_token}",
+        connectable_bind_addr(bind_addr)
+    );
+    let request = websocket_request(
+        url.as_str(),
+        /*bearer_token*/ None,
+        /*origin*/ None,
+    )?;
+    let deadline = Instant::now() + DEFAULT_READ_TIMEOUT;
+    loop {
+        match connect_async(request.clone()).await {
+            Ok((stream, _response)) => return Ok(stream),
+            Err(err) => {
+                if Instant::now() >= deadline {
+                    bail!("failed to connect websocket to {url}: {err}");
+                }
+                sleep(Duration::from_millis(50)).await;
+            }
+        }
+    }
+}
+
+async fn assert_websocket_query_connect_rejected(
+    bind_addr: SocketAddr,
+    query_token: Option<&str>,
+) -> Result<()> {
+    let path_and_query = query_token
+        .map(|token| format!("/?token={token}"))
+        .unwrap_or_default();
+    let url = format!("ws://{}{path_and_query}", connectable_bind_addr(bind_addr));
+    let request = websocket_request(
+        url.as_str(),
+        /*bearer_token*/ None,
+        /*origin*/ None,
+    )?;
+
+    match connect_async(request).await {
+        Ok((_stream, response)) => {
+            bail!(
+                "expected websocket handshake rejection, got {}",
+                response.status()
+            )
+        }
+        // The pinned tungstenite fork can replace an otherwise valid rejected
+        // upgrade response with a synthetic 400 when it normalizes a response
+        // that lacks websocket response headers. Unit coverage verifies that
+        // app-server's authorization result itself is 401.
+        Err(WsError::Http(_response)) => Ok(()),
+        Err(err) => bail!("expected http rejection during websocket handshake: {err}"),
+    }
+}
+
 async fn assert_websocket_connect_rejected(
     bind_addr: SocketAddr,
     bearer_token: Option<&str>,