From b3f5e45cc1de8bcb09d320f3211378db285aa201 Mon Sep 17 00:00:00 2001 From: ostepanian Date: Fri, 4 Sep 2026 14:40:40 +0000 Subject: [PATCH] Add direct SigV4 transport to exec-server (#42781) ## Why Allow remote exec servers to connect directly to AWS-hosted registries that authenticate registry requests and WebSocket handshakes with AWS SigV4. ## What changed - Add `--remote-transport direct` with SigV4 profile, region, and service options while keeping Noise as the default transport. - Register the `direct_jsonrpc_v1` transport and carry plain exec-server JSON-RPC messages over the authenticated WebSocket. - Reuse direct registrations across transient disconnects, refresh them after a `409 Conflict`, and require TLS for non-loopback endpoints. ## Testing - Cover CLI validation and SigV4 request signing. - Exercise direct registration, handshake retry behavior, JSON-RPC interoperability, and process recovery after reconnecting. GitOrigin-RevId: 0755df330ba3abe5db0a516fdaa49338d9bbe2d2 --- codex-rs/Cargo.lock | 1 + codex-rs/cli/Cargo.toml | 1 + codex-rs/cli/src/exec_server_args_tests.rs | 265 +++++++++++ codex-rs/cli/src/exec_server_auth.rs | 77 +++ codex-rs/cli/src/exec_server_auth_tests.rs | 96 ++++ codex-rs/cli/src/main.rs | 210 +++++--- ..._server_help_documents_remote_options.snap | 85 ++++ codex-rs/exec-server/README.md | 35 ++ codex-rs/exec-server/src/client_transport.rs | 140 +++++- codex-rs/exec-server/src/lib.rs | 1 + codex-rs/exec-server/src/remote.rs | 107 ++++- codex-rs/exec-server/src/remote/direct.rs | 277 +++++++++++ .../exec-server/src/remote/direct_tests.rs | 449 ++++++++++++++++++ .../exec-server/tests/accepted_websocket.rs | 313 ++++++++++++ .../exec-server/tests/common/exec_server.rs | 52 +- 15 files changed, 1988 insertions(+), 121 deletions(-) create mode 100644 codex-rs/cli/src/exec_server_args_tests.rs create mode 100644 codex-rs/cli/src/exec_server_auth.rs create mode 100644 codex-rs/cli/src/exec_server_auth_tests.rs create mode 100644 codex-rs/cli/src/snapshots/codex__exec_server_args_tests__exec_server_help_documents_remote_options.snap create mode 100644 codex-rs/exec-server/src/remote/direct.rs create mode 100644 codex-rs/exec-server/src/remote/direct_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 2e0a94c38b..2d2d04dd7f 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2511,6 +2511,7 @@ dependencies = [ "codex-app-server-protocol", "codex-app-server-test-client", "codex-arg0", + "codex-aws-auth", "codex-build-info", "codex-chatgpt", "codex-cloud-config", diff --git a/codex-rs/cli/Cargo.toml b/codex-rs/cli/Cargo.toml index 7e174f8375..197547b6e7 100644 --- a/codex-rs/cli/Cargo.toml +++ b/codex-rs/cli/Cargo.toml @@ -34,6 +34,7 @@ codex-app-server-test-client = { workspace = true } codex-arg0 = { workspace = true } codex-build-info = { workspace = true } codex-api = { workspace = true } +codex-aws-auth = { workspace = true } codex-chatgpt = { workspace = true } codex-cloud-config = { workspace = true } codex-cloud-tasks = { path = "../cloud-tasks" } diff --git a/codex-rs/cli/src/exec_server_args_tests.rs b/codex-rs/cli/src/exec_server_args_tests.rs new file mode 100644 index 0000000000..c58b3b9945 --- /dev/null +++ b/codex-rs/cli/src/exec_server_args_tests.rs @@ -0,0 +1,265 @@ +use clap::error::ErrorKind; +use pretty_assertions::assert_eq; + +use super::*; + +fn exec_server_from_args(args: &[&str]) -> ExecServerCommand { + let cli = MultitoolCli::try_parse_from( + ["codex", "exec-server"] + .into_iter() + .chain(args.iter().copied()), + ) + .expect("parse executor arguments"); + let Some(Subcommand::ExecServer(command)) = cli.subcommand else { + panic!("expected executor command"); + }; + command +} + +#[test] +fn exec_server_help_documents_remote_options() { + let command = MultitoolCli::command() + .term_width(80) + .mut_subcommand("exec-server", |command| { + command.mut_arg("exit_on_stdin_close", |arg| arg.hide_env_values(true)) + }); + let help = command + .try_get_matches_from(["codex", "exec-server", "--help"]) + .expect_err("help should exit before running the executor"); + assert_eq!(help.kind(), ErrorKind::DisplayHelp); + let help_text = help + .to_string() + .lines() + .map(str::trim_end) + .collect::>() + .join("\n"); + insta::assert_snapshot!(help_text); +} + +#[test] +fn exec_server_defaults_preserve_local_and_noise_modes() { + for args in [ + vec![], + vec!["--listen", "stdio"], + vec![ + "--remote", + "https://registry.example.com", + "--environment-id", + "env-1", + ], + vec![ + "--remote", + "https://registry.example.com", + "--environment-id", + "env-1", + "--remote-transport", + "noise", + "--use-agent-identity-auth", + ], + ] { + let command = exec_server_from_args(&args); + assert_eq!(command.remote_transport, ExecServerRemoteTransport::Noise); + assert!(!command.aws_sigv4); + command + .validate_remote_transport() + .expect("valid existing mode"); + } +} + +#[test] +fn exec_server_parses_direct_aws_options() { + for (options, expected) in [ + (vec![], (None, None, "execute-api")), + ( + vec![ + "--aws-profile", + "development", + "--aws-region", + "us-west-2", + "--aws-service", + "bedrock-mantle", + ], + (Some("development"), Some("us-west-2"), "bedrock-mantle"), + ), + ] { + let mut args = vec![ + "--remote", + "https://registry.example.com", + "--environment-id", + "env-1", + "--remote-transport", + "direct", + "--aws-sigv4", + ]; + args.extend(options); + let command = exec_server_from_args(&args); + + assert_eq!( + command.remote.as_deref(), + Some("https://registry.example.com") + ); + assert_eq!(command.environment_id.as_deref(), Some("env-1")); + assert_eq!(command.remote_transport, ExecServerRemoteTransport::Direct); + assert!(command.aws_sigv4); + assert_eq!( + ( + command.aws_profile.as_deref(), + command.aws_region.as_deref(), + command.aws_service.as_str() + ), + expected, + ); + command + .validate_remote_transport() + .expect("valid Direct mode"); + } +} + +#[test] +fn exec_server_rejects_missing_or_conflicting_remote_options() { + for (options, expected) in [ + ( + vec!["--remote-transport", "direct"], + ErrorKind::MissingRequiredArgument, + ), + ( + vec!["--aws-profile", "development"], + ErrorKind::MissingRequiredArgument, + ), + ( + vec!["--aws-region", "us-west-2"], + ErrorKind::MissingRequiredArgument, + ), + ( + vec!["--aws-service", "execute-api"], + ErrorKind::MissingRequiredArgument, + ), + ( + vec![ + "--remote-transport", + "direct", + "--aws-sigv4", + "--use-agent-identity-auth", + ], + ErrorKind::ArgumentConflict, + ), + ( + vec!["--remote-transport", "unsupported"], + ErrorKind::InvalidValue, + ), + ] { + let error = MultitoolCli::try_parse_from( + [ + "codex", + "exec-server", + "--remote", + "https://registry.example.com", + "--environment-id", + "env-1", + ] + .into_iter() + .chain(options), + ) + .expect_err("reject invalid remote arguments"); + assert_eq!(error.kind(), expected); + } +} + +#[test] +fn exec_server_transport_and_aws_options_require_registration_arguments() { + for options in [ + vec!["--remote-transport", "noise"], + vec!["--remote-transport", "direct", "--aws-sigv4"], + vec!["--aws-sigv4"], + vec![ + "--remote", + "https://registry.example.com", + "--remote-transport", + "direct", + "--aws-sigv4", + ], + ] { + let error = + MultitoolCli::try_parse_from(["codex", "exec-server"].into_iter().chain(options)) + .expect_err("require remote URL and environment ID"); + assert_eq!(error.kind(), ErrorKind::MissingRequiredArgument); + } +} + +#[tokio::test] +async fn exec_server_sigv4_does_not_enable_aws_auth_for_noise() { + for options in [vec![], vec!["--remote-transport", "noise"]] { + let mut args = vec![ + "--remote", + "https://registry.example.com", + "--environment-id", + "env-1", + "--aws-sigv4", + ]; + args.extend(options); + let command = exec_server_from_args(&args); + // Unset runtime paths prove validation runs before startup or credential loading. + let error = run_exec_server_command( + command, + &Arg0DispatchPaths::default(), + &CliConfigOverrides::default(), + /*strict_config*/ false, + ) + .await + .expect_err("Noise auth is unchanged"); + assert_eq!( + error.to_string(), + "--aws-sigv4 requires --remote-transport direct" + ); + } +} + +#[tokio::test] +async fn exec_server_direct_forwarding_remains_rejected() { + let command = exec_server_from_args(&[ + "forward", + "--connect", + "ws://127.0.0.1:8765", + "--remote", + "https://registry.example.com", + "--environment-id", + "env-1", + "--remote-transport", + "direct", + "--aws-sigv4", + "--aws-profile", + "development", + "--aws-region", + "us-west-2", + "--aws-service", + "bedrock-mantle", + ]); + assert_eq!( + ( + command.remote_transport, + command.aws_sigv4, + command.aws_profile.as_deref(), + command.aws_region.as_deref(), + command.aws_service.as_str() + ), + ( + ExecServerRemoteTransport::Direct, + true, + Some("development"), + Some("us-west-2"), + "bedrock-mantle" + ), + ); + let error = run_exec_server_command( + command, + &Arg0DispatchPaths::default(), + &CliConfigOverrides::default(), + /*strict_config*/ false, + ) + .await + .expect_err("Direct forwarding is unsupported before startup"); + assert_eq!( + error.to_string(), + "direct exec-server transport does not support forwarding" + ); +} diff --git a/codex-rs/cli/src/exec_server_auth.rs b/codex-rs/cli/src/exec_server_auth.rs new file mode 100644 index 0000000000..ca4a84f37d --- /dev/null +++ b/codex-rs/cli/src/exec_server_auth.rs @@ -0,0 +1,77 @@ +//! CLI-only glue between Codex authentication and generic AWS signing. +//! +//! Keeping this adapter here leaves `codex-api` and `codex-aws-auth` independent. + +use std::sync::Arc; + +use codex_api::AuthError; +use codex_api::AuthProvider; +use codex_api::SharedAuthProvider; +use codex_aws_auth::AwsAuthConfig; +use codex_aws_auth::AwsAuthContext; +use codex_aws_auth::AwsAuthError; +use codex_aws_auth::AwsRequestToSign; +use codex_http_client::Request; +use codex_http_client::RequestBody; +use codex_http_client::RequestCompression; +use http::HeaderMap; + +/// Creates a SigV4 provider, preferring an explicit profile over the default credential chain. +pub(super) async fn aws_sigv4_auth_provider( + mut config: AwsAuthConfig, +) -> Result { + config.profile = config + .profile + .map(|profile| profile.trim().to_string()) + .filter(|profile| !profile.is_empty()); + config.region = config + .region + .map(|region| region.trim().to_string()) + .filter(|region| !region.is_empty()); + let context = if config.profile.is_some() { + AwsAuthContext::load_profile(config).await + } else { + AwsAuthContext::load(config).await + }?; + Ok(Arc::new(AwsSigV4AuthProvider { context })) +} + +#[derive(Debug)] +struct AwsSigV4AuthProvider { + context: AwsAuthContext, +} + +impl AuthProvider for AwsSigV4AuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) {} + + fn apply_auth(&self, mut request: Request) -> codex_api::AuthProviderFuture<'_> { + Box::pin(async move { + let prepared = request.prepare_body_for_send().map_err(AuthError::Build)?; + let signed = self + .context + .sign(AwsRequestToSign { + method: request.method.clone(), + url: request.url.clone(), + headers: prepared.headers.clone(), + body: prepared.body_bytes(), + }) + .await + .map_err(|error| { + if error.is_retryable() { + AuthError::Transient(error.to_string()) + } else { + AuthError::Build(error.to_string()) + } + })?; + request.url = signed.url; + request.headers = signed.headers; + request.body = prepared.body.map(RequestBody::Raw); + request.compression = RequestCompression::None; + Ok(request) + }) + } +} + +#[cfg(test)] +#[path = "exec_server_auth_tests.rs"] +mod tests; diff --git a/codex-rs/cli/src/exec_server_auth_tests.rs b/codex-rs/cli/src/exec_server_auth_tests.rs new file mode 100644 index 0000000000..dd964f8169 --- /dev/null +++ b/codex-rs/cli/src/exec_server_auth_tests.rs @@ -0,0 +1,96 @@ +use std::time::Duration; + +use codex_aws_auth::AwsAccessKeys; +use http::HeaderValue; +use http::Method; +use pretty_assertions::assert_eq; + +use super::*; + +async fn test_provider() -> AwsSigV4AuthProvider { + let context = AwsAuthContext::load_with_access_keys( + AwsAuthConfig { + profile: None, + region: Some("us-east-1".to_string()), + service: "execute-api".to_string(), + }, + AwsAccessKeys { + access_key_id: "test-access-key".to_string(), + secret_access_key: "test-secret-key".to_string(), + session_token: Some("test-session-token".to_string()), + }, + ) + .await + .expect("load fixture signing context"); + AwsSigV4AuthProvider { context } +} + +#[tokio::test] +async fn signs_requests_without_changing_payload_or_metadata() { + let provider = test_provider().await; + let url = "https://executor.example.com/connect?environment_id=environment-1"; + for mut request in [ + Request::new(Method::GET, url.to_string()), + Request::new(Method::POST, url.to_string()) + .with_json(&serde_json::json!({"transport": "direct_jsonrpc_v1"})), + Request::new(Method::POST, url.to_string()) + .with_json(&serde_json::json!({"transport": "direct_jsonrpc_v1"})) + .with_compression(RequestCompression::Zstd), + ] { + request + .headers + .insert("session_id", HeaderValue::from_static("session-1")); + request + .headers + .insert("x-custom-header", HeaderValue::from_static("preserved")); + request.timeout = Some(Duration::from_secs(3)); + let method = request.method.clone(); + let expected = request + .prepare_body_for_send() + .expect("prepare fixture body"); + let signed = provider + .apply_auth(request) + .await + .expect("sign fixture request"); + + assert_eq!(signed.method, method); + assert_eq!(signed.url, url); + assert_eq!(signed.timeout, Some(Duration::from_secs(3))); + assert_eq!(signed.body, expected.body.clone().map(RequestBody::Raw)); + assert_eq!(signed.compression, RequestCompression::None); + for (name, value) in &expected.headers { + assert_eq!(signed.headers.get(name), Some(value)); + } + assert_eq!(signed.headers["x-amz-security-token"], "test-session-token"); + assert!(signed.headers.contains_key("x-amz-date")); + let authorization = signed.headers[http::header::AUTHORIZATION] + .to_str() + .unwrap(); + assert!(authorization.starts_with("AWS4-HMAC-SHA256 ")); + assert!(authorization.contains("/us-east-1/execute-api/aws4_request")); + assert_eq!(signed.prepare_body_for_send().unwrap().body, expected.body); + } +} + +#[tokio::test] +async fn invalid_signing_request_is_a_permanent_auth_error() { + let provider = test_provider().await; + let error = provider + .apply_auth(Request::new(Method::GET, "not a URL".to_string())) + .await + .expect_err("invalid URL should fail signing"); + assert!(matches!(error, AuthError::Build(_))); +} + +#[tokio::test] +async fn invalid_signing_configuration_is_rejected() { + let error = aws_sigv4_auth_provider(AwsAuthConfig { + profile: None, + region: Some("us-east-1".to_string()), + service: " ".to_string(), + }) + .await + .err() + .expect("empty service should fail configuration"); + assert!(matches!(error, AwsAuthError::EmptyService)); +} diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index c70b52d2dd..62f9d2d960 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -24,8 +24,6 @@ use codex_exec::Cli as ExecCli; use codex_exec::Command as ExecCommand; use codex_exec::ReviewArgs; use codex_exec_server::ExecServerRuntimePaths; -use codex_exec_server::ExecServerTelemetry; -use codex_exec_server::RemoteEnvironmentConfig; use codex_execpolicy::ExecPolicyCheckCommand; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; @@ -54,6 +52,10 @@ mod cloud_config; #[cfg(any(target_os = "macos", target_os = "windows"))] mod desktop_app; mod doctor; +#[cfg(test)] +#[path = "exec_server_args_tests.rs"] +mod exec_server_args_tests; +mod exec_server_auth; mod exec_server_telemetry; mod marketplace_cmd; mod mcp_cmd; @@ -639,6 +641,17 @@ struct ExecServerCommand { )] remote: Option, + /// Transport used for the remote executor connection. + #[arg( + long = "remote-transport", + value_enum, + default_value_t = ExecServerRemoteTransport::Noise, + requires = "exec_server_remote", + requires_if("direct", "aws_sigv4"), + global = true + )] + remote_transport: ExecServerRemoteTransport, + /// Environment id to attach to when registering remotely. #[arg(long = "environment-id", value_name = "ID", global = true)] environment_id: Option, @@ -651,10 +664,43 @@ struct ExecServerCommand { #[arg( long = "use-agent-identity-auth", requires = "exec_server_remote", + conflicts_with = "aws_sigv4", global = true )] use_agent_identity_auth: bool, + /// Sign Direct registration and WebSocket handshake requests with AWS SigV4. + #[arg(long = "aws-sigv4", requires = "exec_server_remote", global = true)] + aws_sigv4: bool, + + /// AWS profile used for SigV4 authentication. + #[arg( + long = "aws-profile", + value_name = "PROFILE", + requires = "aws_sigv4", + global = true + )] + aws_profile: Option, + + /// AWS signing region. Uses the SDK region chain when omitted. + #[arg( + long = "aws-region", + value_name = "REGION", + requires = "aws_sigv4", + global = true + )] + aws_region: Option, + + /// AWS signing service. + #[arg( + long = "aws-service", + value_name = "SERVICE", + default_value = "execute-api", + requires = "aws_sigv4", + global = true + )] + aws_service: String, + /// Exit when the parent-owned standard-input pipe closes. #[arg( long = "exit-on-stdin-close", @@ -665,6 +711,37 @@ struct ExecServerCommand { exit_on_stdin_close: bool, } +impl ExecServerCommand { + fn validate_remote_transport(&self) -> anyhow::Result<()> { + match (self.remote_transport, self.aws_sigv4) { + (ExecServerRemoteTransport::Noise, true) => { + anyhow::bail!("--aws-sigv4 requires --remote-transport direct"); + } + (ExecServerRemoteTransport::Direct, false) => { + anyhow::bail!("--remote-transport direct requires --aws-sigv4"); + } + (ExecServerRemoteTransport::Noise, false) + | (ExecServerRemoteTransport::Direct, true) => {} + } + if self.remote_transport == ExecServerRemoteTransport::Direct + && matches!( + self.command.as_ref(), + Some(ExecServerSubcommand::Forward { .. }) + ) + { + anyhow::bail!("direct exec-server transport does not support forwarding"); + } + Ok(()) + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, clap::ValueEnum)] +enum ExecServerRemoteTransport { + #[default] + Noise, + Direct, +} + #[derive(Debug, clap::Subcommand)] enum ExecServerSubcommand { /// Register an existing WebSocket exec-server as a remote environment. @@ -1863,6 +1940,7 @@ async fn run_exec_server_command( root_config_overrides: &CliConfigOverrides, strict_config: bool, ) -> anyhow::Result<()> { + cmd.validate_remote_transport()?; let codex_self_exe = arg0_paths .codex_self_exe .clone() @@ -1880,7 +1958,41 @@ async fn run_exec_server_command( /*enable_workload_identity*/ true, ) .await?; + let direct_transport = cmd.remote_transport == ExecServerRemoteTransport::Direct; let (_otel, telemetry) = exec_server_telemetry::init(Some(&config)); + let auth_provider = if cmd.aws_sigv4 { + exec_server_auth::aws_sigv4_auth_provider(codex_aws_auth::AwsAuthConfig { + profile: cmd.aws_profile, + region: cmd.aws_region, + service: cmd.aws_service, + }) + .await? + } else { + load_exec_server_remote_auth_provider(&config, &base_url, cmd.use_agent_identity_auth) + .await? + }; + let mut remote_config = codex_exec_server::RemoteEnvironmentConfig::new_with_transport( + base_url, + environment_id, + if direct_transport { + codex_exec_server::RemoteEnvironmentTransport::Direct + } else { + codex_exec_server::RemoteEnvironmentTransport::Noise + }, + auth_provider, + config.http_client_factory(), + )?; + if let Some(name) = cmd.name { + remote_config.name = name; + } + remote_config.request_dispatch_mode = cmd.request_dispatch_mode; + let remote_config = remote_config.with_telemetry(telemetry); + let parent_lifetime = if cmd.exit_on_stdin_close { + exec_server_telemetry::ParentLifetime::StdinPipe + } else { + exec_server_telemetry::ParentLifetime::Independent + }; + let (shutdown_sender, shutdown_receiver) = tokio::sync::oneshot::channel(); #[cfg(target_os = "macos")] let runtime_paths = runtime_paths.with_allowed_symlinked_codex_home( codex_config::allowed_symlinked_codex_home( @@ -1888,13 +2000,33 @@ async fn run_exec_server_command( &config.codex_home, ), ); - run_remote_exec_server( - cmd, - base_url, - environment_id, - &config, - runtime_paths, - telemetry, + exec_server_telemetry::run_until_shutdown( + async move { + let shutdown = async move { + let _ = shutdown_receiver.await; + }; + match cmd.command { + Some(ExecServerSubcommand::Forward { connect }) => { + codex_exec_server::run_remote_environment_forward_until_shutdown( + remote_config, + connect, + shutdown, + ) + .await + } + None => { + codex_exec_server::run_remote_environment_until_shutdown( + remote_config, + runtime_paths, + shutdown, + ) + .await + } + } + .map_err(anyhow::Error::new) + }, + parent_lifetime, + exec_server_telemetry::ShutdownBehavior::Graceful(shutdown_sender), ) .await } else { @@ -1940,66 +2072,6 @@ async fn run_exec_server_command( } } -async fn run_remote_exec_server( - cmd: ExecServerCommand, - base_url: String, - environment_id: String, - config: &Config, - runtime_paths: ExecServerRuntimePaths, - telemetry: ExecServerTelemetry, -) -> anyhow::Result<()> { - let auth_provider = - load_exec_server_remote_auth_provider(config, &base_url, cmd.use_agent_identity_auth) - .await?; - let mut remote_config = RemoteEnvironmentConfig::new( - base_url, - environment_id, - auth_provider, - config.http_client_factory(), - )?; - if let Some(name) = cmd.name { - remote_config.name = name; - } - remote_config.request_dispatch_mode = cmd.request_dispatch_mode; - let remote_config = remote_config.with_telemetry(telemetry); - let parent_lifetime = if cmd.exit_on_stdin_close { - exec_server_telemetry::ParentLifetime::StdinPipe - } else { - exec_server_telemetry::ParentLifetime::Independent - }; - let (shutdown_sender, shutdown_receiver) = tokio::sync::oneshot::channel(); - let shutdown = async move { - let _ = shutdown_receiver.await; - }; - let run = async move { - match cmd.command { - Some(ExecServerSubcommand::Forward { connect }) => { - codex_exec_server::run_remote_environment_forward_until_shutdown( - remote_config, - connect, - shutdown, - ) - .await - } - None => { - codex_exec_server::run_remote_environment_until_shutdown( - remote_config, - runtime_paths, - shutdown, - ) - .await - } - } - }; - exec_server_telemetry::run_until_shutdown( - run, - parent_lifetime, - exec_server_telemetry::ShutdownBehavior::Graceful(shutdown_sender), - ) - .await?; - Ok(()) -} - async fn load_exec_server_remote_auth_provider( config: &codex_core::config::Config, base_url: &str, diff --git a/codex-rs/cli/src/snapshots/codex__exec_server_args_tests__exec_server_help_documents_remote_options.snap b/codex-rs/cli/src/snapshots/codex__exec_server_args_tests__exec_server_help_documents_remote_options.snap new file mode 100644 index 0000000000..e2bd9f1278 --- /dev/null +++ b/codex-rs/cli/src/snapshots/codex__exec_server_args_tests__exec_server_help_documents_remote_options.snap @@ -0,0 +1,85 @@ +--- +source: cli/src/exec_server_args_tests.rs +expression: help_text +--- +[EXPERIMENTAL] Run the standalone exec-server service + +Usage: codex exec-server [OPTIONS] [COMMAND] + +Commands: + forward Register an existing WebSocket exec-server as a remote environment + help Print this message or the help of the given subcommand(s) + +Options: + -c, --config + Override a configuration value that would otherwise be loaded from + `~/.codex/config.toml`. Use a dotted path (`foo.bar.baz`) to override + nested values. The `value` portion is parsed as TOML. If it fails to + parse as TOML, the raw string is used as a literal. + + Examples: - `-c model="o3"` - `-c + 'sandbox_permissions=["disk-full-read-access"]'` - `-c + shell_environment_policy.inherit=all` + + --enable + Enable a feature (repeatable). Equivalent to `-c features.=true` + + --strict-config + Error out when config.toml contains fields that are not recognized by + this version of Codex + + --concurrent-requests + Maximum number of requests to process concurrently on each connection + + [default: 1] + + --disable + Disable a feature (repeatable). Equivalent to `-c + features.=false` + + --listen + Transport endpoint URL. Supported values: `ws://IP:PORT` (default), + `stdio`, `stdio://` + + --remote + Register this exec-server as a remote environment using the given base + URL + + --remote-transport + Transport used for the remote executor connection + + [default: noise] + [possible values: noise, direct] + + --environment-id + Environment id to attach to when registering remotely + + --name + Human-readable environment name + + --use-agent-identity-auth + Use Agent Identity auth from CODEX_ACCESS_TOKEN for remote + registration + + --aws-sigv4 + Sign Direct registration and WebSocket handshake requests with AWS + SigV4 + + --aws-profile + AWS profile used for SigV4 authentication + + --aws-region + AWS signing region. Uses the SDK region chain when omitted + + --aws-service + AWS signing service + + [default: execute-api] + + --exit-on-stdin-close + Exit when the parent-owned standard-input pipe closes + + [env: CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE] + + -h, --help + Print help (see a summary with '-h') diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 5a4c3e3f5b..5fbf943fba 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -54,9 +54,44 @@ codex exec-server \ --environment-id "$ENVIRONMENT_ID" ``` +AWS-hosted registries can use SigV4 for registry requests and the executor +WebSocket handshake. Select the transport and authentication with executor +arguments rather than `config.toml` settings: + +```sh +codex exec-server \ + --remote https://example.com \ + --environment-id "$ENVIRONMENT_ID" \ + --remote-transport direct \ + --aws-sigv4 \ + --aws-profile development \ + --aws-region us-west-2 \ + --aws-service bedrock-mantle +``` + +Noise remains the default transport. Direct requires `--aws-sigv4`, which +conflicts with `--use-agent-identity-auth` and is not supported for Noise. +The AWS options require `--aws-sigv4`; Direct forwarding remains unsupported. +The AWS SDK default credential and region chains are used when `--aws-profile` +or `--aws-region` is omitted. The signing service defaults to `execute-api`. +Direct mode registers `direct_jsonrpc_v1` through the AWS-owned +`/cloud/environment/{environment_id}/direct/register` endpoint and carries plain +exec-server JSON-RPC over the authenticated WebSocket. The existing Codex Noise +registration endpoint remains unchanged. Production deployments must use TLS +(`https`/`wss`). + +Direct registration URLs must remain reusable across disconnects and temporary +connection failures. The executor only refreshes its registration when the +WebSocket handshake returns `409 Conflict`. Handshake `408`, `429`, and `5xx` +responses retry with backoff using the current registration; other `4xx` responses +stop the executor. A backend that issues single-use connection URLs must adapt to +this contract. If the initial registration or a registration refresh fails, the +executor returns the error without retrying registration, matching Noise. + Wire framing: - local websocket: one JSON-RPC message per websocket message +- direct remote websocket: one JSON-RPC message per websocket message - Noise remote websocket: binary protobuf relay frames carrying encrypted payloads ## Remote Relay Message Format diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index b1c1ef09ee..bbdf769038 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -14,9 +14,14 @@ use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::debug; use tracing::warn; +use codex_api::AuthError; +use codex_api::AuthProvider; use codex_http_client::HttpClientFactory; +use codex_http_client::Request; +use codex_http_client::RequestCompression; use codex_protocol::shell_environment::scrub_non_inheritable_env_vars; use codex_utils_rustls_provider::ensure_rustls_crypto_provider; +use codex_websocket_client::WebSocketConnection; use codex_websocket_client::WebSocketConnector; use codex_websocket_client::WebSocketTlsMode; use http::HeaderMap; @@ -50,6 +55,113 @@ const INITIAL_REGISTRY_MAX_RETRIES: u32 = 4; const INITIAL_REGISTRY_REQUEST_TIMEOUT: Duration = Duration::from_secs(6); const INITIAL_REGISTRY_OPERATION_TIMEOUT: Duration = Duration::from_secs(14); +pub(crate) async fn connect_websocket_request( + request: http::Request<()>, + diagnostic_url: String, + connector: WebSocketConnector, + connect_timeout: Duration, + use_loopback_direct: bool, +) -> Result { + let websocket_config = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default(); + timeout(connect_timeout, async { + if use_loopback_direct { + connector + .connect_loopback_direct(request, websocket_config) + .await + } else { + connector.connect(request, websocket_config).await + } + }) + .await + .map_err(|_| ExecServerError::WebSocketConnectTimeout { + url: diagnostic_url.clone(), + timeout: connect_timeout, + })? + .map(|(websocket, _)| websocket) + .map_err(|source| ExecServerError::WebSocketConnect { + url: diagnostic_url, + source, + }) +} + +pub(crate) async fn authenticate_websocket_request( + request: &mut http::Request<()>, + auth_provider: &dyn AuthProvider, +) -> Result<(), AuthError> { + let url = request.uri().to_string(); + let signing_url = if let Some(rest) = url.strip_prefix("wss://") { + format!("https://{rest}") + } else if let Some(rest) = url.strip_prefix("ws://") { + format!("http://{rest}") + } else { + url + }; + let mut auth_request = Request::new(request.method().clone(), signing_url); + // Intermediaries may rewrite WebSocket and hop-by-hop headers after signing. + if let Some(host) = request.headers().get(http::header::HOST) { + auth_request + .headers + .insert(http::header::HOST, host.clone()); + } + let authenticated = auth_provider.apply_auth(auth_request).await?; + if authenticated.method != *request.method() { + return Err(AuthError::Build( + "authentication changed the WebSocket request method".to_string(), + )); + } + if authenticated.body.is_some() || authenticated.compression != RequestCompression::None { + return Err(AuthError::Build( + "authentication added a body or compression to the WebSocket request".to_string(), + )); + } + + let authenticated_websocket_url = websocket_url_from_authenticated_url(&authenticated.url)?; + let authenticated_uri = authenticated_websocket_url.parse().map_err(|error| { + AuthError::Build(format!("invalid authenticated WebSocket URL: {error}")) + })?; + let original_host = request.headers().get(http::header::HOST).cloned(); + for (name, value) in &authenticated.headers { + if is_websocket_handshake_header(name) { + if name == http::header::HOST && original_host.as_ref() == Some(value) { + continue; + } + return Err(AuthError::Build(format!( + "authentication changed WebSocket handshake header {name}" + ))); + } + request.headers_mut().insert(name, value.clone()); + } + *request.uri_mut() = authenticated_uri; + Ok(()) +} + +fn websocket_url_from_authenticated_url(url: &str) -> Result { + let mut url = url::Url::parse(url) + .map_err(|error| AuthError::Build(format!("invalid authenticated request URL: {error}")))?; + let websocket_scheme = match url.scheme() { + "https" => "wss", + "http" => "ws", + scheme => { + return Err(AuthError::Build(format!( + "authentication returned unsupported WebSocket URL scheme: {scheme}" + ))); + } + }; + url.set_scheme(websocket_scheme).map_err(|_| { + AuthError::Build("failed to convert authenticated URL to WebSocket scheme".to_string()) + })?; + Ok(url.into()) +} + +fn is_websocket_handshake_header(name: &http::header::HeaderName) -> bool { + name == http::header::HOST + || name == http::header::CONNECTION + || name == http::header::UPGRADE + || name == http::header::CONTENT_LENGTH + || name == http::header::TRANSFER_ENCODING + || name.as_str().starts_with("sec-websocket-") +} + /// Everything the recovery loop needs for one connection attempt. /// /// An attempt may also carry a permit whose lifetime must extend until the @@ -430,26 +542,14 @@ impl ExecServerClient { WebSocketTlsMode::TungsteniteDefault, ) .map_err(|error| ExecServerError::WebSocketConfiguration(error.to_string()))?; - let websocket_config = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default(); - let connect = async { - if !http_headers.is_empty() && request.uri().scheme_str() == Some("ws") { - connector - .connect_loopback_direct(request, websocket_config) - .await - } else { - connector.connect(request, websocket_config).await - } - }; - let (stream, _) = timeout(connect_timeout, connect) - .await - .map_err(|_| ExecServerError::WebSocketConnectTimeout { - url: websocket_url.clone(), - timeout: connect_timeout, - })? - .map_err(|source| ExecServerError::WebSocketConnect { - url: websocket_url.clone(), - source, - })?; + let stream = connect_websocket_request( + request, + websocket_url.clone(), + connector, + connect_timeout, + !http_headers.is_empty() && websocket_url.starts_with("ws://"), + ) + .await?; let connection_label = format!("exec-server websocket {websocket_url}"); let connection = if is_rendezvous_harness_url(&websocket_url) { diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 36bd1cfdd1..633b7b0848 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -201,6 +201,7 @@ pub use protocol::WriteResponse; pub use protocol::WriteStatus; pub use regular_file::read_sensitive_file_to_string; pub use remote::RemoteEnvironmentConfig; +pub use remote::RemoteEnvironmentTransport; pub use remote::run_remote_environment; pub use remote::run_remote_environment_forward_until_shutdown; pub use remote::run_remote_environment_until_shutdown; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index cb4282c04a..0d2711b7a8 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -50,11 +50,25 @@ use crate::server::RequestDispatchMode; use crate::trace_context::current_rendezvous_headers; use crate::trace_context::current_trace_context_headers; +#[path = "remote/direct.rs"] +mod direct; + +use direct::run_direct_environment; + const ERROR_BODY_PREVIEW_BYTES: usize = 4096; const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1"; mod registration_retry; +/// Wire transport used after registering a remote exec-server. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +/// The transport used to connect a remote exec-server environment. +pub enum RemoteEnvironmentTransport { + #[default] + Noise, + Direct, +} + #[derive(Clone)] struct EnvironmentRegistryClient { base_url: String, @@ -515,6 +529,7 @@ pub struct RemoteEnvironmentConfig { pub environment_id: String, pub name: String, pub request_dispatch_mode: RequestDispatchMode, + transport: RemoteEnvironmentTransport, auth_provider: SharedAuthProvider, telemetry: ExecServerTelemetry, http_client_factory: HttpClientFactory, @@ -527,17 +542,36 @@ impl std::fmt::Debug for RemoteEnvironmentConfig { .field("environment_id", &self.environment_id) .field("name", &self.name) .field("request_dispatch_mode", &self.request_dispatch_mode) + .field("transport", &self.transport) .field("auth_provider", &"") .finish() } } impl RemoteEnvironmentConfig { + /// Creates a remote environment configuration using the default Noise transport. pub fn new( base_url: String, environment_id: String, auth_provider: SharedAuthProvider, http_client_factory: HttpClientFactory, + ) -> Result { + Self::new_with_transport( + base_url, + environment_id, + RemoteEnvironmentTransport::Noise, + auth_provider, + http_client_factory, + ) + } + + /// Creates a remote environment configuration using an explicit transport. + pub fn new_with_transport( + base_url: String, + environment_id: String, + transport: RemoteEnvironmentTransport, + auth_provider: SharedAuthProvider, + http_client_factory: HttpClientFactory, ) -> Result { let environment_id = normalize_environment_id(environment_id)?; Ok(Self { @@ -545,6 +579,7 @@ impl RemoteEnvironmentConfig { environment_id, name: "codex-exec-server".to_string(), request_dispatch_mode: RequestDispatchMode::Inline, + transport, auth_provider, telemetry: ExecServerTelemetry::default(), http_client_factory, @@ -557,12 +592,15 @@ impl RemoteEnvironmentConfig { } } -/// Register an exec-server for remote use and serve requests over Noise. +/// Register an exec-server for remote use and serve requests over its configured transport. /// -/// The executor identity is generated once per process and reused across +/// In Noise mode, the executor identity is generated once per process and reused across /// reconnects. The registration and rendezvous URL are also reused until /// rendezvous rejects the URL, at which point the next attempt registers again. /// The websocket carries cleartext routing metadata and encrypted payloads. +/// +/// Direct mode reuses its registration across reconnects. A WebSocket handshake +/// conflict refreshes the registration; other permanent client errors stop the runner. pub async fn run_remote_environment( config: RemoteEnvironmentConfig, runtime_paths: ExecServerRuntimePaths, @@ -588,7 +626,20 @@ where config.request_dispatch_mode, ); - let result = run_remote_transport(config, processor.clone(), shutdown).await; + let result = match config.transport { + RemoteEnvironmentTransport::Noise => { + run_remote_transport(config, shutdown, |config, client| { + run_remote_environment_connections(config, client, processor.clone()) + }) + .await + } + RemoteEnvironmentTransport::Direct => { + run_remote_transport(config, shutdown, |config, client| { + run_direct_environment(config, client, processor.clone()) + }) + .await + } + }; processor.shutdown().await; result } @@ -602,22 +653,35 @@ pub async fn run_remote_environment_forward_until_shutdown( where F: std::future::Future, { + // Forwarder implements the Noise stream bridge. Direct forwarding needs a + // separate Direct-compatible bridge, so reject it rather than use the + // Noise-specific path. + // Remove this guard when a Direct-compatible forwarder is added. + if config.transport == RemoteEnvironmentTransport::Direct { + return Err(ExecServerError::EnvironmentRegistryConfig( + "direct exec-server transport does not support forwarding".to_string(), + )); + } let forwarder = Forwarder::new( websocket_url, &config.http_client_factory, config.telemetry.clone(), )?; - run_remote_transport(config, forwarder, shutdown).await + run_remote_transport(config, shutdown, |config, client| { + run_remote_environment_connections(config, client, forwarder) + }) + .await } -async fn run_remote_transport( +async fn run_remote_transport( config: RemoteEnvironmentConfig, - handler: H, shutdown: F, + run_loop: R, ) -> Result<(), ExecServerError> where F: std::future::Future, - H: NoiseStreamHandler, + R: FnOnce(RemoteEnvironmentConfig, EnvironmentRegistryClient) -> T, + T: std::future::Future>, { ensure_rustls_crypto_provider(); let client = EnvironmentRegistryClient::new_with_telemetry( @@ -626,7 +690,7 @@ where config.telemetry.clone(), config.http_client_factory.clone(), )?; - let run = run_remote_environment_connections(config, client, handler); + let run = run_loop(config, client); tokio::pin!(run, shutdown); tokio::select! { result = &mut run => result, @@ -1205,6 +1269,33 @@ mod tests { ); } + #[test] + fn remote_environment_config_new_defaults_to_noise_transport() { + let config = RemoteEnvironmentConfig::new( + "https://registry.example".to_string(), + "env-1".to_string(), + static_registry_auth_provider(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + .expect("config"); + + assert_eq!(config.transport, RemoteEnvironmentTransport::Noise); + } + + #[test] + fn remote_environment_config_new_with_transport_preserves_direct_transport() { + let config = RemoteEnvironmentConfig::new_with_transport( + "https://registry.example".to_string(), + "env-1".to_string(), + RemoteEnvironmentTransport::Direct, + static_registry_auth_provider(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + .expect("config"); + + assert_eq!(config.transport, RemoteEnvironmentTransport::Direct); + } + #[test] fn debug_output_redacts_auth_provider() { let config = RemoteEnvironmentConfig::new( diff --git a/codex-rs/exec-server/src/remote/direct.rs b/codex-rs/exec-server/src/remote/direct.rs new file mode 100644 index 0000000000..02365f58bb --- /dev/null +++ b/codex-rs/exec-server/src/remote/direct.rs @@ -0,0 +1,277 @@ +use std::time::Instant; + +use codex_api::AuthError; +use codex_api::AuthProvider; +use codex_http_client::Request; +use codex_websocket_client::WebSocketConnection; +use codex_websocket_client::WebSocketConnector; +use http::Method; +use http::StatusCode; +use serde::Deserialize; +use serde::Serialize; +use tokio::time::sleep; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tracing::info; +use tracing::warn; + +use super::EnvironmentRegistryClient; +use super::RemoteEnvironmentConfig; +use crate::ExecServerError; +use crate::client::is_retryable_recovery_error; +use crate::client::registry_recovery_retry_delay; +use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; +use crate::client_transport::authenticate_websocket_request; +use crate::client_transport::connect_websocket_request; +use crate::connection::JsonRpcConnection; +use crate::server::ConnectionProcessor; +use crate::telemetry::ConnectionTransport; +use crate::trace_context::current_trace_context_headers; + +const DIRECT_TRANSPORT: &str = "direct_jsonrpc_v1"; + +#[derive(Debug, Serialize)] +struct DirectRegistrationRequest { + transport: &'static str, +} + +#[derive(Debug, Deserialize)] +struct DirectRegistrationResponse { + environment_id: String, + transport: String, + registration_id: String, + url: String, +} + +impl EnvironmentRegistryClient { + #[tracing::instrument( + name = "codex.exec_server.remote.register", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.register", + result = tracing::field::Empty, + ) + )] + async fn register_direct_environment( + &self, + environment_id: &str, + ) -> Result { + let started_at = Instant::now(); + let response = async { + let url = super::endpoint_url( + &self.base_url, + &format!("/cloud/environment/{environment_id}/direct/register"), + ); + let request = Request::new(Method::POST, url) + .with_json(&DirectRegistrationRequest { + transport: DIRECT_TRANSPORT, + }) + .into_prepared() + .map_err(ExecServerError::EnvironmentRegistryConfig)?; + let request = self + .auth_provider + .apply_auth(request) + .await + .map_err(direct_auth_error)?; + let prepared = request + .prepare_body_for_send() + .map_err(ExecServerError::EnvironmentRegistryConfig)?; + let response = self + .http + .request(request.method, request.url) + .headers(prepared.headers) + .headers(current_trace_context_headers()) + .body(prepared.body.unwrap_or_default()) + .timeout(self.connect_timeout) + .send() + .await?; + let response: DirectRegistrationResponse = self.parse_json_response(response).await?; + if response.environment_id != environment_id { + return Err(ExecServerError::Protocol( + "environment registry returned a different environment id".to_string(), + )); + } + if response.transport != DIRECT_TRANSPORT { + return Err(ExecServerError::Protocol(format!( + "environment registry returned unsupported direct transport `{}`", + response.transport + ))); + } + if response.registration_id.trim().is_empty() || response.url.trim().is_empty() { + return Err(ExecServerError::Protocol( + "environment registry returned incomplete direct connection data".to_string(), + )); + } + require_tls_or_loopback(&response.url, "wss")?; + Ok(response) + } + .await; + let result = if response.is_ok() { "success" } else { "error" }; + tracing::Span::current().record("result", result); + self.telemetry + .remote_registration_completed(result, started_at.elapsed()); + response + } +} + +pub(super) async fn run_direct_environment( + config: RemoteEnvironmentConfig, + client: EnvironmentRegistryClient, + processor: ConnectionProcessor, +) -> Result<(), ExecServerError> { + require_tls_or_loopback(&config.base_url, "https")?; + let mut retry_attempt = 0; + let mut registration = client + .register_direct_environment(&config.environment_id) + .await?; + + loop { + match connect_direct( + ®istration.url, + config.auth_provider.as_ref(), + &config.http_client_factory, + ) + .await + { + Ok(websocket) => { + retry_attempt = 0; + info!( + environment_id = registration.environment_id, + registration_id = registration.registration_id, + "direct exec-server connected" + ); + processor + .run_connection( + JsonRpcConnection::from_websocket( + websocket, + format!( + "direct exec-server websocket {}", + websocket_diagnostic_url(®istration.url) + ), + ), + ConnectionTransport::WebSocket, + ) + .await; + config.telemetry.remote_reconnect("disconnected"); + } + Err(error) + if is_retryable_recovery_error(&error) + && !matches!( + &error, + ExecServerError::WebSocketConnect { + source: tokio_tungstenite::tungstenite::Error::Http(response), + .. + } if response.status().is_client_error() + && !matches!( + response.status(), + StatusCode::REQUEST_TIMEOUT + | StatusCode::CONFLICT + | StatusCode::TOO_MANY_REQUESTS + ) + ) => + { + // A handshake conflict rejects this registration; other transient failures + // reconnect using the existing URL. + if matches!( + &error, + ExecServerError::WebSocketConnect { + source: tokio_tungstenite::tungstenite::Error::Http(response), + .. + } if response.status() == StatusCode::CONFLICT + ) { + registration = client + .register_direct_environment(&config.environment_id) + .await?; + } + warn!("direct exec-server connection failed; retrying"); + config.telemetry.remote_reconnect("connect_failed"); + } + Err(error) => return Err(error), + } + + sleep(registry_recovery_retry_delay( + &config.environment_id, + retry_attempt, + )) + .await; + retry_attempt = retry_attempt.saturating_add(1); + } +} + +async fn connect_direct( + url: &str, + auth_provider: &dyn AuthProvider, + http_client_factory: &codex_http_client::HttpClientFactory, +) -> Result { + let mut request = + url.into_client_request() + .map_err(|source| ExecServerError::WebSocketConnect { + url: websocket_diagnostic_url(url), + source, + })?; + request + .headers_mut() + .extend(current_trace_context_headers()); + authenticate_websocket_request(&mut request, auth_provider) + .await + .map_err(direct_auth_error)?; + let authenticated_url = request.uri().to_string(); + require_tls_or_loopback(&authenticated_url, "wss")?; + let connector = WebSocketConnector::new(http_client_factory) + .map_err(|error| ExecServerError::WebSocketConfiguration(error.to_string()))? + .with_tcp_nodelay(); + connect_websocket_request( + request, + websocket_diagnostic_url(&authenticated_url), + connector, + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + /*use_loopback_direct*/ false, + ) + .await +} + +pub(super) fn require_tls_or_loopback( + url: &str, + secure_scheme: &str, +) -> Result<(), ExecServerError> { + let parsed = url::Url::parse(url).map_err(|error| { + ExecServerError::EnvironmentRegistryConfig(format!("invalid remote endpoint URL: {error}")) + })?; + if parsed.scheme() == secure_scheme { + return Ok(()); + } + + let loopback = match parsed.host() { + Some(url::Host::Domain(host)) => host.eq_ignore_ascii_case("localhost"), + Some(url::Host::Ipv4(host)) => host.is_loopback(), + Some(url::Host::Ipv6(host)) => host.is_loopback(), + None => false, + }; + if loopback + && secure_scheme + .strip_suffix('s') + .is_some_and(|scheme| parsed.scheme() == scheme) + { + return Ok(()); + } + + Err(ExecServerError::EnvironmentRegistryConfig(format!( + "remote transport requires {secure_scheme} for non-loopback endpoints" + ))) +} + +fn websocket_diagnostic_url(url: &str) -> String { + url.split(['?', '#']).next().unwrap_or(url).to_string() +} + +fn direct_auth_error(error: AuthError) -> ExecServerError { + let message = format!("failed to resolve environment registry authentication: {error}"); + match error { + AuthError::Build(_) => ExecServerError::EnvironmentRegistryAuth(message), + AuthError::Transient(_) => ExecServerError::Disconnected(message), + } +} + +#[cfg(test)] +#[path = "direct_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/remote/direct_tests.rs b/codex-rs/exec-server/src/remote/direct_tests.rs new file mode 100644 index 0000000000..d0cbff0d08 --- /dev/null +++ b/codex-rs/exec-server/src/remote/direct_tests.rs @@ -0,0 +1,449 @@ +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Result; +use codex_api::AuthProvider; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use http::HeaderMap; +use http::HeaderValue; +use pretty_assertions::assert_eq; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpListener; +use tokio::time::timeout; +use tokio_tungstenite::accept_hdr_async; +use tokio_tungstenite::tungstenite::handshake::server::Request as HandshakeRequest; +use tokio_tungstenite::tungstenite::handshake::server::Response as HandshakeResponse; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::*; +use crate::ExecServerRuntimePaths; +use crate::RemoteEnvironmentTransport; + +#[derive(Debug)] +struct StaticAuthProvider; + +impl AuthProvider for StaticAuthProvider { + fn add_auth_headers(&self, headers: &mut HeaderMap) { + headers.insert( + http::header::AUTHORIZATION, + HeaderValue::from_static("AWS4-HMAC-SHA256 test-signature"), + ); + } + + fn apply_auth( + &self, + mut request: codex_http_client::Request, + ) -> codex_api::AuthProviderFuture<'_> { + Box::pin(async move { + if request.method == http::Method::GET { + assert_eq!(request.headers.len(), 1); + assert!(request.headers.contains_key(http::header::HOST)); + assert!(request.url.starts_with("http")); + } + self.add_auth_headers(&mut request.headers); + Ok(request) + }) + } +} + +#[derive(Debug)] +struct QueryAuthProvider; + +impl AuthProvider for QueryAuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) {} + + fn apply_auth( + &self, + mut request: codex_http_client::Request, + ) -> codex_api::AuthProviderFuture<'_> { + Box::pin(async move { + let mut url = url::Url::parse(&request.url) + .map_err(|error| codex_api::AuthError::Build(error.to_string()))?; + url.query_pairs_mut().append_pair("auth", "signed-query"); + request.url = url.into(); + Ok(request) + }) + } +} + +#[tokio::test] +async fn direct_websocket_signs_only_host_and_preserves_handshake_headers() -> Result<()> { + let mut request = "wss://executor.example.com/connect".into_client_request()?; + request + .headers_mut() + .insert("traceparent", HeaderValue::from_static("trace-context")); + + authenticate_websocket_request(&mut request, &StaticAuthProvider).await?; + + assert!(request.headers().contains_key("sec-websocket-key")); + assert_eq!(request.headers()["traceparent"], "trace-context"); + assert_eq!( + request.headers()[http::header::AUTHORIZATION], + "AWS4-HMAC-SHA256 test-signature" + ); + Ok(()) +} + +#[tokio::test] +async fn direct_websocket_uses_authenticated_url_query_parameters() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let url = format!("ws://{}/connect?existing=value", listener.local_addr()?); + let acceptor = tokio::spawn(async move { + let (socket, _) = listener.accept().await?; + let callback = |request: &HandshakeRequest, response: HandshakeResponse| { + assert_eq!(request.uri().path(), "/connect"); + assert_eq!( + request.uri().query(), + Some("existing=value&auth=signed-query") + ); + Ok(response) + }; + accept_hdr_async(socket, callback).await.map(|_| ()) + }); + + let connection = connect_direct( + &url, + &QueryAuthProvider, + &HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + .await?; + drop(connection); + acceptor.await??; + Ok(()) +} + +#[test] +fn direct_endpoints_require_tls_except_on_loopback() { + for (url, scheme, allowed) in [ + ("https://registry.example.com", "https", true), + ("wss://executor.example.com/connect", "wss", true), + ("http://127.0.0.1:8080", "https", true), + ("ws://localhost:8080/connect", "wss", true), + ("http://registry.example.com", "https", false), + ("ws://executor.example.com/connect", "wss", false), + ] { + assert_eq!( + require_tls_or_loopback(url, scheme).is_ok(), + allowed, + "{url}" + ); + } +} + +#[test] +fn direct_authentication_retries_only_transient_credential_failures() { + for (error, retryable) in [ + (AuthError::Transient("temporary".to_string()), true), + (AuthError::Build("invalid".to_string()), false), + ] { + assert_eq!( + is_retryable_recovery_error(&direct_auth_error(error)), + retryable + ); + } +} + +#[tokio::test] +async fn direct_registration_validates_connection_data() -> Result<()> { + for (field, invalid_value) in [ + ("environment_id", "different-environment"), + ("transport", "noise_hybrid_ik_v1"), + ("registration_id", ""), + ("url", ""), + ("url", "ws://executor.example.com/connect"), + ] { + let registry = MockServer::start().await; + let mut response = serde_json::json!({ + "environment_id": "environment-requested", + "transport": DIRECT_TRANSPORT, + "registration_id": "registration-1", + "url": "wss://executor.example.com/connect", + }); + response[field] = serde_json::json!(invalid_value); + Mock::given(method("POST")) + .and(path( + "/cloud/environment/environment-requested/direct/register", + )) + .respond_with(ResponseTemplate::new(200).set_body_json(response)) + .expect(1) + .mount(®istry) + .await; + let client = EnvironmentRegistryClient::new(registry.uri(), Arc::new(StaticAuthProvider))?; + let error = client + .register_direct_environment("environment-requested") + .await + .expect_err("invalid registration must be rejected"); + if invalid_value.starts_with("ws://") { + assert!(matches!( + error, + ExecServerError::EnvironmentRegistryConfig(_) + )); + } else { + assert!(matches!(error, ExecServerError::Protocol(_))); + } + registry.verify().await; + } + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn direct_registration_uses_proxy_policy_without_logging_secrets() -> Result<()> { + use tracing_subscriber::Layer; + use tracing_subscriber::layer::SubscriberExt; + + let mut log_file = tempfile::tempfile()?; + let writer = log_file.try_clone()?; + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(move || writer.try_clone().expect("clone log file")) + .with_filter( + tracing_subscriber::filter::Targets::new() + .with_target("codex_http_client", tracing::Level::TRACE) + .with_target("codex_exec_server", tracing::Level::TRACE), + ), + ); + let _guard = tracing::subscriber::set_default(subscriber); + tracing::debug!(target: "codex_exec_server", "direct registry log capture sentinel"); + + let proxy = MockServer::start().await; + let registry_url = "http://direct-registry-proxy.invalid/registry-path-secret"; + let request_url = + format!("{registry_url}/cloud/environment/environment-requested/direct/register"); + codex_http_client::cache_system_proxy_route_for_test(&request_url, proxy.uri()); + Mock::given(method("POST")) + .and(path( + "/registry-path-secret/cloud/environment/environment-requested/direct/register", + )) + .and(header("authorization", "AWS4-HMAC-SHA256 test-signature")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("set-cookie", "session=registry-cookie-secret") + .insert_header( + "location", + "https://registry.example/?token=registry-location-secret", + ) + .set_body_json(serde_json::json!({ + "environment_id": "environment-requested", + "transport": DIRECT_TRANSPORT, + "registration_id": "registration-1", + "url": "wss://executor.example/connect?token=websocket-query-secret", + })), + ) + .expect(1) + .mount(&proxy) + .await; + let client = EnvironmentRegistryClient::new_with_telemetry( + registry_url.to_string(), + Arc::new(StaticAuthProvider), + crate::ExecServerTelemetry::default(), + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + )?; + let response = client + .register_direct_environment("environment-requested") + .await?; + assert_eq!( + response.url, + "wss://executor.example/connect?token=websocket-query-secret" + ); + let requests = proxy + .received_requests() + .await + .expect("record proxy requests"); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].url.as_str(), request_url); + assert_eq!(requests[0].body, br#"{"transport":"direct_jsonrpc_v1"}"#); + proxy.verify().await; + + std::io::Seek::rewind(&mut log_file)?; + let mut logs = String::new(); + std::io::Read::read_to_string(&mut log_file, &mut logs)?; + assert!(logs.contains("direct registry log capture sentinel")); + for secret in [ + "test-signature", + "registry-path-secret", + "registry-cookie-secret", + "registry-location-secret", + "websocket-query-secret", + ] { + assert!(!logs.contains(secret), "registry logs exposed {secret}"); + } + Ok(()) +} + +#[tokio::test] +async fn direct_registration_failure_stops_initial_and_conflict_attempts() -> Result<()> { + for successful_registrations in [0, 1] { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}/connect", listener.local_addr()?); + let registry = MockServer::start().await; + let registration_count = AtomicUsize::new(0); + Mock::given(method("POST")) + .and(path( + "/cloud/environment/environment-requested/direct/register", + )) + .respond_with(move |_: &wiremock::Request| { + let attempt = registration_count.fetch_add(1, Ordering::Relaxed); + if attempt < successful_registrations { + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "environment_id": "environment-requested", + "transport": DIRECT_TRANSPORT, + "registration_id": "registration-1", + "url": websocket_url, + })) + } else { + ResponseTemplate::new(503) + } + }) + .expect((successful_registrations + 1) as u64) + .mount(®istry) + .await; + let config = RemoteEnvironmentConfig::new_with_transport( + registry.uri(), + "environment-requested".to_string(), + RemoteEnvironmentTransport::Direct, + Arc::new(StaticAuthProvider), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + )?; + let runtime_paths = ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?; + let task = tokio::spawn(crate::run_remote_environment(config, runtime_paths)); + if successful_registrations == 1 { + let (mut socket, _) = timeout(Duration::from_secs(5), listener.accept()).await??; + let mut request = [0; 4096]; + let _ = socket.read(&mut request).await?; + socket + .write_all(b"HTTP/1.1 409 Conflict\r\nContent-Length: 0\r\n\r\n") + .await?; + socket.shutdown().await?; + } + let error = timeout(Duration::from_secs(5), task) + .await?? + .expect_err("registration failure should stop the runner"); + assert!(matches!( + error, + ExecServerError::EnvironmentRegistryHttp { + status: StatusCode::SERVICE_UNAVAILABLE, + .. + } + )); + registry.verify().await; + } + Ok(()) +} + +#[tokio::test] +async fn direct_websocket_reuses_registration_and_stops_on_permanent_errors() -> Result<()> { + for (status, should_retry) in [ + (Some(StatusCode::BAD_REQUEST), false), + (Some(StatusCode::UNAUTHORIZED), false), + (Some(StatusCode::FORBIDDEN), false), + (Some(StatusCode::NOT_FOUND), false), + (Some(StatusCode::METHOD_NOT_ALLOWED), false), + (Some(StatusCode::GONE), false), + (Some(StatusCode::REQUEST_TIMEOUT), true), + (Some(StatusCode::CONFLICT), true), + (Some(StatusCode::TOO_MANY_REQUESTS), true), + (Some(StatusCode::INTERNAL_SERVER_ERROR), true), + (Some(StatusCode::SERVICE_UNAVAILABLE), true), + (None, true), // Close the TCP socket without sending an HTTP response. + ] { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let registry = MockServer::start().await; + let expected_calls = if status == Some(StatusCode::CONFLICT) { + 2 + } else { + 1 + }; + let registration_count = AtomicUsize::new(0); + Mock::given(method("POST")) + .and(path( + "/cloud/environment/environment-requested/direct/register", + )) + .and(header("authorization", "AWS4-HMAC-SHA256 test-signature")) + .respond_with(move |_: &wiremock::Request| { + let registration = registration_count.fetch_add(1, Ordering::Relaxed) + 1; + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "environment_id": "environment-requested", + "transport": DIRECT_TRANSPORT, + "registration_id": format!("registration-{registration}"), + "url": format!("{websocket_url}/registration-{registration}"), + })) + }) + .expect(expected_calls) + .mount(®istry) + .await; + let config = RemoteEnvironmentConfig::new_with_transport( + registry.uri(), + "environment-requested".to_string(), + RemoteEnvironmentTransport::Direct, + Arc::new(StaticAuthProvider), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + )?; + let runtime_paths = ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?; + let task = tokio::spawn(crate::run_remote_environment(config, runtime_paths)); + + let (mut socket, _) = timeout(Duration::from_secs(5), listener.accept()).await??; + let mut request = [0; 4096]; + let _ = socket.read(&mut request).await?; + if let Some(status) = status { + let response = format!( + "HTTP/1.1 {} {}\r\nContent-Length: 0\r\n\r\n", + status.as_u16(), + status.canonical_reason().unwrap_or_default() + ); + socket.write_all(response.as_bytes()).await?; + } + socket.shutdown().await?; + drop(socket); + + if should_retry { + let expected_path = format!("/registration-{expected_calls}"); + let check_registration = |request: &HandshakeRequest, response: HandshakeResponse| { + assert_eq!(request.uri().path(), expected_path); + Ok(response) + }; + let (socket, _) = timeout(Duration::from_secs(5), listener.accept()).await??; + let websocket = accept_hdr_async(socket, &check_registration).await?; + let retry_delay = + registry_recovery_retry_delay("environment-requested", /*attempt*/ 0); + let retry_started = Instant::now(); + drop(websocket); + + let (socket, _) = timeout(Duration::from_secs(5), listener.accept()).await??; + assert!(retry_started.elapsed() >= retry_delay); + let _websocket = accept_hdr_async(socket, &check_registration).await?; + registry.verify().await; + task.abort(); + let _ = task.await; + } else { + let error = timeout(Duration::from_secs(5), task) + .await?? + .expect_err("permanent WebSocket client error should be terminal"); + assert!(matches!( + error, + ExecServerError::WebSocketConnect { + source: tokio_tungstenite::tungstenite::Error::Http(response), + .. + } if Some(response.status()) == status + )); + } + } + Ok(()) +} diff --git a/codex-rs/exec-server/tests/accepted_websocket.rs b/codex-rs/exec-server/tests/accepted_websocket.rs index 8ce621fa1a..bd31df7c74 100644 --- a/codex-rs/exec-server/tests/accepted_websocket.rs +++ b/codex-rs/exec-server/tests/accepted_websocket.rs @@ -1,4 +1,7 @@ +mod common; + use std::collections::HashMap; +use std::sync::Arc; use std::time::Duration; use anyhow::Context; @@ -8,19 +11,29 @@ use axum::extract::State; use axum::extract::WebSocketUpgrade; use axum::response::IntoResponse; use axum::routing::any; +use codex_api::AuthProvider; +#[cfg(unix)] +use codex_exec_server::EnvironmentConnectionState; use codex_exec_server::EnvironmentInfo; use codex_exec_server::EnvironmentManager; use codex_exec_server::EnvironmentObservedStatus; use codex_exec_server::EnvironmentStatus; use codex_exec_server::EnvironmentStatusKind; use codex_exec_server::ExecParams; +#[cfg(unix)] +use codex_exec_server::ExecProcessEvent; use codex_exec_server::ExecResponse; use codex_exec_server::ExecServerClientConnectOptions; +use codex_exec_server::ExecServerRuntimePaths; use codex_exec_server::InitializeParams; use codex_exec_server::InitializeResponse; use codex_exec_server::ProcessId; use codex_exec_server::ReadParams; use codex_exec_server::ReadResponse; +use codex_exec_server::RemoteEnvironmentConfig; +use codex_exec_server::RemoteEnvironmentTransport; +#[cfg(unix)] +use codex_exec_server::WriteStatus; use codex_exec_server_protocol::JSONRPCError; use codex_exec_server_protocol::JSONRPCErrorError; use codex_exec_server_protocol::JSONRPCMessage; @@ -30,8 +43,11 @@ use codex_exec_server_protocol::JSONRPCResponse; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; use codex_utils_path_uri::PathUri; +use common::exec_server::DisconnectableWebSocketProxy; use futures::SinkExt; use futures::StreamExt; +use http::HeaderMap; +use http::HeaderValue; use pretty_assertions::assert_eq; use tokio::net::TcpListener; use tokio::sync::mpsc; @@ -41,12 +57,31 @@ use tokio_tungstenite::MaybeTlsStream; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; +use tokio_util::task::AbortOnDropHandle; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; const TEST_TIMEOUT: Duration = Duration::from_secs(5); type AcceptedSocket = axum::extract::ws::WebSocket; const SESSION_ALREADY_ATTACHED_ERROR_CODE: i64 = -32010; +#[derive(Debug)] +struct DirectExecutorAuth; + +impl AuthProvider for DirectExecutorAuth { + fn add_auth_headers(&self, headers: &mut HeaderMap) { + headers.insert( + http::header::AUTHORIZATION, + HeaderValue::from_static("AWS4-HMAC-SHA256 test-signature"), + ); + } +} + #[tokio::test] async fn accepted_websocket_rejects_initial_resume_session_id() -> Result<()> { let (websocket_url, mut accepted_sockets, server_task) = start_acceptor().await?; @@ -97,6 +132,284 @@ async fn accepted_websocket_environment_info_uses_initialization_metadata() -> R Ok(()) } +#[tokio::test] +async fn accepted_websocket_interoperates_and_recovers_with_real_direct_executor() -> Result<()> { + let (websocket_url, mut accepted_sockets, server_task) = start_acceptor().await?; + let proxy = DisconnectableWebSocketProxy::new(&websocket_url).await?; + let registry = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/cloud/environment/environment-1/direct/register")) + .and(header("authorization", "AWS4-HMAC-SHA256 test-signature")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "environment_id": "environment-1", + "transport": "direct_jsonrpc_v1", + "registration_id": "registration-1", + "url": proxy.websocket_url(), + }))) + .expect(1) + .mount(®istry) + .await; + + let http_client_factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + let config = RemoteEnvironmentConfig::new_with_transport( + registry.uri(), + "environment-1".to_string(), + RemoteEnvironmentTransport::Direct, + Arc::new(DirectExecutorAuth), + http_client_factory.clone(), + )?; + let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?; + let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?; + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let executor_task = AbortOnDropHandle::new(tokio::spawn( + codex_exec_server::run_remote_environment_until_shutdown( + config, + runtime_paths, + async move { + let _ = shutdown_rx.await; + }, + ), + )); + let accepted_websocket = timeout(TEST_TIMEOUT, accepted_sockets.recv()) + .await? + .context("direct executor websocket should be accepted")?; + let manager = timeout( + TEST_TIMEOUT, + EnvironmentManager::from_accepted_websocket( + "environment-1".to_string(), + accepted_websocket, + accepted_options(), + http_client_factory, + ), + ) + .await??; + let environment = manager + .default_environment() + .context("direct executor environment should be installed")?; + + assert_eq!( + timeout(TEST_TIMEOUT, environment.force_info()).await??, + EnvironmentInfo::local() + ); + let files = tempfile::tempdir()?; + let large_file_path = files.path().join("large-response.bin"); + let large_file_contents = vec![0x5a; 128 * 1024]; + tokio::fs::write(&large_file_path, &large_file_contents).await?; + assert_eq!( + timeout( + TEST_TIMEOUT, + environment.get_filesystem().read_file( + &PathUri::from_host_native_path(&large_file_path)?, + Default::default(), + /*sandbox*/ None, + ) + ) + .await??, + large_file_contents, + ); + + // The process fixture uses a POSIX shell; metadata and shutdown remain tested on all platforms. + #[cfg(unix)] + { + let mut proxy = proxy; + let backend = environment.get_exec_backend(); + let temp_dir = tempfile::TempDir::new()?; + let gate_path = temp_dir.path().join("release-output"); + let emitted_path = temp_dir.path().join("output-emitted"); + let session = timeout( + TEST_TIMEOUT, + backend.start(ExecParams { + metadata: Default::default(), + process_id: ProcessId::from("proc-recover"), + argv: vec![ + "/bin/sh".to_string(), + "-c".to_string(), + concat!( + "printf 'ready:%s\\n' \"$$\"; ", + "while [ ! -f \"$GATE\" ]; do /bin/sleep 0.01; done; ", + "printf 'during:%s\\n' \"$$\"; ", + ": > \"$EMITTED\"; ", + "IFS= read -r line; ", + "printf 'after:%s:%s\\n' \"$$\" \"$line\"; ", + "exit 7", + ) + .to_string(), + ], + cwd: PathUri::from_host_native_path(std::env::current_dir()?)?, + shell_snapshot: None, + env_policy: /*env_policy*/ None, + env: HashMap::from([ + ( + "GATE".to_string(), + gate_path.to_string_lossy().into_owned(), + ), + ( + "EMITTED".to_string(), + emitted_path.to_string_lossy().into_owned(), + ), + ]), + tty: false, + pipe_stdin: true, + arg0: None, + sandbox: None, + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }), + ) + .await??; + + let process = Arc::clone(&session.process); + let mut events = process.subscribe_events(); + let mut output = Vec::new(); + let mut last_seq = 0; + while !output.ends_with(b"\n") { + match timeout(Duration::from_secs(5), events.recv()).await?? { + ExecProcessEvent::Output(chunk) => { + assert_eq!(chunk.seq, last_seq + 1); + last_seq = chunk.seq; + output.extend_from_slice(&chunk.chunk.into_inner()); + } + event => anyhow::bail!("expected ready output before disconnect, got {event:?}"), + } + } + let ready = String::from_utf8(output.clone())?; + let pid = ready + .strip_prefix("ready:") + .and_then(|line| line.strip_suffix('\n')) + .context("ready output should contain the process id")? + .to_string(); + + let mut connection_state = environment + .subscribe_connection_state() + .context("direct environment connection state")?; + assert_eq!( + *connection_state.borrow_and_update(), + EnvironmentConnectionState::Connected + ); + proxy.pause_and_disconnect().await?; + timeout( + TEST_TIMEOUT, + connection_state.wait_for(|state| *state == EnvironmentConnectionState::Disconnected), + ) + .await??; + tokio::fs::write(&gate_path, b"").await?; + timeout(Duration::from_secs(5), async { + while tokio::fs::metadata(&emitted_path).await.is_err() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("process did not emit output while disconnected")?; + + let process_for_read = Arc::clone(&process); + let mut pending_read = tokio::spawn(async move { + process_for_read + .read( + /*after_seq*/ Some(last_seq), + /*max_bytes*/ None, + /*wait_ms*/ Some(0), + ) + .await + }); + assert!( + timeout(Duration::from_millis(200), &mut pending_read) + .await + .is_err(), + "process reads should wait while recovery is in progress" + ); + proxy.resume()?; + let replacement = timeout(TEST_TIMEOUT, accepted_sockets.recv()) + .await? + .context("real Direct executor should reconnect")?; + timeout( + TEST_TIMEOUT, + manager.replace_accepted_websocket("environment-1", replacement), + ) + .await??; + timeout( + TEST_TIMEOUT, + connection_state.wait_for(|state| *state == EnvironmentConnectionState::Connected), + ) + .await??; + assert!(Arc::ptr_eq( + &environment, + &manager + .default_environment() + .context("same environment should remain installed")?, + )); + assert_eq!( + timeout(TEST_TIMEOUT, environment.force_info()).await??, + EnvironmentInfo::local() + ); + + let recovered_read = timeout(Duration::from_secs(5), pending_read) + .await + .context("timed out waiting for a read after recovery")??; + let recovered_read = recovered_read?; + assert_eq!(recovered_read.failure, None); + let recovered_output = recovered_read + .chunks + .into_iter() + .flat_map(|chunk| chunk.chunk.into_inner()) + .collect::>(); + assert_eq!( + String::from_utf8(recovered_output)?, + format!("during:{pid}\n") + ); + + let write = timeout(Duration::from_secs(5), process.write(b"hello\n".to_vec())) + .await + .context("timed out waiting for a write after recovery")??; + assert_eq!(write.status, WriteStatus::Accepted); + + let mut saw_exit = false; + loop { + match timeout(Duration::from_secs(5), events.recv()).await?? { + ExecProcessEvent::Output(chunk) => { + assert_eq!(chunk.seq, last_seq + 1); + last_seq = chunk.seq; + output.extend_from_slice(&chunk.chunk.into_inner()); + } + ExecProcessEvent::Exited { seq, exit_code, .. } => { + assert_eq!(seq, last_seq + 1); + assert_eq!(exit_code, 7); + last_seq = seq; + saw_exit = true; + } + ExecProcessEvent::Closed { seq } => { + assert!(saw_exit, "closed must be delivered after exit"); + assert_eq!(seq, last_seq + 1); + break; + } + ExecProcessEvent::Failed(message) => { + anyhow::bail!("process recovery failed: {message}"); + } + } + } + assert_eq!( + String::from_utf8(output)?, + format!("ready:{pid}\nduring:{pid}\nafter:{pid}:hello\n") + ); + } + + registry.verify().await; + let registrations = registry + .received_requests() + .await + .context("registration requests")?; + assert_eq!( + registrations[0].body, + br#"{"transport":"direct_jsonrpc_v1"}"# + ); + + let _ = shutdown_tx.send(()); + timeout(TEST_TIMEOUT, executor_task).await???; + server_task.abort(); + let _ = server_task.await; + Ok(()) +} + #[tokio::test] async fn accepted_websocket_environment_is_ready_immediately() -> Result<()> { let (websocket_url, mut accepted_sockets, server_task) = start_acceptor().await?; diff --git a/codex-rs/exec-server/tests/common/exec_server.rs b/codex-rs/exec-server/tests/common/exec_server.rs index ceba4238ba..d0fc3892fc 100644 --- a/codex-rs/exec-server/tests/common/exec_server.rs +++ b/codex-rs/exec-server/tests/common/exec_server.rs @@ -144,30 +144,7 @@ impl ExecServerHarness { pub(crate) async fn disconnectable_websocket_proxy( &self, ) -> anyhow::Result { - let upstream = self - .websocket_url - .strip_prefix("ws://") - .ok_or_else(|| anyhow!("exec-server websocket URL must use ws://"))? - .to_string(); - let listener = TcpListener::bind("127.0.0.1:0").await?; - let websocket_url = format!("ws://{}", listener.local_addr()?); - let (pause_tx, pause_rx) = oneshot::channel(); - let (blocked_connection_tx, blocked_connection_rx) = oneshot::channel(); - let (resume_tx, resume_rx) = oneshot::channel(); - let task = tokio::spawn(run_disconnectable_proxy( - listener, - upstream, - pause_rx, - blocked_connection_tx, - resume_rx, - )); - Ok(DisconnectableWebSocketProxy { - websocket_url, - pause_tx: Some(pause_tx), - blocked_connection_rx: Some(blocked_connection_rx), - resume_tx: Some(resume_tx), - task, - }) + DisconnectableWebSocketProxy::new(&self.websocket_url).await } pub(crate) async fn send_request( @@ -278,6 +255,33 @@ impl ExecServerHarness { } impl DisconnectableWebSocketProxy { + pub(crate) async fn new(websocket_url: &str) -> anyhow::Result { + let upstream = websocket_url + .strip_prefix("ws://") + .ok_or_else(|| anyhow!("exec-server websocket URL must use ws://"))? + .trim_end_matches('/') + .to_string(); + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let (pause_tx, pause_rx) = oneshot::channel(); + let (blocked_connection_tx, blocked_connection_rx) = oneshot::channel(); + let (resume_tx, resume_rx) = oneshot::channel(); + let task = tokio::spawn(run_disconnectable_proxy( + listener, + upstream, + pause_rx, + blocked_connection_tx, + resume_rx, + )); + Ok(DisconnectableWebSocketProxy { + websocket_url, + pause_tx: Some(pause_tx), + blocked_connection_rx: Some(blocked_connection_rx), + resume_tx: Some(resume_tx), + task, + }) + } + pub(crate) fn websocket_url(&self) -> &str { &self.websocket_url }