mirror of
https://github.com/openai/codex.git
synced 2026-09-11 20:36:49 +00:00
9e680a52e7
## What changed - Add `EnvironmentManager::from_accepted_websocket` so embedding hosts can construct a remote environment from an already accepted and authenticated Axum WebSocket. - Add `replace_accepted_websocket` to retire the current transport and resume the same exec-server session on a host-supplied replacement connection. - Serialize replacement handoffs, reject overlapping replacements, and release the handoff claim when a replacement attempt is cancelled or fails. ## Testing - Cover initial connection validation and immediate environment readiness. - Verify replacement retry behavior and recovery of a running process and its output after reconnecting. GitOrigin-RevId: 1f2ab7bcf7b5abbbece5c101801432dc84a8058d
65 lines
2.5 KiB
Rust
65 lines
2.5 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::sync::RwLock;
|
|
|
|
use super::Environment;
|
|
use super::EnvironmentManager;
|
|
use super::validate_environment_id;
|
|
use crate::ExecServerClient;
|
|
use crate::ExecServerClientConnectOptions;
|
|
use crate::ExecServerError;
|
|
use crate::client::LazyRemoteExecServerClient;
|
|
use axum::extract::ws::WebSocket;
|
|
use codex_http_client::HttpClientFactory;
|
|
|
|
impl EnvironmentManager {
|
|
/// Builds a manager around a WebSocket already accepted and authenticated by its host.
|
|
///
|
|
/// The manager owns client construction, session initialization, and later
|
|
/// recovery. The host only supplies the initial socket and authenticated
|
|
/// replacement sockets through [`Self::replace_accepted_websocket`].
|
|
pub async fn from_accepted_websocket(
|
|
environment_id: String,
|
|
websocket: WebSocket,
|
|
options: ExecServerClientConnectOptions,
|
|
http_client_factory: HttpClientFactory,
|
|
) -> Result<Self, ExecServerError> {
|
|
validate_environment_id(&environment_id)?;
|
|
let client = ExecServerClient::connect_accepted_websocket(websocket, options).await?;
|
|
let client =
|
|
LazyRemoteExecServerClient::from_connected(client, http_client_factory.clone());
|
|
let environment = Arc::new(Environment::remote_with_client(
|
|
client, /*local_runtime_paths*/ None,
|
|
));
|
|
Ok(Self {
|
|
default_environment: Some(environment_id.clone()),
|
|
environments: RwLock::new(HashMap::from([(environment_id, environment)])),
|
|
local_environment: None,
|
|
local_runtime_paths: None,
|
|
http_client_factory,
|
|
})
|
|
}
|
|
|
|
/// Hands a replacement WebSocket to an existing accepted environment.
|
|
/// Returns after handoff; recovery continues asynchronously.
|
|
pub async fn replace_accepted_websocket(
|
|
&self,
|
|
environment_id: &str,
|
|
websocket: WebSocket,
|
|
) -> Result<(), ExecServerError> {
|
|
let environment = self.get_environment(environment_id).ok_or_else(|| {
|
|
ExecServerError::Protocol(format!("environment `{environment_id}` is not configured"))
|
|
})?;
|
|
environment
|
|
.remote_client
|
|
.as_ref()
|
|
.ok_or_else(|| {
|
|
ExecServerError::Protocol(
|
|
"local environment does not have a replaceable exec-server client".to_string(),
|
|
)
|
|
})?
|
|
.replace_accepted_websocket(websocket)
|
|
.await
|
|
}
|
|
}
|