Add files using upload-large-folder tool

codex-rs/agent-roles/src/agent_role_config.rs

@@ -0,0 +1,209 @@
+use codex_config::config_toml::ConfigToml;
+use codex_utils_absolute_path::AbsolutePathBufGuard;
+use serde::Deserialize;
+use std::collections::BTreeSet;
+use std::path::Path;
+use std::path::PathBuf;
+use toml::Value as TomlValue;
+
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct AgentRoleConfig {
+    /// Human-facing role documentation used in spawn tool guidance.
+    /// Required for loaded user-defined roles after deprecated/new metadata precedence resolves.
+    pub description: Option<String>,
+    /// Path to a role-specific config layer.
+    pub config_file: Option<PathBuf>,
+    /// Candidate nicknames for agents spawned with this role.
+    pub nickname_candidates: Option<Vec<String>>,
+}
+
+#[derive(Deserialize, Debug, Clone, Default, PartialEq)]
+#[serde(deny_unknown_fields)]
+struct RawAgentRoleFileToml {
+    name: Option<String>,
+    description: Option<String>,
+    nickname_candidates: Option<Vec<String>>,
+    #[serde(flatten)]
+    config: ConfigToml,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct ResolvedAgentRoleFile {
+    pub role_name: String,
+    pub description: Option<String>,
+    pub nickname_candidates: Option<Vec<String>>,
+    pub config: TomlValue,
+}
+
+pub fn parse_agent_role_file_contents(
+    contents: &str,
+    role_file_label: &Path,
+    config_base_dir: &Path,
+    role_name_hint: Option<&str>,
+) -> std::io::Result<ResolvedAgentRoleFile> {
+    let role_file_toml: TomlValue = toml::from_str(contents).map_err(|err| {
+        std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            format!(
+                "failed to parse agent role file at {}: {err}",
+                role_file_label.display()
+            ),
+        )
+    })?;
+    let _guard = AbsolutePathBufGuard::new(config_base_dir);
+    let parsed: RawAgentRoleFileToml = role_file_toml.clone().try_into().map_err(|err| {
+        std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            format!(
+                "failed to deserialize agent role file at {}: {err}",
+                role_file_label.display()
+            ),
+        )
+    })?;
+    let description = normalize_agent_role_description(
+        &format!("agent role file {}.description", role_file_label.display()),
+        parsed.description.as_deref(),
+    )?;
+    validate_agent_role_file_developer_instructions(
+        role_file_label,
+        parsed.config.developer_instructions.as_deref(),
+        role_name_hint.is_none(),
+    )?;
+
+    let role_name = parsed
+        .name
+        .as_deref()
+        .map(str::trim)
+        .filter(|name| !name.is_empty())
+        .map(ToOwned::to_owned)
+        .or_else(|| role_name_hint.map(ToOwned::to_owned))
+        .ok_or_else(|| {
+            std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                format!(
+                    "agent role file at {} must define a non-empty `name`",
+                    role_file_label.display()
+                ),
+            )
+        })?;
+
+    let nickname_candidates = normalize_agent_role_nickname_candidates(
+        &format!(
+            "agent role file {}.nickname_candidates",
+            role_file_label.display()
+        ),
+        parsed.nickname_candidates.as_deref(),
+    )?;
+
+    let mut config = role_file_toml;
+    let Some(config_table) = config.as_table_mut() else {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            format!(
+                "agent role file at {} must contain a TOML table",
+                role_file_label.display()
+            ),
+        ));
+    };
+    config_table.remove("name");
+    config_table.remove("description");
+    config_table.remove("nickname_candidates");
+
+    Ok(ResolvedAgentRoleFile {
+        role_name,
+        description,
+        nickname_candidates,
+        config,
+    })
+}
+
+pub(crate) fn normalize_agent_role_description(
+    field_label: &str,
+    description: Option<&str>,
+) -> std::io::Result<Option<String>> {
+    match description.map(str::trim) {
+        Some("") => Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!("{field_label} cannot be blank"),
+        )),
+        Some(description) => Ok(Some(description.to_string())),
+        None => Ok(None),
+    }
+}
+
+fn validate_agent_role_file_developer_instructions(
+    role_file_label: &Path,
+    developer_instructions: Option<&str>,
+    require_present: bool,
+) -> std::io::Result<()> {
+    match developer_instructions.map(str::trim) {
+        Some("") => Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!(
+                "agent role file at {}.developer_instructions cannot be blank",
+                role_file_label.display()
+            ),
+        )),
+        Some(_) => Ok(()),
+        None if require_present => Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!(
+                "agent role file at {} must define `developer_instructions`",
+                role_file_label.display()
+            ),
+        )),
+        None => Ok(()),
+    }
+}
+
+pub(crate) fn normalize_agent_role_nickname_candidates(
+    field_label: &str,
+    nickname_candidates: Option<&[String]>,
+) -> std::io::Result<Option<Vec<String>>> {
+    let Some(nickname_candidates) = nickname_candidates else {
+        return Ok(None);
+    };
+
+    if nickname_candidates.is_empty() {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!("{field_label} must contain at least one name"),
+        ));
+    }
+
+    let mut normalized_candidates = Vec::with_capacity(nickname_candidates.len());
+    let mut seen_candidates = BTreeSet::new();
+
+    for nickname in nickname_candidates {
+        let normalized_nickname = nickname.trim();
+        if normalized_nickname.is_empty() {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                format!("{field_label} cannot contain blank names"),
+            ));
+        }
+
+        if !seen_candidates.insert(normalized_nickname.to_owned()) {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                format!("{field_label} cannot contain duplicates"),
+            ));
+        }
+
+        if !normalized_nickname
+            .chars()
+            .all(|c| c.is_ascii_alphanumeric() || matches!(c, ' ' | '-' | '_'))
+        {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                format!(
+                    "{field_label} may only contain ASCII letters, digits, spaces, hyphens, and underscores"
+                ),
+            ));
+        }
+
+        normalized_candidates.push(normalized_nickname.to_owned());
+    }
+
+    Ok(Some(normalized_candidates))
+}

codex-rs/agent-roles/src/discovery.rs

@@ -0,0 +1,40 @@
+use codex_file_system::ExecutorFileSystem;
+use codex_utils_absolute_path::AbsolutePathBuf;
+use codex_utils_path_uri::PathUri;
+use std::io;
+use std::io::ErrorKind;
+
+pub(crate) async fn collect_agent_role_files(
+    fs: &dyn ExecutorFileSystem,
+    dir: &AbsolutePathBuf,
+) -> io::Result<Vec<AbsolutePathBuf>> {
+    let mut files = Vec::new();
+    let mut dirs = vec![dir.clone()];
+    while let Some(dir) = dirs.pop() {
+        let dir_uri = PathUri::from_abs_path(&dir);
+        let entries = match fs.read_directory(&dir_uri, /*sandbox*/ None).await {
+            Ok(entries) => entries,
+            Err(err) if err.kind() == ErrorKind::NotFound => continue,
+            Err(err) => return Err(err),
+        };
+
+        for entry in entries {
+            let path = dir.join(entry.file_name);
+            if entry.is_directory {
+                dirs.push(path);
+                continue;
+            }
+            if entry.is_file
+                && path
+                    .as_path()
+                    .extension()
+                    .is_some_and(|extension| extension == "toml")
+            {
+                files.push(path);
+            }
+        }
+    }
+
+    files.sort();
+    Ok(files)
+}

codex-rs/agent-roles/src/lib.rs

@@ -0,0 +1,8 @@
+mod agent_role_config;
+mod discovery;
+mod loader;
+
+pub use agent_role_config::AgentRoleConfig;
+pub use agent_role_config::ResolvedAgentRoleFile;
+pub use agent_role_config::parse_agent_role_file_contents;
+pub use loader::load_agent_roles;

codex-rs/agent-roles/src/loader.rs

@@ -0,0 +1,335 @@
+use crate::AgentRoleConfig;
+use crate::ResolvedAgentRoleFile;
+use crate::agent_role_config::normalize_agent_role_description;
+use crate::agent_role_config::normalize_agent_role_nickname_candidates;
+use crate::discovery::collect_agent_role_files;
+use crate::parse_agent_role_file_contents;
+use codex_config::ConfigLayerStack;
+use codex_config::config_toml::AgentRoleToml;
+use codex_config::config_toml::AgentsToml;
+use codex_config::config_toml::ConfigToml;
+use codex_file_system::ExecutorFileSystem;
+use codex_file_system::GetMetadataOptions;
+use codex_file_system::ReadFileOptions;
+use codex_utils_absolute_path::AbsolutePathBuf;
+use codex_utils_absolute_path::AbsolutePathBufGuard;
+use codex_utils_path_uri::PathUri;
+use std::collections::BTreeMap;
+use std::collections::BTreeSet;
+use std::path::Path;
+use std::path::PathBuf;
+use toml::Value as TomlValue;
+
+pub async fn load_agent_roles(
+    fs: &dyn ExecutorFileSystem,
+    cfg: &ConfigToml,
+    config_layer_stack: &ConfigLayerStack,
+    startup_warnings: &mut Vec<String>,
+) -> std::io::Result<BTreeMap<String, AgentRoleConfig>> {
+    let mut layers = config_layer_stack.layers_low_to_high().peekable();
+    if layers.peek().is_none() {
+        return load_agent_roles_without_layers(fs, cfg).await;
+    }
+
+    let mut roles: BTreeMap<String, AgentRoleConfig> = BTreeMap::new();
+    for layer in layers {
+        let mut layer_roles: BTreeMap<String, AgentRoleConfig> = BTreeMap::new();
+        let mut declared_role_files = BTreeSet::new();
+        let config_folder = layer.config_folder();
+        let agents_toml = match agents_toml_from_layer(&layer.config, config_folder.as_deref()) {
+            Ok(agents_toml) => agents_toml,
+            Err(err) => {
+                push_agent_role_warning(startup_warnings, err);
+                None
+            }
+        };
+        if let Some(agents_toml) = agents_toml {
+            for (declared_role_name, role_toml) in &agents_toml.roles {
+                let (role_name, role) =
+                    match read_declared_role(fs, declared_role_name, role_toml).await {
+                        Ok(role) => role,
+                        Err(err) => {
+                            push_agent_role_warning(startup_warnings, err);
+                            continue;
+                        }
+                    };
+                if let Some(config_file) = role.config_file.clone() {
+                    declared_role_files.insert(config_file);
+                }
+                if layer_roles.contains_key(&role_name) {
+                    push_agent_role_warning(
+                        startup_warnings,
+                        std::io::Error::new(
+                            std::io::ErrorKind::InvalidInput,
+                            format!(
+                                "duplicate agent role name `{role_name}` declared in the same config layer"
+                            ),
+                        ),
+                    );
+                    continue;
+                }
+                layer_roles.insert(role_name, role);
+            }
+        }
+
+        if let Some(config_folder) = layer.config_folder() {
+            for (role_name, role) in discover_agent_roles_in_dir(
+                fs,
+                &config_folder.join("agents"),
+                &declared_role_files,
+                startup_warnings,
+            )
+            .await?
+            {
+                if layer_roles.contains_key(&role_name) {
+                    push_agent_role_warning(
+                        startup_warnings,
+                        std::io::Error::new(
+                            std::io::ErrorKind::InvalidInput,
+                            format!(
+                                "duplicate agent role name `{role_name}` declared in the same config layer"
+                            ),
+                        ),
+                    );
+                    continue;
+                }
+                layer_roles.insert(role_name, role);
+            }
+        }
+
+        for (role_name, role) in layer_roles {
+            let mut merged_role = role;
+            if let Some(existing_role) = roles.get(&role_name) {
+                merge_missing_role_fields(&mut merged_role, existing_role);
+            }
+            if let Err(err) = validate_required_agent_role_description(
+                &role_name,
+                merged_role.description.as_deref(),
+            ) {
+                push_agent_role_warning(startup_warnings, err);
+                continue;
+            }
+            roles.insert(role_name, merged_role);
+        }
+    }
+
+    Ok(roles)
+}
+
+fn push_agent_role_warning(startup_warnings: &mut Vec<String>, err: std::io::Error) {
+    let message = format!("Ignoring malformed agent role definition: {err}");
+    tracing::warn!("{message}");
+    startup_warnings.push(message);
+}
+
+async fn load_agent_roles_without_layers(
+    fs: &dyn ExecutorFileSystem,
+    cfg: &ConfigToml,
+) -> std::io::Result<BTreeMap<String, AgentRoleConfig>> {
+    let mut roles = BTreeMap::new();
+    if let Some(agents_toml) = cfg.agents.as_ref() {
+        for (declared_role_name, role_toml) in &agents_toml.roles {
+            let (role_name, role) = read_declared_role(fs, declared_role_name, role_toml).await?;
+            validate_required_agent_role_description(&role_name, role.description.as_deref())?;
+
+            if roles.insert(role_name.clone(), role).is_some() {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::InvalidInput,
+                    format!("duplicate agent role name `{role_name}` declared in config"),
+                ));
+            }
+        }
+    }
+
+    Ok(roles)
+}
+
+async fn read_declared_role(
+    fs: &dyn ExecutorFileSystem,
+    declared_role_name: &str,
+    role_toml: &AgentRoleToml,
+) -> std::io::Result<(String, AgentRoleConfig)> {
+    let mut role = agent_role_config_from_toml(fs, declared_role_name, role_toml).await?;
+    let mut role_name = declared_role_name.to_string();
+    if let Some(config_file) = role.config_file.as_deref() {
+        let config_file = AbsolutePathBuf::from_absolute_path(config_file)?;
+        let parsed_file =
+            read_resolved_agent_role_file(fs, &config_file, Some(declared_role_name)).await?;
+        role_name = parsed_file.role_name;
+        role.description = parsed_file.description.or(role.description);
+        role.nickname_candidates = parsed_file.nickname_candidates.or(role.nickname_candidates);
+    }
+
+    Ok((role_name, role))
+}
+
+fn merge_missing_role_fields(role: &mut AgentRoleConfig, fallback: &AgentRoleConfig) {
+    role.description = role.description.clone().or(fallback.description.clone());
+    role.config_file = role.config_file.clone().or(fallback.config_file.clone());
+    role.nickname_candidates = role
+        .nickname_candidates
+        .clone()
+        .or(fallback.nickname_candidates.clone());
+}
+
+fn agents_toml_from_layer(
+    layer_toml: &TomlValue,
+    config_base_dir: Option<&Path>,
+) -> std::io::Result<Option<AgentsToml>> {
+    let Some(agents_toml) = layer_toml.get("agents") else {
+        return Ok(None);
+    };
+
+    // AbsolutePathBufGuard resolves relative paths while it remains in scope.
+    let _guard = config_base_dir.map(AbsolutePathBufGuard::new);
+    agents_toml
+        .clone()
+        .try_into()
+        .map(Some)
+        .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))
+}
+
+async fn agent_role_config_from_toml(
+    fs: &dyn ExecutorFileSystem,
+    role_name: &str,
+    role: &AgentRoleToml,
+) -> std::io::Result<AgentRoleConfig> {
+    let config_file = role
+        .config_file
+        .as_ref()
+        .map(AbsolutePathBuf::from_absolute_path)
+        .transpose()?;
+    validate_agent_role_config_file(fs, role_name, config_file.as_ref()).await?;
+    let description = normalize_agent_role_description(
+        &format!("agents.{role_name}.description"),
+        role.description.as_deref(),
+    )?;
+    let nickname_candidates = normalize_agent_role_nickname_candidates(
+        &format!("agents.{role_name}.nickname_candidates"),
+        role.nickname_candidates.as_deref(),
+    )?;
+
+    Ok(AgentRoleConfig {
+        description,
+        config_file: config_file.map(AbsolutePathBuf::into_path_buf),
+        nickname_candidates,
+    })
+}
+
+async fn read_resolved_agent_role_file(
+    fs: &dyn ExecutorFileSystem,
+    path: &AbsolutePathBuf,
+    role_name_hint: Option<&str>,
+) -> std::io::Result<ResolvedAgentRoleFile> {
+    let path_uri = PathUri::from_abs_path(path);
+    let contents = fs
+        .read_file_text(&path_uri, ReadFileOptions::default(), /*sandbox*/ None)
+        .await?;
+    let config_base_dir = path.parent().unwrap_or_else(|| path.clone());
+    parse_agent_role_file_contents(
+        &contents,
+        path.as_path(),
+        config_base_dir.as_path(),
+        role_name_hint,
+    )
+}
+
+fn validate_required_agent_role_description(
+    role_name: &str,
+    description: Option<&str>,
+) -> std::io::Result<()> {
+    if description.is_some() {
+        Ok(())
+    } else {
+        Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!("agent role `{role_name}` must define a description"),
+        ))
+    }
+}
+
+async fn validate_agent_role_config_file(
+    fs: &dyn ExecutorFileSystem,
+    role_name: &str,
+    config_file: Option<&AbsolutePathBuf>,
+) -> std::io::Result<()> {
+    let Some(config_file) = config_file else {
+        return Ok(());
+    };
+
+    let config_file_uri = PathUri::from_abs_path(config_file);
+    let metadata = fs
+        .get_metadata(
+            &config_file_uri,
+            GetMetadataOptions::default(),
+            /*sandbox*/ None,
+        )
+        .await
+        .map_err(|e| {
+            std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                format!(
+                    "agents.{role_name}.config_file must point to an existing file at {}: {e}",
+                    config_file.as_path().display()
+                ),
+            )
+        })?;
+    if metadata.is_file {
+        Ok(())
+    } else {
+        Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!(
+                "agents.{role_name}.config_file must point to a file: {}",
+                config_file.as_path().display()
+            ),
+        ))
+    }
+}
+
+async fn discover_agent_roles_in_dir(
+    fs: &dyn ExecutorFileSystem,
+    agents_dir: &AbsolutePathBuf,
+    declared_role_files: &BTreeSet<PathBuf>,
+    startup_warnings: &mut Vec<String>,
+) -> std::io::Result<BTreeMap<String, AgentRoleConfig>> {
+    let mut roles = BTreeMap::new();
+
+    for agent_file in collect_agent_role_files(fs, agents_dir).await? {
+        if declared_role_files.contains(agent_file.as_path()) {
+            continue;
+        }
+        let parsed_file =
+            match read_resolved_agent_role_file(fs, &agent_file, /*role_name_hint*/ None).await {
+                Ok(parsed_file) => parsed_file,
+                Err(err) => {
+                    push_agent_role_warning(startup_warnings, err);
+                    continue;
+                }
+            };
+        let role_name = parsed_file.role_name;
+        if roles.contains_key(&role_name) {
+            push_agent_role_warning(
+                startup_warnings,
+                std::io::Error::new(
+                    std::io::ErrorKind::InvalidInput,
+                    format!(
+                        "duplicate agent role name `{role_name}` discovered in {}",
+                        agents_dir.as_path().display()
+                    ),
+                ),
+            );
+            continue;
+        }
+        roles.insert(
+            role_name,
+            AgentRoleConfig {
+                description: parsed_file.description,
+                config_file: Some(agent_file.to_path_buf()),
+                nickname_candidates: parsed_file.nickname_candidates,
+            },
+        );
+    }
+
+    Ok(roles)
+}

codex-rs/app-server-protocol-noop-macros/src/lib.rs

@@ -0,0 +1,20 @@
+//! No-op schema derives for production app-server protocol builds.
+//!
+//! The real `ts-rs` and `schemars` derives are only needed when regenerating
+//! the vendored protocol exports. Normal builds retain the annotations so the
+//! protocol definitions stay readable, but use these derives to avoid
+//! generating implementations that cannot be reached at runtime.
+
+use proc_macro::TokenStream;
+
+/// Accepts `#[schemars(...)]` helper attributes without generating an impl.
+#[proc_macro_derive(JsonSchema, attributes(schemars))]
+pub fn derive_json_schema(_input: TokenStream) -> TokenStream {
+    TokenStream::new()
+}
+
+/// Accepts `#[ts(...)]` helper attributes without generating an impl.
+#[proc_macro_derive(TS, attributes(ts))]
+pub fn derive_ts(_input: TokenStream) -> TokenStream {
+    TokenStream::new()
+}

codex-rs/cloud-tasks-mock-client/src/lib.rs

@@ -0,0 +1,3 @@
+mod mock;
+
+pub use mock::MockClient;

codex-rs/cloud-tasks-mock-client/src/mock.rs

@@ -0,0 +1,267 @@
+use chrono::Utc;
+use codex_cloud_tasks_client::ApplyOutcome;
+use codex_cloud_tasks_client::ApplyStatus;
+use codex_cloud_tasks_client::AttemptStatus;
+use codex_cloud_tasks_client::CloudBackend;
+use codex_cloud_tasks_client::CloudBackendFuture;
+use codex_cloud_tasks_client::CloudTaskError;
+use codex_cloud_tasks_client::CreatedTask;
+use codex_cloud_tasks_client::DiffSummary;
+use codex_cloud_tasks_client::Result;
+use codex_cloud_tasks_client::TaskId;
+use codex_cloud_tasks_client::TaskListPage;
+use codex_cloud_tasks_client::TaskStatus;
+use codex_cloud_tasks_client::TaskSummary;
+use codex_cloud_tasks_client::TaskText;
+use codex_cloud_tasks_client::TurnAttempt;
+
+#[derive(Clone, Default)]
+pub struct MockClient;
+
+impl MockClient {
+    async fn list_tasks(
+        &self,
+        _env: Option<&str>,
+        _limit: Option<i64>,
+        _cursor: Option<&str>,
+    ) -> Result<TaskListPage> {
+        // Slightly vary content by env to aid tests that rely on the mock
+        let rows = match _env {
+            Some("env-A") => vec![("T-2000", "A: First", TaskStatus::Ready)],
+            Some("env-B") => vec![
+                ("T-3000", "B: One", TaskStatus::Ready),
+                ("T-3001", "B: Two", TaskStatus::Pending),
+            ],
+            _ => vec![
+                ("T-1000", "Update README formatting", TaskStatus::Ready),
+                ("T-1001", "Fix clippy warnings in core", TaskStatus::Pending),
+                ("T-1002", "Add contributing guide", TaskStatus::Ready),
+            ],
+        };
+        let environment_id = _env.map(str::to_string);
+        let environment_label = match _env {
+            Some("env-A") => Some("Env A".to_string()),
+            Some("env-B") => Some("Env B".to_string()),
+            Some(other) => Some(other.to_string()),
+            None => Some("Global".to_string()),
+        };
+        let mut out = Vec::new();
+        for (id_str, title, status) in rows {
+            let id = TaskId(id_str.to_string());
+            let diff = mock_diff_for(&id);
+            let (a, d) = count_from_unified(&diff);
+            out.push(TaskSummary {
+                id,
+                title: title.to_string(),
+                status,
+                updated_at: Utc::now(),
+                environment_id: environment_id.clone(),
+                environment_label: environment_label.clone(),
+                summary: DiffSummary {
+                    files_changed: 1,
+                    lines_added: a,
+                    lines_removed: d,
+                },
+                is_review: false,
+                attempt_total: Some(if id_str == "T-1000" { 2 } else { 1 }),
+            });
+        }
+        Ok(TaskListPage {
+            tasks: out,
+            cursor: None,
+        })
+    }
+
+    async fn get_task_summary(&self, id: TaskId) -> Result<TaskSummary> {
+        let tasks = self
+            .list_tasks(/*env*/ None, /*limit*/ None, /*cursor*/ None)
+            .await?
+            .tasks;
+        tasks
+            .into_iter()
+            .find(|t| t.id == id)
+            .ok_or_else(|| CloudTaskError::Msg(format!("Task {} not found (mock)", id.0)))
+    }
+
+    async fn get_task_diff(&self, id: TaskId) -> Result<Option<String>> {
+        Ok(Some(mock_diff_for(&id)))
+    }
+
+    async fn get_task_messages(&self, _id: TaskId) -> Result<Vec<String>> {
+        Ok(vec![
+            "Mock assistant output: this task contains no diff.".to_string(),
+        ])
+    }
+
+    async fn get_task_text(&self, _id: TaskId) -> Result<TaskText> {
+        Ok(TaskText {
+            prompt: Some("Why is there no diff?".to_string()),
+            messages: vec!["Mock assistant output: this task contains no diff.".to_string()],
+            turn_id: Some("mock-turn".to_string()),
+            sibling_turn_ids: Vec::new(),
+            attempt_placement: Some(0),
+            attempt_status: AttemptStatus::Completed,
+        })
+    }
+
+    async fn apply_task(&self, id: TaskId, _diff_override: Option<String>) -> Result<ApplyOutcome> {
+        Ok(ApplyOutcome {
+            applied: true,
+            status: ApplyStatus::Success,
+            message: format!("Applied task {} locally (mock)", id.0),
+            skipped_paths: Vec::new(),
+            conflict_paths: Vec::new(),
+        })
+    }
+
+    async fn apply_task_preflight(
+        &self,
+        id: TaskId,
+        _diff_override: Option<String>,
+    ) -> Result<ApplyOutcome> {
+        Ok(ApplyOutcome {
+            applied: false,
+            status: ApplyStatus::Success,
+            message: format!("Preflight passed for task {} (mock)", id.0),
+            skipped_paths: Vec::new(),
+            conflict_paths: Vec::new(),
+        })
+    }
+
+    async fn list_sibling_attempts(
+        &self,
+        task: TaskId,
+        _turn_id: String,
+    ) -> Result<Vec<TurnAttempt>> {
+        if task.0 == "T-1000" {
+            return Ok(vec![TurnAttempt {
+                turn_id: "T-1000-attempt-2".to_string(),
+                attempt_placement: Some(1),
+                created_at: Some(Utc::now()),
+                status: AttemptStatus::Completed,
+                diff: Some(mock_diff_for(&task)),
+                messages: vec!["Mock alternate attempt".to_string()],
+            }]);
+        }
+        Ok(Vec::new())
+    }
+
+    async fn create_task(
+        &self,
+        env_id: &str,
+        prompt: &str,
+        git_ref: &str,
+        qa_mode: bool,
+        best_of_n: usize,
+    ) -> Result<CreatedTask> {
+        let _ = (env_id, prompt, git_ref, qa_mode, best_of_n);
+        let id = format!("task_local_{}", chrono::Utc::now().timestamp_millis());
+        Ok(CreatedTask { id: TaskId(id) })
+    }
+}
+
+impl CloudBackend for MockClient {
+    fn list_tasks<'a>(
+        &'a self,
+        env: Option<&'a str>,
+        limit: Option<i64>,
+        cursor: Option<&'a str>,
+    ) -> CloudBackendFuture<'a, TaskListPage> {
+        Box::pin(MockClient::list_tasks(self, env, limit, cursor))
+    }
+
+    fn get_task_summary(&self, id: TaskId) -> CloudBackendFuture<'_, TaskSummary> {
+        Box::pin(MockClient::get_task_summary(self, id))
+    }
+
+    fn get_task_diff(&self, id: TaskId) -> CloudBackendFuture<'_, Option<String>> {
+        Box::pin(MockClient::get_task_diff(self, id))
+    }
+
+    fn get_task_messages(&self, id: TaskId) -> CloudBackendFuture<'_, Vec<String>> {
+        Box::pin(MockClient::get_task_messages(self, id))
+    }
+
+    fn get_task_text(&self, id: TaskId) -> CloudBackendFuture<'_, TaskText> {
+        Box::pin(MockClient::get_task_text(self, id))
+    }
+
+    fn apply_task(
+        &self,
+        id: TaskId,
+        diff_override: Option<String>,
+    ) -> CloudBackendFuture<'_, ApplyOutcome> {
+        Box::pin(MockClient::apply_task(self, id, diff_override))
+    }
+
+    fn apply_task_preflight(
+        &self,
+        id: TaskId,
+        diff_override: Option<String>,
+    ) -> CloudBackendFuture<'_, ApplyOutcome> {
+        Box::pin(MockClient::apply_task_preflight(self, id, diff_override))
+    }
+
+    fn list_sibling_attempts(
+        &self,
+        task: TaskId,
+        turn_id: String,
+    ) -> CloudBackendFuture<'_, Vec<TurnAttempt>> {
+        Box::pin(MockClient::list_sibling_attempts(self, task, turn_id))
+    }
+
+    fn create_task<'a>(
+        &'a self,
+        env_id: &'a str,
+        prompt: &'a str,
+        git_ref: &'a str,
+        qa_mode: bool,
+        best_of_n: usize,
+    ) -> CloudBackendFuture<'a, CreatedTask> {
+        Box::pin(MockClient::create_task(
+            self, env_id, prompt, git_ref, qa_mode, best_of_n,
+        ))
+    }
+}
+
+fn mock_diff_for(id: &TaskId) -> String {
+    match id.0.as_str() {
+        "T-1000" => {
+            "diff --git a/README.md b/README.md\nindex 000000..111111 100644\n--- a/README.md\n+++ b/README.md\n@@ -1,2 +1,3 @@\n Intro\n-Hello\n+Hello, world!\n+Task: T-1000\n".to_string()
+        }
+        "T-1001" => {
+            "diff --git a/core/src/lib.rs b/core/src/lib.rs\nindex 000000..111111 100644\n--- a/core/src/lib.rs\n+++ b/core/src/lib.rs\n@@ -1,2 +1,1 @@\n-use foo;\n use bar;\n".to_string()
+        }
+        _ => {
+            "diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md\nindex 000000..111111 100644\n--- /dev/null\n+++ b/CONTRIBUTING.md\n@@ -0,0 +1,3 @@\n+## Contributing\n+Please open PRs.\n+Thanks!\n".to_string()
+        }
+    }
+}
+
+fn count_from_unified(diff: &str) -> (usize, usize) {
+    if let Ok(patch) = diffy::Patch::from_str(diff) {
+        patch
+            .hunks()
+            .iter()
+            .flat_map(diffy::Hunk::lines)
+            .fold((0, 0), |(a, d), l| match l {
+                diffy::Line::Insert(_) => (a + 1, d),
+                diffy::Line::Delete(_) => (a, d + 1),
+                _ => (a, d),
+            })
+    } else {
+        let mut a = 0;
+        let mut d = 0;
+        for l in diff.lines() {
+            if l.starts_with("+++") || l.starts_with("---") || l.starts_with("@@") {
+                continue;
+            }
+            match l.as_bytes().first() {
+                Some(b'+') => a += 1,
+                Some(b'-') => d += 1,
+                _ => {}
+            }
+        }
+        (a, d)
+    }
+}

codex-rs/code-mode-protocol/src/description.rs

@@ -0,0 +1,993 @@
+use codex_protocol::ToolName;
+use serde::Deserialize;
+use serde::Serialize;
+use serde_json::Value as JsonValue;
+use std::collections::BTreeMap;
+
+use crate::PUBLIC_TOOL_NAME;
+use crate::json_schema_types::render_json_schema_to_typescript;
+
+const MAX_JS_SAFE_INTEGER: u64 = (1_u64 << 53) - 1;
+const DEFERRED_NESTED_TOOLS_GUIDANCE: &str = r#"Some deferred nested tools may be omitted from this description. They are still available on the global `tools` object and listed in `ALL_TOOLS`.
+To find one, filter `ALL_TOOLS` by `name` and `description`."#;
+const LEGACY_IMAGE_HELPER_DESCRIPTION: &str = r#"`image(imageUrlOrItem: string | { image_url: string; detail?: "auto" | "low" | "high" | "original" | null } | ImageContent, detail?: "auto" | "low" | "high" | "original" | null)`: Appends an image item. `image_url` should be a base64-encoded `data:` URL. To forward an MCP tool image, pass an individual `ImageContent` block from `result.content`, for example `image(result.content[0])`. MCP image blocks may request detail with `_meta: { "codex/imageDetail": "original" }`. When provided, the second `detail` argument overrides any detail embedded in the first argument."#;
+const UNIFIED_IMAGE_HELPER_DESCRIPTION: &str = r#"`image(imageUrlOrItem: string | { image_url: string } | ImageContent)`: Appends an image item. `image_url` should be a base64-encoded `data:` URL. To forward an MCP tool image, pass an individual `ImageContent` block from `result.content`, for example `image(result.content[0])`."#;
+const EXEC_DESCRIPTION_TEMPLATE: &str = r#"Run JavaScript code to orchestrate/compose tool calls
+- Evaluates the provided JavaScript code in a fresh V8 isolate as an async module.
+- All nested tools are available on the global `tools` object, for example `await tools.exec_command(...)`. Tool names are exposed as normalized JavaScript identifiers, for example `await tools.mcp__ologs__get_profile(...)`.
+- Nested tool methods take either a string or an object as their input argument.
+- Nested tools return either an object or a string, based on the description.
+- Runs raw JavaScript -- no Node, no file system, no network access, no console.
+- Accepts raw JavaScript source text, not JSON, quoted strings, or markdown code fences.
+- You may optionally start the tool input with a first-line pragma like `// @exec: {"yield_time_ms": 10000, "max_output_tokens": 1000}`.
+- `yield_time_ms` asks `exec` to yield early if the script is still running. Defaults to 10000 ms.
+- `max_output_tokens` sets the token budget for direct `exec` results. Defaults to 10000 tokens.
+- When the JS code is fully evaluated, the isolate's lifetime ends and unawaited promises are silently discarded.
+
+- Global helpers:
+- `exit()`: Immediately ends the current script successfully (like an early return from the top level).
+- `text(value: string | number | boolean | undefined | null)`: Appends a text item. Non-string values are stringified with `JSON.stringify(...)` when possible.
+- `image(imageUrlOrItem: string | { image_url: string; detail?: "auto" | "low" | "high" | "original" | null } | ImageContent, detail?: "auto" | "low" | "high" | "original" | null)`: Appends an image item. `image_url` should be a base64-encoded `data:` URL. To forward an MCP tool image, pass an individual `ImageContent` block from `result.content`, for example `image(result.content[0])`. MCP image blocks may request detail with `_meta: { "codex/imageDetail": "original" }`. When provided, the second `detail` argument overrides any detail embedded in the first argument.
+- `audio(audioUrlOrItem: string | { audio_url: string } | AudioContent)`: Appends an audio item. `audio_url` should be a base64-encoded `data:` URL. To forward an MCP tool audio block, pass an individual `AudioContent` block from `result.content`, for example `audio(result.content[0])`.
+- `generatedImage(result: { image_url: string; output_hint?: string })`: Appends an image-generation result and its optional output hint. HTTP(S) URLs are not supported.
+- `store(key: string, value: any)`: stores a serializable value under a string key for later `exec` calls in the same session.
+- `load(key: string)`: returns the stored value for a string key, or `undefined` if it is missing.
+- `notify(value: string | number | boolean | undefined | null)`: immediately injects an extra `custom_tool_call_output` for the current `exec` call. Values are stringified like `text(...)`.
+- `setTimeout(callback: () => void, delayMs?: number)`: schedules a callback to run later and returns a timeout id. Pending timeouts do not keep `exec` alive by themselves; await an explicit promise if you need to wait for one.
+- `clearTimeout(timeoutId?: number)`: cancels a timeout created by `setTimeout`.
+- `ALL_TOOLS`: metadata for the enabled nested tools as `{ name, description }` entries.
+- `yield_control()`: yields the accumulated output to the model immediately while the script keeps running."#;
+const WAIT_DESCRIPTION_TEMPLATE: &str = r#"- Use `wait` only after `exec` returns `Script running with cell ID ...`.
+- `cell_id` identifies the running `exec` cell to resume.
+- `yield_time_ms` controls how long to wait for more output before yielding again. Defaults to 10000 ms.
+- `max_tokens` limits how much new output this wait call returns. Defaults to 10000 tokens.
+- `terminate: true` stops the running cell; false or omitted waits for output.
+- `wait` returns only the new output since the last yield, or the final completion or termination result for that cell.
+- If the cell is still running, `wait` may yield again with the same `cell_id`.
+- If the cell has already finished, `wait` returns the completed result and closes the cell."#;
+// Based off of https://modelcontextprotocol.io/specification/draft/schema#calltoolresult
+const MCP_TYPESCRIPT_PREAMBLE: &str = r#"type Role = "user" | "assistant";
+type MetaObject = Record<string, unknown>;
+type Annotations = {
+  audience?: Role[];
+  priority?: number;
+  lastModified?: string;
+};
+type Icon = {
+  src: string;
+  mimeType?: string;
+  sizes?: string[];
+  theme?: "light" | "dark";
+};
+type TextResourceContents = {
+  uri: string;
+  mimeType?: string;
+  _meta?: MetaObject;
+  text: string;
+};
+type BlobResourceContents = {
+  uri: string;
+  mimeType?: string;
+  _meta?: MetaObject;
+  blob: string;
+};
+type TextContent = {
+  type: "text";
+  text: string;
+  annotations?: Annotations;
+  _meta?: MetaObject;
+};
+type ImageContent = {
+  type: "image";
+  data: string;
+  mimeType: string;
+  annotations?: Annotations;
+  _meta?: MetaObject;
+};
+type AudioContent = {
+  type: "audio";
+  data: string;
+  mimeType: string;
+  annotations?: Annotations;
+  _meta?: MetaObject;
+};
+type ResourceLink = {
+  icons?: Icon[];
+  name: string;
+  title?: string;
+  uri: string;
+  description?: string;
+  mimeType?: string;
+  annotations?: Annotations;
+  size?: number;
+  _meta?: MetaObject;
+  type: "resource_link";
+};
+type EmbeddedResource = {
+  type: "resource";
+  resource: TextResourceContents | BlobResourceContents;
+  annotations?: Annotations;
+  _meta?: MetaObject;
+};
+type ContentBlock =
+  | TextContent
+  | ImageContent
+  | AudioContent
+  | ResourceLink
+  | EmbeddedResource;
+type CallToolResult<TStructured = { [key: string]: unknown }> = {
+  _meta?: MetaObject;
+  content: ContentBlock[];
+  isError?: boolean;
+  structuredContent?: TStructured;
+  [key: string]: unknown;
+};"#;
+
+pub const CODE_MODE_PRAGMA_PREFIX: &str = "// @exec:";
+
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum CodeModeToolKind {
+    Function,
+    Freeform,
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+pub struct ToolDefinition {
+    pub name: String,
+    pub tool_name: ToolName,
+    pub description: String,
+    pub kind: CodeModeToolKind,
+    pub input_schema: Option<JsonValue>,
+    pub output_schema: Option<JsonValue>,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct ToolNamespaceDescription {
+    pub name: String,
+    pub description: String,
+}
+
+#[derive(Debug, Default, Deserialize, PartialEq, Eq)]
+#[serde(deny_unknown_fields)]
+struct CodeModeExecPragma {
+    #[serde(default)]
+    yield_time_ms: Option<u64>,
+    #[serde(default)]
+    max_output_tokens: Option<usize>,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub struct ParsedExecSource {
+    pub code: String,
+    pub yield_time_ms: Option<u64>,
+    pub max_output_tokens: Option<usize>,
+}
+
+pub fn parse_exec_source(input: &str) -> Result<ParsedExecSource, String> {
+    if input.trim().is_empty() {
+        return Err(
+            "exec expects raw JavaScript source text (non-empty). Provide JS only, optionally with first-line `// @exec: {\"yield_time_ms\": 10000, \"max_output_tokens\": 1000}`.".to_string(),
+        );
+    }
+
+    let mut args = ParsedExecSource {
+        code: input.to_string(),
+        yield_time_ms: None,
+        max_output_tokens: None,
+    };
+
+    let mut lines = input.splitn(2, '\n');
+    let first_line = lines.next().unwrap_or_default();
+    let rest = lines.next().unwrap_or_default();
+    let trimmed = first_line.trim_start();
+    let Some(pragma) = trimmed.strip_prefix(CODE_MODE_PRAGMA_PREFIX) else {
+        return Ok(args);
+    };
+
+    if rest.trim().is_empty() {
+        return Err(
+            "exec pragma must be followed by JavaScript source on subsequent lines".to_string(),
+        );
+    }
+
+    let directive = pragma.trim();
+    if directive.is_empty() {
+        return Err(
+            "exec pragma must be a JSON object with supported fields `yield_time_ms` and `max_output_tokens`"
+                .to_string(),
+        );
+    }
+
+    let value: serde_json::Value = serde_json::from_str(directive).map_err(|err| {
+        format!(
+            "exec pragma must be valid JSON with supported fields `yield_time_ms` and `max_output_tokens`: {err}"
+        )
+    })?;
+    let object = value.as_object().ok_or_else(|| {
+        "exec pragma must be a JSON object with supported fields `yield_time_ms` and `max_output_tokens`"
+            .to_string()
+    })?;
+    for key in object.keys() {
+        match key.as_str() {
+            "yield_time_ms" | "max_output_tokens" => {}
+            _ => {
+                return Err(format!(
+                    "exec pragma only supports `yield_time_ms` and `max_output_tokens`; got `{key}`"
+                ));
+            }
+        }
+    }
+
+    let pragma: CodeModeExecPragma = serde_json::from_value(value).map_err(|err| {
+        format!(
+            "exec pragma fields `yield_time_ms` and `max_output_tokens` must be non-negative safe integers: {err}"
+        )
+    })?;
+    if pragma
+        .yield_time_ms
+        .is_some_and(|yield_time_ms| yield_time_ms > MAX_JS_SAFE_INTEGER)
+    {
+        return Err(
+            "exec pragma field `yield_time_ms` must be a non-negative safe integer".to_string(),
+        );
+    }
+    if pragma.max_output_tokens.is_some_and(|max_output_tokens| {
+        u64::try_from(max_output_tokens)
+            .map(|max_output_tokens| max_output_tokens > MAX_JS_SAFE_INTEGER)
+            .unwrap_or(true)
+    }) {
+        return Err(
+            "exec pragma field `max_output_tokens` must be a non-negative safe integer".to_string(),
+        );
+    }
+
+    args.code = rest.to_string();
+    args.yield_time_ms = pragma.yield_time_ms;
+    args.max_output_tokens = pragma.max_output_tokens;
+    Ok(args)
+}
+
+pub fn is_code_mode_nested_tool(tool_name: &str) -> bool {
+    tool_name != crate::PUBLIC_TOOL_NAME && tool_name != crate::WAIT_TOOL_NAME
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum ImageDetailVisibility {
+    Visible,
+    Hidden,
+}
+
+pub fn build_exec_tool_description(
+    enabled_tools: &[ToolDefinition],
+    deferred_tools: &[ToolDefinition],
+    namespace_descriptions: &BTreeMap<String, ToolNamespaceDescription>,
+    default_exec_yield_time_ms: u64,
+    code_mode_only: bool,
+    image_detail_visibility: ImageDetailVisibility,
+) -> String {
+    let mut sections = Vec::new();
+    sections.push(EXEC_DESCRIPTION_TEMPLATE.replace(
+        "Defaults to 10000 ms.",
+        &format!("Defaults to {default_exec_yield_time_ms} ms."),
+    ));
+    if image_detail_visibility == ImageDetailVisibility::Hidden {
+        sections[0] = sections[0].replace(
+            LEGACY_IMAGE_HELPER_DESCRIPTION,
+            UNIFIED_IMAGE_HELPER_DESCRIPTION,
+        );
+    }
+    if !deferred_tools.is_empty() {
+        sections.push(DEFERRED_NESTED_TOOLS_GUIDANCE.to_string());
+    }
+    if !code_mode_only {
+        return sections.join("\n\n");
+    }
+
+    let has_mcp_tools = enabled_tools
+        .iter()
+        .chain(deferred_tools)
+        .any(|tool| mcp_structured_content_schema(tool.output_schema.as_ref()).is_some());
+    if has_mcp_tools {
+        sections.push(format!(
+            "Shared MCP Types:\n```ts\n{MCP_TYPESCRIPT_PREAMBLE}\n```"
+        ));
+    }
+
+    if !enabled_tools.is_empty() {
+        let mut current_namespace: Option<&str> = None;
+        let mut nested_tool_sections = Vec::with_capacity(enabled_tools.len());
+
+        for tool in enabled_tools {
+            let name = tool.name.as_str();
+            let nested_description = render_code_mode_sample_for_definition(tool);
+            let namespace_description = tool
+                .tool_name
+                .namespace
+                .as_ref()
+                .and_then(|namespace| namespace_descriptions.get(namespace));
+            let next_namespace = namespace_description
+                .map(|namespace_description| namespace_description.name.as_str());
+            if next_namespace != current_namespace {
+                if let Some(namespace_description) = namespace_description {
+                    let namespace_description_text = namespace_description.description.trim();
+                    if !namespace_description_text.is_empty() {
+                        nested_tool_sections.push(format!(
+                            "## {}\n{namespace_description_text}",
+                            namespace_description.name
+                        ));
+                    }
+                }
+                current_namespace = next_namespace;
+            }
+
+            let global_name = normalize_code_mode_identifier(name);
+            let nested_description = nested_description.trim();
+            if nested_description.is_empty() {
+                nested_tool_sections.push(render_tool_heading(&global_name, name));
+            } else {
+                nested_tool_sections.push(format!(
+                    "{}\n{nested_description}",
+                    render_tool_heading(&global_name, name)
+                ));
+            }
+        }
+
+        sections.push(nested_tool_sections.join("\n\n"));
+    }
+
+    sections.join("\n\n")
+}
+
+pub fn build_wait_tool_description() -> &'static str {
+    WAIT_DESCRIPTION_TEMPLATE
+}
+
+pub fn normalize_code_mode_identifier(tool_key: &str) -> String {
+    let mut identifier = String::new();
+
+    for (index, ch) in tool_key.chars().enumerate() {
+        let is_valid = if index == 0 {
+            ch == '_' || ch == '$' || ch.is_ascii_alphabetic()
+        } else {
+            ch == '_' || ch == '$' || ch.is_ascii_alphanumeric()
+        };
+
+        if is_valid {
+            identifier.push(ch);
+        } else {
+            identifier.push('_');
+        }
+    }
+
+    if identifier.is_empty() {
+        "_".to_string()
+    } else {
+        identifier
+    }
+}
+
+pub fn augment_tool_definition(mut definition: ToolDefinition) -> ToolDefinition {
+    if definition.name != PUBLIC_TOOL_NAME {
+        definition.description = render_code_mode_sample_for_definition(&definition);
+    }
+    definition
+}
+
+pub fn enabled_tool_metadata(definition: &ToolDefinition) -> EnabledToolMetadata {
+    EnabledToolMetadata {
+        tool_name: definition.tool_name.clone(),
+        global_name: normalize_code_mode_identifier(&definition.name),
+        description: definition.description.clone(),
+        kind: definition.kind,
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
+pub struct EnabledToolMetadata {
+    pub tool_name: ToolName,
+    pub global_name: String,
+    pub description: String,
+    pub kind: CodeModeToolKind,
+}
+
+pub fn render_code_mode_sample(
+    description: &str,
+    tool_name: &str,
+    input_name: &str,
+    input_type: String,
+    output_type: String,
+) -> String {
+    let declaration = format!(
+        "declare const tools: {{ {} }};",
+        render_code_mode_tool_declaration(tool_name, input_name, &input_type, &output_type)
+    );
+    format!("{description}\n\nexec tool declaration:\n```ts\n{declaration}\n```")
+}
+
+fn render_code_mode_sample_for_definition(definition: &ToolDefinition) -> String {
+    let input_name = match definition.kind {
+        CodeModeToolKind::Function => "args",
+        CodeModeToolKind::Freeform => "input",
+    };
+    let input_type = match definition.kind {
+        CodeModeToolKind::Function => definition
+            .input_schema
+            .as_ref()
+            .map(render_json_schema_to_typescript)
+            .unwrap_or_else(|| "unknown".to_string()),
+        CodeModeToolKind::Freeform => "string".to_string(),
+    };
+    let output_type = if let Some(structured_content_schema) =
+        mcp_structured_content_schema(definition.output_schema.as_ref())
+    {
+        let structured_content_type = render_json_schema_to_typescript(structured_content_schema);
+        if structured_content_type == "unknown" {
+            "CallToolResult".to_string()
+        } else {
+            format!("CallToolResult<{structured_content_type}>")
+        }
+    } else {
+        definition
+            .output_schema
+            .as_ref()
+            .map(render_json_schema_to_typescript)
+            .unwrap_or_else(|| "unknown".to_string())
+    };
+    render_code_mode_sample(
+        &definition.description,
+        &definition.name,
+        input_name,
+        input_type,
+        output_type,
+    )
+}
+
+fn render_code_mode_tool_declaration(
+    tool_name: &str,
+    input_name: &str,
+    input_type: &str,
+    output_type: &str,
+) -> String {
+    let tool_name = normalize_code_mode_identifier(tool_name);
+    format!("{tool_name}({input_name}: {input_type}): Promise<{output_type}>;")
+}
+
+fn render_tool_heading(global_name: &str, raw_name: &str) -> String {
+    if global_name == raw_name {
+        format!("### `{global_name}`")
+    } else {
+        format!("### `{global_name}` (`{raw_name}`)")
+    }
+}
+
+fn mcp_structured_content_schema(output_schema: Option<&JsonValue>) -> Option<&JsonValue> {
+    let output_schema = output_schema?;
+    let properties = output_schema
+        .get("properties")
+        .and_then(JsonValue::as_object)?;
+    let content_schema = properties.get("content").and_then(JsonValue::as_object)?;
+    if content_schema.get("type").and_then(JsonValue::as_str) != Some("array") {
+        return None;
+    }
+
+    if content_schema
+        .get("items")
+        .and_then(JsonValue::as_object)
+        .is_none_or(|items| items.get("type").and_then(JsonValue::as_str) != Some("object"))
+    {
+        return None;
+    }
+
+    if properties
+        .get("isError")
+        .and_then(JsonValue::as_object)
+        .is_none_or(|schema| schema.get("type").and_then(JsonValue::as_str) != Some("boolean"))
+    {
+        return None;
+    }
+
+    if properties
+        .get("_meta")
+        .and_then(JsonValue::as_object)
+        .is_none_or(|schema| schema.get("type").and_then(JsonValue::as_str) != Some("object"))
+    {
+        return None;
+    }
+
+    Some(
+        properties
+            .get("structuredContent")
+            .unwrap_or(&JsonValue::Bool(true)),
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::CodeModeToolKind;
+    use super::ImageDetailVisibility;
+    use super::ParsedExecSource;
+    use super::ToolDefinition;
+    use super::ToolNamespaceDescription;
+    use super::augment_tool_definition;
+    use super::build_exec_tool_description;
+    use super::normalize_code_mode_identifier;
+    use super::parse_exec_source;
+    use codex_protocol::ToolName;
+    use pretty_assertions::assert_eq;
+    use serde_json::Value as JsonValue;
+    use serde_json::json;
+    use std::collections::BTreeMap;
+
+    fn mcp_call_tool_result_schema(structured_content_schema: JsonValue) -> JsonValue {
+        json!({
+            "type": "object",
+            "properties": {
+                "content": {
+                    "type": "array",
+                    "items": {
+                        "type": "object"
+                    }
+                },
+                "structuredContent": structured_content_schema,
+                "isError": { "type": "boolean" },
+                "_meta": { "type": "object" }
+            },
+            "required": ["content"],
+            "additionalProperties": false
+        })
+    }
+
+    #[test]
+    fn parse_exec_source_without_pragma() {
+        assert_eq!(
+            parse_exec_source("text('hi')").unwrap(),
+            ParsedExecSource {
+                code: "text('hi')".to_string(),
+                yield_time_ms: None,
+                max_output_tokens: None,
+            }
+        );
+    }
+
+    #[test]
+    fn parse_exec_source_with_pragma() {
+        assert_eq!(
+            parse_exec_source("// @exec: {\"yield_time_ms\": 10}\ntext('hi')").unwrap(),
+            ParsedExecSource {
+                code: "text('hi')".to_string(),
+                yield_time_ms: Some(10),
+                max_output_tokens: None,
+            }
+        );
+    }
+
+    #[test]
+    fn normalize_identifier_rewrites_invalid_characters() {
+        assert_eq!(
+            "mcp__ologs__get_profile",
+            normalize_code_mode_identifier("mcp__ologs__get_profile")
+        );
+        assert_eq!(
+            "hidden_dynamic_tool",
+            normalize_code_mode_identifier("hidden-dynamic-tool")
+        );
+    }
+
+    #[test]
+    fn augment_tool_definition_appends_typed_declaration() {
+        let definition = ToolDefinition {
+            name: "hidden_dynamic_tool".to_string(),
+            tool_name: ToolName::plain("hidden_dynamic_tool"),
+            description: "Test tool".to_string(),
+            kind: CodeModeToolKind::Function,
+            input_schema: Some(json!({
+                "type": "object",
+                "properties": { "city": { "type": "string" } },
+                "required": ["city"],
+                "additionalProperties": false
+            })),
+            output_schema: Some(json!({
+                "type": "object",
+                "properties": { "ok": { "type": "boolean" } },
+                "required": ["ok"]
+            })),
+        };
+
+        let description = augment_tool_definition(definition).description;
+        assert!(description.contains("declare const tools"));
+        assert!(
+            description.contains(
+                "hidden_dynamic_tool(args: { city: string; }): Promise<{ ok: boolean; }>;"
+            )
+        );
+    }
+
+    #[test]
+    fn augment_tool_definition_includes_property_descriptions_as_comments() {
+        let definition = ToolDefinition {
+            name: "weather_tool".to_string(),
+            tool_name: ToolName::plain("weather_tool"),
+            description: "Weather tool".to_string(),
+            kind: CodeModeToolKind::Function,
+            input_schema: Some(json!({
+                "type": "object",
+                "properties": {
+                    "weather": {
+                        "type": "array",
+                        "description": "look up weather for a given list of locations",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "location": { "type": "string" }
+                            },
+                            "required": ["location"]
+                        }
+                    }
+                },
+                "required": ["weather"]
+            })),
+            output_schema: Some(json!({
+                "type": "object",
+                "properties": {
+                    "forecast": {
+                        "type": "string",
+                        "description": "human readable weather forecast"
+                    }
+                },
+                "required": ["forecast"]
+            })),
+        };
+
+        let description = augment_tool_definition(definition).description;
+        assert!(description.contains(
+            r#"weather_tool(args: {
+  // look up weather for a given list of locations
+  weather: Array<{ location: string; }>;
+}): Promise<{
+  // human readable weather forecast
+  forecast: string;
+}>;"#
+        ));
+    }
+
+    #[test]
+    fn code_mode_types_structured_content_result_refs() {
+        let definition = ToolDefinition {
+            name: "mcp__sample__search".to_string(),
+            tool_name: ToolName::namespaced("mcp__sample__", "search"),
+            description: "Search".to_string(),
+            kind: CodeModeToolKind::Function,
+            input_schema: Some(json!({
+                "type": "object",
+                "properties": {},
+                "additionalProperties": false
+            })),
+            output_schema: Some(mcp_call_tool_result_schema(json!({
+                "type": "object",
+                "properties": {
+                    "results": {
+                        "type": "array",
+                        "items": { "$ref": "#/definitions/Result~1item~0v1" }
+                    }
+                },
+                "required": ["results"],
+                "additionalProperties": false,
+                "definitions": {
+                    "Result/item~v1": {
+                        "type": "object",
+                        "properties": {
+                            "id": { "type": "string" },
+                            "score": { "type": "number" }
+                        },
+                        "required": ["id", "score"],
+                        "additionalProperties": false
+                    }
+                }
+            }))),
+        };
+
+        let description = augment_tool_definition(definition).description;
+        assert!(description.contains(
+            "mcp__sample__search(args: {}): Promise<CallToolResult<{ results: Array<{ id: string; score: number; }>; }>>;"
+        ));
+    }
+
+    #[test]
+    fn code_mode_only_description_includes_nested_tools() {
+        let description = build_exec_tool_description(
+            &[ToolDefinition {
+                name: "foo".to_string(),
+                tool_name: ToolName::plain("foo"),
+                description: "bar".to_string(),
+                kind: CodeModeToolKind::Function,
+                input_schema: None,
+                output_schema: None,
+            }],
+            &[],
+            &BTreeMap::new(),
+            crate::DEFAULT_EXEC_YIELD_TIME_MS,
+            /*code_mode_only*/ true,
+            ImageDetailVisibility::Visible,
+        );
+        assert!(description.contains(
+            "### `foo`
+bar"
+        ));
+        assert!(!description.contains("do not attempt to use any other tools directly"));
+    }
+
+    #[test]
+    fn exec_description_mentions_timeout_helpers() {
+        let description = build_exec_tool_description(
+            &[],
+            &[],
+            &BTreeMap::new(),
+            crate::DEFAULT_EXEC_YIELD_TIME_MS,
+            /*code_mode_only*/ false,
+            ImageDetailVisibility::Visible,
+        );
+        assert!(description.contains("`audio(audioUrlOrItem:"));
+        assert!(description.contains("`setTimeout(callback: () => void, delayMs?: number)`"));
+        assert!(description.contains("`clearTimeout(timeoutId?: number)`"));
+    }
+
+    #[test]
+    fn code_mode_only_description_groups_namespace_instructions_once() {
+        let namespace_descriptions = BTreeMap::from([(
+            "mcp__sample__".to_string(),
+            ToolNamespaceDescription {
+                name: "mcp__sample".to_string(),
+                description: "Shared namespace guidance.".to_string(),
+            },
+        )]);
+        let description = build_exec_tool_description(
+            &[
+                ToolDefinition {
+                    name: "mcp__sample__alpha".to_string(),
+                    tool_name: ToolName::namespaced("mcp__sample__", "alpha"),
+                    description: "First tool".to_string(),
+                    kind: CodeModeToolKind::Function,
+                    input_schema: Some(json!({
+                        "type": "object",
+                        "properties": {},
+                        "additionalProperties": false
+                    })),
+                    output_schema: Some(mcp_call_tool_result_schema(json!({
+                        "type": "object",
+                        "properties": {},
+                        "additionalProperties": false
+                    }))),
+                },
+                ToolDefinition {
+                    name: "mcp__sample__beta".to_string(),
+                    tool_name: ToolName::namespaced("mcp__sample__", "beta"),
+                    description: "Second tool".to_string(),
+                    kind: CodeModeToolKind::Function,
+                    input_schema: Some(json!({
+                        "type": "object",
+                        "properties": {},
+                        "additionalProperties": false
+                    })),
+                    output_schema: Some(mcp_call_tool_result_schema(json!({
+                        "type": "object",
+                        "properties": {},
+                        "additionalProperties": false
+                    }))),
+                },
+            ],
+            &[],
+            &namespace_descriptions,
+            crate::DEFAULT_EXEC_YIELD_TIME_MS,
+            /*code_mode_only*/ true,
+            ImageDetailVisibility::Visible,
+        );
+        assert_eq!(description.matches("## mcp__sample").count(), 1);
+        assert!(description.contains("## mcp__sample\nShared namespace guidance."));
+        assert!(description.contains(
+            "declare const tools: { mcp__sample__alpha(args: {}): Promise<CallToolResult<{}>>; };"
+        ));
+        assert!(description.contains(
+            "declare const tools: { mcp__sample__beta(args: {}): Promise<CallToolResult<{}>>; };"
+        ));
+    }
+
+    #[test]
+    fn code_mode_only_description_omits_empty_namespace_sections() {
+        let namespace_descriptions = BTreeMap::from([(
+            "mcp__sample__".to_string(),
+            ToolNamespaceDescription {
+                name: "mcp__sample".to_string(),
+                description: String::new(),
+            },
+        )]);
+        let description = build_exec_tool_description(
+            &[ToolDefinition {
+                name: "mcp__sample__alpha".to_string(),
+                tool_name: ToolName::namespaced("mcp__sample__", "alpha"),
+                description: "First tool".to_string(),
+                kind: CodeModeToolKind::Function,
+                input_schema: Some(json!({
+                    "type": "object",
+                    "properties": {},
+                    "additionalProperties": false
+                })),
+                output_schema: Some(mcp_call_tool_result_schema(json!({
+                    "type": "object",
+                    "properties": {},
+                    "additionalProperties": false
+                }))),
+            }],
+            &[],
+            &namespace_descriptions,
+            crate::DEFAULT_EXEC_YIELD_TIME_MS,
+            /*code_mode_only*/ true,
+            ImageDetailVisibility::Visible,
+        );
+
+        assert!(!description.contains("## mcp__sample"));
+        assert!(description.contains("### `mcp__sample__alpha`"));
+    }
+
+    #[test]
+    fn code_mode_only_description_renders_shared_mcp_types_once() {
+        let first_tool = augment_tool_definition(ToolDefinition {
+            name: "mcp__sample__alpha".to_string(),
+            tool_name: ToolName::namespaced("mcp__sample__", "alpha"),
+            description: "First tool".to_string(),
+            kind: CodeModeToolKind::Function,
+            input_schema: Some(json!({
+                "type": "object",
+                "properties": {},
+                "additionalProperties": false
+            })),
+            output_schema: Some(json!({
+                "type": "object",
+                "properties": {
+                    "content": {
+                        "type": "array",
+                        "items": {
+                            "type": "object"
+                        }
+                    },
+                    "structuredContent": {
+                        "type": "object",
+                        "properties": {
+                            "echo": { "type": "string" }
+                        },
+                        "required": ["echo"],
+                        "additionalProperties": false
+                    },
+                    "isError": { "type": "boolean" },
+                    "_meta": { "type": "object" }
+                },
+                "required": ["content"],
+                "additionalProperties": false
+            })),
+        });
+        let second_tool = augment_tool_definition(ToolDefinition {
+            name: "mcp__sample__beta".to_string(),
+            tool_name: ToolName::namespaced("mcp__sample__", "beta"),
+            description: "Second tool".to_string(),
+            kind: CodeModeToolKind::Function,
+            input_schema: Some(json!({
+                "type": "object",
+                "properties": {},
+                "additionalProperties": false
+            })),
+            output_schema: Some(json!({
+                "type": "object",
+                "properties": {
+                    "content": {
+                        "type": "array",
+                        "items": {
+                            "type": "object"
+                        }
+                    },
+                    "structuredContent": {
+                        "type": "object",
+                        "properties": {
+                            "count": { "type": "integer" }
+                        },
+                        "required": ["count"],
+                        "additionalProperties": false
+                    },
+                    "isError": { "type": "boolean" },
+                    "_meta": { "type": "object" }
+                },
+                "required": ["content"],
+                "additionalProperties": false
+            })),
+        });
+
+        let description = build_exec_tool_description(
+            &[
+                ToolDefinition {
+                    name: first_tool.name,
+                    tool_name: first_tool.tool_name,
+                    description: "First tool".to_string(),
+                    kind: first_tool.kind,
+                    input_schema: first_tool.input_schema,
+                    output_schema: first_tool.output_schema,
+                },
+                ToolDefinition {
+                    name: second_tool.name,
+                    tool_name: second_tool.tool_name,
+                    description: "Second tool".to_string(),
+                    kind: second_tool.kind,
+                    input_schema: second_tool.input_schema,
+                    output_schema: second_tool.output_schema,
+                },
+            ],
+            &[],
+            &BTreeMap::new(),
+            crate::DEFAULT_EXEC_YIELD_TIME_MS,
+            /*code_mode_only*/ true,
+            ImageDetailVisibility::Visible,
+        );
+
+        assert_eq!(
+            description
+                .matches("type CallToolResult<TStructured = { [key: string]: unknown }>")
+                .count(),
+            1
+        );
+        assert_eq!(description.matches("Shared MCP Types:").count(), 1);
+    }
+
+    #[test]
+    fn code_mode_only_description_renders_shared_mcp_types_for_deferred_tools() {
+        let deferred_tool = ToolDefinition {
+            name: "mcp__sample__alpha".to_string(),
+            tool_name: ToolName::namespaced("mcp__sample__", "alpha"),
+            description: "Deferred tool".to_string(),
+            kind: CodeModeToolKind::Function,
+            input_schema: Some(json!({
+                "type": "object",
+                "properties": {},
+                "additionalProperties": false
+            })),
+            output_schema: Some(mcp_call_tool_result_schema(json!({
+                "type": "object",
+                "properties": {},
+                "additionalProperties": false
+            }))),
+        };
+
+        let description = build_exec_tool_description(
+            &[],
+            &[deferred_tool],
+            &BTreeMap::new(),
+            crate::DEFAULT_EXEC_YIELD_TIME_MS,
+            /*code_mode_only*/ true,
+            ImageDetailVisibility::Visible,
+        );
+
+        assert!(description.contains("Some deferred nested tools may be omitted"));
+        assert!(description.contains("Shared MCP Types:"));
+        assert!(!description.contains("### `mcp__sample__alpha`"));
+    }
+
+    #[test]
+    fn exec_description_mentions_deferred_nested_tools_when_available() {
+        let description = build_exec_tool_description(
+            &[],
+            &[ToolDefinition {
+                name: "deferred_tool".to_string(),
+                tool_name: ToolName::plain("deferred_tool"),
+                description: "Deferred tool".to_string(),
+                kind: CodeModeToolKind::Function,
+                input_schema: None,
+                output_schema: None,
+            }],
+            &BTreeMap::new(),
+            crate::DEFAULT_EXEC_YIELD_TIME_MS,
+            /*code_mode_only*/ false,
+            ImageDetailVisibility::Visible,
+        );
+
+        assert!(description.contains("Some deferred nested tools may be omitted"));
+        assert!(description.contains("filter `ALL_TOOLS` by `name` and `description`"));
+        assert!(!description.contains("do not print the full `ALL_TOOLS` array"));
+    }
+}

codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto

@@ -0,0 +1,269 @@
+syntax = "proto3";
+
+package codex.code_mode.v1;
+
+// Hosts stateful JavaScript execution and delegates nested tool calls to the
+// session owner. Large tool inputs and outputs stay off the session event
+// stream so independent HTTP/2 streams can make progress concurrently.
+service CodeModeHost {
+  // Opens a session lease. The first event is always SessionOpened; dropping
+  // this stream closes the session and terminates its active cells.
+  rpc OpenSession(OpenSessionRequest) returns (stream SessionEvent);
+  rpc CloseSession(CloseSessionRequest) returns (CloseSessionResponse);
+
+  // Each subscription owns an independent stream of matching invocations. An
+  // empty tool_names filter matches every tool. Each invocation is routed to
+  // exactly one matching subscription, even when filters overlap.
+  rpc SubscribeToToolCalls(SubscribeToToolCallsRequest)
+      returns (stream ToolCall);
+
+  // Each result receives its own HTTP/2 stream, preventing a large response
+  // from blocking unrelated tool completions or session control events.
+  rpc CompleteToolCall(CompleteToolCallRequest)
+      returns (CompleteToolCallResponse);
+  rpc AcknowledgeNotification(AcknowledgeNotificationRequest)
+      returns (AcknowledgeNotificationResponse);
+
+  // Emits ExecutionStarted immediately, followed by one ExecutionOutcome when
+  // the execution yields, completes, or is terminated.
+  rpc Execute(ExecuteRequest) returns (stream ExecuteEvent);
+  rpc Wait(WaitRequest) returns (WaitResponse);
+
+  // Acknowledges that a canceled wait has retired before another wait starts.
+  rpc CancelWait(CancelWaitRequest) returns (CancelWaitResponse);
+  rpc Terminate(TerminateRequest) returns (WaitResponse);
+}
+
+message OpenSessionRequest {
+  optional SessionCellExecutionLimits cell_execution_limits = 1;
+}
+
+message SessionCellExecutionLimits {
+  optional uint64 max_yield_time_ms = 1;
+  optional uint64 max_heap_size_bytes = 2;
+}
+
+message SessionEvent {
+  oneof event {
+    SessionOpened opened = 1;
+    ToolCallCancelled tool_call_cancelled = 2;
+    Notification notification = 3;
+    NotificationCancelled notification_cancelled = 4;
+    CellClosed cell_closed = 5;
+  }
+}
+
+message SessionOpened {
+  string session_id = 1;
+}
+
+message CloseSessionRequest {
+  string session_id = 1;
+}
+
+message CloseSessionResponse {}
+
+message SubscribeToToolCallsRequest {
+  string session_id = 1;
+  repeated ToolName tool_names = 2;
+}
+
+message ToolCall {
+  string session_id = 1;
+
+  // Correlates callbacks with Execute before ExecutionStarted is received.
+  string execution_id = 2;
+  string cell_id = 3;
+  string invocation_id = 4;
+  string runtime_tool_call_id = 5;
+  ToolName tool_name = 6;
+  ToolKind tool_kind = 7;
+  optional bytes input_json = 8;
+
+  // Starts at one and increases independently for each execution.
+  uint64 sequence = 9;
+
+  // The host tool invocation span's W3C parent context for this streamed callback.
+  // gRPC metadata is fixed when the subscription stream opens, so each call
+  // carries its own context in the message.
+  optional string traceparent = 10;
+}
+
+message CompleteToolCallRequest {
+  string session_id = 1;
+  string invocation_id = 2;
+
+  oneof outcome {
+    ToolCallSucceeded succeeded = 3;
+    ToolCallFailed failed = 4;
+  }
+}
+
+message ToolCallSucceeded {
+  bytes output_json = 1;
+}
+
+message ToolCallFailed {
+  string message = 1;
+}
+
+message CompleteToolCallResponse {}
+
+message ToolCallCancelled {
+  string invocation_id = 1;
+
+  // Cancellation can arrive before the corresponding ToolCall because session
+  // control events and tool subscriptions use independent HTTP/2 streams.
+}
+
+message Notification {
+  string notification_id = 1;
+  string execution_id = 2;
+  string cell_id = 3;
+  string call_id = 4;
+  string text = 5;
+}
+
+message NotificationCancelled {
+  string notification_id = 1;
+}
+
+message AcknowledgeNotificationRequest {
+  string session_id = 1;
+  string notification_id = 2;
+}
+
+message AcknowledgeNotificationResponse {}
+
+message CellClosed {
+  string execution_id = 1;
+  string cell_id = 2;
+
+  // Last tool-call sequence issued before closure. Clients may retire the cell
+  // immediately and reject tool calls delivered after its closure.
+  uint64 final_tool_call_sequence = 3;
+}
+
+message ExecuteRequest {
+  string session_id = 1;
+
+  // Chosen by the client so callbacks can be correlated before cell admission.
+  string execution_id = 2;
+  string tool_call_id = 3;
+  string source = 4;
+  repeated ToolDefinition enabled_tools = 5;
+  optional uint64 yield_time_ms = 6;
+  optional uint64 max_output_tokens = 7;
+}
+
+message ExecuteEvent {
+  oneof event {
+    ExecutionStarted started = 1;
+    ExecutionOutcome outcome = 2;
+  }
+}
+
+message ExecutionStarted {
+  string execution_id = 1;
+  string cell_id = 2;
+}
+
+message WaitRequest {
+  string session_id = 1;
+  string cell_id = 2;
+  string wait_id = 3;
+  uint64 yield_time_ms = 4;
+}
+
+message WaitResponse {
+  oneof state {
+    ExecutionOutcome live_cell = 1;
+    ExecutionOutcome missing_cell = 2;
+  }
+}
+
+message CancelWaitRequest {
+  string session_id = 1;
+  string wait_id = 2;
+}
+
+message CancelWaitResponse {}
+
+message TerminateRequest {
+  string session_id = 1;
+  string cell_id = 2;
+}
+
+message ExecutionOutcome {
+  string cell_id = 1;
+  repeated ContentItem content_items = 2;
+
+  // Elapsed monotonic time from receipt of this Execute, Wait, or Terminate
+  // request until its outcome is ready, before serialization and delivery.
+  // Includes nested tool waits, but not background time between requests.
+  // Always supplied by the host; zero is a valid measurement.
+  uint64 code_mode_host_duration_ns = 6;
+
+  oneof outcome {
+    ExecutionYielded yielded = 3;
+    ExecutionTerminated terminated = 4;
+    ExecutionCompleted completed = 5;
+  }
+}
+
+message ExecutionYielded {}
+
+message ExecutionTerminated {}
+
+message ExecutionCompleted {
+  optional string error_text = 1;
+}
+
+message ToolDefinition {
+  string name = 1;
+  ToolName tool_name = 2;
+  string description = 3;
+  ToolKind kind = 4;
+  optional bytes input_schema_json = 5;
+  optional bytes output_schema_json = 6;
+}
+
+message ToolName {
+  string name = 1;
+  optional string namespace = 2;
+}
+
+enum ToolKind {
+  TOOL_KIND_UNSPECIFIED = 0;
+  TOOL_KIND_FUNCTION = 1;
+  TOOL_KIND_FREEFORM = 2;
+}
+
+message ContentItem {
+  oneof item {
+    TextContent text = 1;
+    ImageContent image = 2;
+    AudioContent audio = 3;
+  }
+}
+
+message TextContent {
+  string text = 1;
+}
+
+message ImageContent {
+  string image_url = 1;
+  optional ImageDetail detail = 2;
+}
+
+message AudioContent {
+  string audio_url = 1;
+}
+
+enum ImageDetail {
+  IMAGE_DETAIL_UNSPECIFIED = 0;
+  IMAGE_DETAIL_AUTO = 1;
+  IMAGE_DETAIL_LOW = 2;
+  IMAGE_DETAIL_HIGH = 3;
+  IMAGE_DETAIL_ORIGINAL = 4;
+}

codex-rs/code-mode-protocol/src/grpc/mod.rs

@@ -0,0 +1,7 @@
+#[cfg(codex_bazel)]
+pub use code_mode_proto::codex::code_mode::v1::*;
+
+#[cfg(not(codex_bazel))]
+tonic::include_proto!("codex.code_mode.v1");
+
+pub const MAX_IDENTIFIER_BYTES: usize = 256;

codex-rs/code-mode-protocol/src/host/codec.rs

@@ -0,0 +1,170 @@
+use std::io;
+use std::mem::size_of;
+
+use serde::Serialize;
+use serde::de::DeserializeOwned;
+use tokio::io::AsyncRead;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWrite;
+use tokio::io::AsyncWriteExt;
+
+/// Maximum JSON payload size accepted for one code-mode host frame.
+pub const MAX_FRAME_BYTES: usize = 64 * 1024 * 1024;
+
+/// A serialized IPC frame that has already passed the payload size limit.
+#[derive(Clone, Debug)]
+pub struct EncodedFrame {
+    payload: Vec<u8>,
+}
+
+impl EncodedFrame {
+    pub fn encode<T>(message: &T) -> io::Result<Self>
+    where
+        T: Serialize,
+    {
+        let payload = serde_json::to_vec(message).map_err(|err| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("failed to encode code-mode IPC frame: {err}"),
+            )
+        })?;
+        if payload.len() > MAX_FRAME_BYTES {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!(
+                    "code-mode IPC frame length {} exceeds {MAX_FRAME_BYTES} bytes",
+                    payload.len()
+                ),
+            ));
+        }
+        Ok(Self { payload })
+    }
+
+    /// Returns the complete length-prefixed representation of this frame.
+    pub fn into_framed_bytes(self) -> Vec<u8> {
+        let mut bytes = Vec::with_capacity(size_of::<u32>() + self.payload.len());
+        bytes.extend_from_slice(&(self.payload.len() as u32).to_le_bytes());
+        bytes.extend_from_slice(&self.payload);
+        bytes
+    }
+
+    /// Decodes exactly one complete length-prefixed frame.
+    pub fn decode_framed<T>(bytes: &[u8]) -> io::Result<T>
+    where
+        T: DeserializeOwned,
+    {
+        let length_bytes: [u8; size_of::<u32>()] = bytes
+            .get(..size_of::<u32>())
+            .and_then(|length_bytes| length_bytes.try_into().ok())
+            .ok_or_else(|| {
+                io::Error::new(
+                    io::ErrorKind::InvalidData,
+                    "code-mode IPC frame is missing its length prefix",
+                )
+            })?;
+        let length = u32::from_le_bytes(length_bytes) as usize;
+        if length > MAX_FRAME_BYTES {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("code-mode IPC frame length {length} exceeds {MAX_FRAME_BYTES} bytes"),
+            ));
+        }
+
+        let payload = &bytes[size_of::<u32>()..];
+        if payload.len() != length {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!(
+                    "code-mode IPC frame declares {length} payload bytes but contains {}",
+                    payload.len()
+                ),
+            ));
+        }
+
+        serde_json::from_slice(payload).map_err(|err| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("failed to decode code-mode IPC frame: {err}"),
+            )
+        })
+    }
+}
+
+/// Decodes JSON messages prefixed by a four-byte little-endian payload length.
+pub struct FramedReader<R> {
+    reader: R,
+}
+
+impl<R> FramedReader<R>
+where
+    R: AsyncRead + Unpin,
+{
+    pub fn new(reader: R) -> Self {
+        Self { reader }
+    }
+
+    /// Reads the next frame, returning `None` only for EOF at a frame boundary.
+    pub async fn read<T>(&mut self) -> io::Result<Option<T>>
+    where
+        T: DeserializeOwned,
+    {
+        let mut length_bytes = [0_u8; size_of::<u32>()];
+        if self.reader.read(&mut length_bytes[..1]).await? == 0 {
+            return Ok(None);
+        }
+        self.reader.read_exact(&mut length_bytes[1..]).await?;
+
+        let length = u32::from_le_bytes(length_bytes) as usize;
+        if length > MAX_FRAME_BYTES {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("code-mode IPC frame length {length} exceeds {MAX_FRAME_BYTES} bytes"),
+            ));
+        }
+
+        let mut payload = vec![0; length];
+        self.reader.read_exact(&mut payload).await?;
+        serde_json::from_slice(&payload).map(Some).map_err(|err| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("failed to decode code-mode IPC frame: {err}"),
+            )
+        })
+    }
+}
+
+/// Encodes JSON messages with a four-byte little-endian payload length.
+pub struct FramedWriter<W> {
+    writer: W,
+}
+
+impl<W> FramedWriter<W>
+where
+    W: AsyncWrite + Unpin,
+{
+    pub fn new(writer: W) -> Self {
+        Self { writer }
+    }
+
+    /// Writes and flushes one complete frame.
+    pub async fn write<T>(&mut self, message: &T) -> io::Result<()>
+    where
+        T: Serialize,
+    {
+        self.write_frame(&EncodedFrame::encode(message)?).await
+    }
+
+    /// Writes and flushes a frame encoded before it entered an I/O queue.
+    pub async fn write_frame(&mut self, frame: &EncodedFrame) -> io::Result<()> {
+        let length = u32::try_from(frame.payload.len()).map_err(|_| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                "code-mode IPC frame length exceeds u32",
+            )
+        })?;
+
+        self.writer.write_all(&length.to_le_bytes()).await?;
+        self.writer.write_all(&frame.payload).await?;
+        self.writer.flush().await
+    }
+}

codex-rs/code-mode-protocol/src/host/codec_tests.rs

@@ -0,0 +1,137 @@
+use pretty_assertions::assert_eq;
+use serde_json::json;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+
+use super::EncodedFrame;
+use super::FramedReader;
+use super::FramedWriter;
+use super::MAX_FRAME_BYTES;
+
+#[test]
+fn complete_frame_round_trips_without_a_byte_stream() {
+    let value = json!({"type": "session/open", "sessionId": "session-1"});
+    let bytes = EncodedFrame::encode(&value)
+        .expect("encode frame")
+        .into_framed_bytes();
+
+    assert_eq!(
+        EncodedFrame::decode_framed::<serde_json::Value>(&bytes).expect("decode frame"),
+        value
+    );
+}
+
+#[test]
+fn complete_frame_rejects_truncated_and_trailing_payloads() {
+    let value = json!({"value": 1});
+    let bytes = EncodedFrame::encode(&value)
+        .expect("encode frame")
+        .into_framed_bytes();
+
+    let truncated = &bytes[..bytes.len() - 1];
+    let truncated_error = EncodedFrame::decode_framed::<serde_json::Value>(truncated)
+        .expect_err("truncated frame should fail");
+    assert_eq!(truncated_error.kind(), std::io::ErrorKind::InvalidData);
+
+    let mut trailing = bytes;
+    trailing.push(0);
+    let trailing_error = EncodedFrame::decode_framed::<serde_json::Value>(&trailing)
+        .expect_err("frame with trailing bytes should fail");
+    assert_eq!(trailing_error.kind(), std::io::ErrorKind::InvalidData);
+}
+
+#[tokio::test]
+async fn frame_wire_format_is_little_endian_length_prefixed_json() {
+    let (writer, mut reader) = tokio::io::duplex(/*max_buf_size*/ 128);
+    let write = tokio::spawn(async move {
+        FramedWriter::new(writer)
+            .write(&json!({"value": 1}))
+            .await
+            .expect("write frame");
+    });
+
+    let mut bytes = Vec::new();
+    reader.read_to_end(&mut bytes).await.expect("read bytes");
+    write.await.expect("writer task");
+
+    let payload = br#"{"value":1}"#;
+    let mut expected = (payload.len() as u32).to_le_bytes().to_vec();
+    expected.extend_from_slice(payload);
+    assert_eq!(bytes, expected);
+}
+
+#[tokio::test]
+async fn fragmented_frame_round_trips() {
+    let value = json!({"type": "session/open", "sessionId": "session-1"});
+    let payload = serde_json::to_vec(&value).expect("serialize");
+    let mut bytes = (payload.len() as u32).to_le_bytes().to_vec();
+    bytes.extend(payload);
+
+    let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 128);
+    let write = tokio::spawn(async move {
+        for byte in bytes {
+            writer.write_all(&[byte]).await.expect("write byte");
+            tokio::task::yield_now().await;
+        }
+    });
+
+    assert_eq!(
+        FramedReader::new(reader)
+            .read::<serde_json::Value>()
+            .await
+            .expect("read frame"),
+        Some(value)
+    );
+    write.await.expect("writer task");
+}
+
+#[tokio::test]
+async fn eof_is_clean_only_at_a_frame_boundary() {
+    let (writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16);
+    drop(writer);
+    assert_eq!(
+        FramedReader::new(reader)
+            .read::<serde_json::Value>()
+            .await
+            .expect("clean eof"),
+        None
+    );
+
+    let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16);
+    writer
+        .write_all(&[1, 0])
+        .await
+        .expect("write partial header");
+    drop(writer);
+    let err = FramedReader::new(reader)
+        .read::<serde_json::Value>()
+        .await
+        .expect_err("truncated header");
+    assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
+}
+
+#[tokio::test]
+async fn oversized_and_malformed_frames_are_rejected() {
+    let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16);
+    writer
+        .write_all(&((MAX_FRAME_BYTES as u32) + 1).to_le_bytes())
+        .await
+        .expect("write oversized header");
+    let err = FramedReader::new(reader)
+        .read::<serde_json::Value>()
+        .await
+        .expect_err("oversized frame");
+    assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
+
+    let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16);
+    writer
+        .write_all(&(1_u32).to_le_bytes())
+        .await
+        .expect("write length");
+    writer.write_all(b"{").await.expect("write malformed json");
+    let err = FramedReader::new(reader)
+        .read::<serde_json::Value>()
+        .await
+        .expect_err("malformed frame");
+    assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
+}

codex-rs/code-mode-protocol/src/host/error.rs

@@ -0,0 +1,19 @@
+use serde::Deserialize;
+use serde::Serialize;
+
+use super::Capability;
+use super::SupportedProtocolVersions;
+
+/// Explains why connection negotiation was rejected before any session opened.
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
+pub enum HandshakeRejectReason {
+    #[serde(rename = "noCompatibleVersion")]
+    NoCompatibleVersion {
+        supported_versions: SupportedProtocolVersions,
+    },
+    #[serde(rename = "missingRequiredCapability")]
+    MissingRequiredCapability { capability: Capability },
+    #[serde(rename = "invalidHello")]
+    InvalidHello { message: String },
+}

codex-rs/code-mode-protocol/src/host/host_tests.rs

@@ -0,0 +1,847 @@
+use std::fmt::Debug;
+
+use pretty_assertions::assert_eq;
+use serde::Serialize;
+use serde::de::DeserializeOwned;
+use serde_json::Value;
+use serde_json::json;
+
+use super::Capability;
+use super::CapabilitySet;
+use super::ClientHello;
+use super::ClientToHost;
+use super::DelegateRequest;
+use super::DelegateRequestId;
+use super::DelegateResponse;
+use super::HandshakeRejectReason;
+use super::HostHello;
+use super::HostRequest;
+use super::HostResponse;
+use super::HostToClient;
+use super::ProtocolVersion;
+use super::RequestId;
+use super::SessionId;
+use super::SupportedProtocolVersions;
+use super::WireCellId;
+use super::WireContentItem;
+use super::WireExecuteRequest;
+use super::WireImageDetail;
+use super::WireNestedToolCall;
+use super::WireResult;
+use super::WireRuntimeResponse;
+use super::WireSessionCellExecutionLimits;
+use super::WireToolDefinition;
+use super::WireToolKind;
+use super::WireToolName;
+use super::WireWaitOutcome;
+use super::WireWaitRequest;
+use crate::CodeModeSessionCellExecutionLimits;
+use crate::ExecuteRequest;
+
+fn session_id() -> SessionId {
+    SessionId::new("session-1").expect("valid session ID")
+}
+
+fn cell_id(value: &str) -> WireCellId {
+    WireCellId::new(value)
+}
+
+fn request_id(value: i64) -> RequestId {
+    RequestId::new(value)
+}
+
+fn delegate_request_id(value: i64) -> DelegateRequestId {
+    DelegateRequestId::new(value)
+}
+
+fn capability(value: &str) -> Capability {
+    Capability::new(value).expect("valid capability")
+}
+
+fn supported_versions() -> SupportedProtocolVersions {
+    SupportedProtocolVersions::try_new([ProtocolVersion::V1])
+        .expect("nonempty unique protocol versions")
+}
+
+fn assert_wire_round_trip<T>(message: T, encoded: Value)
+where
+    T: Debug + DeserializeOwned + PartialEq + Serialize,
+{
+    assert_eq!(serde_json::to_value(&message).expect("serialize"), encoded);
+    assert_eq!(
+        serde_json::from_value::<T>(encoded).expect("deserialize"),
+        message
+    );
+}
+
+fn execute_request() -> WireExecuteRequest {
+    WireExecuteRequest {
+        tool_call_id: "call-1".to_string(),
+        enabled_tools: vec![
+            WireToolDefinition {
+                name: "function_tool".to_string(),
+                tool_name: WireToolName {
+                    name: "function_tool".to_string(),
+                    namespace: None,
+                },
+                description: "function tool".to_string(),
+                kind: WireToolKind::Function,
+                input_schema: Some(json!({ "type": "object" })),
+                output_schema: None,
+            },
+            WireToolDefinition {
+                name: "freeform_tool".to_string(),
+                tool_name: WireToolName {
+                    name: "freeform_tool".to_string(),
+                    namespace: Some("mcp__sample__".to_string()),
+                },
+                description: "freeform tool".to_string(),
+                kind: WireToolKind::Freeform,
+                input_schema: None,
+                output_schema: Some(json!({ "type": "string" })),
+            },
+        ],
+        source: "text('hello');".to_string(),
+        yield_time_ms: Some(25),
+        max_output_tokens: Some(100),
+    }
+}
+
+fn content_items() -> Vec<WireContentItem> {
+    vec![
+        WireContentItem::InputText {
+            text: "hello".to_string(),
+        },
+        WireContentItem::InputImage {
+            image_url: "data:image/png;base64,none".to_string(),
+            detail: None,
+        },
+        WireContentItem::InputImage {
+            image_url: "data:image/png;base64,auto".to_string(),
+            detail: Some(WireImageDetail::Auto),
+        },
+        WireContentItem::InputImage {
+            image_url: "data:image/png;base64,low".to_string(),
+            detail: Some(WireImageDetail::Low),
+        },
+        WireContentItem::InputImage {
+            image_url: "data:image/png;base64,high".to_string(),
+            detail: Some(WireImageDetail::High),
+        },
+        WireContentItem::InputImage {
+            image_url: "data:image/png;base64,original".to_string(),
+            detail: Some(WireImageDetail::Original),
+        },
+        WireContentItem::InputAudio {
+            audio_url: "data:audio/wav;base64,YXVkaW8=".to_string(),
+        },
+    ]
+}
+
+fn content_items_json() -> Value {
+    json!([
+        { "type": "input_text", "text": "hello" },
+        { "type": "input_image", "image_url": "data:image/png;base64,none" },
+        {
+            "type": "input_image",
+            "image_url": "data:image/png;base64,auto",
+            "detail": "auto",
+        },
+        {
+            "type": "input_image",
+            "image_url": "data:image/png;base64,low",
+            "detail": "low",
+        },
+        {
+            "type": "input_image",
+            "image_url": "data:image/png;base64,high",
+            "detail": "high",
+        },
+        {
+            "type": "input_image",
+            "image_url": "data:image/png;base64,original",
+            "detail": "original",
+        },
+        {
+            "type": "input_audio",
+            "audio_url": "data:audio/wav;base64,YXVkaW8=",
+        },
+    ])
+}
+
+#[test]
+fn handshake_v1_variants_are_pinned() {
+    assert_wire_round_trip(
+        ClientToHost::ClientHello(
+            ClientHello::new(
+                supported_versions(),
+                CapabilitySet::try_new([capability("required")]).expect("valid required set"),
+                CapabilitySet::try_new([capability("optional")]).expect("valid optional set"),
+            )
+            .expect("disjoint capabilities"),
+        ),
+        json!({
+            "type": "connection/hello",
+            "supportedVersions": [1],
+            "requiredCapabilities": ["required"],
+            "optionalCapabilities": ["optional"],
+        }),
+    );
+    assert_wire_round_trip(
+        HostToClient::HostHello(HostHello::new(
+            ProtocolVersion::V1,
+            CapabilitySet::try_new([capability("required")]).expect("valid capabilities"),
+        )),
+        json!({
+            "type": "connection/ready",
+            "selectedVersion": 1,
+            "capabilities": ["required"],
+        }),
+    );
+    for (reason, encoded) in [
+        (
+            HandshakeRejectReason::NoCompatibleVersion {
+                supported_versions: supported_versions(),
+            },
+            json!({
+                "type": "connection/rejected",
+                "reason": {
+                    "type": "noCompatibleVersion",
+                    "supportedVersions": [1],
+                },
+            }),
+        ),
+        (
+            HandshakeRejectReason::MissingRequiredCapability {
+                capability: capability("required"),
+            },
+            json!({
+                "type": "connection/rejected",
+                "reason": {
+                    "type": "missingRequiredCapability",
+                    "capability": "required",
+                },
+            }),
+        ),
+        (
+            HandshakeRejectReason::InvalidHello {
+                message: "invalid hello".to_string(),
+            },
+            json!({
+                "type": "connection/rejected",
+                "reason": {
+                    "type": "invalidHello",
+                    "message": "invalid hello",
+                },
+            }),
+        ),
+    ] {
+        assert_wire_round_trip(HostToClient::HandshakeRejected { reason }, encoded);
+    }
+}
+
+#[test]
+fn open_session_serializes_optional_cell_execution_limits() {
+    assert_wire_round_trip(
+        HostRequest::OpenSession {
+            session_id: session_id(),
+            cell_execution_limits: Some(WireSessionCellExecutionLimits {
+                max_yield_time_ms: Some(250),
+                max_heap_size_bytes: Some(16 * 1024 * 1024),
+            }),
+        },
+        json!({
+            "method": "session/open",
+            "sessionId": "session-1",
+            "cellExecutionLimits": {
+                "maxYieldTimeMs": 250,
+                "maxHeapSizeBytes": 16 * 1024 * 1024,
+            },
+        }),
+    );
+}
+
+#[test]
+fn session_cell_execution_limits_convert_between_domain_and_wire() {
+    let domain_limits = CodeModeSessionCellExecutionLimits {
+        max_yield_time_ms: Some(250),
+        max_heap_size_bytes: Some(16_usize * 1024 * 1024),
+    };
+    let wire_limits = WireSessionCellExecutionLimits {
+        max_yield_time_ms: Some(250),
+        max_heap_size_bytes: Some(16_u64 * 1024 * 1024),
+    };
+
+    assert_eq!(
+        WireSessionCellExecutionLimits::try_from(domain_limits.clone())
+            .expect("domain limits convert to wire limits"),
+        wire_limits
+    );
+    assert_eq!(
+        CodeModeSessionCellExecutionLimits::try_from(wire_limits)
+            .expect("wire limits convert to domain limits"),
+        domain_limits
+    );
+}
+
+#[cfg(target_pointer_width = "32")]
+#[test]
+fn session_cell_execution_limits_reject_heap_sizes_that_exceed_usize() {
+    let wire_limits = WireSessionCellExecutionLimits {
+        max_yield_time_ms: None,
+        max_heap_size_bytes: Some(u64::from(u32::MAX) + 1),
+    };
+
+    assert!(CodeModeSessionCellExecutionLimits::try_from(wire_limits).is_err());
+}
+
+#[test]
+fn client_to_host_v1_variants_are_pinned() {
+    let execute_request = execute_request();
+    for (id, request, encoded_request) in [
+        (
+            request_id(/*value*/ 1),
+            HostRequest::OpenSession {
+                session_id: session_id(),
+                cell_execution_limits: None,
+            },
+            json!({ "method": "session/open", "sessionId": "session-1" }),
+        ),
+        (
+            request_id(/*value*/ 2),
+            HostRequest::Execute {
+                session_id: session_id(),
+                request: execute_request,
+            },
+            json!({
+                "method": "session/execute",
+                "sessionId": "session-1",
+                "request": {
+                    "tool_call_id": "call-1",
+                    "enabled_tools": [
+                        {
+                            "name": "function_tool",
+                            "tool_name": { "name": "function_tool", "namespace": null },
+                            "description": "function tool",
+                            "kind": "function",
+                            "input_schema": { "type": "object" },
+                            "output_schema": null,
+                        },
+                        {
+                            "name": "freeform_tool",
+                            "tool_name": {
+                                "name": "freeform_tool",
+                                "namespace": "mcp__sample__",
+                            },
+                            "description": "freeform tool",
+                            "kind": "freeform",
+                            "input_schema": null,
+                            "output_schema": { "type": "string" },
+                        },
+                    ],
+                    "source": "text('hello');",
+                    "yield_time_ms": 25,
+                    "max_output_tokens": 100,
+                },
+            }),
+        ),
+        (
+            request_id(/*value*/ 3),
+            HostRequest::Wait {
+                session_id: session_id(),
+                request: WireWaitRequest {
+                    cell_id: cell_id("cell-1"),
+                    yield_time_ms: 50,
+                },
+            },
+            json!({
+                "method": "session/wait",
+                "sessionId": "session-1",
+                "request": { "cell_id": "cell-1", "yield_time_ms": 50 },
+            }),
+        ),
+        (
+            request_id(/*value*/ 4),
+            HostRequest::Terminate {
+                session_id: session_id(),
+                cell_id: cell_id("cell-1"),
+            },
+            json!({
+                "method": "session/terminate",
+                "sessionId": "session-1",
+                "cellId": "cell-1",
+            }),
+        ),
+        (
+            request_id(/*value*/ 5),
+            HostRequest::ShutdownSession {
+                session_id: session_id(),
+            },
+            json!({ "method": "session/shutdown", "sessionId": "session-1" }),
+        ),
+    ] {
+        assert_wire_round_trip(
+            ClientToHost::Request { id, request },
+            json!({
+                "type": "operation/request",
+                "id": id,
+                "request": encoded_request,
+            }),
+        );
+    }
+
+    for (id, result, encoded_result) in [
+        (
+            delegate_request_id(/*value*/ 6),
+            WireResult::Ok {
+                value: DelegateResponse::ToolResult {
+                    result: json!({ "answer": 42 }),
+                },
+            },
+            json!({
+                "status": "ok",
+                "value": { "type": "tool/result", "result": { "answer": 42 } },
+            }),
+        ),
+        (
+            delegate_request_id(/*value*/ 7),
+            WireResult::Ok {
+                value: DelegateResponse::NotificationDelivered,
+            },
+            json!({
+                "status": "ok",
+                "value": { "type": "notification/delivered" },
+            }),
+        ),
+        (
+            delegate_request_id(/*value*/ 8),
+            WireResult::Err {
+                message: "delegate failed".to_string(),
+            },
+            json!({ "status": "error", "message": "delegate failed" }),
+        ),
+    ] {
+        assert_wire_round_trip(
+            ClientToHost::DelegateResponse { id, result },
+            json!({
+                "type": "delegate/response",
+                "id": id,
+                "result": encoded_result,
+            }),
+        );
+    }
+
+    assert_wire_round_trip(
+        ClientToHost::CancelRequest {
+            id: request_id(/*value*/ 9),
+        },
+        json!({
+            "type": "operation/cancel",
+            "id": 9,
+        }),
+    );
+}
+
+#[test]
+fn host_to_client_v1_variants_are_pinned() {
+    for (id, response, encoded_response) in [
+        (
+            request_id(/*value*/ 1),
+            HostResponse::SessionReady {
+                session_id: session_id(),
+            },
+            json!({ "type": "session/ready", "sessionId": "session-1" }),
+        ),
+        (
+            request_id(/*value*/ 2),
+            HostResponse::ExecutionStarted {
+                cell_id: cell_id("cell-1"),
+            },
+            json!({ "type": "execution/started", "cellId": "cell-1" }),
+        ),
+        (
+            request_id(/*value*/ 3),
+            HostResponse::WaitCompleted {
+                outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Yielded {
+                    code_mode_host_duration_ns: 0,
+                    cell_id: cell_id("cell-1"),
+                    content_items: content_items(),
+                }),
+            },
+            json!({
+                "type": "wait/completed",
+                "outcome": {
+                    "LiveCell": {
+                        "Yielded": {
+                            "cell_id": "cell-1",
+                            "content_items": content_items_json(),
+                            "code_mode_host_duration_ns": 0,
+                        },
+                    },
+                },
+            }),
+        ),
+        (
+            request_id(/*value*/ 4),
+            HostResponse::WaitCompleted {
+                outcome: WireWaitOutcome::MissingCell(WireRuntimeResponse::Result {
+                    code_mode_host_duration_ns: 0,
+                    cell_id: cell_id("missing-cell"),
+                    content_items: Vec::new(),
+                    error_text: Some("cell not found".to_string()),
+                }),
+            },
+            json!({
+                "type": "wait/completed",
+                "outcome": {
+                    "MissingCell": {
+                        "Result": {
+                            "cell_id": "missing-cell",
+                            "content_items": [],
+                            "error_text": "cell not found",
+                            "code_mode_host_duration_ns": 0,
+                        },
+                    },
+                },
+            }),
+        ),
+        (
+            request_id(/*value*/ 5),
+            HostResponse::SessionClosed {
+                session_id: session_id(),
+            },
+            json!({ "type": "session/closed", "sessionId": "session-1" }),
+        ),
+    ] {
+        assert_wire_round_trip(
+            HostToClient::Response {
+                id,
+                result: WireResult::Ok { value: response },
+            },
+            json!({
+                "type": "operation/response",
+                "id": id,
+                "result": { "status": "ok", "value": encoded_response },
+            }),
+        );
+    }
+    assert_wire_round_trip(
+        HostToClient::Response {
+            id: request_id(/*value*/ 6),
+            result: WireResult::Err {
+                message: "operation failed".to_string(),
+            },
+        },
+        json!({
+            "type": "operation/response",
+            "id": 6,
+            "result": { "status": "error", "message": "operation failed" },
+        }),
+    );
+
+    assert_wire_round_trip(
+        HostToClient::InitialResponse {
+            id: request_id(/*value*/ 7),
+            result: WireResult::Ok {
+                value: WireRuntimeResponse::Terminated {
+                    code_mode_host_duration_ns: 0,
+                    cell_id: cell_id("cell-1"),
+                    content_items: Vec::new(),
+                },
+            },
+        },
+        json!({
+            "type": "execute/initialResponse",
+            "id": 7,
+            "result": {
+                "status": "ok",
+                "value": {
+                    "Terminated": {
+                        "cell_id": "cell-1",
+                        "content_items": [],
+                        "code_mode_host_duration_ns": 0,
+                    },
+                },
+            },
+        }),
+    );
+    assert_wire_round_trip(
+        HostToClient::InitialResponse {
+            id: request_id(/*value*/ 8),
+            result: WireResult::Err {
+                message: "execution failed".to_string(),
+            },
+        },
+        json!({
+            "type": "execute/initialResponse",
+            "id": 8,
+            "result": { "status": "error", "message": "execution failed" },
+        }),
+    );
+
+    assert_wire_round_trip(
+        HostToClient::DelegateRequest {
+            id: delegate_request_id(/*value*/ 9),
+            session_id: session_id(),
+            request: DelegateRequest::InvokeTool {
+                invocation: WireNestedToolCall {
+                    cell_id: cell_id("cell-1"),
+                    runtime_tool_call_id: "runtime-call-1".to_string(),
+                    tool_name: WireToolName {
+                        name: "freeform_tool".to_string(),
+                        namespace: Some("mcp__sample__".to_string()),
+                    },
+                    tool_kind: WireToolKind::Freeform,
+                    input: Some(json!({ "value": 1 })),
+                },
+            },
+        },
+        json!({
+            "type": "delegate/request",
+            "id": 9,
+            "sessionId": "session-1",
+            "request": {
+                "type": "tool/invoke",
+                "invocation": {
+                    "cell_id": "cell-1",
+                    "runtime_tool_call_id": "runtime-call-1",
+                    "tool_name": {
+                        "name": "freeform_tool",
+                        "namespace": "mcp__sample__",
+                    },
+                    "tool_kind": "freeform",
+                    "input": { "value": 1 },
+                },
+            },
+        }),
+    );
+    assert_wire_round_trip(
+        HostToClient::DelegateRequest {
+            id: delegate_request_id(/*value*/ 10),
+            session_id: session_id(),
+            request: DelegateRequest::Notify {
+                call_id: "call-1".to_string(),
+                cell_id: cell_id("cell-1"),
+                text: "important".to_string(),
+            },
+        },
+        json!({
+            "type": "delegate/request",
+            "id": 10,
+            "sessionId": "session-1",
+            "request": {
+                "type": "notification/send",
+                "callId": "call-1",
+                "cellId": "cell-1",
+                "text": "important",
+            },
+        }),
+    );
+    assert_wire_round_trip(
+        HostToClient::CancelDelegateRequest {
+            id: delegate_request_id(/*value*/ 11),
+        },
+        json!({ "type": "delegate/cancel", "id": 11 }),
+    );
+    assert_wire_round_trip(
+        HostToClient::CellClosed {
+            session_id: session_id(),
+            cell_id: cell_id("cell-1"),
+        },
+        json!({
+            "type": "cell/closed",
+            "sessionId": "session-1",
+            "cellId": "cell-1",
+        }),
+    );
+}
+
+#[test]
+fn execute_request_integer_bounds_are_enforced() {
+    let wire_request = execute_request();
+    let domain_request = ExecuteRequest::try_from(wire_request.clone())
+        .expect("valid wire request converts to the domain");
+    assert_eq!(
+        WireExecuteRequest::try_from(domain_request.clone())
+            .expect("valid domain request converts to the wire"),
+        wire_request
+    );
+
+    let too_large = ExecuteRequest {
+        max_output_tokens: Some(usize::try_from(i32::MAX).expect("i32::MAX fits usize") + 1),
+        ..domain_request
+    };
+    assert!(WireExecuteRequest::try_from(too_large).is_err());
+
+    let negative = WireExecuteRequest {
+        max_output_tokens: Some(-1),
+        ..wire_request
+    };
+    assert!(ExecuteRequest::try_from(negative).is_err());
+}
+
+#[test]
+fn invalid_protocol_states_cannot_be_constructed_or_decoded() {
+    assert!(SessionId::new("").is_err());
+    assert!(Capability::new("   ").is_err());
+    assert!(ProtocolVersion::new(/*value*/ 0).is_none());
+    assert!(SupportedProtocolVersions::try_new([]).is_err());
+    assert!(
+        SupportedProtocolVersions::try_new([ProtocolVersion::V1, ProtocolVersion::V1]).is_err()
+    );
+    assert!(CapabilitySet::try_new([capability("same"), capability("same")]).is_err());
+
+    let version_two = ProtocolVersion::new(/*value*/ 2).expect("valid protocol version");
+    let versions = SupportedProtocolVersions::try_new([ProtocolVersion::V1, version_two])
+        .expect("valid versions");
+    assert!(versions.contains(ProtocolVersion::V1));
+    assert_eq!(
+        versions.iter().collect::<Vec<_>>(),
+        vec![ProtocolVersion::V1, version_two]
+    );
+
+    let overlapping = capability("overlapping");
+    assert!(
+        ClientHello::new(
+            supported_versions(),
+            CapabilitySet::try_new([overlapping.clone()]).expect("valid required set"),
+            CapabilitySet::try_new([overlapping]).expect("valid optional set"),
+        )
+        .is_err()
+    );
+
+    for invalid in [
+        json!({
+            "type": "operation/request",
+            "id": 1,
+            "request": { "method": "session/open", "sessionId": "" },
+        }),
+        json!({
+            "type": "connection/hello",
+            "supportedVersions": [],
+            "requiredCapabilities": [],
+            "optionalCapabilities": [],
+        }),
+        json!({
+            "type": "connection/hello",
+            "supportedVersions": [1],
+            "requiredCapabilities": ["overlapping"],
+            "optionalCapabilities": ["overlapping"],
+        }),
+    ] {
+        assert!(serde_json::from_value::<ClientToHost>(invalid).is_err());
+    }
+}
+
+#[test]
+fn every_nested_v1_object_rejects_unknown_fields() {
+    assert!(
+        serde_json::from_value::<ClientToHost>(json!({
+            "type": "operation/request",
+            "id": 1,
+            "request": { "method": "session/open", "sessionId": "session-1" },
+            "unexpected": true,
+        }))
+        .is_err()
+    );
+    assert!(
+        serde_json::from_value::<HostRequest>(json!({
+            "method": "session/open",
+            "sessionId": "session-1",
+            "unexpected": true,
+        }))
+        .is_err()
+    );
+    assert!(
+        serde_json::from_value::<HostRequest>(json!({
+            "method": "session/open",
+            "sessionId": "session-1",
+            "cellExecutionLimits": {
+                "maxYieldTimeMs": 250,
+                "unexpected": true,
+            },
+        }))
+        .is_err()
+    );
+    assert!(
+        serde_json::from_value::<WireExecuteRequest>(json!({
+            "tool_call_id": "call-1",
+            "enabled_tools": [],
+            "source": "text('hello');",
+            "yield_time_ms": null,
+            "max_output_tokens": null,
+            "unexpected": true,
+        }))
+        .is_err()
+    );
+    assert!(
+        serde_json::from_value::<WireToolDefinition>(json!({
+            "name": "tool",
+            "tool_name": { "name": "tool", "namespace": null },
+            "description": "tool",
+            "kind": "function",
+            "input_schema": null,
+            "output_schema": null,
+            "unexpected": true,
+        }))
+        .is_err()
+    );
+    assert!(
+        serde_json::from_value::<WireToolName>(json!({
+            "name": "tool",
+            "namespace": null,
+            "unexpected": true,
+        }))
+        .is_err()
+    );
+    assert!(
+        serde_json::from_value::<WireWaitRequest>(json!({
+            "cell_id": "cell-1",
+            "yield_time_ms": 50,
+            "unexpected": true,
+        }))
+        .is_err()
+    );
+    assert!(
+        serde_json::from_value::<WireRuntimeResponse>(json!({
+            "Yielded": {
+                "cell_id": "cell-1",
+                "content_items": [],
+                "code_mode_host_duration_ns": 0,
+                "unexpected": true,
+            },
+        }))
+        .is_err()
+    );
+    assert!(
+        serde_json::from_value::<WireContentItem>(json!({
+            "type": "input_text",
+            "text": "hello",
+            "unexpected": true,
+        }))
+        .is_err()
+    );
+    assert!(
+        serde_json::from_value::<WireNestedToolCall>(json!({
+            "cell_id": "cell-1",
+            "runtime_tool_call_id": "runtime-call-1",
+            "tool_name": { "name": "tool", "namespace": null },
+            "tool_kind": "function",
+            "input": null,
+            "unexpected": true,
+        }))
+        .is_err()
+    );
+    assert!(
+        serde_json::from_value::<HostToClient>(json!({
+            "type": "operation/response",
+            "id": 1,
+            "result": {
+                "status": "ok",
+                "value": { "type": "session/ready", "sessionId": "session-1" },
+            },
+            "unexpected": true,
+        }))
+        .is_err()
+    );
+}

codex-rs/code-mode-protocol/src/host/message.rs

@@ -0,0 +1,264 @@
+use std::fmt;
+
+use serde::Deserialize;
+use serde::Serialize;
+use serde_json::Value as JsonValue;
+
+use super::Capability;
+use super::CapabilitySet;
+use super::DelegateRequestId;
+use super::HandshakeRejectReason;
+use super::ProtocolVersion;
+use super::RequestId;
+use super::SessionId;
+use super::SupportedProtocolVersions;
+use super::WireCellId;
+use super::WireExecuteRequest;
+use super::WireNestedToolCall;
+use super::WireRuntimeResponse;
+use super::WireSessionCellExecutionLimits;
+use super::WireWaitOutcome;
+use super::WireWaitRequest;
+
+#[derive(Clone, Debug, PartialEq, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ClientHello {
+    supported_versions: SupportedProtocolVersions,
+    required_capabilities: CapabilitySet,
+    optional_capabilities: CapabilitySet,
+}
+
+impl ClientHello {
+    pub fn new(
+        supported_versions: SupportedProtocolVersions,
+        required_capabilities: CapabilitySet,
+        optional_capabilities: CapabilitySet,
+    ) -> Result<Self, ClientHelloError> {
+        if let Some(capability) = required_capabilities
+            .iter()
+            .find(|capability| optional_capabilities.contains(capability))
+        {
+            return Err(ClientHelloError::OverlappingCapability(capability.clone()));
+        }
+        Ok(Self {
+            supported_versions,
+            required_capabilities,
+            optional_capabilities,
+        })
+    }
+
+    pub fn supported_versions(&self) -> &SupportedProtocolVersions {
+        &self.supported_versions
+    }
+
+    pub fn required_capabilities(&self) -> &CapabilitySet {
+        &self.required_capabilities
+    }
+
+    pub fn optional_capabilities(&self) -> &CapabilitySet {
+        &self.optional_capabilities
+    }
+}
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+struct ClientHelloWire {
+    supported_versions: SupportedProtocolVersions,
+    required_capabilities: CapabilitySet,
+    optional_capabilities: CapabilitySet,
+}
+
+impl<'de> Deserialize<'de> for ClientHello {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let wire = ClientHelloWire::deserialize(deserializer)?;
+        Self::new(
+            wire.supported_versions,
+            wire.required_capabilities,
+            wire.optional_capabilities,
+        )
+        .map_err(serde::de::Error::custom)
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub enum ClientHelloError {
+    OverlappingCapability(Capability),
+}
+
+impl fmt::Display for ClientHelloError {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::OverlappingCapability(capability) => write!(
+                formatter,
+                "capability `{capability}` cannot be both required and optional"
+            ),
+        }
+    }
+}
+
+impl std::error::Error for ClientHelloError {}
+
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+pub struct HostHello {
+    selected_version: ProtocolVersion,
+    capabilities: CapabilitySet,
+}
+
+impl HostHello {
+    pub fn new(selected_version: ProtocolVersion, capabilities: CapabilitySet) -> Self {
+        Self {
+            selected_version,
+            capabilities,
+        }
+    }
+
+    pub fn selected_version(&self) -> ProtocolVersion {
+        self.selected_version
+    }
+
+    pub fn capabilities(&self) -> &CapabilitySet {
+        &self.capabilities
+    }
+}
+
+/// Messages sent from a client to the code-mode host.
+#[derive(Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
+pub enum ClientToHost {
+    #[serde(rename = "connection/hello")]
+    ClientHello(ClientHello),
+    #[serde(rename = "operation/request")]
+    Request { id: RequestId, request: HostRequest },
+    #[serde(rename = "operation/cancel")]
+    CancelRequest { id: RequestId },
+    #[serde(rename = "delegate/response")]
+    DelegateResponse {
+        id: DelegateRequestId,
+        result: WireResult<DelegateResponse>,
+    },
+}
+
+/// Messages sent from the code-mode host to a client.
+#[derive(Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
+pub enum HostToClient {
+    #[serde(rename = "connection/ready")]
+    HostHello(HostHello),
+    #[serde(rename = "connection/rejected")]
+    HandshakeRejected { reason: HandshakeRejectReason },
+    #[serde(rename = "operation/response")]
+    Response {
+        id: RequestId,
+        result: WireResult<HostResponse>,
+    },
+    #[serde(rename = "execute/initialResponse")]
+    InitialResponse {
+        id: RequestId,
+        result: WireResult<WireRuntimeResponse>,
+    },
+    #[serde(rename = "delegate/request")]
+    DelegateRequest {
+        id: DelegateRequestId,
+        session_id: SessionId,
+        request: DelegateRequest,
+    },
+    #[serde(rename = "delegate/cancel")]
+    CancelDelegateRequest { id: DelegateRequestId },
+    #[serde(rename = "cell/closed")]
+    CellClosed {
+        session_id: SessionId,
+        cell_id: WireCellId,
+    },
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, tag = "method", rename_all_fields = "camelCase")]
+pub enum HostRequest {
+    #[serde(rename = "session/open")]
+    OpenSession {
+        session_id: SessionId,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        cell_execution_limits: Option<WireSessionCellExecutionLimits>,
+    },
+    #[serde(rename = "session/execute")]
+    Execute {
+        session_id: SessionId,
+        request: WireExecuteRequest,
+    },
+    #[serde(rename = "session/wait")]
+    Wait {
+        session_id: SessionId,
+        request: WireWaitRequest,
+    },
+    #[serde(rename = "session/terminate")]
+    Terminate {
+        session_id: SessionId,
+        cell_id: WireCellId,
+    },
+    #[serde(rename = "session/shutdown")]
+    ShutdownSession { session_id: SessionId },
+}
+
+#[derive(Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
+pub enum HostResponse {
+    #[serde(rename = "session/ready")]
+    SessionReady { session_id: SessionId },
+    #[serde(rename = "execution/started")]
+    ExecutionStarted { cell_id: WireCellId },
+    #[serde(rename = "wait/completed")]
+    WaitCompleted { outcome: WireWaitOutcome },
+    #[serde(rename = "session/closed")]
+    SessionClosed { session_id: SessionId },
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
+pub enum DelegateRequest {
+    #[serde(rename = "tool/invoke")]
+    InvokeTool { invocation: WireNestedToolCall },
+    #[serde(rename = "notification/send")]
+    Notify {
+        call_id: String,
+        cell_id: WireCellId,
+        text: String,
+    },
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
+pub enum DelegateResponse {
+    #[serde(rename = "tool/result")]
+    ToolResult { result: JsonValue },
+    #[serde(rename = "notification/delivered")]
+    NotificationDelivered,
+}
+
+#[derive(Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, tag = "status", rename_all_fields = "camelCase")]
+pub enum WireResult<T> {
+    #[serde(rename = "ok")]
+    Ok { value: T },
+    #[serde(rename = "error")]
+    Err { message: String },
+}
+
+impl<T> WireResult<T> {
+    pub fn from_result(result: Result<T, String>) -> Self {
+        match result {
+            Ok(value) => Self::Ok { value },
+            Err(message) => Self::Err { message },
+        }
+    }
+
+    pub fn into_result(self) -> Result<T, String> {
+        match self {
+            Self::Ok { value } => Ok(value),
+            Self::Err { message } => Err(message),
+        }
+    }
+}

codex-rs/code-mode-protocol/src/host/mod.rs

@@ -0,0 +1,62 @@
+//! Messages and framing for the code-mode host boundary.
+//!
+//! Protocol version 1 multiplexes session operations and delegate callbacks by
+//! request ID over one ordered connection.
+
+mod codec;
+mod error;
+mod message;
+mod payload;
+mod types;
+
+/// Maximum number of unresolved delegate callbacks allowed per host connection.
+pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024;
+
+/// Negotiated support for cell execution resource limits on `session/open`.
+pub const SESSION_RESOURCE_LIMITS_CAPABILITY: &str = "session-cell-execution-resource-limits";
+
+pub use codec::EncodedFrame;
+pub use codec::FramedReader;
+pub use codec::FramedWriter;
+pub use codec::MAX_FRAME_BYTES;
+pub use error::HandshakeRejectReason;
+pub use message::ClientHello;
+pub use message::ClientHelloError;
+pub use message::ClientToHost;
+pub use message::DelegateRequest;
+pub use message::DelegateResponse;
+pub use message::HostHello;
+pub use message::HostRequest;
+pub use message::HostResponse;
+pub use message::HostToClient;
+pub use message::WireResult;
+pub use payload::WireCellId;
+pub use payload::WireContentItem;
+pub use payload::WireExecuteRequest;
+pub use payload::WireImageDetail;
+pub use payload::WireNestedToolCall;
+pub use payload::WireRuntimeResponse;
+pub use payload::WireSessionCellExecutionLimits;
+pub use payload::WireToolDefinition;
+pub use payload::WireToolKind;
+pub use payload::WireToolName;
+pub use payload::WireWaitOutcome;
+pub use payload::WireWaitRequest;
+pub use types::Capability;
+pub use types::CapabilitySet;
+pub use types::DelegateRequestId;
+pub use types::DuplicateCapability;
+pub use types::InvalidIdentifier;
+pub use types::InvalidSupportedProtocolVersions;
+pub use types::ProtocolVersion;
+pub use types::RequestId;
+pub use types::SessionId;
+pub use types::SupportedProtocolVersions;
+
+#[cfg(test)]
+#[path = "host_tests.rs"]
+mod tests;
+
+#[cfg(test)]
+#[path = "codec_tests.rs"]
+mod codec_tests;

codex-rs/code-mode-protocol/src/host/payload.rs

@@ -0,0 +1,492 @@
+use std::num::TryFromIntError;
+use std::time::Duration;
+
+use codex_protocol::ToolName;
+use serde::Deserialize;
+use serde::Serialize;
+use serde_json::Value as JsonValue;
+
+use crate::CellId;
+use crate::CodeModeNestedToolCall;
+use crate::CodeModeSessionCellExecutionLimits;
+use crate::CodeModeToolKind;
+use crate::ExecuteRequest;
+use crate::FunctionCallOutputContentItem;
+use crate::ImageDetail;
+use crate::MissingCodeModeHostDuration;
+use crate::RuntimeResponse;
+use crate::ToolDefinition;
+use crate::WaitOutcome;
+use crate::WaitRequest;
+
+/// The per-cell execution limits carried by a V1 session-open request.
+#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, rename_all = "camelCase")]
+pub struct WireSessionCellExecutionLimits {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub max_yield_time_ms: Option<u64>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub max_heap_size_bytes: Option<u64>,
+}
+
+impl TryFrom<CodeModeSessionCellExecutionLimits> for WireSessionCellExecutionLimits {
+    type Error = TryFromIntError;
+
+    fn try_from(value: CodeModeSessionCellExecutionLimits) -> Result<Self, Self::Error> {
+        Ok(Self {
+            max_yield_time_ms: value.max_yield_time_ms,
+            max_heap_size_bytes: value.max_heap_size_bytes.map(u64::try_from).transpose()?,
+        })
+    }
+}
+
+impl TryFrom<WireSessionCellExecutionLimits> for CodeModeSessionCellExecutionLimits {
+    type Error = TryFromIntError;
+
+    fn try_from(value: WireSessionCellExecutionLimits) -> Result<Self, Self::Error> {
+        Ok(Self {
+            max_yield_time_ms: value.max_yield_time_ms,
+            max_heap_size_bytes: value.max_heap_size_bytes.map(usize::try_from).transpose()?,
+        })
+    }
+}
+
+/// A cell identifier with a wire representation owned by protocol V1.
+#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
+#[serde(transparent)]
+pub struct WireCellId(String);
+
+impl WireCellId {
+    pub fn new(value: impl Into<String>) -> Self {
+        Self(value.into())
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
+impl From<CellId> for WireCellId {
+    fn from(value: CellId) -> Self {
+        Self(value.as_str().to_string())
+    }
+}
+
+impl From<&CellId> for WireCellId {
+    fn from(value: &CellId) -> Self {
+        Self(value.as_str().to_string())
+    }
+}
+
+impl From<WireCellId> for CellId {
+    fn from(value: WireCellId) -> Self {
+        Self::new(value.0)
+    }
+}
+
+/// The V1 wire representation of a tool's stable name.
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(deny_unknown_fields)]
+pub struct WireToolName {
+    pub name: String,
+    pub namespace: Option<String>,
+}
+
+impl From<ToolName> for WireToolName {
+    fn from(value: ToolName) -> Self {
+        Self {
+            name: value.name,
+            namespace: value.namespace,
+        }
+    }
+}
+
+impl From<WireToolName> for ToolName {
+    fn from(value: WireToolName) -> Self {
+        Self::new(value.namespace, value.name)
+    }
+}
+
+/// The tool invocation shape supported by protocol V1.
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum WireToolKind {
+    Function,
+    Freeform,
+}
+
+impl From<CodeModeToolKind> for WireToolKind {
+    fn from(value: CodeModeToolKind) -> Self {
+        match value {
+            CodeModeToolKind::Function => Self::Function,
+            CodeModeToolKind::Freeform => Self::Freeform,
+        }
+    }
+}
+
+impl From<WireToolKind> for CodeModeToolKind {
+    fn from(value: WireToolKind) -> Self {
+        match value {
+            WireToolKind::Function => Self::Function,
+            WireToolKind::Freeform => Self::Freeform,
+        }
+    }
+}
+
+/// A V1 tool definition embedded in an execute request.
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields)]
+pub struct WireToolDefinition {
+    pub name: String,
+    pub tool_name: WireToolName,
+    pub description: String,
+    pub kind: WireToolKind,
+    pub input_schema: Option<JsonValue>,
+    pub output_schema: Option<JsonValue>,
+}
+
+impl From<ToolDefinition> for WireToolDefinition {
+    fn from(value: ToolDefinition) -> Self {
+        Self {
+            name: value.name,
+            tool_name: value.tool_name.into(),
+            description: value.description,
+            kind: value.kind.into(),
+            input_schema: value.input_schema,
+            output_schema: value.output_schema,
+        }
+    }
+}
+
+impl From<WireToolDefinition> for ToolDefinition {
+    fn from(value: WireToolDefinition) -> Self {
+        Self {
+            name: value.name,
+            tool_name: value.tool_name.into(),
+            description: value.description,
+            kind: value.kind.into(),
+            input_schema: value.input_schema,
+            output_schema: value.output_schema,
+        }
+    }
+}
+
+/// The complete execute request shape supported by protocol V1.
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields)]
+pub struct WireExecuteRequest {
+    pub tool_call_id: String,
+    pub enabled_tools: Vec<WireToolDefinition>,
+    pub source: String,
+    pub yield_time_ms: Option<u64>,
+    pub max_output_tokens: Option<i32>,
+}
+
+impl TryFrom<ExecuteRequest> for WireExecuteRequest {
+    type Error = TryFromIntError;
+
+    fn try_from(value: ExecuteRequest) -> Result<Self, Self::Error> {
+        Ok(Self {
+            tool_call_id: value.tool_call_id,
+            enabled_tools: value.enabled_tools.into_iter().map(Into::into).collect(),
+            source: value.source,
+            yield_time_ms: value.yield_time_ms,
+            max_output_tokens: value.max_output_tokens.map(i32::try_from).transpose()?,
+        })
+    }
+}
+
+impl TryFrom<WireExecuteRequest> for ExecuteRequest {
+    type Error = TryFromIntError;
+
+    fn try_from(value: WireExecuteRequest) -> Result<Self, Self::Error> {
+        Ok(Self {
+            tool_call_id: value.tool_call_id,
+            enabled_tools: value.enabled_tools.into_iter().map(Into::into).collect(),
+            source: value.source,
+            yield_time_ms: value.yield_time_ms,
+            max_output_tokens: value.max_output_tokens.map(usize::try_from).transpose()?,
+        })
+    }
+}
+
+/// The complete wait request shape supported by protocol V1.
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(deny_unknown_fields)]
+pub struct WireWaitRequest {
+    pub cell_id: WireCellId,
+    pub yield_time_ms: u64,
+}
+
+impl From<WaitRequest> for WireWaitRequest {
+    fn from(value: WaitRequest) -> Self {
+        Self {
+            cell_id: value.cell_id.into(),
+            yield_time_ms: value.yield_time_ms,
+        }
+    }
+}
+
+impl From<WireWaitRequest> for WaitRequest {
+    fn from(value: WireWaitRequest) -> Self {
+        Self {
+            cell_id: value.cell_id.into(),
+            yield_time_ms: value.yield_time_ms,
+        }
+    }
+}
+
+/// Image detail values accepted in a V1 runtime response.
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(rename_all = "lowercase")]
+pub enum WireImageDetail {
+    Auto,
+    Low,
+    High,
+    Original,
+}
+
+impl From<ImageDetail> for WireImageDetail {
+    fn from(value: ImageDetail) -> Self {
+        match value {
+            ImageDetail::Auto => Self::Auto,
+            ImageDetail::Low => Self::Low,
+            ImageDetail::High => Self::High,
+            ImageDetail::Original => Self::Original,
+        }
+    }
+}
+
+impl From<WireImageDetail> for ImageDetail {
+    fn from(value: WireImageDetail) -> Self {
+        match value {
+            WireImageDetail::Auto => Self::Auto,
+            WireImageDetail::Low => Self::Low,
+            WireImageDetail::High => Self::High,
+            WireImageDetail::Original => Self::Original,
+        }
+    }
+}
+
+/// One output item emitted by a V1 runtime response.
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields, tag = "type", rename_all = "snake_case")]
+pub enum WireContentItem {
+    InputText {
+        text: String,
+    },
+    InputImage {
+        image_url: String,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        detail: Option<WireImageDetail>,
+    },
+    InputAudio {
+        audio_url: String,
+    },
+}
+
+impl From<FunctionCallOutputContentItem> for WireContentItem {
+    fn from(value: FunctionCallOutputContentItem) -> Self {
+        match value {
+            FunctionCallOutputContentItem::InputText { text } => Self::InputText { text },
+            FunctionCallOutputContentItem::InputImage { image_url, detail } => Self::InputImage {
+                image_url,
+                detail: detail.map(Into::into),
+            },
+            FunctionCallOutputContentItem::InputAudio { audio_url } => {
+                Self::InputAudio { audio_url }
+            }
+        }
+    }
+}
+
+impl From<WireContentItem> for FunctionCallOutputContentItem {
+    fn from(value: WireContentItem) -> Self {
+        match value {
+            WireContentItem::InputText { text } => Self::InputText { text },
+            WireContentItem::InputImage { image_url, detail } => Self::InputImage {
+                image_url,
+                detail: detail.map(Into::into),
+            },
+            WireContentItem::InputAudio { audio_url } => Self::InputAudio { audio_url },
+        }
+    }
+}
+
+/// Runtime output returned over the V1 host connection.
+///
+/// Host time is required and covers this request, not the cell lifetime. The
+/// app-server and host run at the same version; no negotiation is needed.
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields)]
+pub enum WireRuntimeResponse {
+    Yielded {
+        cell_id: WireCellId,
+        content_items: Vec<WireContentItem>,
+        code_mode_host_duration_ns: u64,
+    },
+    Terminated {
+        cell_id: WireCellId,
+        content_items: Vec<WireContentItem>,
+        code_mode_host_duration_ns: u64,
+    },
+    Result {
+        cell_id: WireCellId,
+        content_items: Vec<WireContentItem>,
+        error_text: Option<String>,
+        code_mode_host_duration_ns: u64,
+    },
+}
+
+impl TryFrom<RuntimeResponse> for WireRuntimeResponse {
+    type Error = MissingCodeModeHostDuration;
+
+    /// Preserves the response's timing; the host handler must record it first.
+    fn try_from(value: RuntimeResponse) -> Result<Self, Self::Error> {
+        Ok(match value {
+            RuntimeResponse::Yielded {
+                cell_id,
+                content_items,
+                code_mode_host_duration,
+            } => {
+                let code_mode_host_duration =
+                    code_mode_host_duration.ok_or(MissingCodeModeHostDuration)?;
+                Self::Yielded {
+                    cell_id: cell_id.into(),
+                    content_items: content_items.into_iter().map(Into::into).collect(),
+                    code_mode_host_duration_ns: u64::try_from(code_mode_host_duration.as_nanos())
+                        .unwrap_or(u64::MAX),
+                }
+            }
+            RuntimeResponse::Terminated {
+                cell_id,
+                content_items,
+                code_mode_host_duration,
+            } => {
+                let code_mode_host_duration =
+                    code_mode_host_duration.ok_or(MissingCodeModeHostDuration)?;
+                Self::Terminated {
+                    cell_id: cell_id.into(),
+                    content_items: content_items.into_iter().map(Into::into).collect(),
+                    code_mode_host_duration_ns: u64::try_from(code_mode_host_duration.as_nanos())
+                        .unwrap_or(u64::MAX),
+                }
+            }
+            RuntimeResponse::Result {
+                cell_id,
+                content_items,
+                error_text,
+                code_mode_host_duration,
+            } => {
+                let code_mode_host_duration =
+                    code_mode_host_duration.ok_or(MissingCodeModeHostDuration)?;
+                Self::Result {
+                    cell_id: cell_id.into(),
+                    content_items: content_items.into_iter().map(Into::into).collect(),
+                    error_text,
+                    code_mode_host_duration_ns: u64::try_from(code_mode_host_duration.as_nanos())
+                        .unwrap_or(u64::MAX),
+                }
+            }
+        })
+    }
+}
+
+impl From<WireRuntimeResponse> for RuntimeResponse {
+    fn from(value: WireRuntimeResponse) -> Self {
+        match value {
+            WireRuntimeResponse::Yielded {
+                cell_id,
+                content_items,
+                code_mode_host_duration_ns,
+            } => Self::Yielded {
+                cell_id: cell_id.into(),
+                content_items: content_items.into_iter().map(Into::into).collect(),
+                code_mode_host_duration: Some(Duration::from_nanos(code_mode_host_duration_ns)),
+            },
+            WireRuntimeResponse::Terminated {
+                cell_id,
+                content_items,
+                code_mode_host_duration_ns,
+            } => Self::Terminated {
+                cell_id: cell_id.into(),
+                content_items: content_items.into_iter().map(Into::into).collect(),
+                code_mode_host_duration: Some(Duration::from_nanos(code_mode_host_duration_ns)),
+            },
+            WireRuntimeResponse::Result {
+                cell_id,
+                content_items,
+                error_text,
+                code_mode_host_duration_ns,
+            } => Self::Result {
+                cell_id: cell_id.into(),
+                content_items: content_items.into_iter().map(Into::into).collect(),
+                error_text,
+                code_mode_host_duration: Some(Duration::from_nanos(code_mode_host_duration_ns)),
+            },
+        }
+    }
+}
+
+/// Whether a waited-for cell remained live in protocol V1.
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields)]
+pub enum WireWaitOutcome {
+    LiveCell(WireRuntimeResponse),
+    MissingCell(WireRuntimeResponse),
+}
+
+impl TryFrom<WaitOutcome> for WireWaitOutcome {
+    type Error = MissingCodeModeHostDuration;
+
+    fn try_from(value: WaitOutcome) -> Result<Self, Self::Error> {
+        Ok(match value {
+            WaitOutcome::LiveCell(response) => Self::LiveCell(response.try_into()?),
+            WaitOutcome::MissingCell(response) => Self::MissingCell(response.try_into()?),
+        })
+    }
+}
+
+impl From<WireWaitOutcome> for WaitOutcome {
+    fn from(value: WireWaitOutcome) -> Self {
+        match value {
+            WireWaitOutcome::LiveCell(response) => Self::LiveCell(response.into()),
+            WireWaitOutcome::MissingCell(response) => Self::MissingCell(response.into()),
+        }
+    }
+}
+
+/// A nested tool invocation sent over the V1 host connection.
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+#[serde(deny_unknown_fields)]
+pub struct WireNestedToolCall {
+    pub cell_id: WireCellId,
+    pub runtime_tool_call_id: String,
+    pub tool_name: WireToolName,
+    pub tool_kind: WireToolKind,
+    pub input: Option<JsonValue>,
+}
+
+impl From<CodeModeNestedToolCall> for WireNestedToolCall {
+    fn from(value: CodeModeNestedToolCall) -> Self {
+        Self {
+            cell_id: value.cell_id.into(),
+            runtime_tool_call_id: value.runtime_tool_call_id,
+            tool_name: value.tool_name.into(),
+            tool_kind: value.tool_kind.into(),
+            input: value.input,
+        }
+    }
+}
+
+impl From<WireNestedToolCall> for CodeModeNestedToolCall {
+    fn from(value: WireNestedToolCall) -> Self {
+        Self {
+            cell_id: value.cell_id.into(),
+            runtime_tool_call_id: value.runtime_tool_call_id,
+            tool_name: value.tool_name.into(),
+            tool_kind: value.tool_kind.into(),
+            input: value.input,
+        }
+    }
+}

codex-rs/code-mode-protocol/src/host/types.rs

@@ -0,0 +1,248 @@
+use std::collections::BTreeSet;
+use std::fmt;
+use std::num::NonZeroU32;
+
+use serde::Deserialize;
+use serde::Deserializer;
+use serde::Serialize;
+use serde::Serializer;
+use serde::de::Error as _;
+
+/// Correlates one client operation request with the host's response.
+#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+#[serde(transparent)]
+pub struct RequestId(i64);
+
+impl RequestId {
+    pub const fn new(value: i64) -> Self {
+        Self(value)
+    }
+}
+
+/// Correlates one host delegate request with the client's response.
+#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+#[serde(transparent)]
+pub struct DelegateRequestId(i64);
+
+impl DelegateRequestId {
+    pub const fn new(value: i64) -> Self {
+        Self(value)
+    }
+}
+
+#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+#[serde(transparent)]
+pub struct ProtocolVersion(NonZeroU32);
+
+impl ProtocolVersion {
+    pub const V1: Self = Self(NonZeroU32::MIN);
+
+    pub const fn new(value: u32) -> Option<Self> {
+        match NonZeroU32::new(value) {
+            Some(value) => Some(Self(value)),
+            None => None,
+        }
+    }
+
+    pub const fn get(self) -> u32 {
+        self.0.get()
+    }
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct InvalidIdentifier;
+
+impl fmt::Display for InvalidIdentifier {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter.write_str("identifier must not be empty")
+    }
+}
+
+impl std::error::Error for InvalidIdentifier {}
+
+#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
+struct NonEmptyString(String);
+
+impl NonEmptyString {
+    fn new(value: impl Into<String>) -> Result<Self, InvalidIdentifier> {
+        let value = value.into();
+        if value.trim().is_empty() {
+            Err(InvalidIdentifier)
+        } else {
+            Ok(Self(value))
+        }
+    }
+
+    fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
+impl Serialize for NonEmptyString {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        self.0.serialize(serializer)
+    }
+}
+
+impl<'de> Deserialize<'de> for NonEmptyString {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        Self::new(String::deserialize(deserializer)?).map_err(D::Error::custom)
+    }
+}
+
+/// A named protocol feature advertised during connection negotiation.
+#[derive(Clone, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+#[serde(transparent)]
+pub struct Capability(NonEmptyString);
+
+impl Capability {
+    pub fn new(value: impl Into<String>) -> Result<Self, InvalidIdentifier> {
+        NonEmptyString::new(value).map(Self)
+    }
+
+    pub fn as_str(&self) -> &str {
+        self.0.as_str()
+    }
+}
+
+impl fmt::Display for Capability {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter.write_str(self.as_str())
+    }
+}
+
+/// Identifies one logical code-mode session on a connection.
+#[derive(Clone, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+#[serde(transparent)]
+pub struct SessionId(NonEmptyString);
+
+impl SessionId {
+    pub fn new(value: impl Into<String>) -> Result<Self, InvalidIdentifier> {
+        NonEmptyString::new(value).map(Self)
+    }
+
+    pub fn as_str(&self) -> &str {
+        self.0.as_str()
+    }
+}
+
+impl fmt::Display for SessionId {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter.write_str(self.as_str())
+    }
+}
+
+#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize)]
+#[serde(transparent)]
+pub struct CapabilitySet(BTreeSet<Capability>);
+
+impl CapabilitySet {
+    pub fn empty() -> Self {
+        Self::default()
+    }
+
+    pub fn try_new(
+        capabilities: impl IntoIterator<Item = Capability>,
+    ) -> Result<Self, DuplicateCapability> {
+        let mut unique = BTreeSet::new();
+        for capability in capabilities {
+            if !unique.insert(capability.clone()) {
+                return Err(DuplicateCapability { capability });
+            }
+        }
+        Ok(Self(unique))
+    }
+
+    pub fn contains(&self, capability: &Capability) -> bool {
+        self.0.contains(capability)
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = &Capability> {
+        self.0.iter()
+    }
+}
+
+impl<'de> Deserialize<'de> for CapabilitySet {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        Self::try_new(Vec::<Capability>::deserialize(deserializer)?).map_err(D::Error::custom)
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct DuplicateCapability {
+    capability: Capability,
+}
+
+impl fmt::Display for DuplicateCapability {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(formatter, "duplicate capability `{}`", self.capability)
+    }
+}
+
+impl std::error::Error for DuplicateCapability {}
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
+#[serde(transparent)]
+pub struct SupportedProtocolVersions(BTreeSet<ProtocolVersion>);
+
+impl SupportedProtocolVersions {
+    pub fn try_new(
+        versions: impl IntoIterator<Item = ProtocolVersion>,
+    ) -> Result<Self, InvalidSupportedProtocolVersions> {
+        let mut unique = BTreeSet::new();
+        for version in versions {
+            if !unique.insert(version) {
+                return Err(InvalidSupportedProtocolVersions::Duplicate(version));
+            }
+        }
+        if unique.is_empty() {
+            return Err(InvalidSupportedProtocolVersions::Empty);
+        }
+        Ok(Self(unique))
+    }
+
+    pub fn contains(&self, version: ProtocolVersion) -> bool {
+        self.0.contains(&version)
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = ProtocolVersion> + '_ {
+        self.0.iter().copied()
+    }
+}
+
+impl<'de> Deserialize<'de> for SupportedProtocolVersions {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        Self::try_new(Vec::<ProtocolVersion>::deserialize(deserializer)?).map_err(D::Error::custom)
+    }
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum InvalidSupportedProtocolVersions {
+    Empty,
+    Duplicate(ProtocolVersion),
+}
+
+impl fmt::Display for InvalidSupportedProtocolVersions {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Empty => formatter.write_str("at least one protocol version is required"),
+            Self::Duplicate(version) => {
+                write!(formatter, "duplicate protocol version {}", version.get())
+            }
+        }
+    }
+}
+
+impl std::error::Error for InvalidSupportedProtocolVersions {}

codex-rs/code-mode-protocol/src/json_schema_types.rs

@@ -0,0 +1,538 @@
+use serde_json::Value as JsonValue;
+use std::collections::BTreeMap;
+
+use crate::description::normalize_code_mode_identifier;
+
+// Expose one nested recursive shape, then fall back to `unknown` on the next
+// occurrence so generated tool declarations remain finite.
+const MAX_LOCAL_REF_EXPANSIONS_PER_PATH: usize = 2;
+// Bound repeated refs and DAG fan-out separately from cycle depth so one
+// compact schema cannot expand into an arbitrarily large model-visible item.
+const MAX_TOTAL_LOCAL_REF_EXPANSIONS: usize = 32;
+// Bound individual schema rendering before assembling the final declaration.
+const MAX_RENDERED_SCHEMA_BYTES: usize = 16_000;
+// Charge intermediate render strings as they are built so repeated local refs
+// cannot allocate unbounded expanded copies before the final schema cap runs.
+const MAX_RENDER_WORK_BYTES: usize = MAX_RENDERED_SCHEMA_BYTES * 4;
+
+pub fn render_json_schema_to_typescript(schema: &JsonValue) -> String {
+    let rendered = JsonSchemaTypeRenderer::new(schema).render(schema);
+    if rendered.len() > MAX_RENDERED_SCHEMA_BYTES {
+        "unknown".to_string()
+    } else {
+        rendered
+    }
+}
+
+struct JsonSchemaTypeRenderer<'a> {
+    root: &'a JsonValue,
+    nested_schema_resource_depth: usize,
+    active_local_ref_expansions: BTreeMap<String, usize>,
+    remaining_local_ref_expansions: usize,
+    remaining_render_work_bytes: usize,
+    render_work_budget_exhausted: bool,
+}
+
+impl<'a> JsonSchemaTypeRenderer<'a> {
+    fn new(root: &'a JsonValue) -> Self {
+        Self {
+            root,
+            nested_schema_resource_depth: 0,
+            active_local_ref_expansions: BTreeMap::new(),
+            remaining_local_ref_expansions: MAX_TOTAL_LOCAL_REF_EXPANSIONS,
+            remaining_render_work_bytes: MAX_RENDER_WORK_BYTES,
+            render_work_budget_exhausted: false,
+        }
+    }
+
+    fn render(&mut self, schema: &JsonValue) -> String {
+        if self.render_work_budget_exhausted {
+            return "unknown".to_string();
+        }
+
+        // A nested `$id` starts a new schema resource. Fragment-only refs below
+        // it are scoped to that resource, not to the outer document root.
+        let enters_nested_schema_resource = !std::ptr::eq(schema, self.root)
+            && schema
+                .as_object()
+                .is_some_and(|map| map.contains_key("$id"));
+        if enters_nested_schema_resource {
+            self.nested_schema_resource_depth += 1;
+        }
+
+        let rendered = match schema {
+            JsonValue::Bool(true) => "unknown".to_string(),
+            JsonValue::Bool(false) => "never".to_string(),
+            JsonValue::Object(map) => self.render_map(map),
+            _ => "unknown".to_string(),
+        };
+        if enters_nested_schema_resource {
+            self.nested_schema_resource_depth -= 1;
+        }
+        self.finish_render(rendered)
+    }
+
+    fn render_map(&mut self, map: &serde_json::Map<String, JsonValue>) -> String {
+        if self.render_work_budget_exhausted {
+            return "unknown".to_string();
+        }
+
+        if map.contains_key("$ref") {
+            return self.render_ref(map);
+        }
+
+        if let Some(value) = map.get("const") {
+            return self.render_literal(value);
+        }
+
+        if let Some(values) = map.get("enum").and_then(JsonValue::as_array) {
+            let mut rendered = Vec::new();
+            for value in values {
+                let literal = self.render_literal(value);
+                if self.render_work_budget_exhausted {
+                    return "unknown".to_string();
+                }
+                if !self.consume_render_work(literal.len()) {
+                    return "unknown".to_string();
+                }
+                rendered.push(literal);
+            }
+            if !rendered.is_empty() {
+                return rendered.join(" | ");
+            }
+        }
+
+        for key in ["anyOf", "oneOf"] {
+            if let Some(variants) = map.get(key).and_then(JsonValue::as_array) {
+                let mut rendered = Vec::new();
+                for variant in variants {
+                    if self.render_work_budget_exhausted {
+                        return "unknown".to_string();
+                    }
+                    rendered.push(self.render(variant));
+                }
+                if !rendered.is_empty() {
+                    return rendered.join(" | ");
+                }
+            }
+        }
+
+        if let Some(variants) = map.get("allOf").and_then(JsonValue::as_array) {
+            let mut rendered = Vec::new();
+            for variant in variants {
+                if self.render_work_budget_exhausted {
+                    return "unknown".to_string();
+                }
+                rendered.push(parenthesize_union_for_intersection(self.render(variant)));
+            }
+            if !rendered.is_empty() {
+                return rendered.join(" & ");
+            }
+        }
+
+        if let Some(schema_type) = map.get("type") {
+            if let Some(types) = schema_type.as_array() {
+                let mut rendered = Vec::new();
+                for schema_type in types.iter().filter_map(JsonValue::as_str) {
+                    if self.render_work_budget_exhausted {
+                        return "unknown".to_string();
+                    }
+                    rendered.push(self.render_type_keyword(map, schema_type));
+                }
+                if !rendered.is_empty() {
+                    return rendered.join(" | ");
+                }
+            }
+
+            if let Some(schema_type) = schema_type.as_str() {
+                return self.render_type_keyword(map, schema_type);
+            }
+        }
+
+        if map.contains_key("properties")
+            || map.contains_key("additionalProperties")
+            || map.contains_key("required")
+        {
+            return self.render_object(map);
+        }
+
+        if map.contains_key("items") || map.contains_key("prefixItems") {
+            return self.render_array(map);
+        }
+
+        "unknown".to_string()
+    }
+
+    fn render_ref(&mut self, map: &serde_json::Map<String, JsonValue>) -> String {
+        let referenced_type = if self.nested_schema_resource_depth > 0 {
+            None
+        } else {
+            map.get("$ref")
+                .and_then(JsonValue::as_str)
+                .and_then(local_json_pointer)
+                .and_then(|pointer| {
+                    let active_expansions = self
+                        .active_local_ref_expansions
+                        .get(&pointer)
+                        .copied()
+                        .unwrap_or_default();
+                    if active_expansions >= MAX_LOCAL_REF_EXPANSIONS_PER_PATH {
+                        return Some("unknown".to_string());
+                    }
+                    if self.remaining_local_ref_expansions == 0 {
+                        return Some("unknown".to_string());
+                    }
+
+                    let root = self.root;
+                    let target = if pointer.is_empty() {
+                        Some(root)
+                    } else {
+                        root.pointer(&pointer)
+                    }?;
+                    self.remaining_local_ref_expansions -= 1;
+                    self.active_local_ref_expansions
+                        .insert(pointer.clone(), active_expansions + 1);
+
+                    let rendered = self.render(target);
+                    if active_expansions == 0 {
+                        self.active_local_ref_expansions.remove(&pointer);
+                    } else {
+                        self.active_local_ref_expansions
+                            .insert(pointer, active_expansions);
+                    }
+                    Some(rendered)
+                })
+        }
+        .unwrap_or_else(|| "unknown".to_string());
+        if self.render_work_budget_exhausted {
+            return "unknown".to_string();
+        }
+
+        let siblings = map
+            .iter()
+            .filter(|(key, _)| !matches!(key.as_str(), "$ref" | "$defs" | "definitions"))
+            .map(|(key, value)| (key.clone(), value.clone()))
+            .collect();
+        if !has_renderable_schema_keywords(&siblings) {
+            return referenced_type;
+        }
+
+        let sibling_type = self.render_map(&siblings);
+        match (referenced_type.as_str(), sibling_type.as_str()) {
+            ("unknown", _) => sibling_type,
+            (_, "unknown") => referenced_type,
+            _ => format!("({referenced_type}) & ({sibling_type})"),
+        }
+    }
+
+    fn render_type_keyword(
+        &mut self,
+        map: &serde_json::Map<String, JsonValue>,
+        schema_type: &str,
+    ) -> String {
+        match schema_type {
+            "string" => "string".to_string(),
+            "number" | "integer" => "number".to_string(),
+            "boolean" => "boolean".to_string(),
+            "null" => "null".to_string(),
+            "array" => self.render_array(map),
+            "object" => self.render_object(map),
+            _ => "unknown".to_string(),
+        }
+    }
+
+    fn render_array(&mut self, map: &serde_json::Map<String, JsonValue>) -> String {
+        if let Some(items) = map.get("items") {
+            let item_type = self.render(items);
+            if self.render_work_budget_exhausted {
+                return "unknown".to_string();
+            }
+            return format!("Array<{item_type}>");
+        }
+
+        if let Some(items) = map.get("prefixItems").and_then(JsonValue::as_array) {
+            let mut item_types = Vec::new();
+            for item in items {
+                if self.render_work_budget_exhausted {
+                    return "unknown".to_string();
+                }
+                item_types.push(self.render(item));
+            }
+            if !item_types.is_empty() {
+                return format!("[{}]", item_types.join(", "));
+            }
+        }
+
+        "unknown[]".to_string()
+    }
+
+    fn append_additional_properties_line(
+        &mut self,
+        lines: &mut Vec<String>,
+        map: &serde_json::Map<String, JsonValue>,
+        properties: &serde_json::Map<String, JsonValue>,
+        line_prefix: &str,
+    ) -> bool {
+        if let Some(additional_properties) = map.get("additionalProperties") {
+            let property_type = match additional_properties {
+                JsonValue::Bool(true) => Some("unknown".to_string()),
+                JsonValue::Bool(false) => None,
+                value => Some(self.render(value)),
+            };
+
+            if let Some(property_type) = property_type {
+                return self.push_render_line(
+                    lines,
+                    format!("{line_prefix}[key: string]: {property_type};"),
+                );
+            }
+        } else if properties.is_empty() {
+            return self.push_render_line(lines, format!("{line_prefix}[key: string]: unknown;"));
+        }
+        true
+    }
+
+    fn render_object_property(
+        &mut self,
+        name: &str,
+        value: &JsonValue,
+        required: &[&str],
+    ) -> String {
+        if name.len() > self.remaining_render_work_bytes {
+            self.render_work_budget_exhausted = true;
+            return "unknown".to_string();
+        }
+        let optional = if required.iter().any(|required_name| required_name == &name) {
+            ""
+        } else {
+            "?"
+        };
+        let property_name = render_json_schema_property_name(name);
+        let property_type = self.render(value);
+        if self.render_work_budget_exhausted {
+            return "unknown".to_string();
+        }
+        format!("{property_name}{optional}: {property_type};")
+    }
+
+    fn render_object(&mut self, map: &serde_json::Map<String, JsonValue>) -> String {
+        let required = map
+            .get("required")
+            .and_then(JsonValue::as_array)
+            .map(|items| {
+                items
+                    .iter()
+                    .filter_map(JsonValue::as_str)
+                    .collect::<Vec<_>>()
+            })
+            .unwrap_or_default();
+        let empty_properties = serde_json::Map::new();
+        let properties = map
+            .get("properties")
+            .and_then(JsonValue::as_object)
+            .unwrap_or(&empty_properties);
+
+        let mut sorted_properties = properties.iter().collect::<Vec<_>>();
+        sorted_properties.sort_unstable_by_key(|(name_a, _)| *name_a);
+        if sorted_properties
+            .iter()
+            .any(|(_, value)| has_property_description(value))
+        {
+            let mut lines = Vec::new();
+            if !self.push_render_line(&mut lines, "{".to_string()) {
+                return "unknown".to_string();
+            }
+            for (name, value) in sorted_properties {
+                if let Some(description) = value.get("description").and_then(JsonValue::as_str) {
+                    for description_line in description
+                        .lines()
+                        .map(str::trim)
+                        .filter(|line| !line.is_empty())
+                    {
+                        if description_line.len().saturating_add(5)
+                            > self.remaining_render_work_bytes
+                            || !self
+                                .push_render_line(&mut lines, format!("  // {description_line}"))
+                        {
+                            return "unknown".to_string();
+                        }
+                    }
+                }
+
+                let property = self.render_object_property(name, value, &required);
+                if self.render_work_budget_exhausted
+                    || !self.push_render_line(&mut lines, format!("  {property}"))
+                {
+                    return "unknown".to_string();
+                }
+            }
+
+            if !self.append_additional_properties_line(&mut lines, map, properties, "  ")
+                || !self.push_render_line(&mut lines, "}".to_string())
+            {
+                return "unknown".to_string();
+            }
+            return lines.join("\n");
+        }
+
+        let mut lines = Vec::new();
+        for (name, value) in sorted_properties {
+            let property = self.render_object_property(name, value, &required);
+            if self.render_work_budget_exhausted || !self.push_render_line(&mut lines, property) {
+                return "unknown".to_string();
+            }
+        }
+
+        if !self.append_additional_properties_line(&mut lines, map, properties, "") {
+            return "unknown".to_string();
+        }
+
+        if lines.is_empty() {
+            return "{}".to_string();
+        }
+
+        format!("{{ {} }}", lines.join(" "))
+    }
+
+    fn finish_render(&mut self, rendered: String) -> String {
+        if self.consume_render_work(rendered.len()) {
+            rendered
+        } else {
+            "unknown".to_string()
+        }
+    }
+
+    fn render_literal(&mut self, value: &JsonValue) -> String {
+        if json_literal_serialization_upper_bound(value) > self.remaining_render_work_bytes {
+            self.render_work_budget_exhausted = true;
+            "unknown".to_string()
+        } else {
+            render_json_schema_literal(value)
+        }
+    }
+
+    fn consume_render_work(&mut self, rendered_bytes: usize) -> bool {
+        if rendered_bytes > self.remaining_render_work_bytes {
+            self.render_work_budget_exhausted = true;
+            false
+        } else {
+            self.remaining_render_work_bytes -= rendered_bytes;
+            true
+        }
+    }
+
+    fn push_render_line(&mut self, lines: &mut Vec<String>, line: String) -> bool {
+        if !self.consume_render_work(line.len()) {
+            return false;
+        }
+        lines.push(line);
+        true
+    }
+}
+
+fn parenthesize_union_for_intersection(rendered: String) -> String {
+    if rendered.contains(" | ") {
+        format!("({rendered})")
+    } else {
+        rendered
+    }
+}
+
+fn local_json_pointer(reference: &str) -> Option<String> {
+    let fragment = reference.strip_prefix('#')?;
+    let pointer = percent_decode_uri_fragment(fragment)?;
+    if pointer.is_empty() || pointer.starts_with('/') {
+        Some(pointer)
+    } else {
+        None
+    }
+}
+
+fn percent_decode_uri_fragment(fragment: &str) -> Option<String> {
+    let bytes = fragment.as_bytes();
+    let mut decoded = Vec::with_capacity(bytes.len());
+    let mut index = 0;
+    while index < bytes.len() {
+        if bytes[index] == b'%' {
+            let high = decode_hex_digit(*bytes.get(index + 1)?)?;
+            let low = decode_hex_digit(*bytes.get(index + 2)?)?;
+            decoded.push((high << 4) | low);
+            index += 3;
+        } else {
+            decoded.push(bytes[index]);
+            index += 1;
+        }
+    }
+    String::from_utf8(decoded).ok()
+}
+
+fn decode_hex_digit(digit: u8) -> Option<u8> {
+    match digit {
+        b'0'..=b'9' => Some(digit - b'0'),
+        b'a'..=b'f' => Some(digit - b'a' + 10),
+        b'A'..=b'F' => Some(digit - b'A' + 10),
+        _ => None,
+    }
+}
+
+fn has_renderable_schema_keywords(map: &serde_json::Map<String, JsonValue>) -> bool {
+    [
+        "const",
+        "enum",
+        "anyOf",
+        "oneOf",
+        "allOf",
+        "type",
+        "properties",
+        "additionalProperties",
+        "required",
+        "items",
+        "prefixItems",
+    ]
+    .iter()
+    .any(|key| map.contains_key(*key))
+}
+
+fn has_property_description(value: &JsonValue) -> bool {
+    value
+        .get("description")
+        .and_then(JsonValue::as_str)
+        .is_some_and(|description| !description.is_empty())
+}
+
+fn render_json_schema_property_name(name: &str) -> String {
+    if normalize_code_mode_identifier(name) == name {
+        name.to_string()
+    } else {
+        serde_json::to_string(name).unwrap_or_else(|_| format!("\"{}\"", name.replace('"', "\\\"")))
+    }
+}
+
+fn render_json_schema_literal(value: &JsonValue) -> String {
+    serde_json::to_string(value).unwrap_or_else(|_| "unknown".to_string())
+}
+
+fn json_literal_serialization_upper_bound(value: &JsonValue) -> usize {
+    match value {
+        JsonValue::Null => 4,
+        JsonValue::Bool(false) => 5,
+        JsonValue::Bool(true) => 4,
+        JsonValue::Number(number) => number.to_string().len(),
+        // JSON escaping can expand one UTF-8 byte to at most one six-byte
+        // Unicode escape, so this bounds allocation before serialization.
+        JsonValue::String(string) => string.len().saturating_mul(6).saturating_add(2),
+        JsonValue::Array(values) => values.iter().fold(2, |size, value| {
+            size.saturating_add(1)
+                .saturating_add(json_literal_serialization_upper_bound(value))
+        }),
+        JsonValue::Object(map) => map.iter().fold(2, |size, (key, value)| {
+            size.saturating_add(4)
+                .saturating_add(key.len().saturating_mul(6))
+                .saturating_add(json_literal_serialization_upper_bound(value))
+        }),
+    }
+}
+
+#[cfg(test)]
+#[path = "json_schema_types_tests.rs"]
+mod tests;

codex-rs/code-mode-protocol/src/json_schema_types_tests.rs

@@ -0,0 +1,200 @@
+use super::*;
+use pretty_assertions::assert_eq;
+use serde_json::json;
+
+#[test]
+fn renders_recursive_local_refs_with_escaped_pointer_segments() {
+    let schema = json!({
+        "type": "object",
+        "properties": {
+            "clauses": {
+                "type": "array",
+                "items": { "$ref": "#/$defs/Boolean~1Clause~0v1" }
+            }
+        },
+        "$defs": {
+            "Boolean/Clause~v1": {
+                "type": "object",
+                "properties": {
+                    "query": { "$ref": "#/$defs/Query" }
+                }
+            },
+            "Query": {
+                "oneOf": [
+                    { "type": "string" },
+                    {
+                        "type": "object",
+                        "properties": {
+                            "clauses": {
+                                "type": "array",
+                                "items": { "$ref": "#/$defs/Boolean~1Clause~0v1" }
+                            }
+                        }
+                    }
+                ]
+            }
+        }
+    });
+
+    let rendered = render_json_schema_to_typescript(&schema);
+    assert!(rendered.contains("clauses?: Array<{ query?: string | { clauses?: Array<{"));
+    assert!(rendered.contains("query?: string | { clauses?: Array<unknown>; };"));
+}
+
+#[test]
+fn renders_ref_siblings_uri_fragments_and_all_of_precedence() {
+    assert_eq!(
+        render_json_schema_to_typescript(&json!({
+            "$ref": "#/$defs/Label",
+            "enum": ["A"],
+            "$defs": { "Label": { "type": "string" } }
+        })),
+        r#"(string) & ("A")"#
+    );
+    assert_eq!(
+        render_json_schema_to_typescript(&json!({
+            "$ref": "#/$defs/Foo%20Bar",
+            "$defs": { "Foo Bar": { "type": "string" } }
+        })),
+        "string"
+    );
+    assert_eq!(
+        render_json_schema_to_typescript(&json!({
+            "allOf": [
+                { "$ref": "#/$defs/Choice" },
+                { "type": "object", "properties": { "value": { "type": "string" } } }
+            ],
+            "$defs": {
+                "Choice": { "oneOf": [{ "type": "string" }, { "type": "number" }] }
+            }
+        })),
+        "(string | number) & { value?: string; }"
+    );
+}
+
+#[test]
+fn leaves_local_refs_under_nested_schema_resources_unresolved() {
+    let schema = json!({
+        "$defs": {
+            "Choice": { "type": "string" }
+        },
+        "type": "object",
+        "properties": {
+            "nested": {
+                "$id": "urn:nested",
+                "$defs": {
+                    "Choice": { "type": "number" }
+                },
+                "type": "object",
+                "properties": {
+                    "value": { "$ref": "#/$defs/Choice" }
+                }
+            }
+        }
+    });
+
+    assert_eq!(
+        render_json_schema_to_typescript(&schema),
+        "{ nested?: { value?: unknown; }; }"
+    );
+}
+
+#[test]
+fn bounds_expansions_without_charging_dangling_refs() {
+    let mut properties = (0..MAX_TOTAL_LOCAL_REF_EXPANSIONS)
+        .map(|index| {
+            (
+                format!("a_missing_{index}"),
+                json!({ "$ref": format!("#/$defs/Missing{index}") }),
+            )
+        })
+        .collect::<serde_json::Map<_, _>>();
+    properties.insert("z_valid".to_string(), json!({ "$ref": "#/$defs/Valid" }));
+    let schema = json!({
+        "type": "object",
+        "properties": properties,
+        "$defs": { "Valid": { "type": "string" } }
+    });
+
+    let rendered = render_json_schema_to_typescript(&schema);
+    assert!(rendered.contains("z_valid?: string;"));
+
+    let properties = (0..MAX_TOTAL_LOCAL_REF_EXPANSIONS + 2)
+        .map(|index| {
+            (
+                format!("property_{index}"),
+                json!({ "$ref": "#/$defs/Item" }),
+            )
+        })
+        .collect::<serde_json::Map<_, _>>();
+    let rendered = render_json_schema_to_typescript(&json!({
+        "type": "object",
+        "properties": properties,
+        "$defs": { "Item": { "type": "string" } }
+    }));
+    assert_eq!(
+        rendered.matches("string").count(),
+        MAX_TOTAL_LOCAL_REF_EXPANSIONS
+    );
+    assert_eq!(rendered.matches("unknown").count(), 2);
+}
+
+#[test]
+fn repeated_large_ref_expansions_exhaust_render_work_budget() {
+    let properties = (0..MAX_TOTAL_LOCAL_REF_EXPANSIONS)
+        .map(|index| {
+            (
+                format!("property_{index}"),
+                json!({ "$ref": "#/$defs/Item" }),
+            )
+        })
+        .collect::<serde_json::Map<_, _>>();
+    let schema = json!({
+        "type": "object",
+        "properties": properties,
+        "$defs": {
+            "Item": {
+                "type": "object",
+                "properties": {
+                    "value": {
+                        "type": "string",
+                        "description": "x".repeat(MAX_RENDERED_SCHEMA_BYTES / 2)
+                    }
+                }
+            }
+        }
+    });
+
+    let mut renderer = JsonSchemaTypeRenderer::new(&schema);
+    assert_eq!(renderer.render(&schema), "unknown");
+    assert!(renderer.render_work_budget_exhausted);
+}
+
+#[test]
+fn oversized_ref_literal_exhausts_render_work_budget() {
+    let schema = json!({
+        "$ref": "#/$defs/Value",
+        "$defs": {
+            "Value": {
+                "const": "x".repeat(MAX_RENDER_WORK_BYTES)
+            }
+        }
+    });
+
+    let mut renderer = JsonSchemaTypeRenderer::new(&schema);
+    assert_eq!(renderer.render(&schema), "unknown");
+    assert!(renderer.render_work_budget_exhausted);
+}
+
+#[test]
+fn rendered_schema_has_a_hard_size_cap() {
+    let description = "x".repeat(MAX_RENDERED_SCHEMA_BYTES);
+    let schema = json!({
+        "type": "object",
+        "properties": {
+            "value": { "type": "string", "description": description }
+        }
+    });
+
+    assert_eq!(render_json_schema_to_typescript(&schema), "unknown");
+}

codex-rs/code-mode-protocol/src/lib.rs

@@ -0,0 +1,52 @@
+mod description;
+pub mod grpc;
+pub mod host;
+mod json_schema_types;
+mod response;
+mod runtime;
+mod session;
+
+pub use description::CODE_MODE_PRAGMA_PREFIX;
+pub use description::CodeModeToolKind;
+pub use description::EnabledToolMetadata;
+pub use description::ImageDetailVisibility;
+pub use description::ToolDefinition;
+pub use description::ToolNamespaceDescription;
+pub use description::augment_tool_definition;
+pub use description::build_exec_tool_description;
+pub use description::build_wait_tool_description;
+pub use description::enabled_tool_metadata;
+pub use description::is_code_mode_nested_tool;
+pub use description::normalize_code_mode_identifier;
+pub use description::parse_exec_source;
+pub use description::render_code_mode_sample;
+pub use json_schema_types::render_json_schema_to_typescript;
+pub use response::DEFAULT_IMAGE_DETAIL;
+pub use response::FunctionCallOutputContentItem;
+pub use response::ImageDetail;
+pub use runtime::CodeModeNestedToolCall;
+pub use runtime::DEFAULT_EXEC_YIELD_TIME_MS;
+pub use runtime::DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL;
+pub use runtime::DEFAULT_WAIT_YIELD_TIME_MS;
+pub use runtime::ExecuteRequest;
+pub use runtime::ExecuteToPendingOutcome;
+pub use runtime::MissingCodeModeHostDuration;
+pub use runtime::RuntimeResponse;
+pub use runtime::WaitOutcome;
+pub use runtime::WaitRequest;
+pub use runtime::WaitToPendingOutcome;
+pub use runtime::WaitToPendingRequest;
+pub use session::CellId;
+pub use session::CodeModeSession;
+pub use session::CodeModeSessionCellExecutionLimits;
+pub use session::CodeModeSessionDelegate;
+pub use session::CodeModeSessionProvider;
+pub use session::CodeModeSessionProviderFuture;
+pub use session::CodeModeSessionResultFuture;
+pub use session::NoopCodeModeSessionDelegate;
+pub use session::NotificationFuture;
+pub use session::StartedCell;
+pub use session::ToolInvocationFuture;
+
+pub const PUBLIC_TOOL_NAME: &str = "exec";
+pub const WAIT_TOOL_NAME: &str = "wait";

codex-rs/code-mode-protocol/src/response.rs

@@ -0,0 +1,29 @@
+use serde::Deserialize;
+use serde::Serialize;
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "lowercase")]
+pub enum ImageDetail {
+    Auto,
+    Low,
+    High,
+    Original,
+}
+
+pub const DEFAULT_IMAGE_DETAIL: ImageDetail = ImageDetail::High;
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum FunctionCallOutputContentItem {
+    InputText {
+        text: String,
+    },
+    InputImage {
+        image_url: String,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        detail: Option<ImageDetail>,
+    },
+    InputAudio {
+        audio_url: String,
+    },
+}

codex-rs/code-mode-protocol/src/runtime.rs

@@ -0,0 +1,183 @@
+use std::error::Error;
+use std::fmt;
+use std::time::Duration;
+
+use codex_protocol::ToolName;
+use serde::Deserialize;
+use serde::Serialize;
+use serde_json::Value as JsonValue;
+
+use crate::CellId;
+use crate::CodeModeToolKind;
+use crate::FunctionCallOutputContentItem;
+use crate::ToolDefinition;
+
+pub const DEFAULT_EXEC_YIELD_TIME_MS: u64 = 10_000;
+pub const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000;
+pub const DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL: usize = 10_000;
+
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+pub struct ExecuteRequest {
+    pub tool_call_id: String,
+    pub enabled_tools: Vec<ToolDefinition>,
+    pub source: String,
+    pub yield_time_ms: Option<u64>,
+    pub max_output_tokens: Option<usize>,
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+pub struct WaitRequest {
+    pub cell_id: CellId,
+    pub yield_time_ms: u64,
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+pub struct WaitToPendingRequest {
+    pub cell_id: CellId,
+}
+
+#[derive(Debug, Deserialize, PartialEq, Serialize)]
+pub enum WaitOutcome {
+    LiveCell(RuntimeResponse),
+    MissingCell(RuntimeResponse),
+}
+
+impl WaitOutcome {
+    /// Returns timing for this wait or termination request, when supplied by its host.
+    pub fn code_mode_host_duration(&self) -> Option<Duration> {
+        match self {
+            Self::LiveCell(response) | Self::MissingCell(response) => {
+                response.code_mode_host_duration()
+            }
+        }
+    }
+
+    /// Records the enclosing host request's duration before wire conversion.
+    pub fn with_code_mode_host_duration(self, code_mode_host_duration: Duration) -> Self {
+        match self {
+            Self::LiveCell(response) => {
+                Self::LiveCell(response.with_code_mode_host_duration(code_mode_host_duration))
+            }
+            Self::MissingCell(response) => {
+                Self::MissingCell(response.with_code_mode_host_duration(code_mode_host_duration))
+            }
+        }
+    }
+}
+
+#[derive(Debug, Deserialize, PartialEq, Serialize)]
+pub enum ExecuteToPendingOutcome {
+    Pending {
+        cell_id: CellId,
+        content_items: Vec<FunctionCallOutputContentItem>,
+        pending_tool_call_ids: Vec<String>,
+    },
+    Completed(RuntimeResponse),
+}
+
+#[derive(Debug, Deserialize, PartialEq, Serialize)]
+pub enum WaitToPendingOutcome {
+    LiveCell(ExecuteToPendingOutcome),
+    MissingCell(RuntimeResponse),
+}
+
+impl From<WaitOutcome> for RuntimeResponse {
+    fn from(outcome: WaitOutcome) -> Self {
+        match outcome {
+            WaitOutcome::LiveCell(response) | WaitOutcome::MissingCell(response) => response,
+        }
+    }
+}
+
+/// Runtime output with optional timing for the host request that observed it.
+///
+/// The JavaScript session returns untimed output. The host handler records its
+/// complete request duration before conversion; wire conversions preserve that
+/// field and reject untimed output. Decoded host responses always contain timing,
+/// including measured zero. Raw traces also retain timing when present.
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+pub enum RuntimeResponse {
+    Yielded {
+        cell_id: CellId,
+        content_items: Vec<FunctionCallOutputContentItem>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        code_mode_host_duration: Option<Duration>,
+    },
+    Terminated {
+        cell_id: CellId,
+        content_items: Vec<FunctionCallOutputContentItem>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        code_mode_host_duration: Option<Duration>,
+    },
+    Result {
+        cell_id: CellId,
+        content_items: Vec<FunctionCallOutputContentItem>,
+        error_text: Option<String>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        code_mode_host_duration: Option<Duration>,
+    },
+}
+
+impl RuntimeResponse {
+    /// Returns timing for this observation, excluding background work between requests.
+    pub fn code_mode_host_duration(&self) -> Option<Duration> {
+        match self {
+            Self::Yielded {
+                code_mode_host_duration,
+                ..
+            }
+            | Self::Terminated {
+                code_mode_host_duration,
+                ..
+            }
+            | Self::Result {
+                code_mode_host_duration,
+                ..
+            } => *code_mode_host_duration,
+        }
+    }
+
+    /// Records the enclosing host request's duration before wire conversion.
+    pub fn with_code_mode_host_duration(mut self, code_mode_host_duration: Duration) -> Self {
+        match &mut self {
+            Self::Yielded {
+                code_mode_host_duration: value,
+                ..
+            }
+            | Self::Terminated {
+                code_mode_host_duration: value,
+                ..
+            }
+            | Self::Result {
+                code_mode_host_duration: value,
+                ..
+            } => *value = Some(code_mode_host_duration),
+        }
+        self
+    }
+}
+
+/// An untimed runtime response cannot be encoded for delivery to a host client.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct MissingCodeModeHostDuration;
+
+impl fmt::Display for MissingCodeModeHostDuration {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter.write_str("code-mode response is missing host duration")
+    }
+}
+
+impl Error for MissingCodeModeHostDuration {}
+
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+pub struct CodeModeNestedToolCall {
+    pub cell_id: CellId,
+    pub runtime_tool_call_id: String,
+    pub tool_name: ToolName,
+    pub tool_kind: CodeModeToolKind,
+    pub input: Option<JsonValue>,
+}
+
+#[cfg(test)]
+#[path = "runtime_tests.rs"]
+mod tests;

codex-rs/code-mode-protocol/src/runtime_tests.rs

@@ -0,0 +1,107 @@
+//! Regression coverage for timing metadata at serialization boundaries.
+
+use std::time::Duration;
+
+use pretty_assertions::assert_eq;
+
+use super::MissingCodeModeHostDuration;
+use super::RuntimeResponse;
+use super::WaitOutcome;
+use crate::CellId;
+use crate::FunctionCallOutputContentItem;
+use crate::host::WireRuntimeResponse;
+use crate::host::WireWaitOutcome;
+
+/// Raw runtime serialization preserves absent timing; stdio always carries a
+/// measured duration, including zero, without changing the response payload.
+#[test]
+fn code_mode_host_duration_survives_runtime_and_stdio_serialization() {
+    let content_items = vec![FunctionCallOutputContentItem::InputText {
+        text: "output".to_string(),
+    }];
+    for response in [
+        RuntimeResponse::Yielded {
+            cell_id: CellId::new("yielded-cell".to_string()),
+            content_items: content_items.clone(),
+            code_mode_host_duration: None,
+        },
+        RuntimeResponse::Terminated {
+            cell_id: CellId::new("terminated-cell".to_string()),
+            content_items: content_items.clone(),
+            code_mode_host_duration: None,
+        },
+        RuntimeResponse::Result {
+            cell_id: CellId::new("completed-cell".to_string()),
+            content_items,
+            error_text: Some("execution failed".to_string()),
+            code_mode_host_duration: None,
+        },
+    ] {
+        for duration in [
+            None,
+            Some(Duration::ZERO),
+            Some(Duration::from_nanos(/*nanos*/ 1_234_567_890)),
+            Some(Duration::from_nanos(u64::MAX)),
+        ] {
+            let mut expected = response.clone();
+            match &mut expected {
+                RuntimeResponse::Yielded {
+                    code_mode_host_duration,
+                    ..
+                }
+                | RuntimeResponse::Terminated {
+                    code_mode_host_duration,
+                    ..
+                }
+                | RuntimeResponse::Result {
+                    code_mode_host_duration,
+                    ..
+                } => *code_mode_host_duration = duration,
+            }
+
+            let payload = serde_json::to_value(&expected).expect("serialize response");
+            assert_eq!(
+                serde_json::from_value::<RuntimeResponse>(payload).expect("deserialize response"),
+                expected
+            );
+
+            if duration.is_some() {
+                let wire_payload = serde_json::to_value(
+                    WireRuntimeResponse::try_from(expected.clone()).expect("timed response"),
+                )
+                .expect("serialize response over stdio");
+                assert_eq!(
+                    RuntimeResponse::from(
+                        serde_json::from_value::<WireRuntimeResponse>(wire_payload)
+                            .expect("deserialize stdio response")
+                    ),
+                    expected
+                );
+            }
+        }
+    }
+}
+
+/// Encoding must not turn a missing request measurement into measured zero,
+/// including when no live cell remains to supply output.
+#[test]
+fn stdio_encoding_rejects_untimed_runtime_output() {
+    let response = RuntimeResponse::Terminated {
+        cell_id: CellId::new("cell".to_string()),
+        content_items: Vec::new(),
+        code_mode_host_duration: None,
+    };
+    assert_eq!(
+        WireRuntimeResponse::try_from(response.clone()),
+        Err(MissingCodeModeHostDuration)
+    );
+    for outcome in [
+        WaitOutcome::LiveCell(response.clone()),
+        WaitOutcome::MissingCell(response),
+    ] {
+        assert_eq!(
+            WireWaitOutcome::try_from(outcome),
+            Err(MissingCodeModeHostDuration)
+        );
+    }
+}

codex-rs/code-mode-protocol/src/session.rs

@@ -0,0 +1,200 @@
+use std::fmt;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
+use serde::Deserialize;
+use serde::Serialize;
+use serde_json::Value as JsonValue;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+
+use crate::CodeModeNestedToolCall;
+use crate::ExecuteRequest;
+use crate::RuntimeResponse;
+use crate::WaitOutcome;
+use crate::WaitRequest;
+
+pub type CodeModeSessionResultFuture<'a, T> =
+    Pin<Box<dyn Future<Output = Result<T, String>> + Send + 'a>>;
+pub type CodeModeSessionProviderFuture<'a> =
+    CodeModeSessionResultFuture<'a, Arc<dyn CodeModeSession>>;
+pub type ToolInvocationFuture<'a> =
+    Pin<Box<dyn Future<Output = Result<JsonValue, String>> + Send + 'a>>;
+pub type NotificationFuture<'a> = Pin<Box<dyn Future<Output = Result<(), String>> + Send + 'a>>;
+
+/// Optional resource limits shared by every cell in one code-mode session.
+#[derive(Clone, Debug, Default, Eq, PartialEq)]
+pub struct CodeModeSessionCellExecutionLimits {
+    pub max_yield_time_ms: Option<u64>,
+    pub max_heap_size_bytes: Option<usize>,
+}
+
+#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
+pub struct CellId(String);
+
+impl CellId {
+    pub fn new(value: String) -> Self {
+        Self(value)
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
+impl AsRef<str> for CellId {
+    fn as_ref(&self) -> &str {
+        self.as_str()
+    }
+}
+
+impl fmt::Display for CellId {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter.write_str(self.as_str())
+    }
+}
+
+pub struct StartedCell {
+    pub cell_id: CellId,
+    initial_response: CodeModeSessionResultFuture<'static, RuntimeResponse>,
+}
+
+impl StartedCell {
+    pub fn new(cell_id: CellId, initial_response_rx: oneshot::Receiver<RuntimeResponse>) -> Self {
+        Self::from_future(cell_id, async move {
+            initial_response_rx
+                .await
+                .map_err(|_| "exec runtime ended unexpectedly".to_string())
+        })
+    }
+
+    pub fn from_result_receiver(
+        cell_id: CellId,
+        initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,
+    ) -> Self {
+        Self::from_future(cell_id, async move {
+            initial_response_rx
+                .await
+                .map_err(|_| "exec runtime ended unexpectedly".to_string())?
+        })
+    }
+
+    pub fn from_future(
+        cell_id: CellId,
+        initial_response: impl Future<Output = Result<RuntimeResponse, String>> + Send + 'static,
+    ) -> Self {
+        Self {
+            cell_id,
+            initial_response: Box::pin(initial_response),
+        }
+    }
+
+    pub async fn initial_response(self) -> Result<RuntimeResponse, String> {
+        self.initial_response.await
+    }
+}
+
+/// Host callbacks owned by one code-mode execution.
+///
+/// The session retains the supplied delegate while starting and running the cell,
+/// including across yields, and releases it through its existing close/cancel paths.
+pub trait CodeModeSessionDelegate: Send + Sync {
+    fn invoke_tool<'a>(
+        &'a self,
+        invocation: CodeModeNestedToolCall,
+        cancellation_token: CancellationToken,
+    ) -> ToolInvocationFuture<'a>;
+
+    fn notify<'a>(
+        &'a self,
+        call_id: String,
+        cell_id: CellId,
+        text: String,
+        cancellation_token: CancellationToken,
+    ) -> NotificationFuture<'a>;
+
+    /// Releases delegate state associated with a cell after it reaches a terminal state.
+    fn cell_closed(&self, cell_id: &CellId);
+}
+
+/// A session delegate for clients that do not expose nested tools or notifications.
+pub struct NoopCodeModeSessionDelegate;
+
+impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
+    fn invoke_tool<'a>(
+        &'a self,
+        _invocation: CodeModeNestedToolCall,
+        cancellation_token: CancellationToken,
+    ) -> ToolInvocationFuture<'a> {
+        Box::pin(async move {
+            cancellation_token.cancelled().await;
+            Err("code mode nested tools are unavailable".to_string())
+        })
+    }
+
+    fn notify<'a>(
+        &'a self,
+        _call_id: String,
+        _cell_id: CellId,
+        _text: String,
+        _cancellation_token: CancellationToken,
+    ) -> NotificationFuture<'a> {
+        Box::pin(async { Ok(()) })
+    }
+
+    fn cell_closed(&self, _cell_id: &CellId) {}
+}
+
+/// A durable code-mode session owned by one Codex thread.
+///
+/// Cells executed in the same session share stored values. Separate sessions
+/// must keep those values isolated. Implementations may execute cells
+/// in-process or remotely.
+pub trait CodeModeSession: Send + Sync {
+    fn execute<'a>(
+        &'a self,
+        request: ExecuteRequest,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+    ) -> CodeModeSessionResultFuture<'a, StartedCell>;
+
+    fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
+
+    fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
+
+    fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>;
+}
+
+/// Creates code-mode sessions for Codex threads.
+///
+/// Implementations may share a remote host process across all sessions created
+/// by one provider.
+pub trait CodeModeSessionProvider: Send + Sync {
+    /// Reports whether this provider can execute code without starting its host.
+    fn availability(&self) -> Result<(), String> {
+        Ok(())
+    }
+
+    fn create_session(&self) -> CodeModeSessionProviderFuture<'_>;
+
+    /// Creates a session whose cells share the supplied execution limits.
+    ///
+    /// Existing providers remain compatible with unlimited sessions, but must
+    /// explicitly implement this method before accepting non-default limits.
+    fn create_session_with_limits<'a>(
+        &'a self,
+        limits: CodeModeSessionCellExecutionLimits,
+    ) -> CodeModeSessionProviderFuture<'a> {
+        if limits == CodeModeSessionCellExecutionLimits::default() {
+            self.create_session()
+        } else {
+            Box::pin(async {
+                Err("code-mode session provider does not support resource limits".to_string())
+            })
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "session_tests.rs"]
+mod tests;

codex-rs/code-mode-protocol/src/session_tests.rs

@@ -0,0 +1,19 @@
+use pretty_assertions::assert_eq;
+use tokio::sync::oneshot;
+
+use super::CellId;
+use super::StartedCell;
+
+#[tokio::test]
+async fn started_cell_preserves_remote_initial_response_errors() {
+    let (response_tx, response_rx) = oneshot::channel();
+    response_tx
+        .send(Err("remote runtime failed".to_string()))
+        .expect("initial response receiver should be open");
+    let started = StartedCell::from_result_receiver(CellId::new("1".to_string()), response_rx);
+
+    assert_eq!(
+        started.initial_response().await,
+        Err("remote runtime failed".to_string())
+    );
+}

codex-rs/collaboration-mode-templates/src/lib.rs

@@ -0,0 +1,2 @@
+pub const PLAN: &str = include_str!("../templates/plan.md");
+pub const DEFAULT: &str = include_str!("../templates/default.md");

codex-rs/collaboration-mode-templates/templates/default.md

@@ -0,0 +1,19 @@
+# Collaboration Mode: Default
+
+You are now in Default mode. Any previous instructions for other modes (e.g. Plan mode) are no longer active.
+
+Your active mode changes only when new developer instructions with a different `<collaboration_mode>...</collaboration_mode>` change it; user requests or tool descriptions do not change mode by themselves. Known mode names are Default and Plan.
+
+## request_user_input availability
+
+Use the `request_user_input` tool only when it is listed in the available tools for this turn.
+
+In Default mode, strongly prefer making reasonable assumptions and executing the user's request rather than stopping to ask questions.
+
+Use the `request_user_input` tool only for optional questions where the answer would materially improve the quality of the work.
+
+If `request_user_input` returns no answers, continue with best judgment instead of asking again or treating the turn as blocked.
+
+Never use the `request_user_input` tool for permission requests or permission-related escalations.
+
+If explicit user input is required for another reason before progress can safely continue, do not use the `request_user_input` tool. Ask the user directly with one concise plain-text question instead. Never write a multiple choice question as a textual assistant message.

codex-rs/collaboration-mode-templates/templates/plan.md

@@ -0,0 +1,128 @@
+# Plan Mode (Conversational)
+
+You work in 3 phases, and you should *chat your way* to a great plan before finalizing it. A great plan is very detailed—intent- and implementation-wise—so that it can be handed to another engineer or agent to be implemented right away. It must be **decision complete**, where the implementer does not need to make any decisions.
+
+## Mode rules (strict)
+
+You are in **Plan Mode** until a developer message explicitly ends it.
+
+Plan Mode is not changed by user intent, tone, or imperative language. If a user asks for execution while still in Plan Mode, treat it as a request to **plan the execution**, not perform it.
+
+## Plan Mode vs update_plan tool
+
+Plan Mode is a collaboration mode that can involve requesting user input and eventually issuing a `<proposed_plan>` block.
+
+Separately, `update_plan` is a checklist/progress/TODOs tool; it does not enter or exit Plan Mode. Do not confuse it with Plan mode or try to use it while in Plan mode. If you try to use `update_plan` in Plan mode, it will return an error.
+
+## Execution vs. mutation in Plan Mode
+
+You may explore and execute **non-mutating** actions that improve the plan. You must not perform **mutating** actions.
+
+### Allowed (non-mutating, plan-improving)
+
+Actions that gather truth, reduce ambiguity, or validate feasibility without changing repo-tracked state. Examples:
+
+* Reading or searching files, configs, schemas, types, manifests, and docs
+* Static analysis, inspection, and repo exploration
+* Dry-run style commands when they do not edit repo-tracked files
+* Tests, builds, or checks that may write to caches or build artifacts (for example, `target/`, `.cache/`, or snapshots) so long as they do not edit repo-tracked files
+
+### Not allowed (mutating, plan-executing)
+
+Actions that implement the plan or change repo-tracked state. Examples:
+
+* Editing or writing files
+* Running formatters or linters that rewrite files
+* Applying patches, migrations, or codegen that updates repo-tracked files
+* Side-effectful commands whose purpose is to carry out the plan rather than refine it
+
+When in doubt: if the action would reasonably be described as "doing the work" rather than "planning the work," do not do it.
+
+## PHASE 1 — Ground in the environment (explore first, ask second)
+
+Begin by grounding yourself in the actual environment. Eliminate unknowns in the prompt by discovering facts, not by asking the user. Resolve all questions that can be answered through exploration or inspection. Identify missing or ambiguous details only if they cannot be derived from the environment. Silent exploration between turns is allowed and encouraged.
+
+Before asking the user any question, perform at least one targeted non-mutating exploration pass (for example: search relevant files, inspect likely entrypoints/configs, confirm current implementation shape), unless no local environment/repo is available.
+
+Exception: you may ask clarifying questions about the user's prompt before exploring, ONLY if there are obvious ambiguities or contradictions in the prompt itself. However, if ambiguity might be resolved by exploring, always prefer exploring first.
+
+Do not ask questions that can be answered from the repo or system (for example, "where is this struct?" or "which UI component should we use?" when exploration can make it clear). Only ask once you have exhausted reasonable non-mutating exploration.
+
+## PHASE 2 — Intent chat (what they actually want)
+
+* Keep asking until you can clearly state: goal + success criteria, audience, in/out of scope, constraints, current state, and the key preferences/tradeoffs.
+* Bias toward questions over guessing: if any high-impact ambiguity remains, do NOT plan yet—ask.
+
+## PHASE 3 — Implementation chat (what/how we’ll build)
+
+* Once intent is stable, keep asking until the spec is decision complete: approach, interfaces (APIs/schemas/I/O), data flow, edge cases/failure modes, testing + acceptance criteria, rollout/monitoring, and any migrations/compat constraints.
+
+## Asking questions
+
+Critical rules:
+
+* Strongly prefer using the `request_user_input` tool to ask any questions.
+* Offer only meaningful multiple‑choice options; don’t include filler choices that are obviously wrong or irrelevant.
+* In rare cases where an unavoidable, important question can’t be expressed with reasonable multiple‑choice options (due to extreme ambiguity), you may ask it directly without the tool.
+
+You SHOULD ask many questions, but each question must:
+
+* materially change the spec/plan, OR
+* confirm/lock an assumption, OR
+* choose between meaningful tradeoffs.
+* not be answerable by non-mutating commands.
+
+Use the `request_user_input` tool only for decisions that materially change the plan, for confirming important assumptions, or for information that cannot be discovered via non-mutating exploration.
+
+## Two kinds of unknowns (treat differently)
+
+1. **Discoverable facts** (repo/system truth): explore first.
+
+   * Before asking, run targeted searches and check likely sources of truth (configs/manifests/entrypoints/schemas/types/constants).
+   * Ask only if: multiple plausible candidates; nothing found but you need a missing identifier/context; or ambiguity is actually product intent.
+   * If asking, present concrete candidates (paths/service names) + recommend one.
+   * Never ask questions you can answer from your environment (e.g., “where is this struct”).
+
+2. **Preferences/tradeoffs** (not discoverable): ask early.
+
+   * These are intent or implementation preferences that cannot be derived from exploration.
+   * Provide 2–4 mutually exclusive options + a recommended default.
+   * If unanswered, proceed with the recommended option and record it as an assumption in the final plan.
+
+## Finalization rule
+
+Only output the final plan when it is decision complete and leaves no decisions to the implementer.
+
+When you present the official plan, wrap it in a `<proposed_plan>` block so the client can render it specially:
+
+1) The opening tag must be on its own line.
+2) Start the plan content on the next line (no text on the same line as the tag).
+3) The closing tag must be on its own line.
+4) Use Markdown inside the block.
+5) Keep the tags exactly as `<proposed_plan>` and `</proposed_plan>` (do not translate or rename them), even if the plan content is in another language.
+
+Example:
+
+<proposed_plan>
+plan content
+</proposed_plan>
+
+plan content should be human and agent digestible. The final plan must be plan-only, concise by default, and include:
+
+* A clear title
+* A brief summary section
+* Important changes or additions to public APIs/interfaces/types
+* Test cases and scenarios
+* Explicit assumptions and defaults chosen where needed
+
+When possible, prefer a compact structure with 3-5 short sections, usually: Summary, Key Changes or Implementation Changes, Test Plan, and Assumptions. Do not include a separate Scope section unless scope boundaries are genuinely important to avoid mistakes.
+
+Prefer grouped implementation bullets by subsystem or behavior over file-by-file inventories. Mention files only when needed to disambiguate a non-obvious change, and avoid naming more than 3 paths unless extra specificity is necessary to prevent mistakes. Prefer behavior-level descriptions over symbol-by-symbol removal lists. For v1 feature-addition plans, do not invent detailed schema, validation, precedence, fallback, or wire-shape policy unless the request establishes it or it is needed to prevent a concrete implementation mistake; prefer the intended capability and minimum interface/behavior changes.
+
+Keep bullets short and avoid explanatory sub-bullets unless they are needed to prevent ambiguity. Prefer the minimum detail needed for implementation safety, not exhaustive coverage. Within each section, compress related changes into a few high-signal bullets and omit branch-by-branch logic, repeated invariants, and long lists of unaffected behavior unless they are necessary to prevent a likely implementation mistake. Avoid repeated repo facts and irrelevant edge-case or rollout detail. For straightforward refactors, keep the plan to a compact summary, key edits, tests, and assumptions. If the user asks for more detail, then expand.
+
+Do not ask "should I proceed?" in the final output. The user can easily switch out of Plan mode and request implementation if you have included a `<proposed_plan>` block in your response. Alternatively, they can decide to stay in Plan mode and continue refining the plan.
+
+Only produce at most one `<proposed_plan>` block per turn, and only when you are presenting a complete spec.
+
+If the user stays in Plan mode and asks for revisions after a prior `<proposed_plan>`, any new `<proposed_plan>` must be a complete replacement. If the user indicates that the prior plan is not acceptable but does not provide enough information to produce a complete replacement, address the concern and continue planning without producing a `<proposed_plan>` block. If the follow-up neither requires changes nor calls the plan into question (e.g. clarifying question), answer it before the block, then reproduce the prior `<proposed_plan>` unchanged.

codex-rs/exec-server/src/arg0_exec_helper.rs

@@ -0,0 +1,31 @@
+#[cfg(unix)]
+use std::process::Command;
+
+pub const CODEX_ARG0_EXEC_HELPER_ARG1: &str = "--codex-run-as-arg0-exec-helper";
+
+#[cfg(unix)]
+pub fn main() -> ! {
+    use std::os::unix::process::CommandExt;
+
+    let mut args = std::env::args_os();
+    let _program = args.next();
+    let _helper_mode = args.next();
+    let Some(arg0) = args.next() else {
+        eprintln!("missing arg0 for exec helper");
+        std::process::exit(1);
+    };
+    let Some(program) = args.next() else {
+        eprintln!("missing program for exec helper");
+        std::process::exit(1);
+    };
+
+    let error = Command::new(&program).arg0(arg0).args(args).exec();
+    eprintln!("failed to exec {program:?}: {error}");
+    std::process::exit(1);
+}
+
+#[cfg(not(unix))]
+pub fn main() -> ! {
+    eprintln!("arg0 exec helper is only supported on Unix");
+    std::process::exit(1);
+}

codex-rs/exec-server/src/capability_discovery.rs

@@ -0,0 +1,510 @@
+use std::collections::HashSet;
+use std::io;
+
+use codex_exec_server_protocol::CapabilityRootDiscoverRequest;
+use codex_exec_server_protocol::CapabilityRootDiscovery;
+use codex_exec_server_protocol::CapabilityRootsDiscoverParams;
+use codex_exec_server_protocol::CapabilityRootsDiscoverResponse;
+use codex_exec_server_protocol::CapabilityTextFile;
+use codex_exec_server_protocol::DISCOVERABLE_PLUGIN_MANIFEST_PATHS;
+use codex_exec_server_protocol::DiscoveredPluginFiles;
+use codex_exec_server_protocol::DiscoveredSkillFiles;
+use codex_file_system::ExecutorFileSystem;
+use codex_file_system::FileSystemSandboxContext;
+use codex_file_system::WalkEntryKind;
+use codex_file_system::WalkOptions;
+use codex_utils_path_uri::PathUri;
+use futures::StreamExt;
+use serde::Deserialize;
+use serde_json::Value;
+
+pub(crate) const MAX_ROOTS_PER_REQUEST: usize = 128;
+const MAX_SCAN_DEPTH: usize = 6;
+const MAX_DIRECTORIES_PER_ROOT: usize = 2_000;
+const MAX_ENTRIES_PER_ROOT: usize = 20_000;
+const MAX_FILE_BYTES: usize = 1024 * 1024;
+const MAX_BUNDLE_BYTES_PER_ROOT: usize = 16 * 1024 * 1024;
+const MAX_CONCURRENT_ROOTS: usize = 8;
+const SKILL_FILE_NAME: &str = "SKILL.md";
+const SKILL_METADATA_PATH: &str = "agents/openai.yaml";
+const DEFAULT_MCP_CONFIG_PATH: &str = ".mcp.json";
+
+#[derive(Debug, thiserror::Error)]
+pub enum CapabilityDiscoveryError {
+    #[error("capability root discovery accepts at most {MAX_ROOTS_PER_REQUEST} roots")]
+    TooManyRoots,
+}
+
+/// Discovers and materializes capability manifests using one executor-local filesystem.
+///
+/// Product parsing and policy intentionally remain with the caller. This operation owns the
+/// filesystem-expensive portion: bounded traversal, recognized-file selection, and reads.
+#[tracing::instrument(
+    name = "capability_roots.discover_v1",
+    skip_all,
+    fields(root_count = params.roots.len())
+)]
+pub async fn discover_capability_roots(
+    file_system: &dyn ExecutorFileSystem,
+    params: CapabilityRootsDiscoverParams,
+) -> Result<CapabilityRootsDiscoverResponse, CapabilityDiscoveryError> {
+    if params.roots.len() > MAX_ROOTS_PER_REQUEST {
+        return Err(CapabilityDiscoveryError::TooManyRoots);
+    }
+
+    let roots = futures::stream::iter(params.roots)
+        .map(|root| discover_root(file_system, root))
+        .buffered(MAX_CONCURRENT_ROOTS)
+        .collect()
+        .await;
+    Ok(CapabilityRootsDiscoverResponse { roots })
+}
+
+async fn discover_root(
+    file_system: &dyn ExecutorFileSystem,
+    request: CapabilityRootDiscoverRequest,
+) -> CapabilityRootDiscovery {
+    let CapabilityRootDiscoverRequest { id, path, sandbox } = request;
+    let sandbox = sandbox.as_ref();
+    let mut discovery = CapabilityRootDiscovery {
+        id,
+        path: path.clone(),
+        plugin: None,
+        skills: Vec::new(),
+        namespace_manifests: Vec::new(),
+        warnings: Vec::new(),
+        error: None,
+    };
+
+    #[cfg(target_os = "windows")]
+    if sandbox.is_some_and(|context| {
+        context.should_run_in_sandbox() && !context.windows_sandbox_is_requested()
+    }) {
+        discovery.error = Some("filesystem sandbox is unavailable on this executor".to_string());
+        return discovery;
+    }
+
+    match file_system
+        .get_metadata(&path, Default::default(), sandbox)
+        .await
+    {
+        Ok(metadata) if metadata.is_directory => {}
+        Ok(_) => {
+            discovery.error = Some(format!("capability root {path} is not a directory"));
+            return discovery;
+        }
+        Err(error) => {
+            discovery.error = Some(format!("failed to inspect capability root {path}: {error}"));
+            return discovery;
+        }
+    }
+
+    let walk = match file_system
+        .walk(
+            &path,
+            WalkOptions {
+                max_depth: MAX_SCAN_DEPTH,
+                max_directories: MAX_DIRECTORIES_PER_ROOT,
+                max_entries: MAX_ENTRIES_PER_ROOT,
+                follow_directory_symlinks: true,
+                prune_hidden_directories: false,
+            },
+            sandbox,
+        )
+        .await
+    {
+        Ok(walk) => walk,
+        Err(error) => {
+            discovery.error = Some(format!("failed to scan capability root {path}: {error}"));
+            return discovery;
+        }
+    };
+    discovery
+        .warnings
+        .extend(walk.errors.into_iter().map(|error| {
+            format!(
+                "failed to scan capability path {}: {}",
+                error.path, error.message
+            )
+        }));
+    if walk.truncated {
+        discovery.warnings.push(format!(
+            "capability scan reached its traversal limit (root: {path})"
+        ));
+    }
+
+    let mut skill_paths = Vec::new();
+    let mut namespace_manifest_paths = Vec::new();
+    for entry in walk.entries {
+        if entry.kind != WalkEntryKind::File {
+            continue;
+        }
+        if entry.path.basename().as_deref() == Some(SKILL_FILE_NAME) {
+            skill_paths.push(entry.path.clone());
+        }
+        if is_plugin_manifest_path(&entry.path) {
+            namespace_manifest_paths.push(entry.path);
+        }
+    }
+    skill_paths.sort_unstable_by_key(PathUri::to_string);
+    namespace_manifest_paths.sort_unstable_by(|left, right| {
+        let left_root = plugin_root_for_manifest(left).map(|path| path.to_string());
+        let right_root = plugin_root_for_manifest(right).map(|path| path.to_string());
+        left_root
+            .cmp(&right_root)
+            .then_with(|| plugin_manifest_priority(left).cmp(&plugin_manifest_priority(right)))
+    });
+
+    let mut budget = BundleBudget::default();
+    let root_manifest = read_first_plugin_manifest(
+        file_system,
+        &path,
+        sandbox,
+        &mut budget,
+        &mut discovery.warnings,
+    )
+    .await;
+
+    let inherited_manifest = match root_manifest.as_ref() {
+        Some(manifest) => Some(manifest.clone()),
+        None => {
+            read_nearest_ancestor_manifest(
+                file_system,
+                &path,
+                sandbox,
+                &mut budget,
+                &mut discovery.warnings,
+            )
+            .await
+        }
+    };
+    let mut seen_namespace_roots = HashSet::new();
+    if let Some(manifest) = inherited_manifest {
+        if let Some(plugin_root) = plugin_root_for_manifest(&manifest.path) {
+            seen_namespace_roots.insert(plugin_root);
+        }
+        discovery.namespace_manifests.push(manifest);
+    }
+    for manifest_path in namespace_manifest_paths {
+        let Some(plugin_root) = plugin_root_for_manifest(&manifest_path) else {
+            continue;
+        };
+        if !seen_namespace_roots.insert(plugin_root) {
+            continue;
+        }
+        if let Some(manifest) = read_optional_text_file(
+            file_system,
+            manifest_path,
+            sandbox,
+            &mut budget,
+            &mut discovery.warnings,
+        )
+        .await
+        {
+            discovery.namespace_manifests.push(manifest);
+        }
+    }
+
+    if let Some(manifest) = root_manifest {
+        let declarations = plugin_declaration_paths(&path, &manifest, &mut discovery.warnings);
+        let mcp_path = if declarations.mcp_inline {
+            None
+        } else {
+            declarations
+                .mcp_config
+                .or_else(|| path.join(DEFAULT_MCP_CONFIG_PATH).ok())
+        };
+        let mcp_config = match mcp_path {
+            Some(path) => {
+                read_optional_text_file(
+                    file_system,
+                    path,
+                    sandbox,
+                    &mut budget,
+                    &mut discovery.warnings,
+                )
+                .await
+            }
+            None => None,
+        };
+        let apps_config = match declarations.apps_config {
+            Some(path) => {
+                read_optional_text_file(
+                    file_system,
+                    path,
+                    sandbox,
+                    &mut budget,
+                    &mut discovery.warnings,
+                )
+                .await
+            }
+            None => None,
+        };
+        discovery.plugin = Some(DiscoveredPluginFiles {
+            manifest,
+            mcp_config,
+            apps_config,
+        });
+    }
+
+    for skill_path in skill_paths {
+        let Some(instructions) = read_optional_text_file(
+            file_system,
+            skill_path.clone(),
+            sandbox,
+            &mut budget,
+            &mut discovery.warnings,
+        )
+        .await
+        else {
+            continue;
+        };
+        let metadata = match skill_path
+            .parent()
+            .and_then(|skill_dir| skill_dir.join(SKILL_METADATA_PATH).ok())
+        {
+            Some(metadata_path) => {
+                read_optional_text_file(
+                    file_system,
+                    metadata_path,
+                    sandbox,
+                    &mut budget,
+                    &mut discovery.warnings,
+                )
+                .await
+            }
+            None => None,
+        };
+        discovery.skills.push(DiscoveredSkillFiles {
+            instructions,
+            metadata,
+        });
+    }
+
+    discovery
+}
+
+async fn read_first_plugin_manifest(
+    file_system: &dyn ExecutorFileSystem,
+    root: &PathUri,
+    sandbox: Option<&FileSystemSandboxContext>,
+    budget: &mut BundleBudget,
+    warnings: &mut Vec<String>,
+) -> Option<CapabilityTextFile> {
+    for relative_path in DISCOVERABLE_PLUGIN_MANIFEST_PATHS {
+        let Ok(path) = root.join(relative_path) else {
+            continue;
+        };
+        if let Some(manifest) =
+            read_optional_text_file(file_system, path, sandbox, budget, warnings).await
+        {
+            return Some(manifest);
+        }
+    }
+    None
+}
+
+async fn read_nearest_ancestor_manifest(
+    file_system: &dyn ExecutorFileSystem,
+    root: &PathUri,
+    sandbox: Option<&FileSystemSandboxContext>,
+    budget: &mut BundleBudget,
+    warnings: &mut Vec<String>,
+) -> Option<CapabilityTextFile> {
+    let mut ancestor = root.parent();
+    while let Some(path) = ancestor {
+        if let Some(manifest) =
+            read_first_plugin_manifest(file_system, &path, sandbox, budget, warnings).await
+        {
+            return Some(manifest);
+        }
+        ancestor = path.parent();
+    }
+    None
+}
+
+async fn read_optional_text_file(
+    file_system: &dyn ExecutorFileSystem,
+    path: PathUri,
+    sandbox: Option<&FileSystemSandboxContext>,
+    budget: &mut BundleBudget,
+    warnings: &mut Vec<String>,
+) -> Option<CapabilityTextFile> {
+    let metadata = match file_system
+        .get_metadata(&path, Default::default(), sandbox)
+        .await
+    {
+        Ok(metadata) if metadata.is_file => metadata,
+        Ok(_) => return None,
+        Err(error) if error.kind() == io::ErrorKind::NotFound => return None,
+        Err(error) => {
+            warnings.push(format!("failed to inspect capability file {path}: {error}"));
+            return None;
+        }
+    };
+    let Ok(size) = usize::try_from(metadata.size) else {
+        warnings.push(format!("capability file {path} is too large"));
+        return None;
+    };
+    if size > MAX_FILE_BYTES {
+        warnings.push(format!(
+            "capability file {path} exceeds the {MAX_FILE_BYTES}-byte limit"
+        ));
+        return None;
+    }
+    if !budget.can_add(size) {
+        warnings.push(format!(
+            "capability root bundle exceeds the {MAX_BUNDLE_BYTES_PER_ROOT}-byte limit"
+        ));
+        return None;
+    }
+    let mut stream = match file_system.read_file_stream(&path, sandbox).await {
+        Ok(stream) => stream,
+        Err(error) => {
+            warnings.push(format!("failed to read capability file {path}: {error}"));
+            return None;
+        }
+    };
+    let mut contents = Vec::with_capacity(size);
+    while let Some(chunk) = stream.next().await {
+        let chunk = match chunk {
+            Ok(chunk) => chunk,
+            Err(error) => {
+                warnings.push(format!("failed to read capability file {path}: {error}"));
+                return None;
+            }
+        };
+        let Some(new_len) = contents.len().checked_add(chunk.len()) else {
+            warnings.push(format!("capability file {path} exceeded its read limit"));
+            return None;
+        };
+        if new_len > MAX_FILE_BYTES || !budget.can_add(new_len) {
+            warnings.push(format!("capability file {path} exceeded its read limit"));
+            return None;
+        }
+        contents.extend_from_slice(&chunk);
+    }
+    let contents = match String::from_utf8(contents) {
+        Ok(contents) => contents,
+        Err(error) => {
+            warnings.push(format!("capability file {path} is not UTF-8: {error}"));
+            return None;
+        }
+    };
+    budget.add(contents.len());
+    Some(CapabilityTextFile { path, contents })
+}
+
+fn is_plugin_manifest_path(path: &PathUri) -> bool {
+    plugin_manifest_priority(path).is_some()
+}
+
+fn plugin_manifest_priority(path: &PathUri) -> Option<usize> {
+    if path.basename().as_deref() != Some("plugin.json") {
+        return None;
+    }
+    let manifest_directory = path.parent()?.basename()?;
+    DISCOVERABLE_PLUGIN_MANIFEST_PATHS
+        .iter()
+        .position(|relative_path| {
+            relative_path.strip_suffix("/plugin.json") == Some(manifest_directory.as_str())
+        })
+}
+
+fn plugin_root_for_manifest(path: &PathUri) -> Option<PathUri> {
+    path.parent()?.parent()
+}
+
+#[derive(Default)]
+struct BundleBudget {
+    bytes: usize,
+}
+
+impl BundleBudget {
+    fn can_add(&self, bytes: usize) -> bool {
+        self.bytes
+            .checked_add(bytes)
+            .is_some_and(|total| total <= MAX_BUNDLE_BYTES_PER_ROOT)
+    }
+
+    fn add(&mut self, bytes: usize) {
+        self.bytes += bytes;
+    }
+}
+
+#[derive(Default)]
+struct PluginDeclarationPaths {
+    mcp_config: Option<PathUri>,
+    mcp_inline: bool,
+    apps_config: Option<PathUri>,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct RawPluginDeclarations {
+    #[serde(default)]
+    mcp_servers: Option<Value>,
+    #[serde(default)]
+    apps: Option<Value>,
+}
+
+fn plugin_declaration_paths(
+    root: &PathUri,
+    manifest: &CapabilityTextFile,
+    warnings: &mut Vec<String>,
+) -> PluginDeclarationPaths {
+    let declarations = match serde_json::from_str::<RawPluginDeclarations>(&manifest.contents) {
+        Ok(declarations) => declarations,
+        Err(_) => return PluginDeclarationPaths::default(),
+    };
+    PluginDeclarationPaths {
+        mcp_config: declarations.mcp_servers.as_ref().and_then(|value| {
+            declared_file_path(root, "mcpServers", value, &manifest.path, warnings)
+        }),
+        mcp_inline: declarations
+            .mcp_servers
+            .as_ref()
+            .is_some_and(Value::is_object),
+        apps_config: declarations
+            .apps
+            .as_ref()
+            .and_then(|value| declared_file_path(root, "apps", value, &manifest.path, warnings)),
+    }
+}
+
+fn declared_file_path(
+    root: &PathUri,
+    field: &str,
+    value: &Value,
+    manifest_path: &PathUri,
+    warnings: &mut Vec<String>,
+) -> Option<PathUri> {
+    let Value::String(path) = value else {
+        return None;
+    };
+    let Some(relative_path) = path.strip_prefix("./") else {
+        warnings.push(format!(
+            "ignoring {field} in {manifest_path}: path must start with `./`"
+        ));
+        return None;
+    };
+    if relative_path.is_empty()
+        || relative_path
+            .split(['/', '\\'])
+            .any(|component| component == "..")
+    {
+        warnings.push(format!(
+            "ignoring {field} in {manifest_path}: path must remain below the capability root"
+        ));
+        return None;
+    }
+    match root.join(relative_path) {
+        Ok(path) if path.starts_with(root) => Some(path),
+        Ok(_) | Err(_) => {
+            warnings.push(format!(
+                "ignoring {field} in {manifest_path}: path must remain below the capability root"
+            ));
+            None
+        }
+    }
+}

codex-rs/exec-server/src/capability_discovery_cache.rs

@@ -0,0 +1,246 @@
+use std::collections::BTreeMap;
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+
+use codex_protocol::capabilities::CapabilityRootLocation;
+use codex_protocol::capabilities::SelectedCapabilityRoot;
+use tokio::sync::Mutex;
+
+use crate::CapabilityRootDiscoverRequest;
+use crate::CapabilityRootDiscovery;
+use crate::CapabilityRootsDiscoverParams;
+use crate::EnvironmentManager;
+use crate::ExecutorCapabilityDiscoverySnapshot;
+use crate::FileSystemSandboxContext;
+
+/// Thread-scoped cache shared by capability consumers using the high-level executor API.
+///
+/// A single miss batches every requested root by environment. Successful discoveries and
+/// permanent failures remain cached by root and sandbox; transient failures are retried on the
+/// next request. Recovery is reported so dependent MCP projections can be invalidated.
+pub struct ExecutorCapabilityDiscoveryCache {
+    environment_manager: Arc<EnvironmentManager>,
+    entries: Mutex<Vec<CachedRoot>>,
+    recovered_discovery: AtomicBool,
+}
+
+impl std::fmt::Debug for ExecutorCapabilityDiscoveryCache {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter
+            .debug_struct("ExecutorCapabilityDiscoveryCache")
+            .finish_non_exhaustive()
+    }
+}
+
+struct CachedRoot {
+    selected_root: SelectedCapabilityRoot,
+    sandbox: Option<FileSystemSandboxContext>,
+    result: Result<Arc<CapabilityRootDiscovery>, String>,
+    // Preserve transport classification after the public snapshot reduces errors to strings.
+    retryable: bool,
+}
+
+impl ExecutorCapabilityDiscoveryCache {
+    pub fn new(environment_manager: Arc<EnvironmentManager>) -> Self {
+        Self {
+            environment_manager,
+            entries: Mutex::new(Vec::new()),
+            recovered_discovery: AtomicBool::new(false),
+        }
+    }
+
+    /// Reports whether a previously failed root has recovered since the last observation.
+    pub fn take_recovered_discovery(&self) -> bool {
+        self.recovered_discovery.swap(false, Ordering::AcqRel)
+    }
+
+    /// Returns discoveries in the same order as `selected_roots`.
+    #[tracing::instrument(
+        name = "capability_roots.discovery_cache.resolve",
+        skip_all,
+        fields(root_count = selected_roots.len())
+    )]
+    pub async fn discover(
+        &self,
+        selected_roots: &[SelectedCapabilityRoot],
+        sandbox_contexts: &HashMap<String, FileSystemSandboxContext>,
+    ) -> Vec<Result<Arc<CapabilityRootDiscovery>, String>> {
+        let missing = {
+            let entries = self.entries.lock().await;
+            selected_roots
+                .iter()
+                .filter(|selected_root| {
+                    let CapabilityRootLocation::Environment { environment_id, .. } =
+                        &selected_root.location;
+                    let sandbox = sandbox_contexts.get(environment_id);
+                    !entries.iter().any(|cached| {
+                        cached.selected_root == **selected_root
+                            && cached.sandbox.as_ref() == sandbox
+                            && !cached.retryable
+                    })
+                })
+                .cloned()
+                .collect::<Vec<_>>()
+        };
+        let discovered = self.discover_missing(missing, sandbox_contexts).await;
+        let mut entries = self.entries.lock().await;
+        for discovered_root in discovered {
+            if let Some(cached) = entries
+                .iter_mut()
+                .find(|cached| cached.selected_root == discovered_root.selected_root)
+            {
+                if cached.sandbox != discovered_root.sandbox || cached.result.is_err() {
+                    if cached.result.is_err() && discovered_root.result.is_ok() {
+                        self.recovered_discovery.store(true, Ordering::Release);
+                    }
+                    *cached = discovered_root;
+                }
+            } else {
+                entries.push(discovered_root);
+            }
+        }
+        selected_roots
+            .iter()
+            .map(|selected_root| {
+                let CapabilityRootLocation::Environment { environment_id, .. } =
+                    &selected_root.location;
+                let sandbox = sandbox_contexts.get(environment_id);
+                match entries.iter().find(|cached| {
+                    cached.selected_root == *selected_root && cached.sandbox.as_ref() == sandbox
+                }) {
+                    Some(cached) => cached.result.clone(),
+                    None => Err(format!(
+                        "selected capability root `{}` was not discovered",
+                        selected_root.id
+                    )),
+                }
+            })
+            .collect()
+    }
+
+    /// Resolves the selected roots once and freezes their results for one model step.
+    pub async fn snapshot(
+        &self,
+        selected_roots: &[SelectedCapabilityRoot],
+        sandbox_contexts: &HashMap<String, FileSystemSandboxContext>,
+    ) -> ExecutorCapabilityDiscoverySnapshot {
+        ExecutorCapabilityDiscoverySnapshot::new(
+            selected_roots,
+            self.discover(selected_roots, sandbox_contexts).await,
+            sandbox_contexts.clone(),
+        )
+    }
+
+    async fn discover_missing(
+        &self,
+        missing: Vec<SelectedCapabilityRoot>,
+        sandbox_contexts: &HashMap<String, FileSystemSandboxContext>,
+    ) -> Vec<CachedRoot> {
+        let mut grouped = BTreeMap::<String, Vec<SelectedCapabilityRoot>>::new();
+        for selected_root in missing {
+            let CapabilityRootLocation::Environment { environment_id, .. } =
+                &selected_root.location;
+            grouped
+                .entry(environment_id.clone())
+                .or_default()
+                .push(selected_root);
+        }
+
+        let batches = grouped.into_iter().flat_map(|(environment_id, roots)| {
+            roots
+                .chunks(crate::capability_discovery::MAX_ROOTS_PER_REQUEST)
+                .map(|batch| (environment_id.clone(), batch.to_vec()))
+                .collect::<Vec<_>>()
+        });
+        let discoveries =
+            futures::future::join_all(batches.map(|(environment_id, selected_roots)| async move {
+                let sandbox = sandbox_contexts.get(&environment_id).cloned();
+                let Some(environment) = self.environment_manager.get_environment(&environment_id)
+                else {
+                    let error = format!("environment `{environment_id}` is unavailable");
+                    return selected_roots
+                        .into_iter()
+                        .map(|selected_root| CachedRoot {
+                            selected_root,
+                            sandbox: sandbox.clone(),
+                            result: Err(error.clone()),
+                            retryable: true,
+                        })
+                        .collect::<Vec<_>>();
+                };
+                let params = CapabilityRootsDiscoverParams {
+                    roots: selected_roots
+                        .iter()
+                        .map(|selected_root| {
+                            let CapabilityRootLocation::Environment { path, .. } =
+                                &selected_root.location;
+                            CapabilityRootDiscoverRequest {
+                                id: selected_root.id.clone(),
+                                path: path.clone(),
+                                sandbox: sandbox.clone(),
+                            }
+                        })
+                        .collect(),
+                };
+                let response = match environment.discover_capability_roots(params).await {
+                    Ok(response) => response,
+                    Err(error) => {
+                        let retryable = crate::client::is_retryable_recovery_error(&error);
+                        let error = error.to_string();
+                        return selected_roots
+                            .into_iter()
+                            .map(|selected_root| CachedRoot {
+                                selected_root,
+                                sandbox: sandbox.clone(),
+                                result: Err(error.clone()),
+                                retryable,
+                            })
+                            .collect();
+                    }
+                };
+                if response.roots.len() != selected_roots.len() {
+                    let error = format!(
+                        "exec-server returned {} capability roots for {} requests",
+                        response.roots.len(),
+                        selected_roots.len()
+                    );
+                    return selected_roots
+                        .into_iter()
+                        .map(|selected_root| CachedRoot {
+                            selected_root,
+                            sandbox: sandbox.clone(),
+                            result: Err(error.clone()),
+                            retryable: false,
+                        })
+                        .collect();
+                }
+                selected_roots
+                    .into_iter()
+                    .zip(response.roots)
+                    .map(|(selected_root, discovery)| {
+                        let CapabilityRootLocation::Environment { path, .. } =
+                            &selected_root.location;
+                        let result = if discovery.id == selected_root.id && discovery.path == *path
+                        {
+                            Ok(Arc::new(discovery))
+                        } else {
+                            Err(format!(
+                                "exec-server returned mismatched capability root `{}` at {}",
+                                discovery.id, discovery.path
+                            ))
+                        };
+                        CachedRoot {
+                            selected_root,
+                            sandbox: sandbox.clone(),
+                            result,
+                            retryable: false,
+                        }
+                    })
+                    .collect()
+            }))
+            .await;
+        discoveries.into_iter().flatten().collect()
+    }
+}

codex-rs/exec-server/src/client/accepted.rs

@@ -0,0 +1,266 @@
+use std::sync::Arc;
+
+use axum::extract::ws::WebSocket;
+use futures::lock::Mutex;
+use tokio::sync::OnceCell;
+use tokio::sync::OwnedSemaphorePermit;
+use tokio::sync::Semaphore;
+use tokio::sync::mpsc;
+use tokio::sync::watch;
+
+use super::ConnectionStatus;
+use super::ExecServerClient;
+use super::Inner;
+use super::LazyRemoteExecServerClient;
+use crate::EnvironmentConnectionState;
+use crate::ExecServerClientConnectOptions;
+use crate::ExecServerError;
+use crate::client_transport::ExecServerReconnectStrategy;
+use crate::client_transport::ReconnectAttempt;
+use crate::connection::JsonRpcConnection;
+use codex_http_client::HttpClientFactory;
+
+struct AcceptedReplacement {
+    connection: JsonRpcConnection,
+    permit: OwnedSemaphorePermit,
+}
+
+struct AcceptedConnectionSourceInner {
+    replacements_tx: mpsc::UnboundedSender<AcceptedReplacement>,
+    replacements_rx: Mutex<mpsc::UnboundedReceiver<AcceptedReplacement>>,
+    replacement_slots: Arc<Semaphore>,
+}
+
+/// Receives authenticated connections supplied by an embedding host.
+///
+/// The source owns serialization and cancellation cleanup for replacement
+/// handoffs. The reconnect loop only asks it for the next connection.
+#[derive(Clone)]
+pub(crate) struct AcceptedConnectionSource {
+    inner: Arc<AcceptedConnectionSourceInner>,
+    options: ExecServerClientConnectOptions,
+}
+
+struct AcceptedReplacementSubmission {
+    source: AcceptedConnectionSource,
+    permit: OwnedSemaphorePermit,
+}
+
+impl AcceptedConnectionSource {
+    fn new(options: ExecServerClientConnectOptions) -> Self {
+        let (replacements_tx, replacements_rx) = mpsc::unbounded_channel();
+        Self {
+            inner: Arc::new(AcceptedConnectionSourceInner {
+                replacements_tx,
+                replacements_rx: Mutex::new(replacements_rx),
+                replacement_slots: Arc::new(Semaphore::new(1)),
+            }),
+            options,
+        }
+    }
+
+    fn begin_replacement(&self) -> Result<AcceptedReplacementSubmission, ExecServerError> {
+        let permit = Arc::clone(&self.inner.replacement_slots)
+            .try_acquire_owned()
+            .map_err(|_| {
+                ExecServerError::Protocol(
+                    "an accepted exec-server replacement is already in progress".to_string(),
+                )
+            })?;
+        Ok(AcceptedReplacementSubmission {
+            source: self.clone(),
+            permit,
+        })
+    }
+
+    pub(crate) async fn next_connection(
+        &self,
+        session_id: &str,
+    ) -> Result<ReconnectAttempt, ExecServerError> {
+        let replacement = self
+            .inner
+            .replacements_rx
+            .lock()
+            .await
+            .recv()
+            .await
+            .ok_or_else(|| {
+                ExecServerError::Disconnected(
+                    "accepted exec-server replacement channel closed".to_string(),
+                )
+            })?;
+        let mut options = self.options.clone();
+        options.resume_session_id = Some(session_id.to_string());
+        Ok(ReconnectAttempt::with_attempt_permit(
+            replacement.connection,
+            options,
+            replacement.permit,
+        ))
+    }
+}
+
+impl AcceptedReplacementSubmission {
+    fn submit(self, connection: JsonRpcConnection) -> Result<(), ExecServerError> {
+        self.source
+            .inner
+            .replacements_tx
+            .send(AcceptedReplacement {
+                connection,
+                permit: self.permit,
+            })
+            .map_err(|_| {
+                ExecServerError::Disconnected(
+                    "accepted exec-server connection is no longer awaiting replacements"
+                        .to_string(),
+                )
+            })
+    }
+}
+
+impl ExecServerClient {
+    /// Initializes an exec-server client over a WebSocket accepted by an Axum handler.
+    ///
+    /// The caller owns accepting and authenticating replacement WebSockets.
+    pub(crate) async fn connect_accepted_websocket(
+        websocket: WebSocket,
+        options: ExecServerClientConnectOptions,
+    ) -> Result<Self, ExecServerError> {
+        if options.resume_session_id.is_some() {
+            return Err(ExecServerError::Protocol(
+                "accepted exec-server initial connection cannot resume a session".to_string(),
+            ));
+        }
+        let connection_source = AcceptedConnectionSource::new(options.clone());
+        Self::connect_with_recovery(
+            JsonRpcConnection::from_axum_websocket(
+                websocket,
+                "accepted exec-server websocket".to_string(),
+            ),
+            options,
+            Some(ExecServerReconnectStrategy::Accepted(connection_source)),
+        )
+        .await
+    }
+
+    /// Supplies an authenticated replacement WebSocket for this accepted client.
+    ///
+    /// Retires the old transport before resuming the saved session. Returns
+    /// after handoff; recovery continues asynchronously.
+    pub(crate) async fn replace_accepted_websocket(
+        &self,
+        websocket: WebSocket,
+    ) -> Result<(), ExecServerError> {
+        self.inner
+            .accept_replacement_connection(JsonRpcConnection::from_axum_websocket(
+                websocket,
+                "accepted exec-server replacement websocket".to_string(),
+            ))
+            .await
+    }
+}
+
+impl Inner {
+    /// Hands a replacement connection from the host to the existing accepted client.
+    ///
+    /// This method coordinates the handoff in this order:
+    ///
+    /// 1. Verify that this client uses the accepted connection source.
+    /// 2. Reserve the source so concurrent handoffs are rejected.
+    /// 3. If the old RPC transport is still connected, move the client into
+    ///    recovery and close that transport before attaching the same session to
+    ///    the replacement.
+    /// 4. Queue the raw connection for the recovery task. That task creates the
+    ///    new RPC client, runs the initialize/resume handshake with the saved
+    ///    session ID, and recovers the existing processes.
+    async fn accept_replacement_connection(
+        self: &Arc<Self>,
+        connection: JsonRpcConnection,
+    ) -> Result<(), ExecServerError> {
+        if self.session_id.get().is_none() {
+            return Err(ExecServerError::Protocol(
+                "accepted exec-server connection is missing its session ID".to_string(),
+            ));
+        }
+
+        let Some(ExecServerReconnectStrategy::Accepted(connection_source)) =
+            &self.reconnect_strategy
+        else {
+            return Err(ExecServerError::Protocol(
+                "only an accepted exec-server connection can be replaced directly".to_string(),
+            ));
+        };
+        let (current_rpc_client, replacement_submission) = {
+            let connection = self
+                .connection
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            let current_rpc_client = match &connection.status {
+                ConnectionStatus::Failed(message) => {
+                    return Err(ExecServerError::Disconnected(message.clone()));
+                }
+                ConnectionStatus::Connected(rpc_client) => Some(Arc::clone(rpc_client)),
+                ConnectionStatus::Recovering => None,
+            };
+            let replacement_submission = connection_source.begin_replacement()?;
+            (current_rpc_client, replacement_submission)
+        };
+        if let Some(current_rpc_client) = current_rpc_client {
+            self.request_recovery(
+                Arc::clone(&current_rpc_client),
+                "exec-server connection replaced".to_string(),
+            );
+            current_rpc_client.close_transport().await;
+        }
+        // Synchronize the enqueue with terminal recovery. Recovery can time out
+        // while the handoff is waiting for the old transport to close; in that
+        // case the host must not receive success for a socket that no task will
+        // consume. Holding the connection lock through the synchronous send
+        // makes either the failure or the enqueue win the race unambiguously.
+        let connection_state = self
+            .connection
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        if let ConnectionStatus::Failed(message) = &connection_state.status {
+            return Err(ExecServerError::Disconnected(message.clone()));
+        }
+        replacement_submission.submit(connection)
+    }
+}
+
+#[cfg(test)]
+#[path = "accepted_tests.rs"]
+mod tests;
+
+impl LazyRemoteExecServerClient {
+    pub(crate) fn from_connected(
+        client: ExecServerClient,
+        http_client_factory: HttpClientFactory,
+    ) -> Self {
+        let environment_connection_state_tx =
+            watch::channel(EnvironmentConnectionState::Connected).0;
+        client.attach_environment_connection_state(environment_connection_state_tx.clone());
+        Self {
+            transport_params: None,
+            http_client_factory,
+            recovery_policy: super::RecoveryPolicy::Wait,
+            startup: std::sync::Arc::new(super::ConnectionAttempt {
+                result: OnceCell::new_with(Some(Ok(client.clone()))),
+                ..Default::default()
+            }),
+            current_client: std::sync::Arc::new(std::sync::Mutex::new(Some(client))),
+            reconnect: std::sync::Arc::new(std::sync::Mutex::new(None)),
+            refresh_lock: std::sync::Arc::new(tokio::sync::Mutex::new(())),
+            environment_connection_state_tx,
+        }
+    }
+
+    pub(crate) async fn replace_accepted_websocket(
+        &self,
+        websocket: WebSocket,
+    ) -> Result<(), ExecServerError> {
+        self.get()
+            .await?
+            .replace_accepted_websocket(websocket)
+            .await
+    }
+}

codex-rs/exec-server/src/client/accepted_tests.rs

@@ -0,0 +1,34 @@
+use tokio::sync::oneshot;
+
+use super::AcceptedConnectionSource;
+use crate::ExecServerClientConnectOptions;
+
+#[tokio::test]
+async fn replacement_claim_is_released_when_handoff_is_cancelled() {
+    let source = AcceptedConnectionSource::new(ExecServerClientConnectOptions::default());
+    let (claimed_tx, claimed_rx) = oneshot::channel();
+    let (_release_tx, release_rx) = oneshot::channel::<()>();
+
+    let handoff = tokio::spawn({
+        let source = source.clone();
+        async move {
+            let _submission = source
+                .begin_replacement()
+                .expect("the first replacement should claim the handoff");
+            claimed_tx
+                .send(())
+                .expect("the test should wait for the claim");
+            let _ = release_rx.await;
+        }
+    });
+
+    claimed_rx.await.expect("the handoff should be claimed");
+    assert!(source.begin_replacement().is_err());
+
+    handoff.abort();
+    let _ = handoff.await;
+
+    source
+        .begin_replacement()
+        .expect("a new replacement should be accepted after cancellation");
+}

codex-rs/exec-server/src/client/http_client.rs

@@ -0,0 +1,26 @@
+//! HTTP client capability implementations shared by local and remote environments.
+//!
+//! This module is the facade for the environment-owned [`crate::HttpClient`]
+//! capability:
+//! - [`RouteAwareHttpClient`] executes requests through the shared transport
+//! - [`ExecServerClient`] forwards requests over the JSON-RPC transport
+//! - [`HttpResponseBodyStream`] presents buffered local bodies and streamed
+//!   remote `http/request/bodyDelta` notifications through one byte-stream API
+//!
+//! Runtime split:
+//! - orchestrator process: holds an `Arc<dyn HttpClient>` and chooses local or
+//!   remote execution
+//! - remote runtime: serves the `http/request` RPC and runs the concrete local
+//!   HTTP request there when the orchestrator uses [`ExecServerClient`]
+
+#[path = "http_response_body_stream.rs"]
+pub(crate) mod response_body_stream;
+#[path = "route_aware_http_client.rs"]
+mod route_aware_http_client;
+#[path = "rpc_http_client.rs"]
+mod rpc_http_client;
+
+pub use response_body_stream::HttpResponseBodyStream;
+pub(crate) use route_aware_http_client::PendingRouteAwareHttpBodyStream;
+pub use route_aware_http_client::RouteAwareHttpClient;
+pub(crate) use route_aware_http_client::RouteAwareHttpRequestRunner;

codex-rs/exec-server/src/client/http_response_body_stream.rs

@@ -0,0 +1,446 @@
+//! Shared HTTP response-body stream plumbing for local and remote execution.
+//!
+//! This module owns the byte-stream type exposed by the `HttpClient`
+//! capability plus the remote-side routing table used to turn
+//! `http/request/bodyDelta` notifications back into per-request streams.
+
+use std::collections::HashMap;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::atomic::Ordering;
+
+use bytes::Bytes;
+use codex_http_client::HttpError;
+use codex_http_client::HttpResponse;
+use futures::StreamExt;
+use serde_json::Value;
+use serde_json::from_value;
+use tokio::runtime::Handle;
+use tokio::sync::OwnedSemaphorePermit;
+use tokio::sync::mpsc;
+use tokio::sync::mpsc::error::TrySendError;
+use tracing::debug;
+
+use crate::client::ExecServerError;
+use crate::client::Inner;
+use crate::protocol::HTTP_REQUEST_BODY_DELTA_METHOD;
+use crate::protocol::HttpRequestBodyDeltaNotification;
+use crate::protocol::MAX_HTTP_BODY_DELTA_BYTES;
+use crate::rpc::RpcNotificationSender;
+
+pub(crate) const MAX_QUEUED_HTTP_BODY_BYTES: usize = 16 * 1024 * 1024;
+const MAX_ENCODED_HTTP_BODY_DELTA_BYTES: usize = MAX_HTTP_BODY_DELTA_BYTES.div_ceil(3) * 4;
+
+pub(crate) struct QueuedHttpBodyDelta {
+    notification: HttpRequestBodyDeltaNotification,
+    _byte_permit: Option<OwnedSemaphorePermit>,
+}
+
+impl QueuedHttpBodyDelta {
+    pub(crate) fn new(
+        notification: HttpRequestBodyDeltaNotification,
+        byte_permit: Option<OwnedSemaphorePermit>,
+    ) -> Self {
+        Self {
+            notification,
+            _byte_permit: byte_permit,
+        }
+    }
+}
+
+pub(super) struct HttpBodyStreamRegistration {
+    inner: Arc<Inner>,
+    request_id: String,
+    active: bool,
+}
+
+enum HttpResponseBodyStreamInner {
+    Local {
+        body: Pin<Box<dyn futures::Stream<Item = Result<Bytes, HttpError>> + Send>>,
+    },
+    Remote {
+        inner: Arc<Inner>,
+        request_id: String,
+        next_seq: u64,
+        rx: mpsc::Receiver<QueuedHttpBodyDelta>,
+        pending_eof: bool,
+        closed: bool,
+    },
+}
+
+/// Request-scoped stream of body chunks for an HTTP response.
+///
+/// The initial `http/request` call returns status and headers. This stream then
+/// receives the ordered `http/request/bodyDelta` notifications for that request
+/// id until EOF or a terminal error.
+pub struct HttpResponseBodyStream {
+    inner: HttpResponseBodyStreamInner,
+}
+
+impl HttpResponseBodyStream {
+    /// Creates an in-memory response stream from pre-buffered chunks.
+    ///
+    /// This is useful for [`crate::HttpClient`] implementations that already
+    /// own the response bytes, including lightweight test clients.
+    #[doc(hidden)]
+    pub fn from_chunks(chunks: Vec<Vec<u8>>) -> Self {
+        let body = futures::stream::iter(
+            chunks
+                .into_iter()
+                .map(|chunk| Ok::<Bytes, HttpError>(chunk.into())),
+        );
+        Self {
+            inner: HttpResponseBodyStreamInner::Local {
+                body: Box::pin(body),
+            },
+        }
+    }
+
+    pub(super) fn local(response: HttpResponse) -> Self {
+        Self {
+            inner: HttpResponseBodyStreamInner::Local {
+                body: Box::pin(response.bytes_stream()),
+            },
+        }
+    }
+
+    pub(super) fn remote(
+        inner: Arc<Inner>,
+        request_id: String,
+        rx: mpsc::Receiver<QueuedHttpBodyDelta>,
+    ) -> Self {
+        Self {
+            inner: HttpResponseBodyStreamInner::Remote {
+                inner,
+                request_id,
+                next_seq: 1,
+                rx,
+                pending_eof: false,
+                closed: false,
+            },
+        }
+    }
+
+    /// Receives the next response-body chunk.
+    ///
+    /// Returns `Ok(None)` at EOF and converts sequence gaps or stream-side
+    /// stream errors into protocol errors.
+    pub async fn recv(&mut self) -> Result<Option<Vec<u8>>, ExecServerError> {
+        match &mut self.inner {
+            HttpResponseBodyStreamInner::Local { body } => match body.next().await {
+                Some(chunk) => match chunk {
+                    Ok(bytes) => Ok(Some(bytes.to_vec())),
+                    Err(error) => Err(ExecServerError::HttpRequest(error.to_string())),
+                },
+                None => Ok(None),
+            },
+            HttpResponseBodyStreamInner::Remote {
+                inner,
+                request_id,
+                next_seq,
+                rx,
+                pending_eof,
+                closed,
+            } => {
+                if *pending_eof {
+                    *pending_eof = false;
+                    finish_remote_stream(inner, request_id, closed).await;
+                    return Ok(None);
+                }
+
+                let Some(QueuedHttpBodyDelta {
+                    notification: delta,
+                    ..
+                }) = rx.recv().await
+                else {
+                    finish_remote_stream(inner, request_id, closed).await;
+                    if let Some(error) = inner.take_http_body_stream_failure(request_id).await {
+                        return Err(ExecServerError::Protocol(format!(
+                            "http response stream `{request_id}` failed: {error}",
+                        )));
+                    }
+                    return Ok(None);
+                };
+                if delta.seq != *next_seq {
+                    finish_remote_stream(inner, request_id, closed).await;
+                    return Err(ExecServerError::Protocol(format!(
+                        "http response stream `{request_id}` received seq {}, expected {}",
+                        delta.seq, *next_seq
+                    )));
+                }
+                *next_seq += 1;
+                let chunk = delta.delta.into_inner();
+
+                if let Some(error) = delta.error {
+                    finish_remote_stream(inner, request_id, closed).await;
+                    return Err(ExecServerError::Protocol(format!(
+                        "http response stream `{request_id}` failed: {error}",
+                    )));
+                }
+                if delta.done {
+                    finish_remote_stream(inner, request_id, closed).await;
+                    if chunk.is_empty() {
+                        return Ok(None);
+                    }
+                    *pending_eof = true;
+                }
+                Ok(Some(chunk))
+            }
+        }
+    }
+}
+
+impl Drop for HttpResponseBodyStream {
+    /// Schedules stream-route removal if the consumer drops before EOF.
+    fn drop(&mut self) {
+        if let HttpResponseBodyStreamInner::Remote {
+            inner,
+            request_id,
+            closed,
+            ..
+        } = &mut self.inner
+        {
+            if *closed {
+                return;
+            }
+            *closed = true;
+            spawn_remove_http_body_stream(Arc::clone(inner), request_id.clone());
+        }
+    }
+}
+
+impl HttpBodyStreamRegistration {
+    pub(super) fn new(inner: Arc<Inner>, request_id: String) -> Self {
+        Self {
+            inner,
+            request_id,
+            active: true,
+        }
+    }
+
+    pub(super) fn disarm(&mut self) {
+        self.active = false;
+    }
+}
+
+impl Drop for HttpBodyStreamRegistration {
+    /// Removes the route if the stream request future is cancelled before headers return.
+    fn drop(&mut self) {
+        if self.active {
+            spawn_remove_http_body_stream(Arc::clone(&self.inner), self.request_id.clone());
+        }
+    }
+}
+
+async fn finish_remote_stream(inner: &Arc<Inner>, request_id: &str, closed: &mut bool) {
+    if *closed {
+        return;
+    }
+    *closed = true;
+    inner.remove_http_body_stream(request_id).await;
+}
+
+/// Schedules HTTP body route removal from synchronous drop paths.
+fn spawn_remove_http_body_stream(inner: Arc<Inner>, request_id: String) {
+    if let Ok(handle) = Handle::try_current() {
+        handle.spawn(async move {
+            inner.remove_http_body_stream(&request_id).await;
+        });
+    }
+}
+
+pub(super) async fn send_body_delta(
+    notifications: &RpcNotificationSender,
+    delta: HttpRequestBodyDeltaNotification,
+) -> bool {
+    notifications
+        .notify(HTTP_REQUEST_BODY_DELTA_METHOD, &delta)
+        .await
+        .is_ok()
+}
+
+impl Inner {
+    /// Routes one streamed HTTP body notification into its request-local receiver.
+    pub(crate) async fn handle_http_body_delta_notification(
+        &self,
+        params: Option<Value>,
+    ) -> Result<(), ExecServerError> {
+        let params = params.unwrap_or(Value::Null);
+        if params
+            .get("deltaBase64")
+            .and_then(Value::as_str)
+            .is_some_and(|delta| delta.len() > MAX_ENCODED_HTTP_BODY_DELTA_BYTES)
+        {
+            return Err(ExecServerError::Protocol(format!(
+                "http response body delta exceeds {MAX_HTTP_BODY_DELTA_BYTES} bytes"
+            )));
+        }
+        let params: HttpRequestBodyDeltaNotification = from_value(params)?;
+        if params.delta.0.len() > MAX_HTTP_BODY_DELTA_BYTES {
+            return Err(ExecServerError::Protocol(format!(
+                "http response body delta exceeds {MAX_HTTP_BODY_DELTA_BYTES} bytes"
+            )));
+        }
+        // Unknown request ids are ignored intentionally: a stream may have already
+        // reached EOF and released its route.
+        if let Some(tx) = self
+            .http_body_streams
+            .load()
+            .get(&params.request_id)
+            .cloned()
+        {
+            let request_id = params.request_id.clone();
+            let terminal_delta = params.done || params.error.is_some();
+            let queued_bytes = params
+                .delta
+                .0
+                .len()
+                .saturating_add(params.error.as_deref().map_or(0, str::len));
+            let byte_permit = if queued_bytes == 0 {
+                None
+            } else {
+                u32::try_from(queued_bytes).ok().and_then(|queued_bytes| {
+                    Arc::clone(&self.http_body_stream_byte_budget)
+                        .try_acquire_many_owned(queued_bytes)
+                        .ok()
+                })
+            };
+            if queued_bytes > 0 && byte_permit.is_none() {
+                self.record_http_body_stream_failure(
+                    &request_id,
+                    format!("queued body deltas exceed {MAX_QUEUED_HTTP_BODY_BYTES} bytes"),
+                )
+                .await;
+                self.remove_http_body_stream(&request_id).await;
+                debug!(
+                    "closing http response stream `{request_id}` after exhausting the queued byte budget"
+                );
+                return Ok(());
+            }
+            match tx.try_send(QueuedHttpBodyDelta::new(params, byte_permit)) {
+                Ok(()) => {
+                    if terminal_delta {
+                        self.remove_http_body_stream(&request_id).await;
+                    }
+                }
+                Err(TrySendError::Closed(_)) => {
+                    self.remove_http_body_stream(&request_id).await;
+                    debug!("http response stream receiver dropped before body delta delivery");
+                }
+                Err(TrySendError::Full(_)) => {
+                    self.record_http_body_stream_failure(
+                        &request_id,
+                        "body delta channel filled before delivery".to_string(),
+                    )
+                    .await;
+                    self.remove_http_body_stream(&request_id).await;
+                    debug!(
+                        "closing http response stream `{request_id}` after body delta backpressure"
+                    );
+                }
+            }
+        }
+        Ok(())
+    }
+
+    /// Fails active streamed HTTP bodies so callers do not wait forever after a
+    /// transport disconnect or notification handling failure.
+    pub(crate) async fn fail_all_http_body_streams(&self, message: String) {
+        let _streams_write_guard = self.http_body_streams_write_lock.lock().await;
+        let streams = self.http_body_streams.load();
+        let streams = streams.as_ref().clone();
+        self.http_body_streams.store(Arc::new(HashMap::new()));
+        for (request_id, tx) in streams {
+            // Failure notifications must wake every stream even when no
+            // byte-budget permits remain.
+            if tx
+                .try_send(QueuedHttpBodyDelta::new(
+                    HttpRequestBodyDeltaNotification {
+                        request_id: request_id.clone(),
+                        seq: 1,
+                        delta: Vec::new().into(),
+                        done: true,
+                        error: Some(message.clone()),
+                    },
+                    /*byte_permit*/ None,
+                ))
+                .is_err()
+            {
+                let mut next_failures = self.http_body_stream_failures.load().as_ref().clone();
+                next_failures.insert(request_id, message.clone());
+                self.http_body_stream_failures
+                    .store(Arc::new(next_failures));
+            }
+        }
+    }
+
+    /// Allocates a connection-local streamed HTTP response id.
+    pub(super) fn next_http_body_stream_request_id(&self) -> String {
+        let id = self
+            .http_body_stream_next_id
+            .fetch_add(1, Ordering::Relaxed);
+        format!("http-{id}")
+    }
+
+    /// Registers a request id before issuing a streaming HTTP call.
+    pub(super) async fn insert_http_body_stream(
+        &self,
+        request_id: String,
+        tx: mpsc::Sender<QueuedHttpBodyDelta>,
+    ) -> Result<(), ExecServerError> {
+        let _streams_write_guard = self.http_body_streams_write_lock.lock().await;
+        let streams = self.http_body_streams.load();
+        if streams.contains_key(&request_id) {
+            return Err(ExecServerError::Protocol(format!(
+                "http response stream already registered for request {request_id}"
+            )));
+        }
+        let mut next_streams = streams.as_ref().clone();
+        next_streams.insert(request_id.clone(), tx);
+        self.http_body_streams.store(Arc::new(next_streams));
+        let failures = self.http_body_stream_failures.load();
+        if failures.contains_key(&request_id) {
+            let mut next_failures = failures.as_ref().clone();
+            next_failures.remove(&request_id);
+            self.http_body_stream_failures
+                .store(Arc::new(next_failures));
+        }
+        Ok(())
+    }
+
+    /// Removes a request id after EOF, terminal error, or request failure.
+    pub(super) async fn remove_http_body_stream(
+        &self,
+        request_id: &str,
+    ) -> Option<mpsc::Sender<QueuedHttpBodyDelta>> {
+        let _streams_write_guard = self.http_body_streams_write_lock.lock().await;
+        let streams = self.http_body_streams.load();
+        let stream = streams.get(request_id).cloned();
+        stream.as_ref()?;
+        let mut next_streams = streams.as_ref().clone();
+        next_streams.remove(request_id);
+        self.http_body_streams.store(Arc::new(next_streams));
+        stream
+    }
+
+    async fn record_http_body_stream_failure(&self, request_id: &str, message: String) {
+        let _streams_write_guard = self.http_body_streams_write_lock.lock().await;
+        let failures = self.http_body_stream_failures.load();
+        let mut next_failures = failures.as_ref().clone();
+        next_failures.insert(request_id.to_string(), message);
+        self.http_body_stream_failures
+            .store(Arc::new(next_failures));
+    }
+
+    async fn take_http_body_stream_failure(&self, request_id: &str) -> Option<String> {
+        let _streams_write_guard = self.http_body_streams_write_lock.lock().await;
+        let failures = self.http_body_stream_failures.load();
+        let error = failures.get(request_id).cloned();
+        error.as_ref()?;
+        let mut next_failures = failures.as_ref().clone();
+        next_failures.remove(request_id);
+        self.http_body_stream_failures
+            .store(Arc::new(next_failures));
+        error
+    }
+}

codex-rs/exec-server/src/client/network_policy_audit.rs

@@ -0,0 +1,81 @@
+use super::NetworkPolicyAuditContext;
+use crate::protocol::ExecServerNetworkProtocol;
+use crate::protocol::MAX_NETWORK_POLICY_HOST_BYTES;
+use crate::protocol::MAX_NETWORK_POLICY_PROCESS_ID_BYTES;
+use crate::protocol::MAX_NETWORK_POLICY_REASON_BYTES;
+use crate::protocol::NetworkPolicyDecisionNotification;
+
+const MAX_NETWORK_POLICY_METHOD_BYTES: usize = 32;
+const MAX_NETWORK_POLICY_CLIENT_BYTES: usize = 256;
+const MAX_NETWORK_POLICY_TIMESTAMP_BYTES: usize = 64;
+
+pub(super) fn emit_network_policy_decision(
+    context: &NetworkPolicyAuditContext,
+    decision: &NetworkPolicyDecisionNotification,
+) -> bool {
+    if decision.process_id.is_empty()
+        || decision.process_id.len() > MAX_NETWORK_POLICY_PROCESS_ID_BYTES
+        || decision.host.is_empty()
+        || decision.host.len() > MAX_NETWORK_POLICY_HOST_BYTES
+        || decision.host.chars().any(char::is_control)
+        || decision.host.chars().any(char::is_whitespace)
+        || decision.reason.len() > MAX_NETWORK_POLICY_REASON_BYTES
+        || decision.reason.chars().any(char::is_control)
+        || !matches!(decision.scope.as_str(), "domain" | "non_domain")
+        || !matches!(decision.decision.as_str(), "allow" | "deny" | "ask")
+        || !matches!(
+            decision.source.as_str(),
+            "baseline_policy" | "mode_guard" | "proxy_state" | "decider"
+        )
+        || decision.timestamp.is_empty()
+        || decision.timestamp.len() > MAX_NETWORK_POLICY_TIMESTAMP_BYTES
+        || decision.timestamp.chars().any(char::is_control)
+        || decision.method.as_ref().is_some_and(|method| {
+            method.len() > MAX_NETWORK_POLICY_METHOD_BYTES
+                || method.chars().any(char::is_control)
+                || method.chars().any(char::is_whitespace)
+        })
+        || decision.client.as_ref().is_some_and(|client| {
+            client.len() > MAX_NETWORK_POLICY_CLIENT_BYTES
+                || client.chars().any(char::is_control)
+                || client.chars().any(char::is_whitespace)
+        })
+    {
+        return false;
+    }
+
+    let protocol = match decision.protocol {
+        ExecServerNetworkProtocol::Http => "http",
+        ExecServerNetworkProtocol::HttpsConnect => "https_connect",
+        ExecServerNetworkProtocol::Socks5Tcp => "socks5_tcp",
+        ExecServerNetworkProtocol::Socks5Udp => "socks5_udp",
+    };
+    let metadata = &context.metadata;
+    tracing::event!(
+        target: "codex_otel.log_only",
+        tracing::Level::INFO,
+        event.name = "codex.network_proxy.policy_decision",
+        event.timestamp = decision.timestamp,
+        conversation.id = metadata.conversation_id.as_deref(),
+        app.version = metadata.app_version.as_deref(),
+        auth_mode = metadata.auth_mode.as_deref(),
+        originator = metadata.originator.as_deref(),
+        user.account_id = metadata.user_account_id.as_deref(),
+        user.email = metadata.user_email.as_deref(),
+        terminal.type = metadata.terminal_type.as_deref(),
+        model = metadata.model.as_deref(),
+        slug = metadata.slug.as_deref(),
+        network.policy.scope = decision.scope,
+        network.policy.decision = decision.decision,
+        network.policy.source = decision.source,
+        network.policy.reason = decision.reason,
+        network.transport.protocol = protocol,
+        server.address = decision.host,
+        server.port = decision.port,
+        http.request.method = decision.method.as_deref().unwrap_or("none"),
+        client.address = decision.client.as_deref().unwrap_or("unknown"),
+        execution.id = context.execution_id.as_deref(),
+        network.policy.override = decision.policy_override,
+    );
+    true
+}

codex-rs/exec-server/src/client/route_aware_http_client.rs

@@ -0,0 +1,378 @@
+//! Route-aware local HTTP capability implementation.
+//!
+//! This code runs wherever the real network request should originate:
+//! - in a local environment, that means the orchestrator process
+//! - in a remote environment, that means the remote runtime after the
+//!   orchestrator has forwarded `http/request` over JSON-RPC
+
+use std::time::Duration;
+
+use codex_exec_server_protocol::JSONRPCErrorError;
+use codex_http_client::ClientRouteClass;
+use codex_http_client::HttpClientFactory;
+use codex_http_client::RouteAwareClientPool;
+use codex_http_client::RouteAwareRequestError;
+use codex_protocol::shell_environment::CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR;
+use codex_protocol::shell_environment::OPENAI_FEDERATION_RULE_ID_ENV_VAR;
+use codex_protocol::shell_environment::OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR;
+use codex_protocol::shell_environment::OPENAI_WORKLOAD_IDENTITY_CONTEXT_ENV_VAR;
+use futures::FutureExt;
+use futures::StreamExt;
+use futures::future::BoxFuture;
+use http::HeaderMap;
+use http::HeaderName;
+use http::HeaderValue;
+use http::Method;
+use tracing::Instrument;
+use url::Url;
+
+use super::HttpResponseBodyStream;
+use super::response_body_stream::send_body_delta;
+use crate::HttpClient;
+use crate::client::ExecServerError;
+use crate::protocol::HttpHeader;
+use crate::protocol::HttpRedirectPolicy;
+use crate::protocol::HttpRequestBodyDeltaNotification;
+use crate::protocol::HttpRequestParams;
+use crate::protocol::HttpRequestResponse;
+use crate::protocol::MAX_HTTP_BODY_DELTA_BYTES;
+use crate::rpc::RpcNotificationSender;
+use crate::rpc::internal_error;
+use crate::rpc::invalid_params;
+
+const HTTP_HEADER_ENV_DENYLIST: &[&str] = &[
+    CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR,
+    OPENAI_FEDERATION_RULE_ID_ENV_VAR,
+    OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR,
+    OPENAI_WORKLOAD_IDENTITY_CONTEXT_ENV_VAR,
+    "OPENAI_API_KEY",
+    "CODEX_API_KEY",
+    "CODEX_ACCESS_TOKEN",
+    "CODEX_CONNECTORS_TOKEN",
+    "AWS_ACCESS_KEY_ID",
+    "AWS_SECRET_ACCESS_KEY",
+    "AWS_SESSION_TOKEN",
+    "AZURE_CLIENT_SECRET",
+    "AZURE_FEDERATED_TOKEN_FILE",
+    "GOOGLE_APPLICATION_CREDENTIALS",
+];
+
+/// HTTP capability implementation backed by the shared route-aware transport.
+#[derive(Clone)]
+pub struct RouteAwareHttpClient {
+    follow_redirects: RouteAwareClientPool,
+    stop_redirects: RouteAwareClientPool,
+}
+
+/// Streaming response state held between the initial HTTP response and
+/// downstream body-delta forwarding.
+pub(crate) struct PendingRouteAwareHttpBodyStream {
+    pub(crate) request_id: String,
+    pub(crate) response: codex_http_client::HttpResponse,
+}
+
+/// Validates `http/request` parameters and runs the actual HTTP call used
+/// by the exec-server route and the local [`HttpClient`] backend.
+pub(crate) struct RouteAwareHttpRequestRunner {
+    client: RouteAwareClientPool,
+}
+
+impl RouteAwareHttpClient {
+    pub fn new(http_client_factory: HttpClientFactory) -> Self {
+        Self {
+            follow_redirects: RouteAwareClientPool::with_chatgpt_cloudflare_cookies_without_request_logging(
+                http_client_factory.clone(),
+                // Delegated HTTP targets arbitrary endpoints; route class only labels diagnostics.
+                ClientRouteClass::Other,
+            ),
+            stop_redirects:
+                RouteAwareClientPool::with_chatgpt_cloudflare_cookies_without_redirects_or_request_logging(
+                    http_client_factory,
+                    // Proxy routing comes from the factory, not this diagnostic-only route class.
+                    ClientRouteClass::Other,
+                ),
+        }
+    }
+
+    /// Enables narrowly scoped TLS-backend fallback for both redirect policies.
+    pub fn with_tls_backend_fallback(mut self) -> Self {
+        self.follow_redirects = self.follow_redirects.with_tls_backend_fallback();
+        self.stop_redirects = self.stop_redirects.with_tls_backend_fallback();
+        self
+    }
+
+    pub(crate) fn runner(
+        &self,
+        redirect_policy: HttpRedirectPolicy,
+    ) -> RouteAwareHttpRequestRunner {
+        let client = match redirect_policy {
+            HttpRedirectPolicy::Follow => self.follow_redirects.clone(),
+            HttpRedirectPolicy::Stop => self.stop_redirects.clone(),
+        };
+        RouteAwareHttpRequestRunner { client }
+    }
+}
+
+impl HttpClient for RouteAwareHttpClient {
+    fn http_request(
+        &self,
+        params: HttpRequestParams,
+    ) -> BoxFuture<'_, Result<HttpRequestResponse, ExecServerError>> {
+        async move {
+            let runner = self.runner(params.redirect_policy);
+            let (response, _) = runner
+                .run(HttpRequestParams {
+                    stream_response: false,
+                    ..params
+                })
+                .await
+                .map_err(|error| ExecServerError::HttpRequest(error.message))?;
+            Ok(response)
+        }
+        .boxed()
+    }
+
+    fn http_request_stream(
+        &self,
+        params: HttpRequestParams,
+    ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> {
+        async move {
+            let runner = self.runner(params.redirect_policy);
+            let (response, pending_stream) = runner
+                .run(HttpRequestParams {
+                    stream_response: true,
+                    ..params
+                })
+                .await
+                .map_err(|error| ExecServerError::HttpRequest(error.message))?;
+            let pending_stream = pending_stream.ok_or_else(|| {
+                ExecServerError::Protocol(
+                    "http request stream did not return a response body stream".to_string(),
+                )
+            })?;
+            Ok((
+                response,
+                HttpResponseBodyStream::local(pending_stream.response),
+            ))
+        }
+        .boxed()
+    }
+}
+
+impl RouteAwareHttpRequestRunner {
+    pub(crate) async fn run(
+        &self,
+        params: HttpRequestParams,
+    ) -> Result<(HttpRequestResponse, Option<PendingRouteAwareHttpBodyStream>), JSONRPCErrorError>
+    {
+        let method = Method::from_bytes(params.method.as_bytes())
+            .map_err(|error| invalid_params(format!("http/request method is invalid: {error}")))?;
+        let url = Url::parse(&params.url)
+            .map_err(|error| invalid_params(format!("http/request url is invalid: {error}")))?;
+        match url.scheme() {
+            "http" | "https" => {}
+            scheme => {
+                return Err(invalid_params(format!(
+                    "http/request only supports http and https URLs, got {scheme}"
+                )));
+            }
+        }
+
+        let request_span = tracing::info_span!(
+            "codex.exec_server.http_request",
+            otel.kind = "client",
+            http.request.method = method.as_str(),
+            server.address = url.host_str().unwrap_or_default(),
+            server.port = u64::from(url.port_or_known_default().unwrap_or_default()),
+            http.response.status_code = tracing::field::Empty,
+            error.type = tracing::field::Empty,
+        );
+        let mut headers = Self::build_headers(params.headers)?;
+        codex_otel::inject_span_w3c_trace_headers(&request_span, &mut headers);
+        let mut request = self.client.request(method.clone(), url).headers(headers);
+        if let Some(body) = params.body {
+            request = request.body(body.into_inner());
+        }
+        if let Some(timeout_ms) = params.timeout_ms {
+            request = request.timeout(Duration::from_millis(timeout_ms));
+        }
+
+        let response = match request.send().instrument(request_span.clone()).await {
+            Ok(response) => response,
+            Err(error) => {
+                request_span.record("error.type", "request");
+                let error_message = error.to_string();
+                log_send_error(&method, error);
+                return Err(internal_error(format!(
+                    "http/request failed: {error_message}"
+                )));
+            }
+        };
+        let status = response.status().as_u16();
+        request_span.record("http.response.status_code", u64::from(status));
+        let headers = Self::response_headers(response.headers());
+
+        if params.stream_response {
+            return Ok((
+                HttpRequestResponse {
+                    status,
+                    headers,
+                    body: Vec::new().into(),
+                },
+                Some(PendingRouteAwareHttpBodyStream {
+                    request_id: params.request_id,
+                    response,
+                }),
+            ));
+        }
+
+        let body = response.bytes().await.map_err(|error| {
+            internal_error(format!(
+                "failed to read http/request response body: {error}"
+            ))
+        })?;
+
+        Ok((
+            HttpRequestResponse {
+                status,
+                headers,
+                body: body.to_vec().into(),
+            },
+            None,
+        ))
+    }
+
+    pub(crate) async fn stream_body(
+        pending_stream: PendingRouteAwareHttpBodyStream,
+        notifications: RpcNotificationSender,
+    ) {
+        let PendingRouteAwareHttpBodyStream {
+            request_id,
+            response,
+        } = pending_stream;
+        let mut seq = 1;
+        let mut body = response.bytes_stream();
+        while let Some(chunk) = body.next().await {
+            match chunk {
+                Ok(bytes) => {
+                    for chunk in bytes.chunks(MAX_HTTP_BODY_DELTA_BYTES) {
+                        if !send_body_delta(
+                            &notifications,
+                            HttpRequestBodyDeltaNotification {
+                                request_id: request_id.clone(),
+                                seq,
+                                delta: chunk.to_vec().into(),
+                                done: false,
+                                error: None,
+                            },
+                        )
+                        .await
+                        {
+                            return;
+                        }
+                        seq += 1;
+                    }
+                }
+                Err(error) => {
+                    let _ = send_body_delta(
+                        &notifications,
+                        HttpRequestBodyDeltaNotification {
+                            request_id,
+                            seq,
+                            delta: Vec::new().into(),
+                            done: true,
+                            error: Some(error.to_string()),
+                        },
+                    )
+                    .await;
+                    return;
+                }
+            }
+        }
+
+        let _ = send_body_delta(
+            &notifications,
+            HttpRequestBodyDeltaNotification {
+                request_id,
+                seq,
+                delta: Vec::new().into(),
+                done: true,
+                error: None,
+            },
+        )
+        .await;
+    }
+
+    fn build_headers(headers: Vec<HttpHeader>) -> Result<HeaderMap, JSONRPCErrorError> {
+        let mut header_map = HeaderMap::new();
+        for header in headers {
+            let name = HeaderName::from_bytes(header.name.as_bytes()).map_err(|error| {
+                invalid_params(format!("http/request header name is invalid: {error}"))
+            })?;
+            let value = match header.value_env_var {
+                Some(env_var) => {
+                    if HTTP_HEADER_ENV_DENYLIST
+                        .iter()
+                        .any(|denied| denied.eq_ignore_ascii_case(&env_var))
+                    {
+                        return Err(invalid_params(format!(
+                            "http/request header {} cannot use executor environment variable {env_var}",
+                            header.name
+                        )));
+                    }
+                    let env_value = std::env::var(&env_var).map_err(|_| {
+                        invalid_params(format!(
+                            "http/request header {} requires executor environment variable {env_var}",
+                            header.name
+                        ))
+                    })?;
+                    if env_value.is_empty() {
+                        return Err(invalid_params(format!(
+                            "http/request header {} requires a non-empty executor environment variable {env_var}",
+                            header.name
+                        )));
+                    }
+                    format!("{}{env_value}", header.value)
+                }
+                None => header.value,
+            };
+            let value = HeaderValue::from_str(&value).map_err(|error| {
+                invalid_params(format!(
+                    "http/request header value is invalid for {}: {error}",
+                    header.name
+                ))
+            })?;
+            header_map.append(name, value);
+        }
+        Ok(header_map)
+    }
+
+    fn response_headers(headers: &HeaderMap) -> Vec<HttpHeader> {
+        headers
+            .iter()
+            .filter_map(|(name, value)| {
+                Some(HttpHeader {
+                    name: name.as_str().to_string(),
+                    value: value.to_str().ok()?.to_string(),
+                    value_env_var: None,
+                })
+            })
+            .collect()
+    }
+}
+
+fn log_send_error(method: &Method, error: RouteAwareRequestError) {
+    let error_is_timeout = error.is_timeout();
+    let error_is_connect = error.is_connect();
+    let error = match error {
+        RouteAwareRequestError::Request(error) => error.without_url().to_string(),
+        error => error.to_string(),
+    };
+    tracing::warn!(
+        http_method = method.as_str(),
+        error_is_timeout,
+        error_is_connect,
+        error = %error,
+        "http/request send failed"
+    );
+}

codex-rs/exec-server/src/client/rpc_http_client.rs

@@ -0,0 +1,92 @@
+//! JSON-RPC-backed `HttpClient` implementation.
+//!
+//! This code runs in the orchestrator process. It does not issue network
+//! requests directly; instead it forwards `http/request` to the remote runtime
+//! and then reconstructs streamed bodies from `http/request/bodyDelta`
+//! notifications on the shared connection.
+
+use std::sync::Arc;
+
+use futures::FutureExt;
+use futures::future::BoxFuture;
+use tokio::sync::mpsc;
+
+use super::HttpResponseBodyStream;
+use super::response_body_stream::HttpBodyStreamRegistration;
+use crate::HttpClient;
+use crate::client::ExecServerClient;
+use crate::client::ExecServerError;
+use crate::protocol::HTTP_REQUEST_METHOD;
+use crate::protocol::HttpRequestParams;
+use crate::protocol::HttpRequestResponse;
+
+/// Maximum queued body frames per streamed HTTP response.
+const HTTP_BODY_DELTA_CHANNEL_CAPACITY: usize = 256;
+
+impl ExecServerClient {
+    /// Performs an HTTP request and buffers the response body.
+    pub async fn http_request(
+        &self,
+        mut params: HttpRequestParams,
+    ) -> Result<HttpRequestResponse, ExecServerError> {
+        params.stream_response = false;
+        self.call(HTTP_REQUEST_METHOD, &params).await
+    }
+
+    /// Performs an HTTP request and returns a body stream.
+    ///
+    /// The method sets `stream_response` and replaces any caller-supplied
+    /// `request_id` with a connection-local id, so late deltas from abandoned
+    /// streams cannot be confused with later requests.
+    pub async fn http_request_stream(
+        &self,
+        mut params: HttpRequestParams,
+    ) -> Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError> {
+        let rpc_client = self.rpc_client().await?;
+        params.stream_response = true;
+        let request_id = self.inner.next_http_body_stream_request_id();
+        params.request_id = request_id.clone();
+        let (tx, rx) = mpsc::channel(HTTP_BODY_DELTA_CHANNEL_CAPACITY);
+        self.inner
+            .insert_http_body_stream(request_id.clone(), tx)
+            .await?;
+        let mut registration =
+            HttpBodyStreamRegistration::new(Arc::clone(&self.inner), request_id.clone());
+        let response = match self
+            .call_rpc(&rpc_client, HTTP_REQUEST_METHOD, &params)
+            .await
+        {
+            Ok(response) => response,
+            Err(error) => {
+                self.inner.remove_http_body_stream(&request_id).await;
+                registration.disarm();
+                return Err(error);
+            }
+        };
+        registration.disarm();
+        Ok((
+            response,
+            HttpResponseBodyStream::remote(Arc::clone(&self.inner), request_id, rx),
+        ))
+    }
+}
+
+impl HttpClient for ExecServerClient {
+    /// Orchestrator-side adapter that forwards buffered HTTP requests to the
+    /// remote runtime over the shared JSON-RPC connection.
+    fn http_request(
+        &self,
+        params: HttpRequestParams,
+    ) -> BoxFuture<'_, Result<HttpRequestResponse, ExecServerError>> {
+        async move { ExecServerClient::http_request(self, params).await }.boxed()
+    }
+
+    /// Orchestrator-side adapter that forwards streamed HTTP requests to the
+    /// remote runtime and exposes body deltas as a byte stream.
+    fn http_request_stream(
+        &self,
+        params: HttpRequestParams,
+    ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> {
+        async move { ExecServerClient::http_request_stream(self, params).await }.boxed()
+    }
+}

codex-rs/exec-server/src/client/tests/network_policy_tests.rs

@@ -0,0 +1,544 @@
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::time::Duration;
+
+use codex_exec_server_protocol::JSONRPCMessage;
+use codex_exec_server_protocol::JSONRPCNotification;
+use codex_exec_server_protocol::JSONRPCRequest;
+use codex_exec_server_protocol::JSONRPCResponse;
+use codex_exec_server_protocol::RequestId;
+use codex_http_client::HttpClientFactory;
+use codex_http_client::OutboundProxyPolicy;
+use codex_network_proxy::NetworkDecision;
+use codex_network_proxy::NetworkPolicyDecider;
+use codex_network_proxy::NetworkPolicyRequest;
+use codex_network_proxy::NetworkProxyAuditMetadata;
+use codex_utils_path_uri::PathUri;
+use http::HeaderMap;
+use opentelemetry::trace::TracerProvider as _;
+use opentelemetry_sdk::trace::InMemorySpanExporter;
+use opentelemetry_sdk::trace::SdkTracerProvider;
+use pretty_assertions::assert_eq;
+use tokio::net::TcpListener;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio::time::timeout;
+use tracing::instrument::WithSubscriber;
+use tracing_subscriber::filter::filter_fn;
+use tracing_subscriber::prelude::*;
+
+use super::super::LazyRemoteExecServerClient;
+use super::super::NetworkPolicyAuditContext;
+use super::super::NetworkPolicyDecisionController;
+use super::super::SessionState;
+use super::super::handle_server_notification;
+use super::accept_websocket;
+use super::complete_websocket_initialize;
+use super::read_jsonrpc_websocket;
+use super::write_jsonrpc_websocket;
+use crate::ProcessId;
+use crate::client_api::ExecServerTransportParams;
+use crate::protocol::EXEC_METHOD;
+use crate::protocol::EXEC_TERMINATE_METHOD;
+use crate::protocol::ExecParams;
+use crate::protocol::ExecServerNetworkPolicyDecision;
+use crate::protocol::ExecServerNetworkPolicyRequest;
+use crate::protocol::ExecServerNetworkProtocol;
+use crate::protocol::NETWORK_POLICY_DECISION_METHOD;
+use crate::protocol::NETWORK_POLICY_REQUEST_METHOD;
+use crate::protocol::NetworkPolicyDecisionNotification;
+use crate::protocol::NetworkPolicyRequestParams;
+use crate::protocol::NetworkPolicyRequestResponse;
+use crate::rpc_server_requests::MAX_IN_FLIGHT_SERVER_CALLS;
+
+struct PendingDecisionGuard(mpsc::UnboundedSender<()>);
+
+impl Drop for PendingDecisionGuard {
+    fn drop(&mut self) {
+        let _ = self.0.send(());
+    }
+}
+
+fn policy_request(request_id: i64, process_id: ProcessId, host: &str) -> JSONRPCMessage {
+    JSONRPCMessage::Request(JSONRPCRequest {
+        id: RequestId::Integer(request_id),
+        method: NETWORK_POLICY_REQUEST_METHOD.to_string(),
+        params: Some(
+            serde_json::to_value(NetworkPolicyRequestParams {
+                process_id,
+                request: ExecServerNetworkPolicyRequest {
+                    protocol: ExecServerNetworkProtocol::HttpsConnect,
+                    host: host.to_string(),
+                    port: 443,
+                },
+            })
+            .expect("policy request should serialize"),
+        ),
+        trace: None,
+    })
+}
+
+async fn read_decision(
+    websocket: &mut tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>,
+    request_id: i64,
+) -> ExecServerNetworkPolicyDecision {
+    let JSONRPCMessage::Response(response) = read_jsonrpc_websocket(websocket).await else {
+        panic!("expected network policy response");
+    };
+    assert_eq!(response.id, RequestId::Integer(request_id));
+    serde_json::from_value::<NetworkPolicyRequestResponse>(response.result)
+        .expect("policy response should deserialize")
+        .decision
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn policy_decisions_reject_forged_process_and_use_trusted_controller_metadata() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("listener should bind");
+    let websocket_url = format!("ws://{}", listener.local_addr().expect("listener address"));
+    let (release_tx, release_rx) = oneshot::channel();
+    let (initialized_tx, initialized_rx) = oneshot::channel();
+    let server = tokio::spawn(async move {
+        let mut websocket = accept_websocket(&listener).await;
+        complete_websocket_initialize(
+            &mut websocket,
+            "audit-session",
+            /*expected_resume_session_id*/ None,
+        )
+        .await;
+        initialized_tx
+            .send(())
+            .expect("client should await completed WebSocket initialization");
+        release_rx.await.expect("server should be released");
+    });
+
+    let logs = Arc::new(Mutex::new(Vec::new()));
+    let writer_logs = Arc::clone(&logs);
+    let subscriber = tracing_subscriber::registry().with(
+        tracing_subscriber::fmt::layer()
+            .with_ansi(false)
+            .with_writer(move || AuditLogWriter(Arc::clone(&writer_logs))),
+    );
+    async move {
+        let client = LazyRemoteExecServerClient::new(
+            ExecServerTransportParams::websocket_url(websocket_url, Duration::from_secs(1)),
+            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+        )
+        .get()
+        .await
+        .expect("client should connect");
+        initialized_rx
+            .await
+            .expect("server should complete WebSocket initialization");
+        let mut state = SessionState::new(/*recoverable*/ true);
+        state.network_policy.audit = Some(NetworkPolicyAuditContext {
+            metadata: NetworkProxyAuditMetadata {
+                conversation_id: Some("trusted-conversation".to_string()),
+                user_account_id: Some("trusted-account".to_string()),
+                ..NetworkProxyAuditMetadata::default()
+            },
+            execution_id: Some("trusted-execution".to_string()),
+        });
+        client
+            .inner
+            .insert_session(&ProcessId::from("trusted-process"), Arc::new(state))
+            .expect("trusted process should register");
+        for (process_id, host) in [
+            ("forged-process", "forged.example"),
+            ("trusted-process", "trusted.example"),
+        ] {
+            handle_server_notification(
+                &client.inner,
+                JSONRPCNotification {
+                    method: NETWORK_POLICY_DECISION_METHOD.to_string(),
+                    params: Some(
+                        serde_json::to_value(NetworkPolicyDecisionNotification {
+                            process_id: ProcessId::from(process_id),
+                            timestamp: "2026-08-11T12:00:00.000Z".to_string(),
+                            scope: "domain".to_string(),
+                            decision: "deny".to_string(),
+                            source: "baseline_policy".to_string(),
+                            reason: "not_allowed".to_string(),
+                            protocol: ExecServerNetworkProtocol::HttpsConnect,
+                            host: host.to_string(),
+                            port: 443,
+                            method: None,
+                            client: None,
+                            policy_override: false,
+                        })
+                        .expect("network policy decision should serialize"),
+                    ),
+                },
+            )
+            .await
+            .expect("controller should handle network policy notification");
+        }
+        let output = String::from_utf8(
+            logs.lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .clone(),
+        )
+        .expect("audit log should be UTF-8");
+        assert!(!output.contains("forged.example"));
+        for expected in [
+            "codex_otel.log_only",
+            "trusted-conversation",
+            "trusted-account",
+            "trusted-execution",
+        ] {
+            assert!(
+                output.contains(expected),
+                "missing `{expected}` in {output}"
+            );
+        }
+        release_tx.send(()).expect("server should be released");
+    }
+    .with_subscriber(subscriber)
+    .await;
+    server.await.expect("server should finish");
+}
+
+struct AuditLogWriter(Arc<Mutex<Vec<u8>>>);
+
+impl std::io::Write for AuditLogWriter {
+    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
+        self.0
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .extend_from_slice(bytes);
+        Ok(bytes.len())
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
+}
+
+#[tokio::test]
+async fn abandoned_process_start_unregisters_and_cleans_up() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("listener should bind");
+    let websocket_url = format!("ws://{}", listener.local_addr().expect("listener address"));
+    let (start_seen_tx, start_seen_rx) = oneshot::channel();
+    let (finish_start_tx, finish_start_rx) = oneshot::channel();
+    let server = tokio::spawn(async move {
+        let mut websocket = accept_websocket(&listener).await;
+        complete_websocket_initialize(&mut websocket, "p", Default::default()).await;
+        let JSONRPCMessage::Request(start) = read_jsonrpc_websocket(&mut websocket).await else {
+            panic!("expected process start request");
+        };
+        assert_eq!(start.method, EXEC_METHOD);
+        start_seen_tx.send(()).expect("start should be observed");
+        finish_start_rx.await.expect("start should be released");
+        write_jsonrpc_websocket(
+            &mut websocket,
+            JSONRPCMessage::Response(JSONRPCResponse {
+                id: start.id,
+                result: serde_json::json!({"processId": "pending-start"}),
+            }),
+        )
+        .await;
+        let JSONRPCMessage::Request(terminate) = read_jsonrpc_websocket(&mut websocket).await
+        else {
+            panic!("expected process terminate request");
+        };
+        assert_eq!(terminate.method, EXEC_TERMINATE_METHOD);
+        write_jsonrpc_websocket(
+            &mut websocket,
+            JSONRPCMessage::Response(JSONRPCResponse {
+                id: terminate.id,
+                result: serde_json::json!({"running": true}),
+            }),
+        )
+        .await;
+    });
+
+    let client = LazyRemoteExecServerClient::new(
+        ExecServerTransportParams::websocket_url(websocket_url, Duration::from_secs(1)),
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+    )
+    .get()
+    .await
+    .expect("client should connect");
+    let process_id = ProcessId::from("pending-start");
+    let start_client = client.clone();
+    let start_process_id = process_id.clone();
+    let start = tokio::spawn(async move {
+        let params = ExecParams {
+            metadata: Default::default(),
+            process_id: start_process_id,
+            argv: vec!["true".to_string()],
+            cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd"))
+                .expect("cwd URI"),
+            shell_snapshot: None,
+            env_policy: None,
+            env: Default::default(),
+            tty: false,
+            pipe_stdin: false,
+            arg0: None,
+            sandbox: None,
+            enforce_managed_network: false,
+            managed_network: None,
+            network_proxy: None,
+        };
+        start_client
+            .start_process(params, /*network_policy_decider*/ None)
+            .await
+    });
+    start_seen_rx.await.expect("start should be observed");
+    let state = client
+        .inner
+        .get_session(&process_id)
+        .expect("pending process should be registered");
+    let decider: Arc<dyn NetworkPolicyDecider> =
+        Arc::new(|_request: NetworkPolicyRequest| async { NetworkDecision::Allow });
+    let decider_weak = Arc::downgrade(&decider);
+    state
+        .network_policy
+        .controller
+        .store(Some(Arc::new(NetworkPolicyDecisionController {
+            decider,
+            timeout: Duration::from_secs(30),
+        })));
+
+    start.abort();
+    assert!(start.await.is_err_and(|error| error.is_cancelled()));
+    assert!(state.network_policy.cancelled.is_cancelled());
+    assert!(client.inner.get_session(&process_id).is_none());
+    assert!(decider_weak.upgrade().is_none());
+
+    finish_start_tx.send(()).expect("start should be released");
+    server.await.expect("server task should finish");
+}
+
+#[tokio::test]
+async fn policy_requests_use_process_decider_and_cancel_on_unregister() {
+    let span_exporter = InMemorySpanExporter::default();
+    let tracer_provider = SdkTracerProvider::builder()
+        .with_simple_exporter(span_exporter.clone())
+        .build();
+    let subscriber = tracing_subscriber::registry().with(
+        tracing_opentelemetry::layer()
+            .with_tracer(tracer_provider.tracer("exec-server-test"))
+            .with_filter(filter_fn(codex_otel::OtelProvider::trace_export_filter)),
+    );
+    let _subscriber = tracing::subscriber::set_default(subscriber);
+    tracing::callsite::rebuild_interest_cache();
+
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("listener should bind");
+    let websocket_url = format!(
+        "ws://{}",
+        listener.local_addr().expect("listener should have address")
+    );
+    let process_id = ProcessId::from("policy-process");
+    let server_process_id = process_id.clone();
+    let (ready_tx, ready_rx) = oneshot::channel();
+    let (overflow_checked_tx, overflow_checked_rx) = oneshot::channel();
+    let (unregistered_tx, unregistered_rx) = oneshot::channel();
+    let server = tokio::spawn(async move {
+        let mut websocket = accept_websocket(&listener).await;
+        complete_websocket_initialize(
+            &mut websocket,
+            "policy-session",
+            /*expected_resume_session_id*/ None,
+        )
+        .await;
+        ready_rx.await.expect("process should be registered");
+
+        for (request_id, host, expected) in [
+            (0, "allowed.example", ExecServerNetworkPolicyDecision::Allow),
+            (
+                2,
+                "denied.example",
+                ExecServerNetworkPolicyDecision::Deny {
+                    reason: "blocked".to_string(),
+                },
+            ),
+            (
+                3,
+                "invalid host",
+                ExecServerNetworkPolicyDecision::Deny {
+                    reason: "not_allowed".to_string(),
+                },
+            ),
+        ] {
+            write_jsonrpc_websocket(
+                &mut websocket,
+                policy_request(request_id, server_process_id.clone(), host),
+            )
+            .await;
+            assert_eq!(read_decision(&mut websocket, request_id).await, expected);
+        }
+
+        let first_pending_request_id = 100;
+        for offset in 0..MAX_IN_FLIGHT_SERVER_CALLS {
+            write_jsonrpc_websocket(
+                &mut websocket,
+                policy_request(
+                    first_pending_request_id + offset as i64,
+                    server_process_id.clone(),
+                    "pending.example",
+                ),
+            )
+            .await;
+        }
+        let overflow_request_id = first_pending_request_id + MAX_IN_FLIGHT_SERVER_CALLS as i64;
+        write_jsonrpc_websocket(
+            &mut websocket,
+            policy_request(
+                overflow_request_id,
+                server_process_id.clone(),
+                "pending.example",
+            ),
+        )
+        .await;
+        assert_eq!(
+            read_decision(&mut websocket, overflow_request_id).await,
+            ExecServerNetworkPolicyDecision::Deny {
+                reason: "not_allowed".to_string(),
+            }
+        );
+        overflow_checked_tx.send(()).expect("overflow observed");
+
+        unregistered_rx
+            .await
+            .expect("process should be unregistered");
+
+        let post_unregister_request_id = 900;
+        write_jsonrpc_websocket(
+            &mut websocket,
+            policy_request(
+                post_unregister_request_id,
+                server_process_id,
+                "allowed.example",
+            ),
+        )
+        .await;
+        assert_eq!(
+            read_decision(&mut websocket, post_unregister_request_id).await,
+            ExecServerNetworkPolicyDecision::Deny {
+                reason: "not_allowed".to_string(),
+            }
+        );
+    });
+
+    let client = LazyRemoteExecServerClient::new(
+        ExecServerTransportParams::WebSocketUrl {
+            websocket_url,
+            connect_timeout: Duration::from_secs(1),
+            initialize_timeout: Duration::from_secs(1),
+            http_headers: HeaderMap::new(),
+        },
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+    )
+    .get()
+    .await
+    .expect("client should connect");
+    let session = client
+        .register_session(&process_id)
+        .await
+        .expect("session should register");
+    let (started_tx, mut started_rx) = mpsc::unbounded_channel();
+    let (dropped_tx, mut dropped_rx) = mpsc::unbounded_channel();
+    let decider: Arc<dyn NetworkPolicyDecider> = Arc::new(move |request: NetworkPolicyRequest| {
+        let started_tx = started_tx.clone();
+        let dropped_tx = dropped_tx.clone();
+        async move {
+            assert_eq!(
+                tracing::Span::current()
+                    .metadata()
+                    .map(tracing::Metadata::name),
+                Some("codex.exec_server.request"),
+                "network policy decisions must run inside the inbound request span"
+            );
+            match request.host.as_str() {
+                "allowed.example" => NetworkDecision::Allow,
+                "denied.example" => NetworkDecision::deny("blocked"),
+                "pending.example" => {
+                    started_tx.send(()).expect("decision should start");
+                    let _drop_guard = PendingDecisionGuard(dropped_tx);
+                    std::future::pending().await
+                }
+                host => panic!("unexpected policy host: {host}"),
+            }
+        }
+    });
+    session.state.network_policy.controller.store(Some(Arc::new(
+        NetworkPolicyDecisionController {
+            decider,
+            timeout: Duration::from_secs(30),
+        },
+    )));
+    ready_tx.send(()).expect("server should be waiting");
+    timeout(Duration::from_secs(5), async {
+        for _ in 0..MAX_IN_FLIGHT_SERVER_CALLS {
+            started_rx
+                .recv()
+                .await
+                .expect("pending decision should start");
+        }
+    })
+    .await
+    .expect("pending decisions should start");
+    overflow_checked_rx
+        .await
+        .expect("overflow should be observed");
+    session.unregister().await;
+    timeout(Duration::from_secs(5), async {
+        for _ in 0..MAX_IN_FLIGHT_SERVER_CALLS {
+            dropped_rx
+                .recv()
+                .await
+                .expect("unregistered decision should be dropped");
+        }
+    })
+    .await
+    .expect("unregistered decisions should be cancelled");
+    unregistered_tx
+        .send(())
+        .expect("server should verify late responses");
+    timeout(Duration::from_secs(2), server)
+        .await
+        .expect("policy routing should finish")
+        .expect("server task should finish");
+
+    tracer_provider.force_flush().expect("flush traces");
+    let spans = span_exporter.get_finished_spans().expect("span export");
+    let policy_spans = spans
+        .iter()
+        .filter(|span| span.name.as_ref() == NETWORK_POLICY_REQUEST_METHOD)
+        .collect::<Vec<_>>();
+    assert!(
+        !policy_spans.is_empty(),
+        "network policy requests should export server spans"
+    );
+    let outcomes = policy_spans
+        .iter()
+        .map(|span| {
+            span.attributes
+                .iter()
+                .find(|attribute| attribute.key.as_str() == "result")
+                .map(|attribute| attribute.value.as_str().into_owned())
+        })
+        .collect::<Vec<_>>();
+    assert!(
+        outcomes.iter().all(Option::is_some),
+        "completed, rejected, and cancelled policy requests must all record an outcome"
+    );
+    assert!(
+        outcomes
+            .iter()
+            .any(|outcome| outcome.as_deref() == Some("success")),
+        "completed and capacity-rejected requests should record successful responses"
+    );
+    assert!(
+        outcomes
+            .iter()
+            .any(|outcome| outcome.as_deref() == Some("disconnected")),
+        "cancelled requests should record disconnection"
+    );
+}

codex-rs/exec-server/src/client_api.rs

@@ -0,0 +1,187 @@
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::time::Duration;
+
+use codex_http_client::HttpClientFactory;
+use futures::future::BoxFuture;
+use http::HeaderMap;
+use tokio::sync::watch;
+
+use crate::ExecServerError;
+use crate::HttpRequestParams;
+use crate::HttpRequestResponse;
+use crate::HttpResponseBodyStream;
+use crate::NoiseChannelIdentity;
+use crate::NoiseChannelPublicKey;
+
+pub(crate) const DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
+pub(crate) const DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10);
+
+/// Connection options for any exec-server client transport.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ExecServerClientConnectOptions {
+    pub client_name: String,
+    pub initialize_timeout: Duration,
+    pub resume_session_id: Option<String>,
+}
+
+/// WebSocket connection arguments for a remote exec-server.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct RemoteExecServerConnectArgs {
+    pub websocket_url: String,
+    pub client_name: String,
+    pub connect_timeout: Duration,
+    pub initialize_timeout: Duration,
+    pub resume_session_id: Option<String>,
+    pub http_client_factory: HttpClientFactory,
+}
+
+/// Registry-authorized material for one Noise rendezvous connection attempt.
+///
+/// Treat this as an atomic, single-use bundle. The URL authorization, executor
+/// registration, pinned executor key, and harness-key authorization describe one
+/// physical connection attempt and must not be mixed with values from another
+/// registry response.
+pub struct NoiseRendezvousConnectBundle {
+    pub websocket_url: String,
+    pub environment_id: String,
+    pub executor_registration_id: String,
+    pub executor_public_key: NoiseChannelPublicKey,
+    pub harness_key_authorization: String,
+}
+
+/// Connection arguments for an authenticated Noise rendezvous exec-server.
+///
+/// `harness_identity` identifies the logical harness endpoint and may be reused
+/// across reconnects. In contrast, callers must supply a fresh
+/// [`NoiseRendezvousConnectBundle`] for each physical connection attempt.
+pub struct NoiseRendezvousConnectArgs {
+    pub bundle: NoiseRendezvousConnectBundle,
+    pub harness_identity: NoiseChannelIdentity,
+    pub client_name: String,
+    pub connect_timeout: Duration,
+    pub initialize_timeout: Duration,
+    pub resume_session_id: Option<String>,
+    pub http_client_factory: HttpClientFactory,
+}
+
+/// Supplies fresh registry-authorized material for Noise rendezvous connections.
+pub trait NoiseRendezvousConnectProvider: Send + Sync {
+    /// Fetch a bundle authorizing this harness key for one physical connection.
+    fn connect_bundle(
+        &self,
+        harness_public_key: NoiseChannelPublicKey,
+    ) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>>;
+}
+
+/// Stdio connection arguments for a command-backed exec-server.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct StdioExecServerConnectArgs {
+    pub command: StdioExecServerCommand,
+    pub client_name: String,
+    pub initialize_timeout: Duration,
+    pub resume_session_id: Option<String>,
+}
+
+/// Structured process command used to start an exec-server over stdio.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct StdioExecServerCommand {
+    pub program: String,
+    pub args: Vec<String>,
+    pub env: HashMap<String, String>,
+    pub cwd: Option<PathBuf>,
+}
+
+pub(crate) type DeferredEnvironmentReadiness = watch::Receiver<Option<Result<(), String>>>;
+
+#[derive(Clone)]
+pub(crate) struct Deferred<T> {
+    pub readiness: DeferredEnvironmentReadiness,
+    pub transport: T,
+}
+
+/// Parameters used to connect to a remote exec-server environment.
+#[derive(Clone)]
+pub(crate) enum ExecServerTransportParams {
+    Deferred(Box<Deferred<ExecServerTransportParams>>),
+    WebSocketUrl {
+        websocket_url: String,
+        connect_timeout: Duration,
+        initialize_timeout: Duration,
+        http_headers: HeaderMap,
+    },
+    NoiseRendezvous {
+        provider: Arc<dyn NoiseRendezvousConnectProvider>,
+        identity: NoiseChannelIdentity,
+    },
+    #[allow(dead_code)]
+    StdioCommand {
+        command: StdioExecServerCommand,
+        initialize_timeout: Duration,
+    },
+}
+
+impl std::fmt::Debug for ExecServerTransportParams {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Deferred(deferred) => f
+                .debug_struct("Deferred")
+                .field("transport", &deferred.transport)
+                .finish_non_exhaustive(),
+            Self::WebSocketUrl {
+                websocket_url,
+                connect_timeout,
+                initialize_timeout,
+                ..
+            } => f
+                .debug_struct("WebSocketUrl")
+                .field("websocket_url", websocket_url)
+                .field("connect_timeout", connect_timeout)
+                .field("initialize_timeout", initialize_timeout)
+                .field("http_headers", &"<redacted>")
+                .finish(),
+            Self::NoiseRendezvous { .. } => {
+                f.debug_struct("NoiseRendezvous").finish_non_exhaustive()
+            }
+            Self::StdioCommand {
+                command,
+                initialize_timeout,
+            } => f
+                .debug_struct("StdioCommand")
+                .field("command", command)
+                .field("initialize_timeout", initialize_timeout)
+                .finish(),
+        }
+    }
+}
+
+impl ExecServerTransportParams {
+    pub(crate) fn websocket_url(websocket_url: String, connect_timeout: Duration) -> Self {
+        Self::WebSocketUrl {
+            websocket_url,
+            connect_timeout,
+            initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
+            http_headers: HeaderMap::new(),
+        }
+    }
+}
+
+/// Sends HTTP requests through a runtime-selected transport.
+///
+/// This is the HTTP capability counterpart to [`crate::ExecBackend`]. Callers
+/// use it when they need environment-owned network requests but should not
+/// depend on the concrete connection type or how that connection is established.
+pub trait HttpClient: Send + Sync {
+    /// Perform an HTTP request and buffer the response body.
+    fn http_request(
+        &self,
+        params: HttpRequestParams,
+    ) -> BoxFuture<'_, Result<HttpRequestResponse, ExecServerError>>;
+
+    /// Perform an HTTP request and return a streamed body handle.
+    fn http_request_stream(
+        &self,
+        params: HttpRequestParams,
+    ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>>;
+}

codex-rs/exec-server/src/client_recovery.rs

@@ -0,0 +1,900 @@
+use std::collections::hash_map::DefaultHasher;
+use std::hash::Hash;
+use std::hash::Hasher;
+use std::sync::Arc;
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+
+use codex_network_proxy::NetworkDecision;
+use codex_network_proxy::NetworkPolicyDecision;
+use codex_network_proxy::NetworkPolicyRequest;
+use codex_network_proxy::NetworkProtocol;
+use codex_network_proxy::NetworkRequestCancellation;
+use codex_network_proxy::NetworkRequestCancellationReason;
+use serde_json::Value;
+use tokio::sync::mpsc;
+use tokio::time::Instant;
+use tokio::time::sleep;
+use tokio::time::timeout;
+use tokio::time::timeout_at;
+use tokio_util::sync::CancellationToken;
+use tracing::Instrument;
+use tracing::debug;
+
+use super::ConnectionStatus;
+use super::ExecServerClient;
+use super::ExecServerError;
+use super::Inner;
+use super::OrderedSessionEvents;
+use super::RecoveryPolicy;
+use super::SessionState;
+use super::disconnected_message;
+use super::fail_all_in_flight_work;
+use super::handle_server_notification;
+use super::is_transport_closed_error;
+use crate::client_transport::ExecServerReconnectStrategy;
+use crate::process::ExecProcessEvent;
+use crate::protocol::EXEC_READ_METHOD;
+use crate::protocol::EXEC_TERMINATE_METHOD;
+use crate::protocol::ExecServerNetworkPolicyDecision;
+use crate::protocol::ExecServerNetworkProtocol;
+use crate::protocol::MAX_NETWORK_POLICY_HOST_BYTES;
+use crate::protocol::MAX_NETWORK_POLICY_PROCESS_ID_BYTES;
+use crate::protocol::MAX_NETWORK_POLICY_REASON_BYTES;
+use crate::protocol::NETWORK_POLICY_REQUEST_METHOD;
+use crate::protocol::NetworkPolicyRequestParams;
+use crate::protocol::NetworkPolicyRequestResponse;
+use crate::protocol::ReadParams;
+use crate::protocol::ReadResponse;
+use crate::protocol::TerminateParams;
+use crate::protocol::TerminateResponse;
+use crate::rpc::RpcClient;
+use crate::rpc::RpcClientEvent;
+use crate::rpc::RpcInboundRequestAdmissionError;
+use crate::rpc::SESSION_ALREADY_ATTACHED_ERROR_CODE;
+use crate::rpc::invalid_params;
+use crate::rpc::method_not_found;
+
+#[cfg(test)]
+const SESSION_RECOVERY_TIMEOUT: Duration = Duration::from_millis(500);
+#[cfg(not(test))]
+// Leave margin inside the server's 30-second retention windows because the
+// client and server start their disconnect clocks independently.
+const SESSION_RECOVERY_TIMEOUT: Duration = Duration::from_secs(25);
+const SESSION_RECOVERY_RETRY_INTERVAL: Duration = Duration::from_millis(100);
+const REGISTRY_RECOVERY_INITIAL_RETRY_INTERVAL: Duration = Duration::from_millis(500);
+const REGISTRY_RECOVERY_MAX_RETRY_INTERVAL: Duration = Duration::from_secs(5);
+const NETWORK_POLICY_DENIAL_REASON: &str = "not_allowed";
+
+struct ClientRequestOutcome {
+    span: tracing::Span,
+    result: &'static str,
+}
+
+impl ClientRequestOutcome {
+    fn complete(&mut self, result: &'static str) {
+        self.result = result;
+    }
+}
+
+impl Drop for ClientRequestOutcome {
+    fn drop(&mut self) {
+        self.span.record("result", self.result);
+    }
+}
+
+impl SessionState {
+    fn last_published_seq(&self) -> u64 {
+        self.ordered_events
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .last_published_seq
+    }
+
+    fn recover_events(&self, response: ReadResponse) -> Result<bool, ExecServerError> {
+        let ReadResponse {
+            chunks,
+            next_seq,
+            exited,
+            exit_code,
+            closed,
+            failure,
+            sandbox_denied,
+        } = response;
+        if let Some(message) = failure {
+            return Err(ExecServerError::Protocol(format!(
+                "process failed while recovering: {message}"
+            )));
+        }
+
+        let target_seq = next_seq.saturating_sub(1);
+        let published_closed = {
+            let mut ordered_events = self
+                .ordered_events
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            if ordered_events.failure.is_some()
+                || ordered_events.closed_published
+                || target_seq <= ordered_events.last_published_seq
+            {
+                return Ok(false);
+            }
+            let pending_exit = ordered_events.pending.range_mut(..=target_seq).find_map(
+                |(_, event)| match event {
+                    ExecProcessEvent::Exited {
+                        sandbox_denied: pending_sandbox_denied,
+                        ..
+                    } => Some(pending_sandbox_denied),
+                    _ => None,
+                },
+            );
+            let exit_pending = pending_exit.is_some();
+            if let Some(pending_sandbox_denied) = pending_exit {
+                *pending_sandbox_denied =
+                    Some(pending_sandbox_denied.unwrap_or(false) || sandbox_denied);
+            }
+            let mut exit_known = ordered_events.exit_published || exit_pending;
+            if closed
+                && (matches!(
+                    ordered_events.pending.get(&target_seq),
+                    Some(event) if !matches!(event, ExecProcessEvent::Closed { .. })
+                ) || chunks.iter().any(|chunk| chunk.seq == target_seq))
+            {
+                return Err(ExecServerError::Protocol(format!(
+                    "process close sequence {target_seq} conflicts with recovered output"
+                )));
+            }
+            let mut published_closed = false;
+            for chunk in chunks {
+                if chunk.seq > target_seq {
+                    return Err(ExecServerError::Protocol(format!(
+                        "recovered process output sequence {} exceeds target sequence {target_seq}",
+                        chunk.seq
+                    )));
+                }
+                let next_seq = ordered_events.last_published_seq.saturating_add(1);
+                if exited && !exit_known && chunk.seq > next_seq {
+                    let exit_code = exit_code.ok_or_else(|| {
+                        ExecServerError::Protocol(
+                            "recovering exited process did not include its exit code".to_string(),
+                        )
+                    })?;
+                    ordered_events
+                        .insert_pending(ExecProcessEvent::Exited {
+                            seq: next_seq,
+                            exit_code,
+                            sandbox_denied: Some(sandbox_denied),
+                        })
+                        .map_err(ExecServerError::Protocol)?;
+                    published_closed |= self.publish_ready(&mut ordered_events);
+                    exit_known = true;
+                }
+                if chunk.seq > ordered_events.last_published_seq {
+                    ordered_events
+                        .insert_pending(ExecProcessEvent::Output(chunk))
+                        .map_err(ExecServerError::Protocol)?;
+                    published_closed |= self.publish_ready(&mut ordered_events);
+                }
+            }
+            if closed
+                && !ordered_events.closed_published
+                && !matches!(
+                    ordered_events.pending.get(&target_seq),
+                    Some(ExecProcessEvent::Closed { .. })
+                )
+            {
+                ordered_events
+                    .insert_pending(ExecProcessEvent::Closed { seq: target_seq })
+                    .map_err(ExecServerError::Protocol)?;
+            }
+
+            let event_count = target_seq.saturating_sub(ordered_events.last_published_seq);
+            let first_unpublished_seq = ordered_events.last_published_seq.saturating_add(1);
+            let retained_count = if first_unpublished_seq <= target_seq {
+                ordered_events
+                    .pending
+                    .range(first_unpublished_seq..=target_seq)
+                    .count() as u64
+            } else {
+                0
+            };
+            let missing_count = event_count.saturating_sub(retained_count);
+            if exited && !exit_known {
+                if missing_count != 1 {
+                    return Err(recovery_gap_error(target_seq));
+                }
+                let seq = first_missing_seq(&ordered_events, target_seq);
+                let exit_code = exit_code.ok_or_else(|| {
+                    ExecServerError::Protocol(
+                        "recovering exited process did not include its exit code".to_string(),
+                    )
+                })?;
+                ordered_events
+                    .insert_pending(ExecProcessEvent::Exited {
+                        seq,
+                        exit_code,
+                        sandbox_denied: Some(sandbox_denied),
+                    })
+                    .map_err(ExecServerError::Protocol)?;
+            } else if missing_count != 0 {
+                return Err(recovery_gap_error(target_seq));
+            }
+            published_closed |= self.publish_ready(&mut ordered_events);
+            published_closed
+        };
+
+        self.note_change(target_seq);
+        Ok(published_closed)
+    }
+}
+
+fn first_missing_seq(events: &OrderedSessionEvents, target_seq: u64) -> u64 {
+    let mut expected = events.last_published_seq.saturating_add(1);
+    for seq in events
+        .pending
+        .range(expected..=target_seq)
+        .map(|(seq, _)| *seq)
+    {
+        if seq != expected {
+            break;
+        }
+        expected = expected.saturating_add(1);
+    }
+    expected
+}
+
+fn recovery_gap_error(target_seq: u64) -> ExecServerError {
+    ExecServerError::Protocol(format!(
+        "process events are no longer retained while recovering through sequence {target_seq}"
+    ))
+}
+
+impl Inner {
+    pub(super) async fn rpc_client(self: &Arc<Self>) -> Result<Arc<RpcClient>, ExecServerError> {
+        let mut connection_changed = self.connection_changed.subscribe();
+        loop {
+            if let Some(message) = self.failure_message() {
+                return Err(ExecServerError::Disconnected(message));
+            }
+
+            let rpc_client = {
+                let connection = self
+                    .connection
+                    .lock()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner);
+                match &connection.status {
+                    ConnectionStatus::Connected(rpc_client) => Some(Arc::clone(rpc_client)),
+                    ConnectionStatus::Recovering | ConnectionStatus::Failed(_) => None,
+                }
+            };
+            let Some(rpc_client) = rpc_client else {
+                let _ = connection_changed.changed().await;
+                continue;
+            };
+            if !rpc_client.is_disconnected() {
+                return Ok(rpc_client);
+            }
+
+            let _ = connection_changed.changed().await;
+        }
+    }
+
+    pub(super) fn begin_process_start(&self, expected: &Arc<RpcClient>) -> bool {
+        let mut connection = self
+            .connection
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let ConnectionStatus::Connected(current) = &connection.status else {
+            return false;
+        };
+        if !Arc::ptr_eq(current, expected) || expected.is_disconnected() {
+            return false;
+        }
+        connection.active_process_starts += 1;
+        true
+    }
+
+    pub(super) fn finish_process_start(&self) {
+        {
+            let mut connection = self
+                .connection
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            if connection.active_process_starts == 0 {
+                tracing::error!("finished an exec-server process start that was not active");
+                return;
+            }
+            connection.active_process_starts -= 1;
+        }
+        self.notify_connection_changed();
+    }
+
+    pub(super) fn is_failed(&self) -> bool {
+        self.failure_message().is_some()
+    }
+
+    pub(super) fn failure_message(&self) -> Option<String> {
+        let connection = self
+            .connection
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        match &connection.status {
+            ConnectionStatus::Failed(message) => Some(message.clone()),
+            ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => None,
+        }
+    }
+
+    pub(super) fn request_recovery(
+        self: &Arc<Self>,
+        failed_rpc_client: Arc<RpcClient>,
+        disconnect_message: String,
+    ) {
+        let should_recover = {
+            let mut connection = self
+                .connection
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            match &connection.status {
+                ConnectionStatus::Connected(current)
+                    if Arc::ptr_eq(current, &failed_rpc_client) =>
+                {
+                    connection.set_status(ConnectionStatus::Recovering);
+                    true
+                }
+                ConnectionStatus::Connected(_)
+                | ConnectionStatus::Recovering
+                | ConnectionStatus::Failed(_) => false,
+            }
+        };
+        if !should_recover {
+            return;
+        }
+
+        self.notify_connection_changed();
+        let inner = Arc::clone(self);
+        tokio::spawn(async move {
+            tokio::select! {
+                biased;
+                _ = inner.retired.cancelled() => {},
+                _ = inner.recover(disconnect_message) => {},
+            }
+        });
+    }
+
+    async fn recover(self: &Arc<Self>, disconnect_message: String) {
+        let deadline = Instant::now() + SESSION_RECOVERY_TIMEOUT;
+        self.fail_all_http_body_streams(disconnect_message.clone())
+            .await;
+        if timeout_at(deadline, self.wait_for_process_starts())
+            .await
+            .is_err()
+        {
+            let message = format!(
+                "{disconnect_message}; failed to resume exec-server session: recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}"
+            );
+            self.fail(message).await;
+            return;
+        }
+        if self.reconnect_strategy.is_none() {
+            self.fail(disconnect_message).await;
+            return;
+        }
+
+        let Some(session_id) = self.session_id.get().cloned() else {
+            let message = format!(
+                "{disconnect_message}; failed to resume exec-server session: missing session id"
+            );
+            self.fail(message).await;
+            return;
+        };
+        let uses_registry_backoff = matches!(
+            self.reconnect_strategy.as_ref(),
+            Some(ExecServerReconnectStrategy::NoiseRendezvous { .. })
+        );
+        let mut registry_retry_attempt = 0;
+        let last_error = loop {
+            match timeout_at(deadline, self.resume_once(&session_id)).await {
+                Ok(Ok((rpc_client, _attempt))) => {
+                    if !rpc_client.is_disconnected() && self.install_recovered_client(rpc_client) {
+                        return;
+                    }
+                }
+                Ok(Err(error)) if !is_retryable_recovery_error(&error) => {
+                    break error.to_string();
+                }
+                Ok(Err(_)) => {}
+                Err(_) => {
+                    break format!("recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}");
+                }
+            }
+
+            let retry_delay = if uses_registry_backoff {
+                let delay = registry_recovery_retry_delay(&session_id, registry_retry_attempt);
+                registry_retry_attempt = registry_retry_attempt.saturating_add(1);
+                delay
+            } else {
+                SESSION_RECOVERY_RETRY_INTERVAL
+            };
+
+            let now = Instant::now();
+            if now >= deadline {
+                break format!("recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}");
+            }
+            sleep(retry_delay.min(deadline - now)).await;
+        };
+
+        let message =
+            format!("{disconnect_message}; failed to resume exec-server session: {last_error}");
+        self.fail(message).await;
+    }
+
+    async fn wait_for_process_starts(&self) {
+        let mut connection_changed = self.connection_changed.subscribe();
+        loop {
+            let starts_are_done = self
+                .connection
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .active_process_starts
+                == 0;
+            if starts_are_done {
+                return;
+            }
+            let _ = connection_changed.changed().await;
+        }
+    }
+
+    fn install_recovered_client(&self, rpc_client: Arc<RpcClient>) -> bool {
+        let installed = {
+            let mut connection = self
+                .connection
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            if !matches!(connection.status, ConnectionStatus::Recovering)
+                || rpc_client.is_disconnected()
+            {
+                false
+            } else {
+                connection.set_status(ConnectionStatus::Connected(rpc_client));
+                true
+            }
+        };
+        if installed {
+            self.notify_connection_changed();
+        }
+        installed
+    }
+
+    fn notify_connection_changed(&self) {
+        self.connection_changed.send_replace(());
+    }
+
+    async fn resume_once(
+        self: &Arc<Self>,
+        session_id: &str,
+    ) -> Result<(Arc<RpcClient>, Option<tokio::sync::OwnedSemaphorePermit>), ExecServerError> {
+        let reconnect_strategy = self
+            .reconnect_strategy
+            .as_ref()
+            .ok_or_else(|| ExecServerError::Protocol("missing reconnect strategy".to_string()))?;
+        let attempt = reconnect_strategy.resume(session_id).await?;
+        let (connection, options, attempt_permit, noise_context) = attempt.into_parts();
+        let (rpc_client, events_rx) = RpcClient::new(connection);
+        let rpc_client = Arc::new(rpc_client);
+        let client = ExecServerClient {
+            inner: Arc::clone(self),
+            recovery_policy: RecoveryPolicy::Wait,
+        };
+        // Resuming a session redirects notifications from its running processes
+        // to this connection during initialize. Drain them immediately so a
+        // burst cannot fill the bounded event channel and block the initialize
+        // response behind it.
+        client.spawn_rpc_reader(&rpc_client, events_rx);
+        client
+            .initialize_rpc(&rpc_client, options, noise_context)
+            .await?;
+
+        self.recover_processes(&rpc_client).await?;
+        Ok((rpc_client, attempt_permit))
+    }
+
+    async fn recover_processes(
+        self: &Arc<Self>,
+        rpc_client: &RpcClient,
+    ) -> Result<(), ExecServerError> {
+        let sessions = self.sessions.load_full();
+        for (process_id, session) in sessions.iter() {
+            if !session.recoverable.load(Ordering::Acquire) {
+                continue;
+            }
+            let response = rpc_client
+                .call::<_, ReadResponse>(
+                    EXEC_READ_METHOD,
+                    &ReadParams {
+                        process_id: process_id.clone(),
+                        after_seq: Some(session.last_published_seq()),
+                        max_bytes: None,
+                        wait_ms: Some(0),
+                    },
+                )
+                .await
+                .map_err(ExecServerError::from);
+            let recovered = match response {
+                Ok(response) => session.recover_events(response),
+                Err(error) if is_transport_closed_error(&error) => return Err(error),
+                Err(error) => Err(error),
+            };
+            match recovered {
+                Ok(true) => self.remove_session_if(process_id, session),
+                Ok(false) => {}
+                Err(error) => {
+                    session
+                        .network_policy
+                        .cancellation
+                        .record(NetworkRequestCancellationReason::ProcessCancelled);
+                    let terminated: Result<TerminateResponse, ExecServerError> = rpc_client
+                        .call_for_cleanup(
+                            EXEC_TERMINATE_METHOD,
+                            &TerminateParams {
+                                process_id: process_id.clone(),
+                            },
+                        )
+                        .await
+                        .map_err(ExecServerError::from);
+                    if let Err(terminate_error) = terminated
+                        && is_transport_closed_error(&terminate_error)
+                    {
+                        return Err(terminate_error);
+                    }
+                    self.remove_session_if(process_id, session);
+                    session.set_failure(format!("failed to recover process {process_id}: {error}"));
+                }
+            }
+        }
+        Ok(())
+    }
+
+    async fn fail(self: &Arc<Self>, message: String) {
+        let (message, newly_failed) = {
+            let mut connection = self
+                .connection
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            match &connection.status {
+                ConnectionStatus::Failed(existing) => (existing.clone(), false),
+                ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => {
+                    connection.set_status(ConnectionStatus::Failed(message.clone()));
+                    (message, true)
+                }
+            }
+        };
+        if newly_failed {
+            self.notify_connection_changed();
+            fail_all_in_flight_work(self, message.clone()).await;
+        }
+    }
+}
+
+impl ExecServerClient {
+    pub(super) fn spawn_rpc_reader(
+        &self,
+        rpc_client: &Arc<RpcClient>,
+        mut events_rx: mpsc::Receiver<RpcClientEvent>,
+    ) {
+        let inner = Arc::downgrade(&self.inner);
+        let rpc_inbound_request_slots = Arc::clone(&self.inner.rpc_inbound_request_slots);
+        let rpc_client = Arc::downgrade(rpc_client);
+        let connection_cancelled = CancellationToken::new();
+        let connection_cancel_guard = connection_cancelled.clone().drop_guard();
+        tokio::spawn(async move {
+            let _connection_cancel_guard = connection_cancel_guard;
+            while let Some(event) = events_rx.recv().await {
+                let (Some(inner), Some(rpc_client)) = (inner.upgrade(), rpc_client.upgrade())
+                else {
+                    return;
+                };
+                match event {
+                    RpcClientEvent::Request {
+                        request,
+                        request_span,
+                    } => {
+                        let mut request_outcome = ClientRequestOutcome {
+                            span: request_span,
+                            result: "disconnected",
+                        };
+                        if request.method != NETWORK_POLICY_REQUEST_METHOD {
+                            let error = method_not_found(format!(
+                                "exec-server client does not implement `{}` yet",
+                                request.method
+                            ));
+                            if rpc_client.respond_error(request.id, error).await.is_err() {
+                                inner.request_recovery(
+                                    rpc_client,
+                                    disconnected_message(/*reason*/ None),
+                                );
+                                return;
+                            }
+                            request_outcome.complete("error");
+                            continue;
+                        }
+                        request_outcome
+                            .span
+                            .record("otel.name", NETWORK_POLICY_REQUEST_METHOD);
+
+                        let request_guard = match rpc_client
+                            .admit_inbound_request(&request.id, &rpc_inbound_request_slots)
+                        {
+                            Ok(request_guard) => request_guard,
+                            Err(RpcInboundRequestAdmissionError::InvalidRequestId) => {
+                                rpc_client.close_transport().await;
+                                inner.request_recovery(
+                                    rpc_client,
+                                    "exec-server sent an invalid request ID".to_string(),
+                                );
+                                return;
+                            }
+                            Err(RpcInboundRequestAdmissionError::DuplicateRequestId) => {
+                                rpc_client.close_transport().await;
+                                inner.request_recovery(
+                                    rpc_client,
+                                    "exec-server reused an in-flight request ID".to_string(),
+                                );
+                                return;
+                            }
+                            Err(RpcInboundRequestAdmissionError::AtCapacity) => {
+                                let response = NetworkPolicyRequestResponse {
+                                    decision: ExecServerNetworkPolicyDecision::Deny {
+                                        reason: NETWORK_POLICY_DENIAL_REASON.to_string(),
+                                    },
+                                };
+                                if rpc_client.respond(request.id, &response).await.is_err() {
+                                    inner.request_recovery(
+                                        rpc_client,
+                                        disconnected_message(/*reason*/ None),
+                                    );
+                                    return;
+                                }
+                                request_outcome.complete("success");
+                                continue;
+                            }
+                        };
+                        let request_id = request.id;
+                        let params: NetworkPolicyRequestParams =
+                            match serde_json::from_value(request.params.unwrap_or(Value::Null)) {
+                                Ok(params) => params,
+                                Err(_) => {
+                                    let error = invalid_params(
+                                        "invalid network policy request params".to_string(),
+                                    );
+                                    if rpc_client.respond_error(request_id, error).await.is_err() {
+                                        inner.request_recovery(
+                                            rpc_client,
+                                            disconnected_message(/*reason*/ None),
+                                        );
+                                        return;
+                                    }
+                                    request_outcome.complete("error");
+                                    continue;
+                                }
+                            };
+                        let process_id = params.process_id;
+                        let request = params.request;
+                        let process_id_valid = !process_id.is_empty()
+                            && process_id.len() <= MAX_NETWORK_POLICY_PROCESS_ID_BYTES;
+                        let host_valid = !request.host.is_empty()
+                            && request.host.len() <= MAX_NETWORK_POLICY_HOST_BYTES
+                            && !request.host.chars().any(char::is_control)
+                            && !request.host.chars().any(char::is_whitespace);
+                        let session = (process_id_valid && host_valid)
+                            .then(|| inner.get_session(&process_id))
+                            .flatten();
+                        let controller = session
+                            .as_ref()
+                            .and_then(|session| session.network_policy.controller.load_full());
+                        let process_cancelled = session
+                            .as_ref()
+                            .map(|session| session.network_policy.cancelled.clone());
+                        let process_cancellation = session
+                            .as_ref()
+                            .map(|session| session.network_policy.cancellation.clone());
+                        let cancellation = NetworkRequestCancellation::default();
+                        let expected_session = session.as_ref().map(Arc::downgrade);
+                        let policy_request =
+                            (process_id_valid && host_valid).then_some(NetworkPolicyRequest {
+                                protocol: match request.protocol {
+                                    ExecServerNetworkProtocol::Http => NetworkProtocol::Http,
+                                    ExecServerNetworkProtocol::HttpsConnect => {
+                                        NetworkProtocol::HttpsConnect
+                                    }
+                                    ExecServerNetworkProtocol::Socks5Tcp => {
+                                        NetworkProtocol::Socks5Tcp
+                                    }
+                                    ExecServerNetworkProtocol::Socks5Udp => {
+                                        NetworkProtocol::Socks5Udp
+                                    }
+                                },
+                                host: request.host,
+                                port: request.port,
+                                environment_id: None,
+                                client_addr: None,
+                                method: None,
+                                command: None,
+                                exec_policy_hint: None,
+                                execution_id: None,
+                                disconnect: None,
+                                cancellation: Some(cancellation.clone()),
+                            });
+                        let inner = Arc::downgrade(&inner);
+                        let rpc_client = Arc::downgrade(&rpc_client);
+                        let connection_cancelled = connection_cancelled.clone();
+                        let task_span = request_outcome.span.clone();
+                        let task = async move {
+                            let _request_guard = request_guard;
+                            let decision = match (controller, policy_request, process_cancelled) {
+                                (Some(controller), Some(request), Some(process_cancelled)) => {
+                                    // Keep the decision future outside select/timeout so its
+                                    // guard sees the cancellation cause before it is dropped.
+                                    let mut decision = controller.decider.decide(request);
+                                    tokio::select! {
+                                        biased;
+                                        _ = connection_cancelled.cancelled() => {
+                                            cancellation.record(NetworkRequestCancellationReason::ConnectionClosed);
+                                            return;
+                                        },
+                                        _ = process_cancelled.cancelled() => {
+                                            cancellation.record(process_cancellation.as_ref()
+                                                .and_then(NetworkRequestCancellation::reason)
+                                                .unwrap_or(NetworkRequestCancellationReason::ProcessCancelled));
+                                            NetworkDecision::deny(NETWORK_POLICY_DENIAL_REASON)
+                                        }
+                                        result = timeout(
+                                            controller.timeout,
+                                            &mut decision,
+                                        ) => result.unwrap_or_else(|_| {
+                                            cancellation.record(NetworkRequestCancellationReason::TimedOut);
+                                            NetworkDecision::deny(NETWORK_POLICY_DENIAL_REASON)
+                                        }),
+                                    }
+                                }
+                                (None, _, _) | (_, None, _) | (_, _, None) => {
+                                    NetworkDecision::deny(NETWORK_POLICY_DENIAL_REASON)
+                                }
+                            };
+                            if let Some(expected_session) = expected_session {
+                                let (Some(inner), Some(expected_session)) =
+                                    (inner.upgrade(), expected_session.upgrade())
+                                else {
+                                    return;
+                                };
+                                if !inner
+                                    .get_session(&process_id)
+                                    .is_some_and(|session| Arc::ptr_eq(&session, &expected_session))
+                                {
+                                    return;
+                                }
+                            }
+                            let Some(rpc_client) = rpc_client.upgrade() else {
+                                return;
+                            };
+                            let decision = match decision {
+                                NetworkDecision::Allow => ExecServerNetworkPolicyDecision::Allow,
+                                NetworkDecision::Deny {
+                                    reason, decision, ..
+                                } if reason.len() <= MAX_NETWORK_POLICY_REASON_BYTES
+                                    && !reason.chars().any(char::is_control) =>
+                                {
+                                    match decision {
+                                        NetworkPolicyDecision::Deny => {
+                                            ExecServerNetworkPolicyDecision::Deny { reason }
+                                        }
+                                        NetworkPolicyDecision::Ask => {
+                                            ExecServerNetworkPolicyDecision::Ask { reason }
+                                        }
+                                    }
+                                }
+                                NetworkDecision::Deny { .. } => {
+                                    ExecServerNetworkPolicyDecision::Deny {
+                                        reason: NETWORK_POLICY_DENIAL_REASON.to_string(),
+                                    }
+                                }
+                            };
+                            if let Err(error) = rpc_client
+                                .respond(request_id, &NetworkPolicyRequestResponse { decision })
+                                .await
+                            {
+                                debug!(
+                                    ?error,
+                                    "failed to send network policy decision to exec-server"
+                                );
+                            } else {
+                                request_outcome.complete("success");
+                            }
+                        };
+                        tokio::spawn(task.instrument(task_span));
+                    }
+                    RpcClientEvent::Notification(notification) => {
+                        if let Err(error) = handle_server_notification(&inner, notification).await {
+                            rpc_client.close_transport().await;
+                            inner.request_recovery(
+                                rpc_client,
+                                format!("exec-server notification handling failed: {error}"),
+                            );
+                            return;
+                        }
+                    }
+                    RpcClientEvent::Disconnected { reason } => {
+                        inner.request_recovery(rpc_client, disconnected_message(reason.as_deref()));
+                        return;
+                    }
+                }
+            }
+        });
+    }
+}
+
+pub(crate) fn is_retryable_recovery_error(error: &ExecServerError) -> bool {
+    if let ExecServerError::ConnectionAttempt(error) = error {
+        return is_retryable_recovery_error(error.as_ref());
+    }
+    is_transport_closed_error(error)
+        || matches!(
+            error,
+            ExecServerError::ProvisioningFailed(_)
+                | ExecServerError::WebSocketConnectTimeout { .. }
+                | ExecServerError::WebSocketConnect { .. }
+                | ExecServerError::InitializeTimedOut { .. }
+        )
+        || is_retryable_registry_error(error)
+        || matches!(
+            error,
+            ExecServerError::Server { code, .. }
+                if *code == SESSION_ALREADY_ATTACHED_ERROR_CODE
+        )
+}
+
+pub(crate) fn is_retryable_registry_error(error: &ExecServerError) -> bool {
+    matches!(
+        error,
+        ExecServerError::EnvironmentRegistryRequest(error)
+            if error.is_connect()
+                || error.is_timeout()
+                || error.is_body()
+                || matches!(
+                    error,
+                    codex_http_client::RouteAwareRequestError::Request(error)
+                        if error.is_decode()
+                )
+    ) || matches!(
+        error,
+        ExecServerError::EnvironmentRegistryHttp { status, .. }
+            if status.is_server_error()
+                || *status == http::StatusCode::REQUEST_TIMEOUT
+                || *status == http::StatusCode::TOO_MANY_REQUESTS
+    ) || is_environment_offline_error(error)
+}
+
+pub(crate) fn is_environment_offline_error(error: &ExecServerError) -> bool {
+    matches!(
+        error,
+        ExecServerError::EnvironmentRegistryHttp { status, code, .. }
+            if *status == http::StatusCode::CONFLICT
+                && code.as_deref() == Some("environment_offline")
+    )
+}
+
+pub(crate) fn registry_recovery_retry_delay(retry_key: &str, attempt: u32) -> Duration {
+    let multiplier = 1_u32.checked_shl(attempt.min(4)).unwrap_or(u32::MAX);
+    let base_delay = REGISTRY_RECOVERY_INITIAL_RETRY_INTERVAL
+        .saturating_mul(multiplier)
+        .min(REGISTRY_RECOVERY_MAX_RETRY_INTERVAL);
+    let base_millis = base_delay.as_millis() as u64;
+    let mut hasher = DefaultHasher::new();
+    retry_key.hash(&mut hasher);
+    attempt.hash(&mut hasher);
+
+    Duration::from_millis(base_millis + hasher.finish() % (base_millis / 2 + 1))
+}
+
+#[cfg(test)]
+#[path = "client_recovery_tests.rs"]
+mod tests;

codex-rs/exec-server/src/client_recovery_tests.rs

@@ -0,0 +1,262 @@
+use std::time::Duration;
+
+use pretty_assertions::assert_eq;
+
+use super::*;
+use crate::protocol::ExecOutputStream;
+use crate::protocol::ProcessOutputChunk;
+
+fn registry_error(status: http::StatusCode, code: Option<&str>) -> ExecServerError {
+    ExecServerError::EnvironmentRegistryHttp {
+        status,
+        code: code.map(str::to_string),
+        message: "registry unavailable".to_string(),
+    }
+}
+
+#[test]
+fn registry_recovery_retry_delay_exponentially_backs_off_and_caps() {
+    let cases = [
+        (0, Duration::from_millis(500)),
+        (1, Duration::from_secs(1)),
+        (2, Duration::from_secs(2)),
+        (3, Duration::from_secs(4)),
+        (4, Duration::from_secs(5)),
+        (20, Duration::from_secs(5)),
+    ];
+
+    for (attempt, base) in cases {
+        let delay = registry_recovery_retry_delay("session-1", attempt);
+        assert!(delay >= base, "delay {delay:?} for attempt {attempt}");
+        assert!(
+            delay <= base + base / 2,
+            "delay {delay:?} for attempt {attempt}"
+        );
+    }
+}
+
+#[test]
+fn recovery_retries_transient_registry_errors() {
+    for status in [
+        http::StatusCode::REQUEST_TIMEOUT,
+        http::StatusCode::TOO_MANY_REQUESTS,
+        http::StatusCode::INTERNAL_SERVER_ERROR,
+        http::StatusCode::BAD_GATEWAY,
+        http::StatusCode::SERVICE_UNAVAILABLE,
+    ] {
+        let error = registry_error(status, /*code*/ None);
+
+        assert!(is_retryable_registry_error(&error));
+        assert!(is_retryable_recovery_error(&error));
+        assert!(is_retryable_recovery_error(
+            &ExecServerError::ConnectionAttempt(Arc::new(error))
+        ));
+    }
+}
+
+#[test]
+fn recovery_retries_registry_request_timeouts() {
+    let error = ExecServerError::EnvironmentRegistryRequest(
+        codex_http_client::RouteAwareRequestError::Timeout,
+    );
+
+    assert!(is_retryable_registry_error(&error));
+    assert!(is_retryable_recovery_error(&error));
+}
+
+#[test]
+fn recovery_retries_environment_offline_conflicts() {
+    let error = registry_error(http::StatusCode::CONFLICT, Some("environment_offline"));
+
+    assert!(is_retryable_registry_error(&error));
+    assert!(is_retryable_recovery_error(&error));
+}
+
+#[test]
+fn recovery_does_not_retry_other_registry_conflicts() {
+    let error = registry_error(http::StatusCode::CONFLICT, Some("registration_conflict"));
+
+    assert!(!is_retryable_registry_error(&error));
+    assert!(!is_retryable_recovery_error(&error));
+    assert!(!is_retryable_recovery_error(
+        &ExecServerError::ConnectionAttempt(Arc::new(error))
+    ));
+}
+
+#[test]
+fn process_event_reorder_rejects_oversized_output() {
+    let state = SessionState::new(/*recoverable*/ true);
+
+    let error = state
+        .publish_ordered_event(ExecProcessEvent::Output(ProcessOutputChunk {
+            seq: 1,
+            stream: ExecOutputStream::Stdout,
+            chunk: vec![0; super::super::MAX_PENDING_PROCESS_EVENT_BYTES + 1].into(),
+        }))
+        .expect_err("oversized pending process output should be rejected");
+
+    assert!(error.contains("bytes"));
+}
+
+#[test]
+fn process_event_reorder_accepts_gap_closing_event_at_limits() {
+    let state = SessionState::new(/*recoverable*/ true);
+    let chunk_size =
+        super::super::MAX_PENDING_PROCESS_EVENT_BYTES / super::super::MAX_PENDING_PROCESS_EVENTS;
+    let last_seq = super::super::MAX_PENDING_PROCESS_EVENTS as u64 + 1;
+
+    for seq in 2..=last_seq {
+        assert!(
+            !state
+                .publish_ordered_event(ExecProcessEvent::Output(ProcessOutputChunk {
+                    seq,
+                    stream: ExecOutputStream::Stdout,
+                    chunk: vec![0; chunk_size].into(),
+                }))
+                .expect("future output should fit within reorder limits")
+        );
+    }
+    assert!(
+        !state
+            .publish_ordered_event(ExecProcessEvent::Output(ProcessOutputChunk {
+                seq: 1,
+                stream: ExecOutputStream::Stdout,
+                chunk: b"x".to_vec().into(),
+            }))
+            .expect("gap-closing output should drain the reorder buffer")
+    );
+
+    let ordered_events = state
+        .ordered_events
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner);
+    assert_eq!(
+        (
+            ordered_events.last_published_seq,
+            ordered_events.pending.len(),
+            ordered_events.pending_bytes,
+        ),
+        (last_seq, 0, 0)
+    );
+}
+
+#[test]
+fn recovery_handles_dense_tail_output_and_newer_notification() {
+    let state = SessionState::new(/*recoverable*/ true);
+    let last_seq = super::super::MAX_PENDING_PROCESS_EVENTS as u64 + 2;
+    let live_seq = last_seq + 1;
+    assert!(
+        !state
+            .publish_ordered_event(ExecProcessEvent::Output(ProcessOutputChunk {
+                seq: live_seq,
+                stream: ExecOutputStream::Stdout,
+                chunk: b"live".to_vec().into(),
+            }))
+            .expect("live output should remain bounded while recovery fills the gap")
+    );
+    let chunks = (2..=last_seq)
+        .map(|seq| ProcessOutputChunk {
+            seq,
+            stream: ExecOutputStream::Stdout,
+            chunk: b"x".to_vec().into(),
+        })
+        .collect();
+
+    assert!(
+        !state
+            .recover_events(ReadResponse {
+                chunks,
+                next_seq: last_seq + 1,
+                exited: true,
+                exit_code: Some(17),
+                closed: false,
+                failure: None,
+                sandbox_denied: false,
+            })
+            .expect("dense retained output should recover")
+    );
+
+    let ordered_events = state
+        .ordered_events
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner);
+    assert_eq!(
+        (
+            ordered_events.last_published_seq,
+            ordered_events.pending.len(),
+            ordered_events.pending_bytes,
+        ),
+        (live_seq, 0, 0)
+    );
+}
+
+#[test]
+fn recovery_rejects_output_at_closed_sequence() {
+    let state = SessionState::new(/*recoverable*/ true);
+
+    let error = state
+        .recover_events(ReadResponse {
+            chunks: vec![ProcessOutputChunk {
+                seq: 1,
+                stream: ExecOutputStream::Stdout,
+                chunk: b"output".to_vec().into(),
+            }],
+            next_seq: 2,
+            exited: false,
+            exit_code: None,
+            closed: true,
+            failure: None,
+            sandbox_denied: false,
+        })
+        .expect_err("output should not occupy the closed sequence");
+
+    assert!(
+        error
+            .to_string()
+            .contains("conflicts with recovered output")
+    );
+}
+
+#[tokio::test]
+async fn recovery_adds_sandbox_denial_to_pending_exit_event() {
+    let state = SessionState::new(/*recoverable*/ true);
+    assert!(
+        !state
+            .publish_ordered_event(ExecProcessEvent::Exited {
+                seq: 2,
+                exit_code: 1,
+                sandbox_denied: None,
+            })
+            .expect("pending exit should fit within reorder limits")
+    );
+
+    state
+        .recover_events(ReadResponse {
+            chunks: vec![ProcessOutputChunk {
+                seq: 1,
+                stream: ExecOutputStream::Stderr,
+                chunk: b"sandbox denied".to_vec().into(),
+            }],
+            next_seq: 3,
+            exited: true,
+            exit_code: Some(1),
+            closed: false,
+            failure: None,
+            sandbox_denied: true,
+        })
+        .expect("recovery should publish the pending exit");
+
+    let mut events = state.subscribe_events();
+    assert!(matches!(
+        events.recv().await,
+        Ok(ExecProcessEvent::Output(_))
+    ));
+    assert_eq!(
+        events.recv().await,
+        Ok(ExecProcessEvent::Exited {
+            seq: 2,
+            exit_code: 1,
+            sandbox_denied: Some(true),
+        })
+    );
+}

codex-rs/exec-server/src/client_refresh.rs

@@ -0,0 +1,269 @@
+//! Explicit connection refresh after a planned executor replacement.
+//!
+//! Ordinary recovery tries to resume the same executor session after a transient
+//! disconnect. Replacement needs a fresh session, without waiting for old recovery
+//! to give up. The caller supplies the ordering: register the replacement first,
+//! then refresh. Executor identity stays inside this connection layer.
+//!
+//! Flow: fresh registry lookup -> reuse or retire session -> connect if needed ->
+//! live status probe. The lazy client and its `Environment` remain the same objects;
+//! only the underlying `ExecServerClient` may change. The public caller contract is on
+//! `Environment::refresh_connection`.
+//!
+//! Two races determine the synchronization here. A client installed during the
+//! lookup makes that lookup stale, so refresh checks again. A connection attempt
+//! cancelled by refresh must never install later. Cancellation and installation
+//! synchronize on the `current_client` lock; acquire `reconnect` first when both are needed.
+//! `refresh_lock` serializes only explicit refreshes; ordinary connection and recovery
+//! work can continue concurrently. Retired sessions cannot publish environment state.
+
+use std::sync::Arc;
+use std::sync::Mutex as StdMutex;
+
+use futures::future::BoxFuture;
+use tokio::sync::OnceCell;
+use tokio::sync::watch;
+use tokio::time::timeout;
+use tokio_util::sync::CancellationToken;
+
+use super::ConnectionResult;
+use super::ConnectionStatus;
+use super::ExecServerClient;
+use super::ExecServerError;
+use super::Inner;
+use super::LazyRemoteExecServerClient;
+use super::fail_all_in_flight_work;
+use crate::EnvironmentConnectionState;
+use crate::NoiseChannelPublicKey;
+use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
+use crate::client_api::ExecServerTransportParams;
+use crate::client_api::NoiseRendezvousConnectBundle;
+use crate::client_api::NoiseRendezvousConnectProvider;
+use crate::client_transport::ExecServerReconnectStrategy;
+
+/// Shared startup/reconnect result plus cancellation for work superseded by refresh.
+/// The optional transport carries a refresh lookup's bundle into the normal connector.
+#[derive(Default)]
+pub(super) struct ConnectionAttempt {
+    pub(super) result: OnceCell<ConnectionResult>,
+    pub(super) cancelled: CancellationToken,
+    pub(super) transport: Option<ExecServerTransportParams>,
+}
+
+// Use the compared bundle intact for the first connection: address, key and authorization
+// belong together. Later lookups, including authorization refresh, use the real provider.
+struct PrefetchedConnectProvider {
+    bundle: StdMutex<Option<NoiseRendezvousConnectBundle>>,
+    provider: Arc<dyn NoiseRendezvousConnectProvider>,
+}
+
+impl NoiseRendezvousConnectProvider for PrefetchedConnectProvider {
+    fn connect_bundle(
+        &self,
+        harness_public_key: NoiseChannelPublicKey,
+    ) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
+        Box::pin(async move {
+            let bundle = self
+                .bundle
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .take();
+            match bundle {
+                Some(bundle) => Ok(bundle),
+                None => self.provider.connect_bundle(harness_public_key).await,
+            }
+        })
+    }
+}
+
+impl LazyRemoteExecServerClient {
+    #[expect(
+        clippy::await_holding_invalid_type,
+        reason = "serialize explicit refreshes, not ordinary connection or recovery attempts"
+    )]
+    pub(crate) async fn refresh_connection(&self) -> Result<(), ExecServerError> {
+        let _refresh = self.refresh_lock.lock().await;
+        let (previous, attempt) = loop {
+            let observed = self.cached_client();
+            let mut transport = self.transport_params.clone().ok_or_else(|| {
+                ExecServerError::Protocol(
+                    "connection refresh requires a Noise registry".to_string(),
+                )
+            })?;
+            let target = match &mut transport {
+                ExecServerTransportParams::Deferred(deferred) => &mut deferred.transport,
+                transport => transport,
+            };
+            let ExecServerTransportParams::NoiseRendezvous { provider, identity } = target else {
+                return Err(ExecServerError::Protocol(
+                    "connection refresh requires a Noise registry".to_string(),
+                ));
+            };
+            // This lookup is independent of the old session and its recovery deadline.
+            let bundle = timeout(
+                DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
+                provider.connect_bundle(identity.public_key()),
+            )
+            .await
+            .map_err(|_| {
+                ExecServerError::EnvironmentRegistryRequest(
+                    codex_http_client::RouteAwareRequestError::Timeout,
+                )
+            })??;
+            let executor_public_key = bundle.executor_public_key.clone();
+            *provider = Arc::new(PrefetchedConnectProvider {
+                bundle: StdMutex::new(Some(bundle)),
+                provider: Arc::clone(provider),
+            });
+
+            let mut reconnect = self
+                .reconnect
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            let current = self
+                .current_client
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            // An ordinary connection may have finished during the lookup. Re-read the
+            // registry rather than retire a newer client using a superseded response.
+            if !match (&observed, &*current) {
+                (Some(observed), Some(current)) => Arc::ptr_eq(&observed.inner, &current.inner),
+                (None, None) => true,
+                _ => false,
+            } {
+                continue;
+            }
+            // is_disconnected means terminally failed, not temporarily recovering.
+            // Preserve same-executor recovery; the final probe fails fast if still recovering.
+            if current.as_ref().is_some_and(|client| {
+                !client.is_disconnected()
+                    && matches!(
+                        client.inner.reconnect_strategy.as_ref(),
+                        Some(ExecServerReconnectStrategy::NoiseRendezvous {
+                            executor_public_key: key, ..
+                        }) if key == &executor_public_key
+                    )
+            }) {
+                break (current.clone(), None);
+            }
+            // Cancellation and connection installation use the same lock. A late
+            // handshake cannot install a client after its attempt has been superseded.
+            self.startup.cancelled.cancel();
+            if let Some(attempt) = reconnect.as_ref() {
+                attempt.cancelled.cancel();
+            }
+            self.environment_connection_state_tx
+                .send_replace(EnvironmentConnectionState::Disconnected);
+            let attempt = Arc::new(ConnectionAttempt {
+                transport: Some(transport),
+                ..Default::default()
+            });
+            *reconnect = Some(Arc::clone(&attempt));
+            break (current.clone(), Some(attempt));
+        };
+        let client = match attempt {
+            Some(attempt) => {
+                if let Some(previous) = previous {
+                    previous.inner.retire().await;
+                }
+                let result = attempt
+                    .result
+                    .get_or_init(|| self.connect_once(&attempt))
+                    .await
+                    .clone();
+                let mut reconnect = self
+                    .reconnect
+                    .lock()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner);
+                if reconnect
+                    .as_ref()
+                    .is_some_and(|current| Arc::ptr_eq(current, &attempt))
+                {
+                    *reconnect = None;
+                }
+                result.map_err(ExecServerError::ConnectionAttempt)?
+            }
+            None => previous.ok_or_else(|| {
+                ExecServerError::Protocol("current executor session is missing".to_string())
+            })?,
+        };
+        // Metadata may be cached; readiness requires a live, non-recovering probe.
+        client.environment_status().await.map(drop)
+    }
+
+    #[tracing::instrument(name = "codex.exec_server.remote.connect", skip_all)]
+    pub(super) fn connect_once<'a>(
+        &'a self,
+        attempt: &'a ConnectionAttempt,
+    ) -> BoxFuture<'a, ConnectionResult> {
+        // Keep the transport future out of every caller's async layout, including
+        // the CLI entry point, which otherwise exceeds rustc's query-depth limit.
+        Box::pin(async move {
+            let transport = attempt
+                .transport
+                .as_ref()
+                .or(self.transport_params.as_ref())
+                .ok_or_else(|| {
+                    Arc::new(ExecServerError::Protocol(
+                        "missing transport params for lazy exec-server connection".to_string(),
+                    ))
+                })?;
+            let client = tokio::select! {
+                biased;
+                _ = attempt.cancelled.cancelled() => return Err(Arc::new(ExecServerError::Disconnected("connection attempt was superseded".to_string()))),
+                result = ExecServerClient::connect_for_transport(transport.clone(), self.http_client_factory.clone()) => result.map_err(Arc::new)?,
+            };
+            // Cancellation can race with a completed handshake. Recheck before attaching
+            // state or installing the client, under the same lock used by refresh.
+            {
+                let mut current = self
+                    .current_client
+                    .lock()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner);
+                if !attempt.cancelled.is_cancelled() {
+                    client.attach_environment_connection_state(
+                        self.environment_connection_state_tx.clone(),
+                    );
+                    *current = Some(client.clone());
+                    return Ok(client);
+                }
+            }
+            client.inner.retire().await;
+            Err(Arc::new(ExecServerError::Disconnected(
+                "connection attempt was superseded".to_string(),
+            )))
+        })
+    }
+}
+
+impl Inner {
+    async fn retire(self: &Arc<Self>) {
+        let message = "exec-server executor was replaced".to_string();
+        let rpc_client = {
+            let mut connection = self
+                .connection
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            // Detach before a later transport completion can publish stale state.
+            connection.environment_connection_state_tx =
+                watch::channel(EnvironmentConnectionState::Disconnected).0;
+            let rpc_client = match &connection.status {
+                ConnectionStatus::Connected(client) => Some(Arc::clone(client)),
+                ConnectionStatus::Recovering | ConnectionStatus::Failed(_) => None,
+            };
+            self.retired.cancel();
+            connection.set_status(ConnectionStatus::Failed(message.clone()));
+            rpc_client
+        };
+        self.connection_changed.send_replace(());
+        // Drain pending RPCs before stream cleanup, which may wait for other work.
+        if let Some(rpc_client) = rpc_client {
+            rpc_client.close_transport().await;
+        }
+        fail_all_in_flight_work(self, message).await;
+    }
+}
+
+#[cfg(test)]
+#[path = "client_refresh_tests.rs"]
+mod tests;

codex-rs/exec-server/src/client_refresh_tests.rs

@@ -0,0 +1,621 @@
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::time::Duration;
+
+use anyhow::Result;
+use futures::future::BoxFuture;
+use pretty_assertions::assert_eq;
+use tokio::net::TcpListener;
+use tokio::sync::Notify;
+use tokio::sync::oneshot;
+use tokio::task::JoinSet;
+use tokio_util::task::AbortOnDropHandle;
+
+use super::*;
+use crate::ExecServerRuntimePaths;
+use crate::NoiseChannelIdentity;
+use crate::ProcessId;
+use crate::relay::HarnessKeyValidator;
+use crate::relay::run_multiplexed_environment;
+use crate::server::ConnectionProcessor;
+use codex_http_client::HttpClientFactory;
+use codex_http_client::OutboundProxyPolicy;
+
+// These tests exercise real sockets, not keepalive deadlines. The shared unit-test
+// Pong timeout is only 100 ms and can expire during Noise handshakes under load.
+// A blocking task prevents paused time from auto-advancing while socket I/O is
+// pending; dropping the returned sender releases it without polling or sleeping.
+fn freeze_clock() -> std::sync::mpsc::Sender<()> {
+    tokio::time::pause();
+    let (guard, dropped) = std::sync::mpsc::channel();
+    tokio::task::spawn_blocking(move || {
+        let _ = dropped.recv();
+    });
+    guard
+}
+
+#[derive(Clone)]
+struct Target {
+    url: String,
+    identity: NoiseChannelIdentity,
+    registration: String,
+}
+
+struct Registry {
+    target: Mutex<Target>,
+    next_lookup: Mutex<Option<oneshot::Receiver<()>>>,
+    lookup_started: Notify,
+}
+
+impl Registry {
+    fn block_next_lookup(&self) -> oneshot::Sender<()> {
+        let (tx, rx) = oneshot::channel();
+        *self.next_lookup.lock().unwrap() = Some(rx);
+        tx
+    }
+}
+
+impl NoiseRendezvousConnectProvider for Registry {
+    fn connect_bundle(
+        &self,
+        _: NoiseChannelPublicKey,
+    ) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
+        Box::pin(async move {
+            let target = self.target.lock().unwrap().clone();
+            let block = self.next_lookup.lock().unwrap().take();
+            if let Some(block) = block {
+                self.lookup_started.notify_one();
+                block
+                    .await
+                    .map_err(|_| ExecServerError::Protocol("test lookup failed".to_owned()))?;
+            }
+            Ok(NoiseRendezvousConnectBundle {
+                websocket_url: target.url,
+                environment_id: "environment".to_owned(),
+                executor_registration_id: target.registration,
+                executor_public_key: target.identity.public_key(),
+                harness_key_authorization: "authorization".to_owned(),
+            })
+        })
+    }
+}
+
+#[derive(Clone, Default)]
+struct Validator {
+    handshake: Option<Arc<Notify>>,
+    started: Arc<Notify>,
+}
+
+impl HarnessKeyValidator for Validator {
+    async fn validate_harness_key(
+        &self,
+        _: &NoiseChannelPublicKey,
+        _: &str,
+    ) -> Result<(), ExecServerError> {
+        if let Some(handshake) = &self.handshake {
+            self.started.notify_one();
+            handshake.notified().await;
+        }
+        Ok(())
+    }
+}
+
+struct Executor {
+    target: Target,
+    _server: AbortOnDropHandle<()>,
+}
+
+impl Executor {
+    async fn start(validator: Validator) -> Result<Self> {
+        let listener = TcpListener::bind("127.0.0.1:0").await?;
+        let target = Target {
+            url: format!("ws://{}", listener.local_addr()?),
+            identity: NoiseChannelIdentity::generate()?,
+            registration: uuid::Uuid::new_v4().to_string(),
+        };
+        let executor = target.clone();
+        let processor = ConnectionProcessor::new(ExecServerRuntimePaths::new(
+            std::env::current_exe()?,
+            /*codex_linux_sandbox_exe*/ None,
+        )?);
+        let server = tokio::spawn(async move {
+            let mut connections = JoinSet::new();
+            loop {
+                tokio::select! {
+                    socket = listener.accept() => {
+                        let (socket, _) = socket.unwrap();
+                        let executor = executor.clone();
+                        let processor = processor.clone();
+                        let validator = validator.clone();
+                        connections.spawn(async move {
+                            let socket = tokio_tungstenite::accept_async(socket).await.unwrap();
+                            run_multiplexed_environment(socket, processor, "environment".to_owned(), executor.registration, executor.identity, validator).await;
+                        });
+                    }
+                    _ = connections.join_next(), if !connections.is_empty() => {}
+                }
+            }
+        });
+        Ok(Self {
+            target,
+            _server: AbortOnDropHandle::new(server),
+        })
+    }
+
+    fn client(&self) -> Result<(LazyRemoteExecServerClient, Arc<Registry>)> {
+        let registry = Arc::new(Registry {
+            target: Mutex::new(self.target.clone()),
+            next_lookup: Mutex::new(None),
+            lookup_started: Notify::new(),
+        });
+        let client = LazyRemoteExecServerClient::new(
+            ExecServerTransportParams::NoiseRendezvous {
+                provider: registry.clone(),
+                identity: NoiseChannelIdentity::generate()?,
+            },
+            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+        );
+        Ok((client, registry))
+    }
+}
+
+async fn disconnect(client: &ExecServerClient) {
+    let rpc_client = {
+        let connection = client.inner.connection.lock().unwrap();
+        let ConnectionStatus::Connected(rpc_client) = &connection.status else {
+            panic!("expected connected session")
+        };
+        Arc::clone(rpc_client)
+    };
+    rpc_client.close_transport().await;
+}
+
+#[tokio::test]
+async fn refresh_cancels_old_recovery_and_connects_without_resuming_old_session() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let original = client.get().await?;
+    let process = original
+        .register_session(&ProcessId::from("old-process"))
+        .await?;
+    let blocked_recovery = registry.block_next_lookup();
+    disconnect(&original).await;
+    registry.lookup_started.notified().await;
+    *registry.target.lock().unwrap() = new.target.clone();
+    let refresh = client.refresh_connection();
+    tokio::pin!(refresh);
+    let started = tokio::time::Instant::now();
+    assert!(futures::poll!(refresh.as_mut()).is_pending());
+    assert!(original.inner.retired.is_cancelled());
+    assert!(original.is_disconnected());
+    assert_eq!(started.elapsed(), Duration::ZERO);
+    refresh.await?;
+    let replacement = client.get().await?;
+    assert_ne!(original.session_id(), replacement.session_id());
+    assert_eq!(
+        *client.environment_connection_state_tx.borrow(),
+        EnvironmentConnectionState::Connected
+    );
+    assert!(matches!(
+        process.write(b"never replay".to_vec()).await,
+        Err(ExecServerError::Disconnected(_))
+    ));
+    // Releasing an old registry response cannot reinstall or disconnect the replacement.
+    let _ = blocked_recovery.send(());
+    tokio::task::yield_now().await;
+    assert!(Arc::ptr_eq(&client.get().await?.inner, &replacement.inner));
+    replacement.environment_status().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_retires_a_still_connected_old_executor() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let original = client.get().await?;
+    *registry.target.lock().unwrap() = new.target.clone();
+    client.refresh_connection().await?;
+    assert!(original.is_disconnected());
+    assert_ne!(original.session_id(), client.get().await?.session_id());
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_preserves_a_current_session_across_registration_renewal() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (client, registry) = executor.client()?;
+    let original = client.get().await?;
+    registry.target.lock().unwrap().registration = "renewed-registration".to_owned();
+    let concurrent = client.clone();
+    let concurrent = tokio::spawn(async move { concurrent.refresh_connection().await });
+    client.refresh_connection().await?;
+    concurrent.await??;
+    assert!(Arc::ptr_eq(&original.inner, &client.get().await?.inner));
+    original.environment_status().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_cancels_a_stalled_initial_lookup() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let release = registry.block_next_lookup();
+    let initial = client.get();
+    tokio::pin!(initial);
+    assert!(futures::poll!(initial.as_mut()).is_pending());
+    *registry.target.lock().unwrap() = new.target.clone();
+    client.refresh_connection().await?;
+    assert!(initial.await.is_err());
+    let _ = release.send(());
+    client.get().await?.environment_status().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_cancels_a_stalled_noise_handshake() -> Result<()> {
+    let _clock = freeze_clock();
+    let validator = Validator {
+        handshake: Some(Arc::new(Notify::new())),
+        ..Default::default()
+    };
+    let old = Executor::start(validator.clone()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let connecting = client.clone();
+    let initial = tokio::spawn(async move { connecting.get().await });
+    validator.started.notified().await;
+    *registry.target.lock().unwrap() = new.target.clone();
+    client.refresh_connection().await?;
+    assert!(initial.await?.is_err());
+    validator.handshake.unwrap().notify_one();
+    client.get().await?.environment_status().await?;
+    assert_eq!(
+        *client.environment_connection_state_tx.borrow(),
+        EnvironmentConnectionState::Connected
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn failed_refresh_lookup_leaves_the_existing_session_usable() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (client, registry) = executor.client()?;
+    let original = client.get().await?;
+    drop(registry.block_next_lookup());
+    assert!(client.refresh_connection().await.is_err());
+    assert!(Arc::ptr_eq(&original.inner, &client.get().await?.inner));
+    original.environment_status().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn failed_replacement_connection_keeps_old_handles_retired_and_get_retries() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let original = client.get().await?;
+    let process = original
+        .register_session(&ProcessId::from("old-process"))
+        .await?;
+
+    // The registry knows the replacement, but its endpoint drops the new connection.
+    let unavailable = TcpListener::bind("127.0.0.1:0").await?;
+    let target = Target {
+        url: format!("ws://{}", unavailable.local_addr()?),
+        ..new.target.clone()
+    };
+    let _rejected_connection = AbortOnDropHandle::new(tokio::spawn(async move {
+        drop(unavailable.accept().await.unwrap());
+    }));
+    *registry.target.lock().unwrap() = target;
+    assert!(matches!(
+        client.refresh_connection().await,
+        Err(ExecServerError::ConnectionAttempt(_))
+    ));
+    assert!(original.inner.retired.is_cancelled());
+    assert!(matches!(
+        original.environment_status().await,
+        Err(ExecServerError::Disconnected(_))
+    ));
+    assert!(matches!(
+        process.write(b"never replay".to_vec()).await,
+        Err(ExecServerError::Disconnected(_))
+    ));
+    assert_eq!(
+        *client.environment_connection_state_tx.borrow(),
+        EnvironmentConnectionState::Disconnected
+    );
+
+    // A later caller retries against the registry instead of reviving the retired client.
+    *registry.target.lock().unwrap() = new.target.clone();
+    let replacement = client.get().await?;
+    assert_ne!(original.session_id(), replacement.session_id());
+    replacement.environment_status().await?;
+    assert_eq!(
+        *client.environment_connection_state_tx.borrow(),
+        EnvironmentConnectionState::Connected
+    );
+    assert!(matches!(
+        process.write(b"still retired".to_vec()).await,
+        Err(ExecServerError::Disconnected(_))
+    ));
+    Ok(())
+}
+
+#[tokio::test]
+async fn superseded_refresh_lookup_does_not_retire_a_newer_session() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let original = client.get().await?;
+    let release = registry.block_next_lookup();
+    let refreshing = client.refresh_connection();
+    tokio::pin!(refreshing);
+    assert!(futures::poll!(refreshing.as_mut()).is_pending());
+    *registry.target.lock().unwrap() = new.target.clone();
+    original.inner.retire().await;
+    let replacement = client.get().await?;
+    release.send(()).unwrap();
+    refreshing.await?;
+    assert!(Arc::ptr_eq(&replacement.inner, &client.get().await?.inner));
+    assert!(!replacement.inner.retired.is_cancelled());
+    replacement.environment_status().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn environment_refresh_preserves_environment_and_filesystem_handles() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (_, registry) = old.client()?;
+    let manager = crate::EnvironmentManager::from_snapshot(
+        crate::environment_provider::EnvironmentProviderSnapshot {
+            environments: Vec::new(),
+            default: crate::environment_provider::EnvironmentDefault::Disabled,
+            include_local: false,
+        },
+        /*local_runtime_paths*/ None,
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+    )?;
+    let environment = manager
+        .materialize_pending_noise_environment("environment".to_owned(), registry.clone())?;
+    manager.report_environment_provisioning_status(
+        "environment".to_owned(),
+        Ok(crate::EnvironmentReadyInfo {
+            selected_capability_roots: Vec::new(),
+        }),
+        registry.clone(),
+    )?;
+    environment.info().await?;
+    let filesystem = environment.get_filesystem();
+    *registry.target.lock().unwrap() = new.target.clone();
+    environment.refresh_connection().await?;
+    assert!(Arc::ptr_eq(
+        &environment,
+        &manager.get_environment("environment").unwrap()
+    ));
+    assert!(Arc::ptr_eq(&filesystem, &environment.get_filesystem()));
+    environment.info().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_does_not_accept_cached_metadata_during_recovery() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (client, registry) = executor.client()?;
+    let original = client.get().await?;
+    original.environment_info().await?;
+    let blocked_recovery = registry.block_next_lookup();
+    disconnect(&original).await;
+    registry.lookup_started.notified().await;
+    let started = tokio::time::Instant::now();
+    assert!(matches!(
+        client.refresh_connection().await,
+        Err(ExecServerError::Disconnected(_))
+    ));
+    assert_eq!(started.elapsed(), Duration::ZERO);
+    assert!(!original.inner.retired.is_cancelled());
+    drop(blocked_recovery);
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_before_startup_marks_startup_finished() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (client, _) = executor.client()?;
+    assert!(!client.startup_finished());
+    client.refresh_connection().await?;
+    assert!(client.startup_finished());
+    assert!(matches!(client.readiness_result(), Some(Ok(()))));
+    Ok(())
+}
+
+struct ControlledRpc {
+    client: ExecServerClient,
+    requests: tokio::sync::mpsc::Receiver<codex_exec_server_protocol::JSONRPCMessage>,
+    responses: tokio::sync::mpsc::Sender<crate::connection::JsonRpcConnectionEvent>,
+}
+
+async fn controlled_rpc() -> Result<ControlledRpc> {
+    use crate::connection::JsonRpcConnection;
+    use crate::connection::JsonRpcConnectionEvent;
+    use crate::connection::JsonRpcTransport;
+    use codex_exec_server_protocol::JSONRPCMessage;
+    use codex_exec_server_protocol::JSONRPCResponse;
+    let (outgoing_tx, mut requests) = tokio::sync::mpsc::channel(/*buffer*/ 8);
+    let (responses, incoming_rx) = tokio::sync::mpsc::channel(/*buffer*/ 8);
+    let connection = JsonRpcConnection {
+        outgoing_tx,
+        incoming_rx,
+        disconnected_rx: tokio::sync::watch::channel(/*init*/ false).1,
+        task_handles: Vec::new(),
+        transport: JsonRpcTransport::Plain,
+    };
+    let connecting = ExecServerClient::connect(connection, /*options*/ Default::default());
+    tokio::pin!(connecting);
+    assert!(futures::poll!(connecting.as_mut()).is_pending());
+    let Some(JSONRPCMessage::Request(initialize)) = requests.recv().await else {
+        anyhow::bail!("expected initialize request");
+    };
+    responses
+        .send(JsonRpcConnectionEvent::message(JSONRPCMessage::Response(
+            JSONRPCResponse {
+                id: initialize.id,
+                result: serde_json::json!({"sessionId": "controlled-session"}),
+            },
+        )))
+        .await?;
+    let client = connecting.await?;
+    assert!(matches!(
+        requests.recv().await,
+        Some(JSONRPCMessage::Notification(_))
+    ));
+    Ok(ControlledRpc {
+        client,
+        requests,
+        responses,
+    })
+}
+
+#[tokio::test]
+#[expect(
+    clippy::await_holding_invalid_type,
+    reason = "hold stream cleanup pending to exercise retirement ordering"
+)]
+async fn retirement_rejects_pending_mutation_before_stream_cleanup() -> Result<()> {
+    use crate::connection::JsonRpcConnectionEvent;
+    use codex_exec_server_protocol::JSONRPCMessage;
+    use codex_exec_server_protocol::JSONRPCResponse;
+    for response_queued in [false, true] {
+        let mut rpc = controlled_rpc().await?;
+        let call = rpc.client.fs_remove(crate::protocol::FsRemoveParams {
+            path: "file:///retired-file".parse()?,
+            recursive: None,
+            force: None,
+            follow_symlinks: None,
+            sandbox: None,
+        });
+        tokio::pin!(call);
+        assert!(futures::poll!(call.as_mut()).is_pending());
+        let Some(JSONRPCMessage::Request(request)) = rpc.requests.recv().await else {
+            anyhow::bail!("expected filesystem request");
+        };
+        let response = JSONRPCMessage::Response(JSONRPCResponse {
+            id: request.id,
+            result: serde_json::json!({}),
+        });
+        if response_queued {
+            rpc.responses
+                .send(JsonRpcConnectionEvent::message(response.clone()))
+                .await?;
+            let transport = rpc.client.rpc_client_without_recovery()?;
+            tokio::time::timeout(Duration::from_secs(5), async {
+                while transport.pending_request_count().await != 0 {
+                    tokio::task::yield_now().await;
+                }
+            })
+            .await?;
+        }
+        // Hold stream cleanup. Cover both a late response and one already queued
+        // for the caller when retirement begins; closing the socket alone misses the latter.
+        let streams = rpc.client.inner.http_body_streams_write_lock.lock().await;
+        let retirement = rpc.client.inner.retire();
+        tokio::pin!(retirement);
+        assert!(futures::poll!(retirement.as_mut()).is_pending());
+        if !response_queued {
+            rpc.responses
+                .send(JsonRpcConnectionEvent::message(response))
+                .await?;
+        }
+        assert!(matches!(call.await, Err(ExecServerError::Disconnected(_))));
+        drop(streams);
+        retirement.await;
+    }
+    Ok(())
+}
+
+#[tokio::test]
+#[expect(
+    clippy::await_holding_invalid_type,
+    reason = "hold stream cleanup pending to exercise retirement ordering"
+)]
+async fn retirement_rejects_pending_process_start_before_stream_cleanup() -> Result<()> {
+    use crate::connection::JsonRpcConnectionEvent;
+    use codex_exec_server_protocol::JSONRPCMessage;
+    use codex_exec_server_protocol::JSONRPCResponse;
+    for response_queued in [false, true] {
+        let mut rpc = controlled_rpc().await?;
+        let process_id = ProcessId::from("retired-process");
+        let call = rpc.client.start_process(
+            crate::protocol::ExecParams {
+                metadata: Default::default(),
+                process_id: process_id.clone(),
+                argv: vec!["unused".to_owned()],
+                cwd: "file:///".parse()?,
+                shell_snapshot: None,
+                env_policy: None,
+                env: Default::default(),
+                tty: false,
+                pipe_stdin: false,
+                arg0: None,
+                sandbox: None,
+                enforce_managed_network: false,
+                managed_network: None,
+                network_proxy: None,
+            },
+            /*network_policy_decider*/ None,
+        );
+        tokio::pin!(call);
+        assert!(futures::poll!(call.as_mut()).is_pending());
+        let Some(JSONRPCMessage::Request(request)) = rpc.requests.recv().await else {
+            anyhow::bail!("expected process start request");
+        };
+        let response = JSONRPCMessage::Response(JSONRPCResponse {
+            id: request.id,
+            result: serde_json::json!({"processId": "retired-process"}),
+        });
+        if response_queued {
+            rpc.responses
+                .send(JsonRpcConnectionEvent::message(response.clone()))
+                .await?;
+            let state = rpc
+                .client
+                .inner
+                .get_session(&process_id)
+                .expect("pending process");
+            // The start task has a second response channel to the original caller.
+            tokio::time::timeout(Duration::from_secs(5), async {
+                while !state.recoverable.load(std::sync::atomic::Ordering::Acquire) {
+                    tokio::task::yield_now().await;
+                }
+            })
+            .await?;
+        }
+        let streams = rpc.client.inner.http_body_streams_write_lock.lock().await;
+        let retirement = rpc.client.inner.retire();
+        tokio::pin!(retirement);
+        assert!(futures::poll!(retirement.as_mut()).is_pending());
+        if !response_queued {
+            rpc.responses
+                .send(JsonRpcConnectionEvent::message(response))
+                .await?;
+        }
+        assert!(matches!(call.await, Err(ExecServerError::Disconnected(_))));
+        drop(streams);
+        retirement.await;
+    }
+    Ok(())
+}

codex-rs/exec-server/src/client_telemetry.rs

@@ -0,0 +1,23 @@
+//! Caller-side RPC attempt counts using the client's protocol method names, independent of tracing.
+
+use codex_otel::EXEC_SERVER_CLIENT_REQUEST_COUNT_METRIC;
+use codex_otel::MetricsClient;
+
+pub(crate) fn record_client_request(metrics: Option<&MetricsClient>, method: &str) {
+    let Some(metrics) = metrics else {
+        return;
+    };
+    // Record before local admission so failures and cancelled calls still count
+    // as attempts. Notifications and responses never enter these call paths.
+    if metrics
+        .counter_with_description(
+            EXEC_SERVER_CLIENT_REQUEST_COUNT_METRIC,
+            "Total number of client-side exec-server RPC attempts, including local failures.",
+            /*inc*/ 1,
+            &[("method", method)],
+        )
+        .is_err()
+    {
+        tracing::warn!("failed to emit exec-server client request counter");
+    }
+}

codex-rs/exec-server/src/client_transport.rs

@@ -0,0 +1,814 @@
+use std::process::Stdio;
+use std::sync::Arc;
+use std::time::Duration;
+
+use tokio::io::AsyncBufReadExt;
+use tokio::io::BufReader;
+use tokio::process::Command;
+use tokio::sync::OwnedSemaphorePermit;
+use tokio::time::Instant;
+use tokio::time::sleep;
+use tokio::time::timeout;
+use tokio::time::timeout_at;
+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;
+
+use crate::ExecServerClient;
+use crate::ExecServerError;
+use crate::client::NoiseInitializeContext;
+use crate::client::accepted::AcceptedConnectionSource;
+use crate::client::is_retryable_registry_error;
+use crate::client::registry_recovery_retry_delay;
+use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
+use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT;
+use crate::client_api::ExecServerClientConnectOptions;
+use crate::client_api::ExecServerTransportParams;
+use crate::client_api::NoiseRendezvousConnectArgs;
+use crate::client_api::NoiseRendezvousConnectBundle;
+use crate::client_api::NoiseRendezvousConnectProvider;
+use crate::client_api::RemoteExecServerConnectArgs;
+use crate::client_api::StdioExecServerCommand;
+use crate::client_api::StdioExecServerConnectArgs;
+use crate::connection::JsonRpcConnection;
+use crate::noise_channel::NoiseChannelIdentity;
+use crate::noise_relay::NoiseHarnessConnectionArgs;
+use crate::noise_relay::noise_harness_connection_from_websocket_with_readiness;
+use crate::noise_relay::noise_relay_websocket_config;
+use crate::relay::harness_connection_from_websocket;
+use crate::trace_context::current_rendezvous_headers;
+
+const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment";
+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<WebSocketConnection, ExecServerError> {
+    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<String, AuthError> {
+    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
+/// attempt finishes.
+pub(crate) struct ReconnectAttempt {
+    connection: JsonRpcConnection,
+    options: ExecServerClientConnectOptions,
+    attempt_permit: Option<OwnedSemaphorePermit>,
+    noise_context: Option<NoiseInitializeContext>,
+}
+
+struct OpenNoiseRendezvousConnection {
+    connection: JsonRpcConnection,
+    options: ExecServerClientConnectOptions,
+    handshake_ready: tokio::sync::oneshot::Receiver<()>,
+}
+
+struct ReadyNoiseRendezvousConnection {
+    connection: JsonRpcConnection,
+    options: ExecServerClientConnectOptions,
+    noise_context: NoiseInitializeContext,
+}
+
+impl ReconnectAttempt {
+    pub(crate) fn new(
+        connection: JsonRpcConnection,
+        options: ExecServerClientConnectOptions,
+    ) -> Self {
+        Self {
+            connection,
+            options,
+            attempt_permit: None,
+            noise_context: None,
+        }
+    }
+
+    fn with_noise_context(
+        connection: JsonRpcConnection,
+        options: ExecServerClientConnectOptions,
+        noise_context: NoiseInitializeContext,
+    ) -> Self {
+        Self {
+            connection,
+            options,
+            attempt_permit: None,
+            noise_context: Some(noise_context),
+        }
+    }
+
+    pub(crate) fn with_attempt_permit(
+        connection: JsonRpcConnection,
+        options: ExecServerClientConnectOptions,
+        attempt_permit: OwnedSemaphorePermit,
+    ) -> Self {
+        Self {
+            connection,
+            options,
+            attempt_permit: Some(attempt_permit),
+            noise_context: None,
+        }
+    }
+
+    pub(crate) fn into_parts(
+        self,
+    ) -> (
+        JsonRpcConnection,
+        ExecServerClientConnectOptions,
+        Option<OwnedSemaphorePermit>,
+        Option<NoiseInitializeContext>,
+    ) {
+        (
+            self.connection,
+            self.options,
+            self.attempt_permit,
+            self.noise_context,
+        )
+    }
+}
+
+/// Reopens the transport for one logical exec-server client session.
+///
+/// URL connections reuse their configured endpoint. Noise connections retain
+/// the harness identity but fetch a fresh single-use authorization bundle for
+/// every physical connection attempt.
+#[derive(Clone)]
+pub(crate) enum ExecServerReconnectStrategy {
+    Accepted(AcceptedConnectionSource),
+    WebSocket {
+        args: RemoteExecServerConnectArgs,
+        http_headers: HeaderMap,
+    },
+    NoiseRendezvous {
+        // The executor that created the session, not the latest recovery lookup.
+        executor_public_key: crate::NoiseChannelPublicKey,
+        provider: Arc<dyn NoiseRendezvousConnectProvider>,
+        identity: NoiseChannelIdentity,
+        client_name: String,
+        connect_timeout: Duration,
+        initialize_timeout: Duration,
+        http_client_factory: HttpClientFactory,
+    },
+}
+
+impl ExecServerReconnectStrategy {
+    pub(crate) async fn resume(
+        &self,
+        session_id: &str,
+    ) -> Result<ReconnectAttempt, ExecServerError> {
+        match self {
+            Self::Accepted(source) => source.next_connection(session_id).await,
+            Self::WebSocket { args, http_headers } => {
+                let mut args = args.clone();
+                args.resume_session_id = Some(session_id.to_string());
+                let connection =
+                    ExecServerClient::open_websocket_connection(&args, http_headers).await?;
+                Ok(ReconnectAttempt::new(connection, args.into()))
+            }
+            Self::NoiseRendezvous {
+                executor_public_key: _,
+                provider,
+                identity,
+                client_name,
+                connect_timeout,
+                initialize_timeout,
+                http_client_factory,
+            } => {
+                let bundle = provider.connect_bundle(identity.public_key()).await?;
+                let opened = ExecServerClient::open_noise_rendezvous_connection(
+                    NoiseRendezvousConnectArgs {
+                        bundle,
+                        harness_identity: identity.clone(),
+                        client_name: client_name.clone(),
+                        connect_timeout: *connect_timeout,
+                        initialize_timeout: *initialize_timeout,
+                        resume_session_id: Some(session_id.to_string()),
+                        http_client_factory: http_client_factory.clone(),
+                    },
+                )
+                .await?;
+                let ready = ExecServerClient::finish_noise_rendezvous_connection(opened).await?;
+                Ok(ReconnectAttempt::with_noise_context(
+                    ready.connection,
+                    ready.options,
+                    ready.noise_context,
+                ))
+            }
+        }
+    }
+}
+
+impl ExecServerClient {
+    /// Open the selected transport and run the common JSON-RPC initialization.
+    /// Noise connection details are fetched here so reconnects get a fresh URL
+    /// and authorization without replacing the harness identity.
+    pub(crate) async fn connect_for_transport(
+        transport_params: ExecServerTransportParams,
+        http_client_factory: HttpClientFactory,
+    ) -> Result<Self, ExecServerError> {
+        let (transport_params, deferred_readiness) = match transport_params {
+            ExecServerTransportParams::Deferred(deferred) => {
+                (deferred.transport, Some(deferred.readiness))
+            }
+            transport_params => (transport_params, None),
+        };
+
+        if let Some(mut readiness) = deferred_readiness {
+            let provisioning_result = readiness
+                .wait_for(Option::is_some)
+                .await
+                .map_err(|_| {
+                    ExecServerError::Disconnected(
+                        "environment unavailable: environment provisioning ended before completion"
+                            .to_string(),
+                    )
+                })?
+                .clone()
+                .ok_or_else(|| {
+                    ExecServerError::Disconnected(
+                        "environment unavailable: provisioning remained pending after completion"
+                            .to_string(),
+                    )
+                })?;
+            provisioning_result.map_err(ExecServerError::ProvisioningFailed)?;
+        }
+
+        let websocket = match transport_params {
+            ExecServerTransportParams::Deferred(_) => {
+                return Err(ExecServerError::Protocol(
+                    "nested deferred exec-server transports are unsupported".to_string(),
+                ));
+            }
+            ExecServerTransportParams::WebSocketUrl {
+                websocket_url,
+                connect_timeout,
+                initialize_timeout,
+                http_headers,
+            } => (
+                websocket_url,
+                connect_timeout,
+                initialize_timeout,
+                http_headers,
+            ),
+            ExecServerTransportParams::NoiseRendezvous { provider, identity } => {
+                let (ready, executor_public_key) = Self::open_initial_noise_rendezvous_connection(
+                    &provider,
+                    &identity,
+                    http_client_factory.clone(),
+                )
+                .await?;
+                let reconnect_strategy = ExecServerReconnectStrategy::NoiseRendezvous {
+                    executor_public_key,
+                    provider,
+                    identity,
+                    client_name: ENVIRONMENT_CLIENT_NAME.to_string(),
+                    connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
+                    initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
+                    http_client_factory,
+                };
+                return Self::connect_with_recovery_and_noise_context(
+                    ready.connection,
+                    ready.options,
+                    Some(reconnect_strategy),
+                    ready.noise_context,
+                )
+                .await;
+            }
+            ExecServerTransportParams::StdioCommand {
+                command,
+                initialize_timeout,
+            } => {
+                return Self::connect_stdio_command(StdioExecServerConnectArgs {
+                    command,
+                    client_name: ENVIRONMENT_CLIENT_NAME.to_string(),
+                    initialize_timeout,
+                    resume_session_id: None,
+                })
+                .await;
+            }
+        };
+        let (websocket_url, connect_timeout, initialize_timeout, http_headers) = websocket;
+        Self::connect_websocket_with_headers(
+            RemoteExecServerConnectArgs {
+                websocket_url,
+                client_name: ENVIRONMENT_CLIENT_NAME.to_string(),
+                connect_timeout,
+                initialize_timeout,
+                resume_session_id: None,
+                http_client_factory,
+            },
+            http_headers,
+        )
+        .await
+    }
+
+    #[tracing::instrument(name = "codex.exec_server.remote.noise.connect", skip_all)]
+    async fn open_initial_noise_rendezvous_connection(
+        provider: &Arc<dyn NoiseRendezvousConnectProvider>,
+        identity: &NoiseChannelIdentity,
+        http_client_factory: HttpClientFactory,
+    ) -> Result<(ReadyNoiseRendezvousConnection, crate::NoiseChannelPublicKey), ExecServerError>
+    {
+        let open_connection = |bundle: NoiseRendezvousConnectBundle| {
+            Self::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs {
+                bundle,
+                harness_identity: identity.clone(),
+                client_name: ENVIRONMENT_CLIENT_NAME.to_string(),
+                connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
+                initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
+                resume_session_id: None,
+                http_client_factory: http_client_factory.clone(),
+            })
+        };
+        let mut deadline = Instant::now() + INITIAL_REGISTRY_OPERATION_TIMEOUT;
+        let retry_key = uuid::Uuid::new_v4().to_string();
+        let mut retries = 0;
+        let mut refreshed_unauthorized_bundle = false;
+        let connect_bundle = || async {
+            timeout(
+                INITIAL_REGISTRY_REQUEST_TIMEOUT,
+                provider.connect_bundle(identity.public_key()),
+            )
+            .await
+            .unwrap_or_else(|_| {
+                Err(ExecServerError::EnvironmentRegistryRequest(
+                    codex_http_client::RouteAwareRequestError::Timeout,
+                ))
+            })
+        };
+        let mut result = connect_bundle().await;
+        loop {
+            let bundle = match result {
+                Ok(bundle) => bundle,
+                Err(error)
+                    if is_retryable_registry_error(&error)
+                        && retries < INITIAL_REGISTRY_MAX_RETRIES =>
+                {
+                    // Session resumption owns its separate recovery deadline.
+                    let delay = registry_recovery_retry_delay(&retry_key, retries);
+                    retries += 1;
+                    result = match timeout_at(deadline, async {
+                        sleep(delay).await;
+                        connect_bundle().await
+                    })
+                    .await
+                    {
+                        Ok(result) => result,
+                        Err(_) => return Err(error),
+                    };
+                    continue;
+                }
+                Err(error) => return Err(error),
+            };
+            let executor_public_key = bundle.executor_public_key.clone();
+            match open_connection(bundle).await {
+                Err(error)
+                    if !refreshed_unauthorized_bundle
+                        && matches!(
+                            &error,
+                            ExecServerError::WebSocketConnect { source, .. }
+                                if matches!(
+                                    source,
+                                    tokio_tungstenite::tungstenite::Error::Http(response)
+                                        if response.status().as_u16() == 401
+                                )
+                        ) =>
+                {
+                    refreshed_unauthorized_bundle = true;
+                    deadline = Instant::now() + INITIAL_REGISTRY_OPERATION_TIMEOUT;
+                    retries = 0;
+                    result = connect_bundle().await;
+                }
+                result => {
+                    let opened = result?;
+                    let ready = Self::finish_noise_rendezvous_connection(opened).await?;
+                    return Ok((ready, executor_public_key));
+                }
+            }
+        }
+    }
+
+    pub async fn connect_websocket(
+        args: RemoteExecServerConnectArgs,
+    ) -> Result<Self, ExecServerError> {
+        Self::connect_websocket_with_headers(args, HeaderMap::new()).await
+    }
+
+    async fn connect_websocket_with_headers(
+        args: RemoteExecServerConnectArgs,
+        http_headers: HeaderMap,
+    ) -> Result<Self, ExecServerError> {
+        let connection = Self::open_websocket_connection(&args, &http_headers).await?;
+        let options = args.clone().into();
+        Self::connect_with_recovery(
+            connection,
+            options,
+            Some(ExecServerReconnectStrategy::WebSocket { args, http_headers }),
+        )
+        .await
+    }
+
+    pub(crate) async fn open_websocket_connection(
+        args: &RemoteExecServerConnectArgs,
+        http_headers: &HeaderMap,
+    ) -> Result<JsonRpcConnection, ExecServerError> {
+        ensure_rustls_crypto_provider();
+        let websocket_url = args.websocket_url.clone();
+        let connect_timeout = args.connect_timeout;
+        let mut request = websocket_url
+            .as_str()
+            .into_client_request()
+            .map_err(|source| ExecServerError::WebSocketConnect {
+                url: websocket_url.clone(),
+                source,
+            })?;
+        request.headers_mut().extend(http_headers.clone());
+        let connector = WebSocketConnector::new_with_tls_mode(
+            &args.http_client_factory,
+            WebSocketTlsMode::TungsteniteDefault,
+        )
+        .map_err(|error| ExecServerError::WebSocketConfiguration(error.to_string()))?;
+        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) {
+            harness_connection_from_websocket(stream, connection_label)
+        } else {
+            JsonRpcConnection::from_websocket(stream, connection_label)
+        };
+        Ok(connection)
+    }
+
+    /// Connect to one exec-server through an authenticated rendezvous stream
+    /// using a caller-supplied single-use authorization bundle.
+    ///
+    /// The executor key is pinned before JSON-RPC starts; the websocket carries
+    /// only ciphertext after that. Environment-managed connections use a
+    /// retained [`NoiseRendezvousConnectProvider`] so recovery can fetch a fresh
+    /// bundle for each reconnect.
+    #[tracing::instrument(
+        name = "codex.exec_server.remote.harness.connect",
+        skip_all,
+        fields(
+            otel.kind = "client",
+            otel.name = "codex.exec_server.remote.harness.connect",
+        )
+    )]
+    pub async fn connect_noise_rendezvous(
+        args: NoiseRendezvousConnectArgs,
+    ) -> Result<Self, ExecServerError> {
+        let opened = Self::open_noise_rendezvous_connection(args).await?;
+        let ready = Self::finish_noise_rendezvous_connection(opened).await?;
+        Self::connect_with_recovery_and_noise_context(
+            ready.connection,
+            ready.options,
+            /*reconnect_strategy*/ None,
+            ready.noise_context,
+        )
+        .await
+    }
+
+    #[tracing::instrument(
+        name = "codex.exec_server.remote.noise.websocket_connect",
+        skip_all,
+        fields(
+            otel.kind = "client",
+            otel.name = "codex.exec_server.remote.noise.websocket_connect",
+            environment_id = %args.bundle.environment_id,
+            executor_registration_id = %args.bundle.executor_registration_id,
+        )
+    )]
+    async fn open_noise_rendezvous_connection(
+        args: NoiseRendezvousConnectArgs,
+    ) -> Result<OpenNoiseRendezvousConnection, ExecServerError> {
+        ensure_rustls_crypto_provider();
+        // Keep the registry-issued URL, key, and authorization together for this
+        // connection attempt.
+        let NoiseRendezvousConnectArgs {
+            bundle,
+            harness_identity,
+            client_name,
+            connect_timeout,
+            initialize_timeout,
+            resume_session_id,
+            http_client_factory,
+        } = args;
+        let NoiseRendezvousConnectBundle {
+            websocket_url,
+            environment_id,
+            executor_registration_id,
+            executor_public_key,
+            harness_key_authorization,
+        } = bundle;
+        let diagnostic_url = websocket_url
+            .split(['?', '#'])
+            .next()
+            .unwrap_or(websocket_url.as_str())
+            .to_string();
+        let mut request = websocket_url
+            .as_str()
+            .into_client_request()
+            .map_err(|source| ExecServerError::WebSocketConnect {
+                url: diagnostic_url.clone(),
+                source,
+            })?;
+        request.headers_mut().extend(current_rendezvous_headers());
+        let (stream, _) = timeout(
+            connect_timeout,
+            WebSocketConnector::new_with_tls_mode(
+                &http_client_factory,
+                WebSocketTlsMode::TungsteniteDefault,
+            )
+            .map_err(|error| ExecServerError::WebSocketConfiguration(error.to_string()))?
+            .with_tcp_nodelay()
+            .connect(request, noise_relay_websocket_config()),
+        )
+        .await
+        .map_err(|_| ExecServerError::WebSocketConnectTimeout {
+            url: diagnostic_url.clone(),
+            timeout: connect_timeout,
+        })?
+        .map_err(|source| ExecServerError::WebSocketConnect {
+            url: diagnostic_url.clone(),
+            source,
+        })?;
+
+        let connection_label = format!("Noise exec-server rendezvous websocket {diagnostic_url}");
+        let connection = noise_harness_connection_from_websocket_with_readiness(
+            stream,
+            NoiseHarnessConnectionArgs {
+                connection_label,
+                environment_id,
+                executor_registration_id,
+                identity: harness_identity,
+                responder_public_key: executor_public_key,
+                harness_key_authorization,
+            },
+        );
+        Ok(OpenNoiseRendezvousConnection {
+            connection: connection.connection,
+            options: ExecServerClientConnectOptions {
+                client_name,
+                initialize_timeout,
+                resume_session_id,
+            },
+            handshake_ready: connection.handshake_ready,
+        })
+    }
+
+    #[tracing::instrument(
+        name = "codex.exec_server.remote.noise.handshake",
+        skip_all,
+        parent = initialize_span,
+        fields(
+            otel.kind = "client",
+            otel.name = "codex.exec_server.remote.noise.handshake",
+        )
+    )]
+    async fn wait_for_noise_handshake(
+        handshake_ready: &mut tokio::sync::oneshot::Receiver<()>,
+        deadline: Instant,
+        initialize_timeout: Duration,
+        initialize_span: &tracing::Span,
+    ) -> Result<(), ExecServerError> {
+        match timeout_at(deadline, handshake_ready).await {
+            Ok(Ok(())) => Ok(()),
+            Ok(Err(_)) => Err(ExecServerError::Disconnected(
+                "Noise harness handshake failed before connection became ready".to_string(),
+            )),
+            Err(_) => Err(ExecServerError::InitializeTimedOut {
+                timeout: initialize_timeout,
+            }),
+        }
+    }
+
+    async fn finish_noise_rendezvous_connection(
+        mut connection: OpenNoiseRendezvousConnection,
+    ) -> Result<ReadyNoiseRendezvousConnection, ExecServerError> {
+        // Preserve the legacy initialize request span as the post-WebSocket
+        // startup parent while making its two child operations visible.
+        let initialize_timeout = connection.options.initialize_timeout;
+        let noise_context = NoiseInitializeContext {
+            span: tracing::info_span!(
+                "codex.exec_server.request",
+                otel.kind = "client",
+                otel.name = "initialize",
+                method = "initialize",
+            ),
+            timeout_for_error: initialize_timeout,
+        };
+        let deadline = Instant::now() + initialize_timeout;
+        let readiness = Self::wait_for_noise_handshake(
+            &mut connection.handshake_ready,
+            deadline,
+            initialize_timeout,
+            &noise_context.span,
+        )
+        .await;
+        if let Err(error) = readiness {
+            // Unlike the normal connect path, the connection has not reached
+            // RpcClient yet, so its Drop implementation cannot abort the
+            // transport task for us.
+            connection.connection.transport.terminate();
+            for task in &connection.connection.task_handles {
+                task.abort();
+            }
+            return Err(error);
+        }
+        let mut options = connection.options;
+        options.initialize_timeout = deadline.saturating_duration_since(Instant::now());
+        Ok(ReadyNoiseRendezvousConnection {
+            connection: connection.connection,
+            options,
+            noise_context,
+        })
+    }
+
+    pub(crate) async fn connect_stdio_command(
+        args: StdioExecServerConnectArgs,
+    ) -> Result<Self, ExecServerError> {
+        let mut child = stdio_command_process(&args.command)
+            .stdin(Stdio::piped())
+            .stdout(Stdio::piped())
+            .stderr(Stdio::piped())
+            .spawn()
+            .map_err(ExecServerError::Spawn)?;
+
+        let stdin = child.stdin.take().ok_or_else(|| {
+            ExecServerError::Protocol("spawned exec-server command has no stdin".to_string())
+        })?;
+        let stdout = child.stdout.take().ok_or_else(|| {
+            ExecServerError::Protocol("spawned exec-server command has no stdout".to_string())
+        })?;
+        if let Some(stderr) = child.stderr.take() {
+            tokio::spawn(async move {
+                let mut lines = BufReader::new(stderr).lines();
+                loop {
+                    match lines.next_line().await {
+                        Ok(Some(line)) => debug!("exec-server stdio stderr: {line}"),
+                        Ok(None) => break,
+                        Err(err) => {
+                            warn!("failed to read exec-server stdio stderr: {err}");
+                            break;
+                        }
+                    }
+                }
+            });
+        }
+
+        Self::connect(
+            JsonRpcConnection::from_stdio(stdout, stdin, "exec-server stdio command".to_string())
+                .with_child_process(child),
+            args.into(),
+        )
+        .await
+    }
+}
+
+fn is_rendezvous_harness_url(websocket_url: &str) -> bool {
+    let Some((_path, query)) = websocket_url.split_once('?') else {
+        return false;
+    };
+    query
+        .split('&')
+        .filter_map(|pair| pair.split_once('='))
+        .any(|(key, value)| key == "role" && value == "harness")
+}
+
+fn stdio_command_process(stdio_command: &StdioExecServerCommand) -> Command {
+    let mut command = Command::new(&stdio_command.program);
+    command.args(&stdio_command.args);
+    command.envs(&stdio_command.env);
+    scrub_non_inheritable_env_vars(command.as_std_mut());
+    if let Some(cwd) = &stdio_command.cwd {
+        command.current_dir(cwd);
+    }
+    #[cfg(unix)]
+    command.process_group(0);
+    command
+}
+
+#[cfg(test)]
+#[path = "client_transport_tests.rs"]
+mod tests;

codex-rs/exec-server/src/client_transport_tests.rs

@@ -0,0 +1,581 @@
+use std::collections::VecDeque;
+use std::future::Future;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+use anyhow::Result;
+use codex_exec_server_protocol::JSONRPCMessage;
+use futures::FutureExt;
+use futures::SinkExt;
+use futures::StreamExt;
+use futures::future::BoxFuture;
+use pretty_assertions::assert_eq;
+use tokio::io::AsyncBufReadExt;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+use tokio::io::BufReader;
+use tokio::io::duplex;
+use tokio::net::TcpListener;
+use tokio_tungstenite::accept_async;
+use tokio_tungstenite::tungstenite::Message;
+
+use super::ExecServerClient;
+use super::ExecServerReconnectStrategy;
+use super::INITIAL_REGISTRY_MAX_RETRIES;
+use super::INITIAL_REGISTRY_OPERATION_TIMEOUT;
+use super::INITIAL_REGISTRY_REQUEST_TIMEOUT;
+use crate::ExecServerError;
+use crate::NoiseChannelIdentity;
+use crate::NoiseChannelPublicKey;
+use crate::NoiseRendezvousConnectArgs;
+use crate::NoiseRendezvousConnectBundle;
+use crate::NoiseRendezvousConnectProvider;
+use crate::client::NoiseInitializeContext;
+use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
+use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT;
+use crate::client_api::ExecServerClientConnectOptions;
+use crate::connection::JsonRpcConnection;
+use crate::noise_channel::PendingResponderHandshake;
+use crate::noise_channel::noise_channel_prologue;
+use crate::protocol::INITIALIZE_METHOD;
+use crate::relay::RelayFrameBodyKind;
+use crate::relay::decode_relay_message_frame;
+use crate::relay::encode_relay_message_frame;
+use crate::relay_proto::RelayMessageFrame;
+
+#[derive(Default)]
+struct SequenceNoiseConnectProvider {
+    bundles:
+        Mutex<VecDeque<BoxFuture<'static, Result<NoiseRendezvousConnectBundle, ExecServerError>>>>,
+    returned_urls: Mutex<Vec<String>>,
+    requested_keys: Mutex<Vec<NoiseChannelPublicKey>>,
+}
+
+impl SequenceNoiseConnectProvider {
+    fn push_response(
+        &self,
+        response: impl Future<Output = Result<NoiseRendezvousConnectBundle, ExecServerError>>
+        + Send
+        + 'static,
+    ) {
+        self.bundles.lock().unwrap().push_back(response.boxed());
+    }
+
+    fn push_error(&self, error: ExecServerError) {
+        self.push_response(futures::future::ready(Err(error)));
+    }
+
+    fn push_pending(&self) {
+        self.push_response(futures::future::pending());
+    }
+
+    fn requested_keys(&self) -> Vec<NoiseChannelPublicKey> {
+        self.requested_keys.lock().unwrap().clone()
+    }
+
+    fn assert_requested_identity(&self, identity: &NoiseChannelIdentity, requests: usize) {
+        assert_eq!(self.requested_keys(), vec![identity.public_key(); requests]);
+    }
+
+    fn returned_urls(&self) -> Vec<String> {
+        self.returned_urls
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .clone()
+    }
+
+    async fn connect(
+        self: &Arc<Self>,
+        identity: &NoiseChannelIdentity,
+    ) -> Result<
+        (
+            super::JsonRpcConnection,
+            super::ExecServerClientConnectOptions,
+        ),
+        ExecServerError,
+    > {
+        let provider: Arc<dyn NoiseRendezvousConnectProvider> = self.clone();
+        ExecServerClient::open_initial_noise_rendezvous_connection(
+            &provider,
+            identity,
+            codex_http_client::HttpClientFactory::new(
+                codex_http_client::OutboundProxyPolicy::ReqwestDefault,
+            ),
+        )
+        .await
+        .map(|(ready, _)| (ready.connection, ready.options))
+    }
+}
+
+impl NoiseRendezvousConnectProvider for SequenceNoiseConnectProvider {
+    fn connect_bundle(
+        &self,
+        harness_public_key: NoiseChannelPublicKey,
+    ) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
+        self.requested_keys.lock().unwrap().push(harness_public_key);
+        let response = self
+            .bundles
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .pop_front()
+            .expect("test Noise provider exhausted");
+        Box::pin(async move {
+            let result = response.await;
+            if let Ok(bundle) = &result {
+                self.returned_urls
+                    .lock()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner)
+                    .push(bundle.websocket_url.clone());
+            }
+            result
+        })
+    }
+}
+
+fn test_bundle(websocket_url: String) -> Result<NoiseRendezvousConnectBundle> {
+    Ok(NoiseRendezvousConnectBundle {
+        websocket_url,
+        environment_id: "environment".to_string(),
+        executor_registration_id: "registration".to_string(),
+        executor_public_key: NoiseChannelIdentity::generate()?.public_key(),
+        harness_key_authorization: "authorization".to_string(),
+    })
+}
+
+fn registry_error(status: http::StatusCode, code: &str) -> ExecServerError {
+    ExecServerError::EnvironmentRegistryHttp {
+        status,
+        code: Some(code.to_string()),
+        message: "registry unavailable".to_string(),
+    }
+}
+
+#[tokio::test]
+async fn noise_handshake_uses_initialize_timeout() -> Result<()> {
+    let listener = TcpListener::bind("127.0.0.1:0").await?;
+    let websocket_url = format!("ws://{}", listener.local_addr()?);
+    let server = tokio::spawn(async move {
+        let (socket, _) = listener.accept().await?;
+        let mut websocket = accept_async(socket).await?;
+        // Drain the frames sent before the harness waits for the responder,
+        // then verify that a timed-out readiness wait closes the socket.
+        assert!(websocket.next().await.is_some());
+        assert!(websocket.next().await.is_some());
+        let closed =
+            tokio::time::timeout(std::time::Duration::from_secs(1), websocket.next()).await?;
+        assert!(
+            matches!(closed, None | Some(Ok(Message::Close(_))) | Some(Err(_))),
+            "timed-out Noise handshake must close its websocket"
+        );
+        anyhow::Ok(())
+    });
+    let initialize_timeout = std::time::Duration::from_millis(1);
+    let opened = ExecServerClient::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs {
+        bundle: NoiseRendezvousConnectBundle {
+            websocket_url,
+            environment_id: "environment".to_string(),
+            executor_registration_id: "registration".to_string(),
+            executor_public_key: NoiseChannelIdentity::generate()?.public_key(),
+            harness_key_authorization: "authorization".to_string(),
+        },
+        harness_identity: NoiseChannelIdentity::generate()?,
+        client_name: "test".to_string(),
+        connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
+        initialize_timeout,
+        resume_session_id: None,
+        http_client_factory: codex_http_client::HttpClientFactory::new(
+            codex_http_client::OutboundProxyPolicy::ReqwestDefault,
+        ),
+    })
+    .await?;
+
+    let error = ExecServerClient::finish_noise_rendezvous_connection(opened)
+        .await
+        .err()
+        .expect("stalled Noise handshake must time out");
+    assert!(matches!(
+        error,
+        ExecServerError::InitializeTimedOut { timeout } if timeout == initialize_timeout
+    ));
+
+    server.await??;
+    Ok(())
+}
+
+#[tokio::test]
+async fn deferred_initialize_timeout_reports_configured_budget() {
+    let (client_stdin, server_reader) = duplex(1 << 20);
+    let (server_writer, client_stdout) = duplex(1 << 20);
+    let server = tokio::spawn(async move {
+        let _server_writer = server_writer;
+        let mut lines = BufReader::new(server_reader).lines();
+        let line = lines
+            .next_line()
+            .await
+            .expect("initialize read should succeed")
+            .expect("initialize request should arrive");
+        let request: JSONRPCMessage =
+            serde_json::from_str(&line).expect("initialize request should parse");
+        assert!(
+            matches!(
+                request,
+                JSONRPCMessage::Request(ref request) if request.method == INITIALIZE_METHOD
+            ),
+            "expected initialize request, got {request:?}"
+        );
+        futures::future::pending::<()>().await;
+    });
+    let configured_timeout = std::time::Duration::from_secs(10);
+    let error = ExecServerClient::connect_with_recovery_and_noise_context(
+        JsonRpcConnection::from_stdio(
+            client_stdout,
+            client_stdin,
+            "timeout-test-client".to_string(),
+        ),
+        ExecServerClientConnectOptions {
+            client_name: "timeout-test-client".to_string(),
+            initialize_timeout: std::time::Duration::from_millis(1),
+            resume_session_id: None,
+        },
+        /*reconnect_strategy*/ None,
+        NoiseInitializeContext {
+            span: tracing::info_span!("codex.exec_server.request"),
+            timeout_for_error: configured_timeout,
+        },
+    )
+    .await
+    .err()
+    .expect("initialize RPC must time out");
+    assert!(matches!(
+        error,
+        ExecServerError::InitializeTimedOut { timeout } if timeout == configured_timeout
+    ));
+    server.abort();
+    let _ = server.await;
+}
+
+#[tokio::test(start_paused = true)]
+async fn initial_noise_connection_bounds_offline_retries() -> Result<()> {
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    for _ in 0..=INITIAL_REGISTRY_MAX_RETRIES {
+        sequence.push_error(registry_error(
+            http::StatusCode::CONFLICT,
+            "environment_offline",
+        ));
+    }
+    let identity = NoiseChannelIdentity::generate()?;
+    let started = tokio::time::Instant::now();
+    let error = sequence
+        .connect(&identity)
+        .await
+        .err()
+        .expect("offline retries must end");
+
+    assert!(crate::client::is_environment_offline_error(&error));
+    let requests = sequence.requested_keys().len();
+    assert!((4..=INITIAL_REGISTRY_MAX_RETRIES as usize + 1).contains(&requests));
+    sequence.assert_requested_identity(&identity, requests);
+    assert!(started.elapsed() <= INITIAL_REGISTRY_OPERATION_TIMEOUT);
+    Ok(())
+}
+
+#[tokio::test(start_paused = true)]
+async fn initial_noise_connection_bounds_a_stalled_retry_request() -> Result<()> {
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    sequence.push_error(registry_error(
+        http::StatusCode::CONFLICT,
+        "environment_offline",
+    ));
+    for _ in 0..INITIAL_REGISTRY_MAX_RETRIES {
+        sequence.push_pending();
+    }
+    let identity = NoiseChannelIdentity::generate()?;
+    let started = tokio::time::Instant::now();
+    let error = sequence
+        .connect(&identity)
+        .await
+        .err()
+        .expect("stalled retry must time out");
+
+    assert!(matches!(
+        error,
+        ExecServerError::EnvironmentRegistryRequest(error) if error.is_timeout()
+    ));
+    assert_eq!(started.elapsed(), INITIAL_REGISTRY_OPERATION_TIMEOUT);
+    let requests = sequence.requested_keys().len();
+    assert!((2..=3).contains(&requests));
+    sequence.assert_requested_identity(&identity, requests);
+    Ok(())
+}
+
+#[tokio::test(start_paused = true)]
+async fn initial_noise_connection_bounds_a_stalled_initial_request() -> Result<()> {
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    for _ in 0..=INITIAL_REGISTRY_MAX_RETRIES {
+        sequence.push_pending();
+    }
+    let identity = NoiseChannelIdentity::generate()?;
+    let started = tokio::time::Instant::now();
+
+    let error = sequence
+        .connect(&identity)
+        .await
+        .err()
+        .expect("stalled initial request must time out");
+
+    assert!(matches!(
+        error,
+        ExecServerError::EnvironmentRegistryRequest(error) if error.is_timeout()
+    ));
+    assert_eq!(started.elapsed(), INITIAL_REGISTRY_OPERATION_TIMEOUT);
+    let requests = sequence.requested_keys().len();
+    assert!((2..=3).contains(&requests));
+    sequence.assert_requested_identity(&identity, requests);
+    Ok(())
+}
+
+#[tokio::test(start_paused = true)]
+async fn initial_noise_connection_retries_a_stalled_initial_request() -> Result<()> {
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    sequence.push_pending();
+    sequence.push_error(registry_error(http::StatusCode::FORBIDDEN, "forbidden"));
+    let identity = NoiseChannelIdentity::generate()?;
+    let started = tokio::time::Instant::now();
+
+    let error = sequence
+        .connect(&identity)
+        .await
+        .err()
+        .expect("terminal response must stop the retry sequence");
+
+    assert!(matches!(
+        error,
+        ExecServerError::EnvironmentRegistryHttp {
+            status: http::StatusCode::FORBIDDEN,
+            ..
+        }
+    ));
+    assert!(started.elapsed() >= INITIAL_REGISTRY_REQUEST_TIMEOUT);
+    assert!(started.elapsed() < INITIAL_REGISTRY_OPERATION_TIMEOUT);
+    sequence.assert_requested_identity(&identity, /*requests*/ 2);
+    Ok(())
+}
+
+#[tokio::test(start_paused = true)]
+async fn initial_noise_connection_retries_transient_registry_statuses() -> Result<()> {
+    for status in [
+        http::StatusCode::REQUEST_TIMEOUT,
+        http::StatusCode::TOO_MANY_REQUESTS,
+        http::StatusCode::INTERNAL_SERVER_ERROR,
+        http::StatusCode::BAD_GATEWAY,
+        http::StatusCode::SERVICE_UNAVAILABLE,
+    ] {
+        let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+        sequence.push_error(registry_error(status, "temporarily_unavailable"));
+        sequence.push_error(registry_error(http::StatusCode::FORBIDDEN, "forbidden"));
+        let identity = NoiseChannelIdentity::generate()?;
+
+        let error = sequence
+            .connect(&identity)
+            .await
+            .err()
+            .expect("terminal response must stop the retry sequence");
+
+        assert!(matches!(
+            error,
+            ExecServerError::EnvironmentRegistryHttp {
+                status: http::StatusCode::FORBIDDEN,
+                ..
+            }
+        ));
+        sequence.assert_requested_identity(&identity, /*requests*/ 2);
+    }
+    Ok(())
+}
+
+#[tokio::test(start_paused = true)]
+async fn initial_noise_connection_retries_registry_request_timeouts() -> Result<()> {
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    sequence.push_error(ExecServerError::EnvironmentRegistryRequest(
+        codex_http_client::RouteAwareRequestError::Timeout,
+    ));
+    sequence.push_error(registry_error(http::StatusCode::FORBIDDEN, "forbidden"));
+    let identity = NoiseChannelIdentity::generate()?;
+
+    let error = sequence
+        .connect(&identity)
+        .await
+        .err()
+        .expect("terminal response must stop the retry sequence");
+
+    assert!(matches!(
+        error,
+        ExecServerError::EnvironmentRegistryHttp {
+            status: http::StatusCode::FORBIDDEN,
+            ..
+        }
+    ));
+    sequence.assert_requested_identity(&identity, /*requests*/ 2);
+    Ok(())
+}
+
+#[tokio::test(start_paused = true)]
+async fn initial_noise_connection_does_not_retry_permanent_registry_errors() -> Result<()> {
+    for (status, code) in [
+        (http::StatusCode::UNAUTHORIZED, "unauthorized"),
+        (http::StatusCode::FORBIDDEN, "forbidden"),
+        (http::StatusCode::BAD_REQUEST, "bad_request"),
+        (http::StatusCode::NOT_FOUND, "environment_not_found"),
+        (http::StatusCode::CONFLICT, "registration_conflict"),
+        (http::StatusCode::CONFLICT, "route_unavailable"),
+    ] {
+        // A terminal error must also stop a retry sequence already in progress.
+        for initial_offline in [false, true] {
+            let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+            if initial_offline {
+                sequence.push_error(registry_error(
+                    http::StatusCode::CONFLICT,
+                    "environment_offline",
+                ));
+            }
+            sequence.push_error(registry_error(status, code));
+            let identity = NoiseChannelIdentity::generate()?;
+            let error = sequence
+                .connect(&identity)
+                .await
+                .err()
+                .expect("other errors must propagate");
+            assert!(
+                matches!(error, ExecServerError::EnvironmentRegistryHttp { status: actual_status, code: Some(actual_code), .. } if actual_status == status && actual_code == code)
+            );
+            sequence.assert_requested_identity(&identity, 1 + usize::from(initial_offline));
+        }
+    }
+    Ok(())
+}
+
+#[tokio::test(start_paused = true)]
+async fn noise_session_resume_leaves_offline_retries_to_recovery() -> Result<()> {
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    sequence.push_error(registry_error(
+        http::StatusCode::CONFLICT,
+        "environment_offline",
+    ));
+    let identity = NoiseChannelIdentity::generate()?;
+    let strategy = ExecServerReconnectStrategy::NoiseRendezvous {
+        executor_public_key: NoiseChannelIdentity::generate()?.public_key(),
+        provider: sequence.clone(),
+        identity: identity.clone(),
+        client_name: "test".to_string(),
+        connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
+        initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
+        http_client_factory: codex_http_client::HttpClientFactory::new(
+            codex_http_client::OutboundProxyPolicy::ReqwestDefault,
+        ),
+    };
+    let started = tokio::time::Instant::now();
+    let error = strategy
+        .resume("session")
+        .await
+        .err()
+        .expect("resume must return the offline error");
+    assert!(crate::client::is_environment_offline_error(&error));
+    assert_eq!(started.elapsed(), std::time::Duration::ZERO);
+    sequence.assert_requested_identity(&identity, /*requests*/ 1);
+    Ok(())
+}
+
+#[tokio::test]
+async fn initial_noise_connection_refreshes_bundle_after_exhausting_initial_retries() -> Result<()>
+{
+    let unauthorized_listener = TcpListener::bind("127.0.0.1:0").await?;
+    let unauthorized_url = format!("ws://{}", unauthorized_listener.local_addr()?);
+    let unauthorized_server = tokio::spawn(async move {
+        let (mut socket, _) = unauthorized_listener.accept().await?;
+        let mut request = [0_u8; 4096];
+        let _ = socket.read(&mut request).await?;
+        socket
+            .write_all(
+                b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
+            )
+            .await?;
+        socket.shutdown().await?;
+        anyhow::Ok(())
+    });
+    let accepted_listener = TcpListener::bind("127.0.0.1:0").await?;
+    let accepted_url = format!("ws://{}", accepted_listener.local_addr()?);
+    let executor_identity = NoiseChannelIdentity::generate()?;
+    let executor_public_key = executor_identity.public_key();
+    let accepted_server = tokio::spawn(async move {
+        let (socket, _) = accepted_listener.accept().await?;
+        let mut websocket = accept_async(socket).await?;
+        let Message::Binary(resume_payload) = websocket.next().await.unwrap()? else {
+            anyhow::bail!("expected Noise relay resume frame");
+        };
+        let resume = decode_relay_message_frame(resume_payload.as_ref())?;
+        assert_eq!(resume.validate()?, RelayFrameBodyKind::Resume);
+        let Message::Binary(handshake_payload) = websocket.next().await.unwrap()? else {
+            anyhow::bail!("expected Noise relay handshake frame");
+        };
+        let handshake = decode_relay_message_frame(handshake_payload.as_ref())?;
+        let stream_id = handshake.stream_id.clone();
+        let prologue = noise_channel_prologue("environment", "registration", &stream_id);
+        let pending = PendingResponderHandshake::read_request(
+            &executor_identity,
+            &prologue,
+            &handshake.into_handshake_payload()?,
+        )?;
+        let (_transport, response) = pending.complete()?;
+        websocket
+            .send(Message::Binary(
+                encode_relay_message_frame(&RelayMessageFrame::handshake(stream_id, response))
+                    .into(),
+            ))
+            .await?;
+        anyhow::Ok(())
+    });
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    let unauthorized_bundle = test_bundle(unauthorized_url.clone())?;
+    let mut accepted_bundle = test_bundle(accepted_url.clone())?;
+    accepted_bundle.executor_public_key = executor_public_key;
+    sequence.push_response(async {
+        tokio::time::pause();
+        Err(registry_error(
+            http::StatusCode::CONFLICT,
+            "environment_offline",
+        ))
+    });
+    for _ in 1..INITIAL_REGISTRY_MAX_RETRIES {
+        sequence.push_error(registry_error(
+            http::StatusCode::CONFLICT,
+            "environment_offline",
+        ));
+    }
+    sequence.push_response(async move {
+        tokio::time::resume();
+        Ok(unauthorized_bundle)
+    });
+    sequence.push_response(async {
+        tokio::time::pause();
+        Err(registry_error(
+            http::StatusCode::CONFLICT,
+            "environment_offline",
+        ))
+    });
+    sequence.push_response(async move {
+        tokio::time::resume();
+        Ok(accepted_bundle)
+    });
+    let identity = NoiseChannelIdentity::generate()?;
+
+    let _connection = sequence.connect(&identity).await?;
+
+    assert_eq!(
+        sequence.returned_urls(),
+        vec![unauthorized_url, accepted_url]
+    );
+    sequence.assert_requested_identity(&identity, INITIAL_REGISTRY_MAX_RETRIES as usize + 3);
+    unauthorized_server.await??;
+    accepted_server.await??;
+    Ok(())
+}

codex-rs/exec-server/src/connection.rs

@@ -0,0 +1,1042 @@
+#[cfg(windows)]
+use std::process::Stdio;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+use std::time::Instant;
+
+use axum::extract::ws::Message as AxumWebSocketMessage;
+use axum::extract::ws::WebSocket as AxumWebSocket;
+use codex_exec_server_protocol::JSONRPCMessage;
+use codex_exec_server_protocol::JSONRPCRequest;
+use futures::Sink;
+use futures::SinkExt;
+use futures::Stream;
+use futures::StreamExt;
+use tokio::io::AsyncRead;
+use tokio::io::AsyncWrite;
+use tokio::process::Child;
+use tokio::sync::mpsc;
+use tokio::sync::watch;
+use tokio::time::timeout;
+#[cfg(test)]
+use tokio_tungstenite::WebSocketStream;
+use tokio_tungstenite::tungstenite::Message;
+use tracing::debug;
+use tracing::warn;
+
+use tokio::io::AsyncBufReadExt;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+use tokio::io::BufReader;
+use tokio::io::BufWriter;
+
+pub(crate) const CHANNEL_CAPACITY: usize = 128;
+// Match the existing serialized JSON-RPC message ceiling used by Noise and
+// WebSocket transports so stdio has the same per-message bound.
+const MAX_STDIO_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024;
+const STDIO_TERMINATION_GRACE_PERIOD: Duration = Duration::from_secs(2);
+#[cfg(test)]
+pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_millis(25);
+#[cfg(not(test))]
+pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
+
+#[derive(Debug)]
+pub(crate) enum JsonRpcConnectionEvent {
+    Message(JSONRPCMessage),
+    QueuedRequest {
+        request: JSONRPCRequest,
+        request_span: tracing::Span,
+        queued_at: Instant,
+    },
+    MalformedMessage {
+        reason: String,
+    },
+    Disconnected {
+        reason: Option<String>,
+    },
+}
+
+impl JsonRpcConnectionEvent {
+    pub(crate) fn message(message: JSONRPCMessage) -> Self {
+        let JSONRPCMessage::Request(request) = message else {
+            return Self::Message(message);
+        };
+
+        let queued_at = Instant::now();
+        let request_span = tracing::info_span!(
+            "codex.exec_server.request",
+            otel.kind = "server",
+            otel.name = "unknown",
+            method = request.method.as_str(),
+            result = tracing::field::Empty,
+        );
+        if let Some(trace) = &request.trace
+            && !codex_otel::set_parent_from_w3c_trace_context(&request_span, trace)
+        {
+            warn!(
+                method = request.method.as_str(),
+                "ignoring invalid inbound exec-server trace carrier"
+            );
+        }
+
+        Self::QueuedRequest {
+            request,
+            request_span,
+            queued_at,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub(crate) enum JsonRpcTransport {
+    // Plain means no child process; transport bytes may still be encrypted.
+    Plain,
+    Stdio { transport: StdioTransport },
+}
+
+impl JsonRpcTransport {
+    fn from_child_process(child_process: Child) -> Self {
+        Self::Stdio {
+            transport: StdioTransport::spawn(child_process),
+        }
+    }
+
+    pub(crate) fn terminate(&self) {
+        match self {
+            Self::Plain => {}
+            Self::Stdio { transport } => transport.terminate(),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub(crate) struct StdioTransport {
+    handle: Arc<StdioTransportHandle>,
+}
+
+struct StdioTransportHandle {
+    terminate_tx: watch::Sender<bool>,
+    terminate_requested: AtomicBool,
+}
+
+impl StdioTransport {
+    fn spawn(child_process: Child) -> Self {
+        let (terminate_tx, terminate_rx) = watch::channel(false);
+        let handle = Arc::new(StdioTransportHandle {
+            terminate_tx,
+            terminate_requested: AtomicBool::new(false),
+        });
+        spawn_stdio_child_supervisor(child_process, terminate_rx);
+        Self { handle }
+    }
+
+    fn terminate(&self) {
+        self.handle.terminate();
+    }
+}
+
+impl StdioTransportHandle {
+    fn terminate(&self) {
+        if !self.terminate_requested.swap(true, Ordering::AcqRel) {
+            let _ = self.terminate_tx.send(true);
+        }
+    }
+}
+
+impl Drop for StdioTransportHandle {
+    fn drop(&mut self) {
+        self.terminate();
+    }
+}
+
+fn spawn_stdio_child_supervisor(mut child_process: Child, mut terminate_rx: watch::Receiver<bool>) {
+    let process_group_id = child_process.id();
+    tokio::spawn(async move {
+        tokio::select! {
+            result = child_process.wait() => {
+                log_stdio_child_wait_result(result);
+                kill_process_tree(&mut child_process, process_group_id);
+            }
+            () = wait_for_stdio_termination(&mut terminate_rx) => {
+                terminate_stdio_child(&mut child_process, process_group_id).await;
+            }
+        }
+    });
+}
+
+async fn wait_for_stdio_termination(terminate_rx: &mut watch::Receiver<bool>) {
+    loop {
+        if *terminate_rx.borrow() {
+            return;
+        }
+        if terminate_rx.changed().await.is_err() {
+            return;
+        }
+    }
+}
+
+async fn terminate_stdio_child(child_process: &mut Child, process_group_id: Option<u32>) {
+    terminate_process_tree(child_process, process_group_id);
+    match timeout(STDIO_TERMINATION_GRACE_PERIOD, child_process.wait()).await {
+        Ok(result) => {
+            log_stdio_child_wait_result(result);
+        }
+        Err(_) => {
+            kill_process_tree(child_process, process_group_id);
+            log_stdio_child_wait_result(child_process.wait().await);
+        }
+    }
+}
+
+fn terminate_process_tree(child_process: &mut Child, process_group_id: Option<u32>) {
+    let Some(process_group_id) = process_group_id else {
+        kill_direct_child(child_process, "terminate");
+        return;
+    };
+
+    #[cfg(unix)]
+    if let Err(err) = codex_utils_pty::process_group::terminate_process_group(process_group_id) {
+        warn!("failed to terminate exec-server stdio process group {process_group_id}: {err}");
+        kill_direct_child(child_process, "terminate");
+    }
+
+    #[cfg(windows)]
+    if !kill_windows_process_tree(process_group_id) {
+        kill_direct_child(child_process, "terminate");
+    }
+
+    #[cfg(not(any(unix, windows)))]
+    {
+        let _ = process_group_id;
+        kill_direct_child(child_process, "terminate");
+    }
+}
+
+fn kill_process_tree(child_process: &mut Child, process_group_id: Option<u32>) {
+    let Some(process_group_id) = process_group_id else {
+        kill_direct_child(child_process, "kill");
+        return;
+    };
+
+    #[cfg(unix)]
+    if let Err(err) = codex_utils_pty::process_group::kill_process_group(process_group_id) {
+        warn!("failed to kill exec-server stdio process group {process_group_id}: {err}");
+    }
+
+    #[cfg(windows)]
+    if !kill_windows_process_tree(process_group_id) {
+        kill_direct_child(child_process, "kill");
+    }
+
+    #[cfg(not(any(unix, windows)))]
+    {
+        let _ = process_group_id;
+        kill_direct_child(child_process, "kill");
+    }
+}
+
+fn kill_direct_child(child_process: &mut Child, action: &str) {
+    if let Err(err) = child_process.start_kill() {
+        debug!("failed to {action} exec-server stdio child: {err}");
+    }
+}
+
+#[cfg(windows)]
+fn kill_windows_process_tree(pid: u32) -> bool {
+    let pid = pid.to_string();
+    match std::process::Command::new("taskkill")
+        .args(["/PID", pid.as_str(), "/T", "/F"])
+        .stdin(Stdio::null())
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .status()
+    {
+        Ok(status) => status.success(),
+        Err(err) => {
+            warn!("failed to run taskkill for exec-server stdio process tree {pid}: {err}");
+            false
+        }
+    }
+}
+
+fn log_stdio_child_wait_result(result: std::io::Result<std::process::ExitStatus>) {
+    if let Err(err) = result {
+        debug!("failed to wait for exec-server stdio child: {err}");
+    }
+}
+
+pub(crate) struct JsonRpcConnection {
+    pub(crate) outgoing_tx: mpsc::Sender<JSONRPCMessage>,
+    pub(crate) incoming_rx: mpsc::Receiver<JsonRpcConnectionEvent>,
+    pub(crate) disconnected_rx: watch::Receiver<bool>,
+    pub(crate) task_handles: Vec<tokio::task::JoinHandle<()>>,
+    pub(crate) transport: JsonRpcTransport,
+}
+
+impl JsonRpcConnection {
+    pub(crate) fn from_stdio<R, W>(reader: R, writer: W, connection_label: String) -> Self
+    where
+        R: AsyncRead + Unpin + Send + 'static,
+        W: AsyncWrite + Unpin + Send + 'static,
+    {
+        Self::from_stdio_with_max_message_len(
+            reader,
+            writer,
+            connection_label,
+            MAX_STDIO_JSONRPC_MESSAGE_LEN,
+        )
+    }
+
+    fn from_stdio_with_max_message_len<R, W>(
+        reader: R,
+        writer: W,
+        connection_label: String,
+        max_message_len: usize,
+    ) -> Self
+    where
+        R: AsyncRead + Unpin + Send + 'static,
+        W: AsyncWrite + Unpin + Send + 'static,
+    {
+        let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
+        let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
+        let (disconnected_tx, disconnected_rx) = watch::channel(false);
+
+        let reader_label = connection_label.clone();
+        let incoming_tx_for_reader = incoming_tx.clone();
+        let disconnected_tx_for_reader = disconnected_tx.clone();
+        // Read one byte past the payload limit so an unterminated oversized
+        // message fails promptly. A trailing CR gets one more byte of lookahead
+        // because it may be the first half of a valid CRLF terminator.
+        let read_limit = u64::try_from(max_message_len.saturating_add(1)).unwrap_or(u64::MAX);
+        let reader_task = tokio::spawn(async move {
+            let mut reader = BufReader::new(reader);
+            let mut line = String::new();
+            loop {
+                line.clear();
+                let read_result = (&mut reader).take(read_limit).read_line(&mut line).await;
+                match read_result {
+                    Ok(0) => {
+                        send_disconnected(
+                            &incoming_tx_for_reader,
+                            &disconnected_tx_for_reader,
+                            /*reason*/ None,
+                        )
+                        .await;
+                        break;
+                    }
+                    Ok(_) => {
+                        if line.ends_with('\n') {
+                            line.pop();
+                            if line.ends_with('\r') {
+                                line.pop();
+                            }
+                        } else if line.len() > max_message_len && line.ends_with('\r') {
+                            match reader.read_u8().await {
+                                Ok(b'\n') => {
+                                    line.pop();
+                                }
+                                Ok(_) => {}
+                                Err(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => {}
+                                Err(err) => {
+                                    send_disconnected(
+                                        &incoming_tx_for_reader,
+                                        &disconnected_tx_for_reader,
+                                        Some(format!(
+                                            "failed to read JSON-RPC message from {reader_label}: {err}"
+                                        )),
+                                    )
+                                    .await;
+                                    break;
+                                }
+                            }
+                        }
+                        if line.len() > max_message_len {
+                            send_disconnected(
+                                &incoming_tx_for_reader,
+                                &disconnected_tx_for_reader,
+                                Some(format!(
+                                    "JSON-RPC message from {reader_label} exceeds maximum length of {max_message_len} bytes"
+                                )),
+                            )
+                            .await;
+                            break;
+                        }
+                        if line.trim().is_empty() {
+                            continue;
+                        }
+                        match serde_json::from_str::<JSONRPCMessage>(&line) {
+                            Ok(message) => {
+                                if incoming_tx_for_reader
+                                    .send(JsonRpcConnectionEvent::message(message))
+                                    .await
+                                    .is_err()
+                                {
+                                    break;
+                                }
+                            }
+                            Err(err) => {
+                                send_malformed_message(
+                                    &incoming_tx_for_reader,
+                                    Some(format!(
+                                        "failed to parse JSON-RPC message from {reader_label}: {err}"
+                                    )),
+                                )
+                                .await;
+                            }
+                        }
+                    }
+                    Err(err) => {
+                        send_disconnected(
+                            &incoming_tx_for_reader,
+                            &disconnected_tx_for_reader,
+                            Some(format!(
+                                "failed to read JSON-RPC message from {reader_label}: {err}"
+                            )),
+                        )
+                        .await;
+                        break;
+                    }
+                }
+            }
+        });
+
+        let writer_task = tokio::spawn(async move {
+            let mut writer = BufWriter::new(writer);
+            while let Some(message) = outgoing_rx.recv().await {
+                if let Err(err) = write_jsonrpc_line_message(&mut writer, &message).await {
+                    send_disconnected(
+                        &incoming_tx,
+                        &disconnected_tx,
+                        Some(format!(
+                            "failed to write JSON-RPC message to {connection_label}: {err}"
+                        )),
+                    )
+                    .await;
+                    break;
+                }
+            }
+        });
+
+        Self {
+            outgoing_tx,
+            incoming_rx,
+            disconnected_rx,
+            task_handles: vec![reader_task, writer_task],
+            transport: JsonRpcTransport::Plain,
+        }
+    }
+
+    pub(crate) fn from_websocket<T, E>(stream: T, connection_label: String) -> Self
+    where
+        T: Sink<Message, Error = E> + Stream<Item = Result<Message, E>> + Unpin + Send + 'static,
+        E: std::fmt::Display + Send + 'static,
+    {
+        Self::from_websocket_stream(stream, connection_label, /*ping_interval*/ None)
+    }
+
+    pub(crate) fn from_axum_websocket(stream: AxumWebSocket, connection_label: String) -> Self {
+        Self::from_websocket_stream(stream, connection_label, Some(WEBSOCKET_KEEPALIVE_INTERVAL))
+    }
+
+    fn from_websocket_stream<T, M, E>(
+        mut websocket: T,
+        connection_label: String,
+        ping_interval: Option<Duration>,
+    ) -> Self
+    where
+        T: Sink<M, Error = E> + Stream<Item = Result<M, E>> + Unpin + Send + 'static,
+        M: JsonRpcWebSocketMessage,
+        E: std::fmt::Display + Send + 'static,
+    {
+        let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
+        let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
+        let (disconnected_tx, disconnected_rx) = watch::channel(false);
+
+        let websocket_task = tokio::spawn(async move {
+            let mut ping_interval = ping_interval.map(|ping_interval| {
+                let mut interval = tokio::time::interval_at(
+                    tokio::time::Instant::now() + ping_interval,
+                    ping_interval,
+                );
+                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+                interval
+            });
+
+            loop {
+                tokio::select! {
+                    maybe_message = outgoing_rx.recv() => {
+                        let Some(message) = maybe_message else {
+                            break;
+                        };
+                        if let Err(reason) = send_websocket_jsonrpc_message(
+                            &mut websocket,
+                            &connection_label,
+                            &message,
+                        )
+                        .await
+                        {
+                            send_disconnected(&incoming_tx, &disconnected_tx, Some(reason)).await;
+                            break;
+                        }
+                    }
+                    _ = async {
+                        match ping_interval.as_mut() {
+                            Some(interval) => interval.tick().await,
+                            None => std::future::pending().await,
+                        }
+                    } => {
+                        if let Err(err) = websocket.send(M::ping()).await {
+                            send_disconnected(
+                                &incoming_tx,
+                                &disconnected_tx,
+                                Some(format!(
+                                    "failed to write websocket ping to {connection_label}: {err}"
+                                )),
+                            )
+                            .await;
+                            break;
+                        }
+                    }
+                    incoming_message = websocket.next() => {
+                        match incoming_message {
+                            Some(Ok(message)) => match message.parse_jsonrpc_frame() {
+                                Ok(JsonRpcWebSocketFrame::Message(message)) => {
+                                    if incoming_tx
+                                        .send(JsonRpcConnectionEvent::message(message))
+                                        .await
+                                        .is_err()
+                                    {
+                                        break;
+                                    }
+                                }
+                                Ok(JsonRpcWebSocketFrame::Close) => {
+                                    send_disconnected(
+                                        &incoming_tx,
+                                        &disconnected_tx,
+                                        /*reason*/ None,
+                                    )
+                                    .await;
+                                    break;
+                                }
+                                Ok(JsonRpcWebSocketFrame::Ignore) => {}
+                                Err(err) => {
+                                    send_malformed_message(
+                                        &incoming_tx,
+                                        Some(format!(
+                                            "failed to parse websocket JSON-RPC message from {connection_label}: {err}"
+                                        )),
+                                    )
+                                    .await;
+                                }
+                            },
+                            Some(Err(err)) => {
+                                send_disconnected(
+                                    &incoming_tx,
+                                    &disconnected_tx,
+                                    Some(format!(
+                                        "failed to read websocket JSON-RPC message from {connection_label}: {err}"
+                                    )),
+                                )
+                                .await;
+                                break;
+                            }
+                            None => {
+                                send_disconnected(
+                                    &incoming_tx,
+                                    &disconnected_tx,
+                                    /*reason*/ None,
+                                )
+                                .await;
+                                break;
+                            }
+                        }
+                    }
+                }
+            }
+        });
+
+        Self {
+            outgoing_tx,
+            incoming_rx,
+            disconnected_rx,
+            task_handles: vec![websocket_task],
+            transport: JsonRpcTransport::Plain,
+        }
+    }
+
+    pub(crate) fn with_child_process(mut self, child_process: Child) -> Self {
+        self.transport = JsonRpcTransport::from_child_process(child_process);
+        self
+    }
+}
+
+enum JsonRpcWebSocketFrame {
+    Message(JSONRPCMessage),
+    Close,
+    Ignore,
+}
+
+trait JsonRpcWebSocketMessage: Send + 'static {
+    fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error>;
+    fn from_text(text: String) -> Self;
+    fn ping() -> Self;
+}
+
+impl JsonRpcWebSocketMessage for Message {
+    fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
+        match self {
+            Message::Text(text) => {
+                serde_json::from_str(text.as_ref()).map(JsonRpcWebSocketFrame::Message)
+            }
+            Message::Binary(bytes) => {
+                serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message)
+            }
+            Message::Close(_) => Ok(JsonRpcWebSocketFrame::Close),
+            Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => {
+                Ok(JsonRpcWebSocketFrame::Ignore)
+            }
+        }
+    }
+
+    fn from_text(text: String) -> Self {
+        Self::Text(text.into())
+    }
+
+    fn ping() -> Self {
+        Self::Ping(Vec::new().into())
+    }
+}
+
+impl JsonRpcWebSocketMessage for AxumWebSocketMessage {
+    fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
+        match self {
+            AxumWebSocketMessage::Text(text) => {
+                serde_json::from_str(text.as_ref()).map(JsonRpcWebSocketFrame::Message)
+            }
+            AxumWebSocketMessage::Binary(bytes) => {
+                serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message)
+            }
+            AxumWebSocketMessage::Close(_) => Ok(JsonRpcWebSocketFrame::Close),
+            AxumWebSocketMessage::Ping(_) | AxumWebSocketMessage::Pong(_) => {
+                Ok(JsonRpcWebSocketFrame::Ignore)
+            }
+        }
+    }
+
+    fn from_text(text: String) -> Self {
+        Self::Text(text.into())
+    }
+
+    fn ping() -> Self {
+        Self::Ping(Vec::new().into())
+    }
+}
+
+async fn send_disconnected(
+    incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>,
+    disconnected_tx: &watch::Sender<bool>,
+    reason: Option<String>,
+) {
+    let _ = disconnected_tx.send(true);
+    let _ = incoming_tx
+        .send(JsonRpcConnectionEvent::Disconnected { reason })
+        .await;
+}
+
+async fn send_malformed_message(
+    incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>,
+    reason: Option<String>,
+) {
+    let _ = incoming_tx
+        .send(JsonRpcConnectionEvent::MalformedMessage {
+            reason: reason.unwrap_or_else(|| "malformed JSON-RPC message".to_string()),
+        })
+        .await;
+}
+
+async fn write_jsonrpc_line_message<W>(
+    writer: &mut BufWriter<W>,
+    message: &JSONRPCMessage,
+) -> std::io::Result<()>
+where
+    W: AsyncWrite + Unpin,
+{
+    let encoded =
+        serialize_jsonrpc_message(message).map_err(|err| std::io::Error::other(err.to_string()))?;
+    writer.write_all(encoded.as_bytes()).await?;
+    writer.write_all(b"\n").await?;
+    writer.flush().await
+}
+
+async fn send_websocket_jsonrpc_message<W, M, E>(
+    websocket_writer: &mut W,
+    connection_label: &str,
+    message: &JSONRPCMessage,
+) -> Result<(), String>
+where
+    W: Sink<M, Error = E> + Unpin,
+    M: JsonRpcWebSocketMessage,
+    E: std::fmt::Display,
+{
+    match serialize_jsonrpc_message(message) {
+        Ok(encoded) => websocket_writer
+            .send(M::from_text(encoded))
+            .await
+            .map_err(|err| {
+                format!("failed to write websocket JSON-RPC message to {connection_label}: {err}")
+            }),
+        Err(err) => Err(format!(
+            "failed to serialize JSON-RPC message for {connection_label}: {err}"
+        )),
+    }
+}
+
+fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result<String, serde_json::Error> {
+    serde_json::to_string(message)
+}
+
+#[cfg(test)]
+mod tests {
+    use std::pin::Pin;
+    use std::sync::Arc;
+    use std::sync::atomic::AtomicBool;
+    use std::sync::atomic::Ordering;
+    use std::task::Context;
+    use std::task::Poll;
+
+    use codex_exec_server_protocol::JSONRPCRequest;
+    use codex_exec_server_protocol::RequestId;
+    use futures::channel::mpsc as futures_mpsc;
+    use futures::task::AtomicWaker;
+    use pretty_assertions::assert_eq;
+    use tokio::net::TcpListener;
+    use tokio::time::timeout;
+    use tokio_tungstenite::accept_async;
+    use tokio_tungstenite::connect_async;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn stdio_connection_accepts_message_at_size_limit() -> anyhow::Result<()> {
+        let message = test_jsonrpc_message();
+        let encoded = serde_json::to_string(&message)?;
+        let max_message_len = encoded.len();
+
+        for line_ending in [b"\n".as_slice(), b"\r\n".as_slice()] {
+            let (reader, mut peer) =
+                tokio::io::duplex(max_message_len.saturating_add(line_ending.len()));
+            let mut connection = JsonRpcConnection::from_stdio_with_max_message_len(
+                reader,
+                tokio::io::sink(),
+                "test stdio peer".to_string(),
+                max_message_len,
+            );
+
+            peer.write_all(encoded.as_bytes()).await?;
+            peer.write_all(line_ending).await?;
+            let event = timeout(Duration::from_secs(1), connection.incoming_rx.recv())
+                .await?
+                .expect("stdio connection should report the message");
+            match event {
+                JsonRpcConnectionEvent::QueuedRequest { request, .. } => {
+                    assert_eq!(JSONRPCMessage::Request(request), message)
+                }
+                event => anyhow::bail!("expected JSON-RPC message, got {event:?}"),
+            }
+
+            drop(peer);
+            drop(connection);
+        }
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn stdio_connection_rejects_overlong_unterminated_message() -> anyhow::Result<()> {
+        let max_message_len: usize = 32;
+        let (reader, mut peer) = tokio::io::duplex(max_message_len.saturating_add(1));
+        let mut connection = JsonRpcConnection::from_stdio_with_max_message_len(
+            reader,
+            tokio::io::sink(),
+            "hostile stdio peer".to_string(),
+            max_message_len,
+        );
+        let overlong_message = vec![b'x'; max_message_len + 1];
+
+        peer.write_all(&overlong_message).await?;
+        let event = timeout(Duration::from_secs(1), connection.incoming_rx.recv())
+            .await?
+            .expect("stdio connection should report the framing violation");
+        match event {
+            JsonRpcConnectionEvent::Disconnected { reason } => assert_eq!(
+                reason,
+                Some(
+                    "JSON-RPC message from hostile stdio peer exceeds maximum length of 32 bytes"
+                        .to_string()
+                )
+            ),
+            event => anyhow::bail!("expected stdio disconnect, got {event:?}"),
+        }
+
+        drop(peer);
+        drop(connection);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn websocket_connection_sends_configured_ping() -> anyhow::Result<()> {
+        let (client_websocket, mut server_websocket) = websocket_pair().await?;
+        let connection = JsonRpcConnection::from_websocket_stream(
+            client_websocket,
+            "test".into(),
+            Some(WEBSOCKET_KEEPALIVE_INTERVAL),
+        );
+
+        let message = timeout(Duration::from_secs(1), server_websocket.next())
+            .await?
+            .expect("websocket should stay open")?;
+        assert!(matches!(message, Message::Ping(_)));
+
+        drop(connection);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn websocket_connection_ignores_server_pong() -> anyhow::Result<()> {
+        let (client_websocket, mut server_websocket) = websocket_pair().await?;
+        let mut connection = JsonRpcConnection::from_websocket(client_websocket, "test".into());
+
+        server_websocket
+            .send(Message::Pong(b"check".to_vec().into()))
+            .await?;
+        assert!(
+            timeout(Duration::from_millis(50), connection.incoming_rx.recv())
+                .await
+                .is_err()
+        );
+
+        drop(connection);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn websocket_connection_reports_server_close() -> anyhow::Result<()> {
+        let (client_websocket, mut server_websocket) = websocket_pair().await?;
+        let mut connection = JsonRpcConnection::from_websocket(client_websocket, "test".into());
+
+        server_websocket.close(None).await?;
+        assert!(matches!(
+            timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
+            Some(JsonRpcConnectionEvent::Disconnected { reason: None })
+        ));
+
+        drop(connection);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn websocket_connection_accepts_binary_jsonrpc_message() -> anyhow::Result<()> {
+        let (client_websocket, mut server_websocket) = websocket_pair().await?;
+        let mut connection = JsonRpcConnection::from_websocket(client_websocket, "test".into());
+        let message = JSONRPCMessage::Request(JSONRPCRequest {
+            id: RequestId::Integer(1),
+            method: "test".to_string(),
+            params: None,
+            trace: None,
+        });
+
+        server_websocket
+            .send(Message::Binary(serde_json::to_vec(&message)?.into()))
+            .await?;
+        let Some(JsonRpcConnectionEvent::QueuedRequest { request, .. }) =
+            timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?
+        else {
+            anyhow::bail!("expected a queued JSON-RPC request");
+        };
+        assert_eq!(JSONRPCMessage::Request(request), message);
+
+        drop(connection);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn websocket_connection_keeps_outbound_message_while_send_is_backpressured()
+    -> anyhow::Result<()> {
+        let (websocket, control, mut outbound_rx) =
+            ControlledWebSocket::new(/*write_ready*/ false);
+        let mut connection = JsonRpcConnection::from_websocket_stream(
+            websocket,
+            "test".into(),
+            /*ping_interval*/ None,
+        );
+        let message = test_jsonrpc_message();
+
+        connection.outgoing_tx.send(message.clone()).await?;
+        control.wait_for_blocked_write().await?;
+        control.send_inbound(Message::Pong(b"check".to_vec().into()))?;
+        assert!(
+            timeout(Duration::from_millis(50), connection.incoming_rx.recv())
+                .await
+                .is_err()
+        );
+
+        control.set_write_ready();
+        assert!(matches!(
+            timeout(Duration::from_secs(1), outbound_rx.next()).await?,
+            Some(Message::Text(text)) if serde_json::from_str::<JSONRPCMessage>(&text)? == message
+        ));
+        drop(connection);
+        Ok(())
+    }
+
+    async fn websocket_pair() -> anyhow::Result<(
+        WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
+        WebSocketStream<tokio::net::TcpStream>,
+    )> {
+        let listener = TcpListener::bind("127.0.0.1:0").await?;
+        let websocket_url = format!("ws://{}", listener.local_addr()?);
+        let server_task = tokio::spawn(async move {
+            let (stream, _) = listener.accept().await?;
+            accept_async(stream).await.map_err(anyhow::Error::from)
+        });
+        let (client_websocket, _) = connect_async(websocket_url).await?;
+        let server_websocket = server_task.await??;
+        Ok((client_websocket, server_websocket))
+    }
+
+    fn test_jsonrpc_message() -> JSONRPCMessage {
+        JSONRPCMessage::Request(JSONRPCRequest {
+            id: RequestId::Integer(1),
+            method: "test".to_string(),
+            params: None,
+            trace: None,
+        })
+    }
+
+    struct ControlledWebSocket {
+        inbound_rx: futures_mpsc::UnboundedReceiver<Result<Message, std::convert::Infallible>>,
+        outbound_tx: futures_mpsc::UnboundedSender<Message>,
+        write_ready: Arc<AtomicBool>,
+        write_blocked: Arc<AtomicBool>,
+        write_blocked_waker: Arc<AtomicWaker>,
+        write_waker: Arc<AtomicWaker>,
+    }
+
+    struct ControlledWebSocketHandle {
+        inbound_tx: futures_mpsc::UnboundedSender<Result<Message, std::convert::Infallible>>,
+        write_ready: Arc<AtomicBool>,
+        write_blocked: Arc<AtomicBool>,
+        write_blocked_waker: Arc<AtomicWaker>,
+        write_waker: Arc<AtomicWaker>,
+    }
+
+    impl ControlledWebSocket {
+        fn new(
+            write_ready: bool,
+        ) -> (
+            Self,
+            ControlledWebSocketHandle,
+            futures_mpsc::UnboundedReceiver<Message>,
+        ) {
+            let (inbound_tx, inbound_rx) = futures_mpsc::unbounded();
+            let (outbound_tx, outbound_rx) = futures_mpsc::unbounded();
+            let write_ready = Arc::new(AtomicBool::new(write_ready));
+            let write_blocked = Arc::new(AtomicBool::new(false));
+            let write_blocked_waker = Arc::new(AtomicWaker::new());
+            let write_waker = Arc::new(AtomicWaker::new());
+            (
+                Self {
+                    inbound_rx,
+                    outbound_tx,
+                    write_ready: Arc::clone(&write_ready),
+                    write_blocked: Arc::clone(&write_blocked),
+                    write_blocked_waker: Arc::clone(&write_blocked_waker),
+                    write_waker: Arc::clone(&write_waker),
+                },
+                ControlledWebSocketHandle {
+                    inbound_tx,
+                    write_ready,
+                    write_blocked,
+                    write_blocked_waker,
+                    write_waker,
+                },
+                outbound_rx,
+            )
+        }
+    }
+
+    impl ControlledWebSocketHandle {
+        fn send_inbound(&self, message: Message) -> anyhow::Result<()> {
+            self.inbound_tx
+                .unbounded_send(Ok(message))
+                .map_err(anyhow::Error::from)
+        }
+
+        fn set_write_ready(&self) {
+            self.write_ready.store(true, Ordering::Release);
+            self.write_waker.wake();
+        }
+
+        async fn wait_for_blocked_write(&self) -> anyhow::Result<()> {
+            timeout(
+                Duration::from_secs(1),
+                futures::future::poll_fn(|cx| {
+                    if self.write_blocked.load(Ordering::Acquire) {
+                        Poll::Ready(())
+                    } else {
+                        self.write_blocked_waker.register(cx.waker());
+                        Poll::Pending
+                    }
+                }),
+            )
+            .await?;
+            Ok(())
+        }
+    }
+
+    impl Sink<Message> for ControlledWebSocket {
+        type Error = std::convert::Infallible;
+
+        fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            if self.write_ready.load(Ordering::Acquire) {
+                Poll::Ready(Ok(()))
+            } else {
+                self.write_blocked.store(true, Ordering::Release);
+                self.write_blocked_waker.wake();
+                self.write_waker.register(cx.waker());
+                Poll::Pending
+            }
+        }
+
+        fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
+            self.outbound_tx
+                .unbounded_send(item)
+                .expect("test outbound receiver should stay open");
+            Ok(())
+        }
+
+        fn poll_flush(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_close(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+    }
+
+    impl Stream for ControlledWebSocket {
+        type Item = Result<Message, std::convert::Infallible>;
+
+        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+            Pin::new(&mut self.inbound_rx).poll_next(cx)
+        }
+    }
+}