codex_exec_server/
remote.rs

1use std::sync::Arc;
2use std::time::Duration;
3use std::time::Instant;
4
5use codex_api::AuthProvider;
6use codex_api::SharedAuthProvider;
7use futures::FutureExt;
8use http::HeaderMap;
9use http::HeaderName;
10use http::HeaderValue;
11use reqwest::StatusCode;
12use serde::Deserialize;
13use tokio::time::sleep;
14use tokio_tungstenite::MaybeTlsStream;
15use tokio_tungstenite::WebSocketStream;
16use tokio_tungstenite::connect_async_with_config;
17use tokio_tungstenite::tungstenite::client::IntoClientRequest;
18use tracing::debug;
19use tracing::info;
20use tracing::warn;
21
22use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
23
24use crate::EnvironmentRegistryConnectRequest;
25use crate::EnvironmentRegistryConnectResponse;
26use crate::EnvironmentRegistryHarnessKeyValidationRequest;
27use crate::EnvironmentRegistryHarnessKeyValidationResponse;
28use crate::EnvironmentRegistryRegistrationRequest;
29use crate::EnvironmentRegistryRegistrationResponse;
30use crate::ExecServerError;
31use crate::ExecServerRuntimePaths;
32use crate::ExecServerTelemetry;
33use crate::NoiseChannelIdentity;
34use crate::NoiseChannelPublicKey;
35use crate::NoiseRendezvousConnectBundle;
36use crate::NoiseRendezvousConnectProvider;
37use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
38use crate::noise_relay::noise_relay_websocket_config;
39use crate::relay::HarnessKeyValidator;
40use crate::relay::run_multiplexed_environment;
41use crate::server::ConnectionProcessor;
42use crate::trace_context::current_trace_context_headers;
43
44const ERROR_BODY_PREVIEW_BYTES: usize = 4096;
45const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1";
46
47#[derive(Clone)]
48struct EnvironmentRegistryClient {
49    base_url: String,
50    auth_provider: SharedAuthProvider,
51    http: reqwest::Client,
52    connect_timeout: Duration,
53    telemetry: ExecServerTelemetry,
54}
55
56impl std::fmt::Debug for EnvironmentRegistryClient {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        f.debug_struct("EnvironmentRegistryClient")
59            .field("base_url", &self.base_url)
60            .field("auth_provider", &"<redacted>")
61            .finish_non_exhaustive()
62    }
63}
64
65impl EnvironmentRegistryClient {
66    #[cfg(test)]
67    fn new(base_url: String, auth_provider: SharedAuthProvider) -> Result<Self, ExecServerError> {
68        Self::new_with_telemetry(base_url, auth_provider, ExecServerTelemetry::default())
69    }
70
71    fn new_with_telemetry(
72        base_url: String,
73        auth_provider: SharedAuthProvider,
74        telemetry: ExecServerTelemetry,
75    ) -> Result<Self, ExecServerError> {
76        let base_url = normalize_base_url(base_url)?;
77        Ok(Self {
78            base_url,
79            auth_provider,
80            http: reqwest::Client::builder()
81                .redirect(reqwest::redirect::Policy::none())
82                .build()?,
83            connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
84            telemetry,
85        })
86    }
87
88    /// Register the executor public key and obtain the rendezvous allocation.
89    /// The returned registration ID is included in each stream's Noise prologue.
90    #[tracing::instrument(
91        name = "codex.exec_server.remote.register",
92        skip_all,
93        fields(
94            otel.kind = "client",
95            otel.name = "codex.exec_server.remote.register",
96            result = tracing::field::Empty,
97        )
98    )]
99    async fn register_environment(
100        &self,
101        environment_id: &str,
102        executor_public_key: &NoiseChannelPublicKey,
103    ) -> Result<EnvironmentRegistryRegistrationResponse, ExecServerError> {
104        let started_at = Instant::now();
105        let response = self
106            .register_environment_inner(environment_id, executor_public_key)
107            .await;
108        let result = if response.is_ok() { "success" } else { "error" };
109        tracing::Span::current().record("result", result);
110        self.telemetry
111            .remote_registration_completed(result, started_at.elapsed());
112        response
113    }
114
115    async fn register_environment_inner(
116        &self,
117        environment_id: &str,
118        executor_public_key: &NoiseChannelPublicKey,
119    ) -> Result<EnvironmentRegistryRegistrationResponse, ExecServerError> {
120        let response = self
121            .http
122            .post(endpoint_url(
123                &self.base_url,
124                &format!("/cloud/environment/{environment_id}/register"),
125            ))
126            .headers(self.auth_provider.to_auth_headers())
127            .headers(current_trace_context_headers())
128            .json(&EnvironmentRegistryRegistrationRequest {
129                security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(),
130                executor_public_key: executor_public_key.clone(),
131            })
132            .send()
133            .await?;
134        let response: EnvironmentRegistryRegistrationResponse =
135            self.parse_json_response(response).await?;
136        if response.environment_id != environment_id {
137            return Err(ExecServerError::Protocol(
138                "environment registry returned a different environment id".to_string(),
139            ));
140        }
141        if response.security_profile != NOISE_RELAY_SECURITY_PROFILE {
142            return Err(ExecServerError::Protocol(format!(
143                "environment registry returned unsupported security profile `{}`",
144                response.security_profile
145            )));
146        }
147        info!(
148            noise_event = "registration",
149            noise_outcome = "ok",
150            security_profile = NOISE_RELAY_SECURITY_PROFILE,
151            "Noise executor registration completed"
152        );
153        debug!(
154            environment_id = response.environment_id,
155            executor_registration_id = response.executor_registration_id,
156            "Noise executor registration details"
157        );
158        Ok(response)
159    }
160
161    /// Authorize one Noise harness key and obtain the full rendezvous bundle.
162    async fn connect_environment(
163        &self,
164        environment_id: &str,
165        harness_public_key: NoiseChannelPublicKey,
166    ) -> Result<NoiseRendezvousConnectBundle, ExecServerError> {
167        let response = self
168            .http
169            .post(endpoint_url(
170                &self.base_url,
171                &format!("/cloud/environment/{environment_id}/connect"),
172            ))
173            .headers(self.auth_provider.to_auth_headers())
174            .json(&EnvironmentRegistryConnectRequest { harness_public_key })
175            .timeout(self.connect_timeout)
176            .send()
177            .await?;
178        let response: EnvironmentRegistryConnectResponse =
179            self.parse_json_response(response).await?;
180        if response.environment_id != environment_id {
181            return Err(ExecServerError::Protocol(
182                "environment registry returned a different environment id".to_string(),
183            ));
184        }
185        if response.security_profile != NOISE_RELAY_SECURITY_PROFILE {
186            return Err(ExecServerError::Protocol(format!(
187                "environment registry returned unsupported security profile `{}`",
188                response.security_profile
189            )));
190        }
191        if response.url.trim().is_empty()
192            || response.executor_registration_id.trim().is_empty()
193            || response.harness_key_authorization.trim().is_empty()
194        {
195            return Err(ExecServerError::Protocol(
196                "environment registry returned incomplete Noise connection data".to_string(),
197            ));
198        }
199        Ok(NoiseRendezvousConnectBundle {
200            websocket_url: response.url,
201            environment_id: response.environment_id,
202            executor_registration_id: response.executor_registration_id,
203            executor_public_key: response.executor_public_key,
204            harness_key_authorization: response.harness_key_authorization,
205        })
206    }
207
208    async fn parse_json_response<R>(
209        &self,
210        response: reqwest::Response,
211    ) -> Result<R, ExecServerError>
212    where
213        R: for<'de> Deserialize<'de>,
214    {
215        if response.status().is_success() {
216            return response.json::<R>().await.map_err(ExecServerError::from);
217        }
218
219        let status = response.status();
220        let body = response.text().await.unwrap_or_default();
221        if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) {
222            return Err(environment_registry_auth_error(status, &body));
223        }
224
225        Err(environment_registry_http_error(status, &body))
226    }
227}
228
229#[derive(Clone)]
230struct RegistryHarnessKeyValidator {
231    client: EnvironmentRegistryClient,
232    environment_id: String,
233    executor_registration_id: String,
234}
235
236impl HarnessKeyValidator for RegistryHarnessKeyValidator {
237    /// Authorize the harness key recovered from the first IK message.
238    /// Noise proves key possession; the registry decides whether that key may use
239    /// this executor. The authorization token and public key are checked together.
240    async fn validate_harness_key(
241        &self,
242        harness_public_key: &NoiseChannelPublicKey,
243        authorization: &str,
244    ) -> Result<(), ExecServerError> {
245        let environment_id = &self.environment_id;
246        let response = self
247            .client
248            .http
249            .post(endpoint_url(
250                &self.client.base_url,
251                &format!("/cloud/environment/{environment_id}/validate"),
252            ))
253            .headers(self.client.auth_provider.to_auth_headers())
254            .json(&EnvironmentRegistryHarnessKeyValidationRequest {
255                executor_registration_id: self.executor_registration_id.clone(),
256                harness_public_key: harness_public_key.clone(),
257                harness_key_authorization: authorization.to_string(),
258            })
259            .send()
260            .await?;
261        let status = response.status();
262        if !status.is_success() {
263            // The request contains the short-lived authorization. Do not include
264            // a response body that might echo it in logs or error chains.
265            if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) {
266                return Err(ExecServerError::EnvironmentRegistryAuth(format!(
267                    "environment registry harness key validation authentication failed ({status})"
268                )));
269            }
270            return Err(ExecServerError::EnvironmentRegistryHttp {
271                status,
272                code: None,
273                message: "environment registry harness key validation failed".to_string(),
274            });
275        }
276        let response = response
277            .json::<EnvironmentRegistryHarnessKeyValidationResponse>()
278            .await?;
279        if !response.valid {
280            return Err(ExecServerError::Protocol(
281                "environment registry rejected Noise relay harness key".to_string(),
282            ));
283        }
284        Ok(())
285    }
286}
287
288/// Noise connection configuration for a Codex harness.
289///
290/// The provider holds the authenticated registry client so every reconnect
291/// receives fresh URL and harness-key authorization material.
292#[derive(Clone)]
293pub(crate) struct NoiseRendezvousEnvironmentConfig {
294    provider: Arc<dyn NoiseRendezvousConnectProvider>,
295}
296
297impl std::fmt::Debug for NoiseRendezvousEnvironmentConfig {
298    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
299        f.debug_struct("NoiseRendezvousEnvironmentConfig")
300            .field("provider", &"<redacted>")
301            .finish()
302    }
303}
304
305impl NoiseRendezvousEnvironmentConfig {
306    pub(crate) fn new(
307        base_url: String,
308        environment_id: String,
309        bearer_token: String,
310        chatgpt_account_id: Option<String>,
311    ) -> Result<Self, ExecServerError> {
312        let environment_id = normalize_environment_id(environment_id)?;
313        let auth_provider = static_bearer_auth_provider(bearer_token, chatgpt_account_id)?;
314        let client = EnvironmentRegistryClient::new_with_telemetry(
315            base_url,
316            auth_provider,
317            ExecServerTelemetry::default(),
318        )?;
319        Ok(Self {
320            provider: Arc::new(EnvironmentRegistryNoiseConnectProvider {
321                client,
322                environment_id,
323            }),
324        })
325    }
326
327    pub(crate) fn connect_provider(&self) -> Arc<dyn NoiseRendezvousConnectProvider> {
328        Arc::clone(&self.provider)
329    }
330}
331
332#[derive(Clone, Debug)]
333struct EnvironmentRegistryNoiseConnectProvider {
334    client: EnvironmentRegistryClient,
335    environment_id: String,
336}
337
338impl NoiseRendezvousConnectProvider for EnvironmentRegistryNoiseConnectProvider {
339    fn connect_bundle(
340        &self,
341        harness_public_key: NoiseChannelPublicKey,
342    ) -> futures::future::BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
343        async move {
344            self.client
345                .connect_environment(&self.environment_id, harness_public_key)
346                .await
347        }
348        .boxed()
349    }
350}
351
352#[derive(Clone)]
353struct StaticBearerAuthProvider {
354    authorization: HeaderValue,
355    chatgpt_account_id: Option<HeaderValue>,
356}
357
358impl std::fmt::Debug for StaticBearerAuthProvider {
359    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
360        f.debug_struct("StaticBearerAuthProvider")
361            .field("authorization", &"<redacted>")
362            .field(
363                "chatgpt_account_id",
364                &self.chatgpt_account_id.as_ref().map(|_| "<redacted>"),
365            )
366            .finish()
367    }
368}
369
370impl AuthProvider for StaticBearerAuthProvider {
371    fn add_auth_headers(&self, headers: &mut HeaderMap) {
372        headers.insert(http::header::AUTHORIZATION, self.authorization.clone());
373        if let Some(chatgpt_account_id) = &self.chatgpt_account_id {
374            headers.insert(
375                HeaderName::from_static("chatgpt-account-id"),
376                chatgpt_account_id.clone(),
377            );
378        }
379    }
380}
381
382fn static_bearer_auth_provider(
383    bearer_token: String,
384    chatgpt_account_id: Option<String>,
385) -> Result<SharedAuthProvider, ExecServerError> {
386    let bearer_token = bearer_token.trim();
387    if bearer_token.is_empty() {
388        return Err(ExecServerError::EnvironmentRegistryConfig(
389            "environment registry bearer token is required".to_string(),
390        ));
391    }
392    let authorization =
393        HeaderValue::try_from(format!("Bearer {bearer_token}")).map_err(|error| {
394            ExecServerError::EnvironmentRegistryConfig(format!(
395                "environment registry bearer token is not a valid HTTP header: {error}"
396            ))
397        })?;
398    let chatgpt_account_id = chatgpt_account_id
399        .as_deref()
400        .map(str::trim)
401        .filter(|account_id| !account_id.is_empty())
402        .map(|account_id| {
403            HeaderValue::try_from(account_id).map_err(|error| {
404                ExecServerError::EnvironmentRegistryConfig(format!(
405                    "ChatGPT account id is not a valid HTTP header: {error}"
406                ))
407            })
408        })
409        .transpose()?;
410    Ok(Arc::new(StaticBearerAuthProvider {
411        authorization,
412        chatgpt_account_id,
413    }))
414}
415
416/// Configuration for registering an exec-server for remote use.
417#[derive(Clone)]
418pub struct RemoteEnvironmentConfig {
419    pub base_url: String,
420    pub environment_id: String,
421    pub name: String,
422    auth_provider: SharedAuthProvider,
423    telemetry: ExecServerTelemetry,
424}
425
426impl std::fmt::Debug for RemoteEnvironmentConfig {
427    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
428        f.debug_struct("RemoteEnvironmentConfig")
429            .field("base_url", &self.base_url)
430            .field("environment_id", &self.environment_id)
431            .field("name", &self.name)
432            .field("auth_provider", &"<redacted>")
433            .finish()
434    }
435}
436
437impl RemoteEnvironmentConfig {
438    pub fn new(
439        base_url: String,
440        environment_id: String,
441        auth_provider: SharedAuthProvider,
442    ) -> Result<Self, ExecServerError> {
443        let environment_id = normalize_environment_id(environment_id)?;
444        Ok(Self {
445            base_url,
446            environment_id,
447            name: "codex-exec-server".to_string(),
448            auth_provider,
449            telemetry: ExecServerTelemetry::default(),
450        })
451    }
452
453    pub fn with_telemetry(mut self, telemetry: ExecServerTelemetry) -> Self {
454        self.telemetry = telemetry;
455        self
456    }
457}
458
459/// Register an exec-server for remote use and serve requests over Noise.
460///
461/// The executor identity is generated once per process and reused across
462/// reconnects. The registration and rendezvous URL are also reused until
463/// rendezvous rejects the URL, at which point the next attempt registers again.
464/// The websocket carries cleartext routing metadata and encrypted payloads.
465pub async fn run_remote_environment(
466    config: RemoteEnvironmentConfig,
467    runtime_paths: ExecServerRuntimePaths,
468) -> Result<(), ExecServerError> {
469    ensure_rustls_crypto_provider();
470    let client = EnvironmentRegistryClient::new_with_telemetry(
471        config.base_url.clone(),
472        config.auth_provider.clone(),
473        config.telemetry.clone(),
474    )?;
475    let processor =
476        ConnectionProcessor::new_with_telemetry(runtime_paths, config.telemetry.clone());
477    let identity = NoiseChannelIdentity::generate().map_err(|error| {
478        ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}"))
479    })?;
480    let mut backoff = Duration::from_secs(1);
481    let mut response = client
482        .register_environment(&config.environment_id, &identity.public_key())
483        .await?;
484
485    loop {
486        match connect_rendezvous(&response.url, &config.telemetry).await {
487            Ok(websocket) => {
488                backoff = Duration::from_secs(1);
489                let executor_registration_id = response.executor_registration_id.clone();
490                info!(
491                    noise_event = "rendezvous_connection",
492                    noise_outcome = "ok",
493                    "Noise executor connected to rendezvous"
494                );
495                let disconnect_reason = run_multiplexed_environment(
496                    websocket,
497                    processor.clone(),
498                    response.environment_id.clone(),
499                    executor_registration_id.clone(),
500                    identity.clone(),
501                    RegistryHarnessKeyValidator {
502                        client: client.clone(),
503                        environment_id: config.environment_id.clone(),
504                        executor_registration_id,
505                    },
506                )
507                .await;
508                info!(
509                    noise_event = "rendezvous_connection",
510                    noise_outcome = "disconnected",
511                    noise_reason = disconnect_reason.as_str(),
512                    "Noise executor disconnected from rendezvous"
513                );
514                config
515                    .telemetry
516                    .remote_reconnect(disconnect_reason.as_str());
517            }
518            Err(error) => {
519                let registration_rejected = matches!(
520                    &error,
521                    tokio_tungstenite::tungstenite::Error::Http(response)
522                        if response.status().is_client_error()
523                );
524                warn!(
525                    noise_event = "rendezvous_connection",
526                    noise_outcome = "error",
527                    noise_reason = "websocket_error",
528                    "Noise executor failed to connect to rendezvous"
529                );
530                debug!(error = %error, "Noise executor rendezvous connection error");
531                if registration_rejected {
532                    config.telemetry.remote_reconnect("registration_rejected");
533                    response = client
534                        .register_environment(&config.environment_id, &identity.public_key())
535                        .await?;
536                } else {
537                    config.telemetry.remote_reconnect("connect_failed");
538                }
539            }
540        }
541
542        sleep(backoff).await;
543        backoff = (backoff * 2).min(Duration::from_secs(30));
544    }
545}
546
547#[tracing::instrument(
548    name = "codex.exec_server.remote.rendezvous.connect",
549    skip_all,
550    fields(
551        otel.kind = "client",
552        otel.name = "codex.exec_server.remote.rendezvous.connect",
553        result = tracing::field::Empty,
554    )
555)]
556async fn connect_rendezvous(
557    url: &str,
558    telemetry: &ExecServerTelemetry,
559) -> Result<
560    WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>,
561    tokio_tungstenite::tungstenite::Error,
562> {
563    let started_at = Instant::now();
564    let result = async {
565        let mut request = url.into_client_request()?;
566        request
567            .headers_mut()
568            .extend(current_trace_context_headers());
569        connect_async_with_config(
570            request,
571            Some(noise_relay_websocket_config()),
572            // Rendezvous sends small, latency-sensitive frames, so avoid Nagle's coalescing delay.
573            /*disable_nagle*/
574            true,
575        )
576        .await
577        .map(|(websocket, _)| websocket)
578    }
579    .await;
580    let result_name = if result.is_ok() { "success" } else { "error" };
581    tracing::Span::current().record("result", result_name);
582    telemetry.remote_rendezvous_completed(result_name, started_at.elapsed());
583    result
584}
585
586fn normalize_environment_id(environment_id: String) -> Result<String, ExecServerError> {
587    let environment_id = environment_id.trim().to_string();
588    if environment_id.is_empty() {
589        return Err(ExecServerError::EnvironmentRegistryConfig(
590            "environment id is required for remote exec-server registration".to_string(),
591        ));
592    }
593    Ok(environment_id)
594}
595
596#[derive(Deserialize)]
597struct RegistryErrorBody {
598    error: Option<RegistryError>,
599}
600
601#[derive(Deserialize)]
602struct RegistryError {
603    code: Option<String>,
604    message: Option<String>,
605}
606
607fn normalize_base_url(base_url: String) -> Result<String, ExecServerError> {
608    let trimmed = base_url.trim().trim_end_matches('/').to_string();
609    if trimmed.is_empty() {
610        return Err(ExecServerError::EnvironmentRegistryConfig(
611            "environment registry base URL is required".to_string(),
612        ));
613    }
614    Ok(trimmed)
615}
616
617fn endpoint_url(base_url: &str, path: &str) -> String {
618    format!("{base_url}/{}", path.trim_start_matches('/'))
619}
620
621fn environment_registry_auth_error(status: StatusCode, body: &str) -> ExecServerError {
622    let message = registry_error_message(body).unwrap_or_else(|| "empty error body".to_string());
623    ExecServerError::EnvironmentRegistryAuth(format!(
624        "environment registry authentication failed ({status}): {message}"
625    ))
626}
627
628fn environment_registry_http_error(status: StatusCode, body: &str) -> ExecServerError {
629    let parsed = serde_json::from_str::<RegistryErrorBody>(body).ok();
630    let (code, message) = parsed
631        .and_then(|body| body.error)
632        .map(|error| {
633            (
634                error.code,
635                error.message.unwrap_or_else(|| {
636                    preview_error_body(body).unwrap_or_else(|| "empty error body".to_string())
637                }),
638            )
639        })
640        .unwrap_or_else(|| {
641            (
642                None,
643                preview_error_body(body)
644                    .unwrap_or_else(|| "empty or malformed error body".to_string()),
645            )
646        });
647    ExecServerError::EnvironmentRegistryHttp {
648        status,
649        code,
650        message,
651    }
652}
653
654fn registry_error_message(body: &str) -> Option<String> {
655    serde_json::from_str::<RegistryErrorBody>(body)
656        .ok()
657        .and_then(|body| body.error)
658        .and_then(|error| error.message)
659        .or_else(|| preview_error_body(body))
660}
661
662fn preview_error_body(body: &str) -> Option<String> {
663    let trimmed = body.trim();
664    if trimmed.is_empty() {
665        return None;
666    }
667    Some(trimmed.chars().take(ERROR_BODY_PREVIEW_BYTES).collect())
668}
669
670#[cfg(test)]
671mod tests {
672    use std::sync::Arc;
673
674    use codex_api::AuthProvider;
675    use http::HeaderMap;
676    use http::HeaderValue;
677    use opentelemetry::trace::TracerProvider as _;
678    use opentelemetry_sdk::trace::SdkTracerProvider;
679    use pretty_assertions::assert_eq;
680    use tracing::Instrument;
681    use tracing_subscriber::prelude::*;
682    use wiremock::Mock;
683    use wiremock::MockServer;
684    use wiremock::ResponseTemplate;
685    use wiremock::matchers::body_partial_json;
686    use wiremock::matchers::header;
687    use wiremock::matchers::header_regex;
688    use wiremock::matchers::method;
689    use wiremock::matchers::path;
690
691    use super::*;
692
693    #[derive(Debug)]
694    struct StaticRegistryAuthProvider;
695
696    impl AuthProvider for StaticRegistryAuthProvider {
697        fn add_auth_headers(&self, headers: &mut HeaderMap) {
698            let _ = headers.insert(
699                http::header::AUTHORIZATION,
700                HeaderValue::from_static("Bearer registry-token"),
701            );
702            let _ = headers.insert(
703                "ChatGPT-Account-ID",
704                HeaderValue::from_static("workspace-123"),
705            );
706        }
707    }
708
709    fn static_registry_auth_provider() -> SharedAuthProvider {
710        Arc::new(StaticRegistryAuthProvider)
711    }
712
713    #[tokio::test(flavor = "current_thread")]
714    async fn register_environment_posts_with_auth_provider_headers() {
715        let provider = SdkTracerProvider::builder().build();
716        let tracer = provider.tracer("exec-server-test");
717        let subscriber =
718            tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer));
719        let _guard = subscriber.set_default();
720        tracing::callsite::rebuild_interest_cache();
721        let server = MockServer::start().await;
722        let executor_public_key = NoiseChannelIdentity::generate()
723            .expect("identity")
724            .public_key();
725        Mock::given(method("POST"))
726            .and(path("/cloud/environment/environment-requested/register"))
727            .and(header("authorization", "Bearer registry-token"))
728            .and(header("chatgpt-account-id", "workspace-123"))
729            .and(header_regex(
730                "traceparent",
731                "^00-[0-9a-f]{32}-[0-9a-f]{16}-0[01]$",
732            ))
733            .and(body_partial_json(serde_json::json!({
734                "security_profile": NOISE_RELAY_SECURITY_PROFILE,
735                "executor_public_key": executor_public_key.clone(),
736            })))
737            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
738                "environment_id": "environment-requested",
739                "url": "wss://rendezvous.test/cloud-agent/default/ws/environment/environment-requested?role=environment&sig=abc",
740                "security_profile": NOISE_RELAY_SECURITY_PROFILE,
741                "executor_registration_id": "registration-1",
742            })))
743            .mount(&server)
744            .await;
745        let client = EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider())
746            .expect("client");
747
748        let response = client
749            .register_environment("environment-requested", &executor_public_key)
750            .instrument(tracing::info_span!("remote-operation"))
751            .await
752            .expect("register environment");
753
754        assert_eq!(
755            response,
756            EnvironmentRegistryRegistrationResponse {
757                environment_id: "environment-requested".to_string(),
758                url: "wss://rendezvous.test/cloud-agent/default/ws/environment/environment-requested?role=environment&sig=abc".to_string(),
759                security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(),
760                executor_registration_id: "registration-1".to_string(),
761            }
762        );
763    }
764
765    #[tokio::test]
766    async fn noise_connect_provider_requests_and_validates_a_full_bundle() {
767        let server = MockServer::start().await;
768        let harness_public_key = NoiseChannelIdentity::generate()
769            .expect("identity")
770            .public_key();
771        let executor_public_key = NoiseChannelIdentity::generate()
772            .expect("identity")
773            .public_key();
774        Mock::given(method("POST"))
775            .and(path("/cloud/environment/environment-requested/connect"))
776            .and(header("authorization", "Bearer registry-token"))
777            .and(header("chatgpt-account-id", "workspace-123"))
778            .and(body_partial_json(serde_json::json!({
779                "harness_public_key": harness_public_key.clone(),
780            })))
781            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
782                "environment_id": "environment-requested",
783                "url": "wss://rendezvous.test/cloud-agent/default/ws/environment/environment-requested?role=harness&sig=abc",
784                "security_profile": NOISE_RELAY_SECURITY_PROFILE,
785                "executor_registration_id": "registration-1",
786                "executor_public_key": executor_public_key.clone(),
787                "harness_key_authorization": "authorization-1",
788            })))
789            .mount(&server)
790            .await;
791        let config = NoiseRendezvousEnvironmentConfig::new(
792            server.uri(),
793            "environment-requested".to_string(),
794            "registry-token".to_string(),
795            Some("workspace-123".to_string()),
796        )
797        .expect("noise configuration");
798
799        let bundle = config
800            .connect_provider()
801            .connect_bundle(harness_public_key)
802            .await
803            .expect("Noise connect bundle");
804
805        assert_eq!(
806            bundle.websocket_url,
807            "wss://rendezvous.test/cloud-agent/default/ws/environment/environment-requested?role=harness&sig=abc"
808        );
809        assert_eq!(bundle.environment_id, "environment-requested");
810        assert_eq!(bundle.executor_registration_id, "registration-1");
811        assert_eq!(bundle.executor_public_key, executor_public_key);
812        assert_eq!(bundle.harness_key_authorization, "authorization-1");
813    }
814
815    #[tokio::test]
816    async fn connect_environment_times_out_when_registry_stalls() {
817        let server = MockServer::start().await;
818        Mock::given(method("POST"))
819            .and(path("/cloud/environment/environment-requested/connect"))
820            .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(1)))
821            .mount(&server)
822            .await;
823        let mut client =
824            EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider())
825                .expect("client");
826        client.connect_timeout = Duration::from_millis(50);
827        let harness_public_key = NoiseChannelIdentity::generate()
828            .expect("identity")
829            .public_key();
830
831        let error = match client
832            .connect_environment("environment-requested", harness_public_key)
833            .await
834        {
835            Ok(_) => panic!("stalled connect response should time out"),
836            Err(error) => error,
837        };
838
839        assert!(matches!(
840            error,
841            ExecServerError::EnvironmentRegistryRequest(error) if error.is_timeout()
842        ));
843    }
844
845    #[tokio::test]
846    async fn register_environment_does_not_follow_redirects_with_auth_headers() {
847        let server = MockServer::start().await;
848        let executor_public_key = NoiseChannelIdentity::generate()
849            .expect("identity")
850            .public_key();
851        Mock::given(method("POST"))
852            .and(path("/cloud/environment/environment-requested/register"))
853            .and(header("authorization", "Bearer registry-token"))
854            .respond_with(
855                ResponseTemplate::new(302)
856                    .insert_header("location", format!("{}/redirect-target", server.uri())),
857            )
858            .mount(&server)
859            .await;
860        Mock::given(path("/redirect-target"))
861            .and(header("authorization", "Bearer registry-token"))
862            .respond_with(ResponseTemplate::new(200))
863            .expect(0)
864            .mount(&server)
865            .await;
866        let client = EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider())
867            .expect("client");
868
869        let error = client
870            .register_environment("environment-requested", &executor_public_key)
871            .await
872            .expect_err("redirect response should not be followed");
873
874        assert!(matches!(
875            error,
876            ExecServerError::EnvironmentRegistryHttp {
877                status: StatusCode::FOUND,
878                ..
879            }
880        ));
881    }
882
883    #[test]
884    fn debug_output_redacts_auth_provider() {
885        let config = RemoteEnvironmentConfig::new(
886            "https://registry.example".to_string(),
887            "env-1".to_string(),
888            static_registry_auth_provider(),
889        )
890        .expect("config");
891
892        let debug = format!("{config:?}");
893
894        assert!(debug.contains("<redacted>"));
895        assert!(!debug.contains("workspace-123"));
896    }
897}
898
899#[cfg(test)]
900#[path = "remote/noise_tests.rs"]
901mod noise_tests;
codex_exec_server/remote.rs

codex_exec_server/
remote.rs
