openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-22 17:58:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into etraut/rollout_creation_strategy

Browse Source
This commit is contained in:
Eric Traut
2026-02-06 11:29:43 -08:00
parent 71edf960e0 e416e578bb
commit c268d55764
27 changed files with 1671 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json
+40
View File
@@ -9825,6 +9825,32 @@
},
"type": "object"
},
"AppConfig": {
"properties": {
"disabled_reason": {
"anyOf": [
{
"$ref": "#/definitions/v2/AppDisabledReason"
},
{
"type": "null"
}
]
},
"enabled": {
"default": true,
"type": "boolean"
}
},
"type": "object"
},
"AppDisabledReason": {
"enum": [
"unknown",
"user"
],
"type": "string"
},
"AppInfo": {
"properties": {
"description": {
@@ -9874,6 +9900,9 @@
],
"type": "object"
},
"AppsConfig": {
"type": "object"
},
"AppsListParams": {
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
@@ -10489,6 +10518,17 @@
}
]
},
"apps": {
"anyOf": [
{
"$ref": "#/definitions/v2/AppsConfig"
},
{
"type": "null"
}
],
"default": null
},
"compact_prompt": {
"type": [
"string",
codex-rs/app-server-protocol/schema/json/v2/ConfigReadResponse.json
+40
View File
@@ -17,6 +17,35 @@
},
"type": "object"
},
"AppConfig": {
"properties": {
"disabled_reason": {
"anyOf": [
{
"$ref": "#/definitions/AppDisabledReason"
},
{
"type": "null"
}
]
},
"enabled": {
"default": true,
"type": "boolean"
}
},
"type": "object"
},
"AppDisabledReason": {
"enum": [
"unknown",
"user"
],
"type": "string"
},
"AppsConfig": {
"type": "object"
},
"AskForApproval": {
"enum": [
"untrusted",
@@ -49,6 +78,17 @@
}
]
},
"apps": {
"anyOf": [
{
"$ref": "#/definitions/AppsConfig"
},
{
"type": "null"
}
],
"default": null
},
"compact_prompt": {
"type": [
"string",
codex-rs/app-server-protocol/schema/typescript/v2/AppDisabledReason.ts
+5
View File
@@ -0,0 +1,5 @@
// GENERATED CODE! DO NOT MODIFY BY HAND!
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
export type AppDisabledReason = "unknown" | "user";
codex-rs/app-server-protocol/schema/typescript/v2/AppsConfig.ts
+6
View File
@@ -0,0 +1,6 @@
// GENERATED CODE! DO NOT MODIFY BY HAND!
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
import type { AppDisabledReason } from "./AppDisabledReason";
export type AppsConfig = { [key in string]?: { enabled: boolean, disabled_reason: AppDisabledReason | null, } };
codex-rs/app-server-protocol/schema/typescript/v2/Config.ts
+2 -1
View File
@@ -8,10 +8,11 @@ import type { Verbosity } from "../Verbosity";
import type { WebSearchMode } from "../WebSearchMode";
import type { JsonValue } from "../serde_json/JsonValue";
import type { AnalyticsConfig } from "./AnalyticsConfig";
import type { AppsConfig } from "./AppsConfig";
import type { AskForApproval } from "./AskForApproval";
import type { ProfileV2 } from "./ProfileV2";
import type { SandboxMode } from "./SandboxMode";
import type { SandboxWorkspaceWrite } from "./SandboxWorkspaceWrite";
import type { ToolsV2 } from "./ToolsV2";
export type Config = { model: string | null, review_model: string | null, model_context_window: bigint | null, model_auto_compact_token_limit: bigint | null, model_provider: string | null, approval_policy: AskForApproval | null, sandbox_mode: SandboxMode | null, sandbox_workspace_write: SandboxWorkspaceWrite | null, forced_chatgpt_workspace_id: string | null, forced_login_method: ForcedLoginMethod | null, web_search: WebSearchMode | null, tools: ToolsV2 | null, profile: string | null, profiles: { [key in string]?: ProfileV2 }, instructions: string | null, developer_instructions: string | null, compact_prompt: string | null, model_reasoning_effort: ReasoningEffort | null, model_reasoning_summary: ReasoningSummary | null, model_verbosity: Verbosity | null, analytics: AnalyticsConfig | null, } & ({ [key in string]?: number | string | boolean | Array<JsonValue> | { [key in string]?: JsonValue } | null });
export type Config = { model: string | null, review_model: string | null, model_context_window: bigint | null, model_auto_compact_token_limit: bigint | null, model_provider: string | null, approval_policy: AskForApproval | null, sandbox_mode: SandboxMode | null, sandbox_workspace_write: SandboxWorkspaceWrite | null, forced_chatgpt_workspace_id: string | null, forced_login_method: ForcedLoginMethod | null, web_search: WebSearchMode | null, tools: ToolsV2 | null, profile: string | null, profiles: { [key in string]?: ProfileV2 }, instructions: string | null, developer_instructions: string | null, compact_prompt: string | null, model_reasoning_effort: ReasoningEffort | null, model_reasoning_summary: ReasoningSummary | null, model_verbosity: Verbosity | null, analytics: AnalyticsConfig | null, apps: AppsConfig | null, } & ({ [key in string]?: number | string | boolean | Array<JsonValue> | { [key in string]?: JsonValue } | null });
codex-rs/app-server-protocol/schema/typescript/v2/index.ts
+2
View File
@@ -6,7 +6,9 @@ export type { AccountRateLimitsUpdatedNotification } from "./AccountRateLimitsUp
export type { AccountUpdatedNotification } from "./AccountUpdatedNotification";
export type { AgentMessageDeltaNotification } from "./AgentMessageDeltaNotification";
export type { AnalyticsConfig } from "./AnalyticsConfig";
export type { AppDisabledReason } from "./AppDisabledReason";
export type { AppInfo } from "./AppInfo";
export type { AppsConfig } from "./AppsConfig";
export type { AppsListParams } from "./AppsListParams";
export type { AppsListResponse } from "./AppsListResponse";
export type { AskForApproval } from "./AskForApproval";
codex-rs/app-server-protocol/src/protocol/v2.rs
+32
View File
@@ -374,6 +374,36 @@ pub struct AnalyticsConfig {
pub additional: HashMap<String, JsonValue>,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export_to = "v2/")]
pub enum AppDisabledReason {
Unknown,
User,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export_to = "v2/")]
pub struct AppConfig {
#[serde(default = "default_enabled")]
pub enabled: bool,
pub disabled_reason: Option<AppDisabledReason>,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export_to = "v2/")]
pub struct AppsConfig {
#[serde(default, flatten)]
#[schemars(with = "HashMap<String, AppConfig>")]
pub apps: HashMap<String, AppConfig>,
}
const fn default_enabled() -> bool {
true
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export_to = "v2/")]
@@ -400,6 +430,8 @@ pub struct Config {
pub model_reasoning_summary: Option<ReasoningSummary>,
pub model_verbosity: Option<Verbosity>,
pub analytics: Option<AnalyticsConfig>,
#[serde(default)]
pub apps: Option<AppsConfig>,
#[serde(default, flatten)]
pub additional: HashMap<String, JsonValue>,
}
codex-rs/app-server/tests/suite/v2/config_rpc.rs
+71
View File
@@ -3,6 +3,9 @@ use app_test_support::McpProcess;
use app_test_support::test_path_buf_with_windows;
use app_test_support::test_tmp_path_buf;
use app_test_support::to_response;
use codex_app_server_protocol::AppConfig;
use codex_app_server_protocol::AppDisabledReason;
use codex_app_server_protocol::AppsConfig;
use codex_app_server_protocol::AskForApproval;
use codex_app_server_protocol::ConfigBatchWriteParams;
use codex_app_server_protocol::ConfigEdit;
@@ -146,6 +149,74 @@ view_image = false
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn config_read_includes_apps() -> Result<()> {
let codex_home = TempDir::new()?;
write_config(
&codex_home,
r#"
[apps.app1]
enabled = false
disabled_reason = "user"
"#,
)?;
let codex_home_path = codex_home.path().canonicalize()?;
let user_file = AbsolutePathBuf::try_from(codex_home_path.join("config.toml"))?;
let mut mcp = McpProcess::new(codex_home.path()).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
let request_id = mcp
.send_config_read_request(ConfigReadParams {
include_layers: true,
cwd: None,
})
.await?;
let resp: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
let ConfigReadResponse {
config,
origins,
layers,
} = to_response(resp)?;
assert_eq!(
config.apps,
Some(AppsConfig {
apps: std::collections::HashMap::from([(
"app1".to_string(),
AppConfig {
enabled: false,
disabled_reason: Some(AppDisabledReason::User),
},
)]),
})
);
assert_eq!(
origins.get("apps.app1.enabled").expect("origin").name,
ConfigLayerSource::User {
file: user_file.clone(),
}
);
assert_eq!(
origins
.get("apps.app1.disabled_reason")
.expect("origin")
.name,
ConfigLayerSource::User {
file: user_file.clone(),
}
);
let layers = layers.expect("layers present");
assert_layers_user_then_optional_system(&layers, user_file)?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn config_read_includes_project_layers_for_cwd() -> Result<()> {
let codex_home = TempDir::new()?;
codex-rs/core/config.schema.json
+43
View File
@@ -55,6 +55,40 @@
},
"type": "object"
},
"AppConfig": {
"additionalProperties": false,
"description": "Config values for a single app/connector.",
"properties": {
"disabled_reason": {
"allOf": [
{
"$ref": "#/definitions/AppDisabledReason"
}
],
"description": "Reason this app was disabled."
},
"enabled": {
"default": true,
"description": "When `false`, Codex does not surface this app.",
"type": "boolean"
}
},
"type": "object"
},
"AppDisabledReason": {
"enum": [
"unknown",
"user"
],
"type": "string"
},
"AppsConfigToml": {
"additionalProperties": {
"$ref": "#/definitions/AppConfig"
},
"description": "App/connector settings loaded from `config.toml`.",
"type": "object"
},
"AskForApproval": {
"description": "Determines the conditions under which the user is consulted to approve running the command proposed by Codex.",
"oneOf": [
@@ -1135,6 +1169,15 @@
],
"description": "Default approval policy for executing commands."
},
"apps": {
"allOf": [
{
"$ref": "#/definitions/AppsConfigToml"
}
],
"default": null,
"description": "Settings for app-specific controls."
},
"chatgpt_base_url": {
"description": "Base URL for requests to ChatGPT (as opposed to the OpenAI API).",
"type": "string"
codex-rs/core/src/client.rs
+501 -104
View File
@@ -9,10 +9,35 @@
//! call site.
//!
//! A [`ModelClientSession`] is created per turn and is used to stream one or more Responses API
//! requests during that turn. It caches a Responses WebSocket connection (opened lazily) and
//! stores per-turn state such as the `x-codex-turn-state` token used for sticky routing.
//! requests during that turn. It caches a Responses WebSocket connection (opened lazily, or reused
//! from a session-level preconnect) and stores per-turn state such as the `x-codex-turn-state`
//! token used for sticky routing.
//!
//! Preconnect is intentionally handshake-only: it may warm a socket and capture sticky-routing
//! state, but the first `response.create` payload is still sent only when a turn starts.
//!
//! Internally, startup preconnect and warmed-socket adoption share one session-level lifecycle:
//! `Idle` (no task/socket), `InFlight` (startup preconnect task running), and `Ready` (one-shot
//! warmed socket available). On first use in a turn, the session tries to adopt `Ready`; if not
//! ready, it awaits `InFlight` and retries adoption before opening a new websocket. This prevents
//! racing duplicate first-turn handshakes while keeping preconnect best-effort.
//!
//! ## Retry-Budget Tradeoff
//!
//! `stream_max_retries` applies to retryable turn stream failures, not to background startup
//! preconnect handshakes. In failure cases this can produce two websocket handshakes on the first
//! turn (startup preconnect, then turn-time connect) before HTTP fallback becomes sticky. We keep
//! this split intentionally so opportunistic preconnect cannot consume the user-visible stream
//! retry budget before any turn payload is sent.
//!
//! If this policy needs to change later, preconnect can be modeled as an explicit first connection
//! attempt in the same retry budget as turn streaming. That would require plumbing websocket
//! attempt accounting from connection acquisition into the turn retry loop and updating fallback
//! expectations/tests accordingly.
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
@@ -64,8 +89,10 @@ use reqwest::StatusCode;
use serde_json::Value;
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::Error;
use tokio_tungstenite::tungstenite::Message;
use tracing::debug;
use tracing::warn;
use crate::AuthManager;
@@ -81,6 +108,8 @@ use crate::flags::CODEX_RS_SSE_FIXTURE;
use crate::model_provider_info::ModelProviderInfo;
use crate::model_provider_info::WireApi;
use crate::tools::spec::create_tools_json_for_responses_api;
use crate::turn_metadata::build_turn_metadata_header;
use crate::turn_metadata::resolve_turn_metadata_header_with_timeout;
pub const OPENAI_BETA_HEADER: &str = "OpenAI-Beta";
pub const OPENAI_BETA_RESPONSES_WEBSOCKETS: &str = "responses_websockets=2026-02-04";
@@ -93,7 +122,6 @@ pub const X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER: &str =
///
/// This is intentionally kept minimal so `ModelClient` does not need to hold a full `Config`. Most
/// configuration is per turn and is passed explicitly to streaming/unary methods.
#[derive(Debug)]
struct ModelClientState {
auth_manager: Option<Arc<AuthManager>>,
conversation_id: ThreadId,
@@ -105,6 +133,70 @@ struct ModelClientState {
include_timing_metrics: bool,
beta_features_header: Option<String>,
disable_websockets: AtomicBool,
/// Session-scoped preconnect lifecycle state.
///
/// This keeps startup preconnect task tracking and warmed-socket adoption in one lock so
/// turn-time websocket setup observes a single, coherent state.
preconnect: Mutex<PreconnectState>,
}
impl std::fmt::Debug for ModelClientState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ModelClientState")
.field("auth_manager", &self.auth_manager)
.field("conversation_id", &self.conversation_id)
.field("provider", &self.provider)
.field("session_source", &self.session_source)
.field("model_verbosity", &self.model_verbosity)
.field(
"enable_responses_websockets",
&self.enable_responses_websockets,
)
.field(
"enable_request_compression",
&self.enable_request_compression,
)
.field("include_timing_metrics", &self.include_timing_metrics)
.field("beta_features_header", &self.beta_features_header)
.field(
"disable_websockets",
&self.disable_websockets.load(Ordering::Relaxed),
)
.field("preconnect", &"<opaque>")
.finish()
}
}
/// Resolved API client setup for a single request attempt.
///
/// Keeping this as a single bundle ensures preconnect and normal request paths
/// share the same auth/provider setup flow.
struct CurrentClientSetup {
auth: Option<CodexAuth>,
api_provider: codex_api::Provider,
api_auth: CoreAuthProvider,
}
/// One-shot preconnected websocket slot consumed by the next turn.
///
/// This bundles the socket with optional sticky-routing state captured during
/// handshake so they are taken and cleared atomically.
struct PreconnectedWebSocket {
connection: ApiWebSocketConnection,
turn_state: Option<String>,
}
/// Session-level lifecycle of startup websocket preconnect.
///
/// `InFlight` tracks the startup task so the first turn can await it and reuse the same socket.
/// `Ready` stores a one-shot warmed socket for turn adoption.
enum PreconnectState {
/// No startup preconnect task is active and no warmed socket is available.
Idle,
/// Startup preconnect is currently running; first turn may await this task.
InFlight(JoinHandle<()>),
/// Startup preconnect finished and produced a one-shot warmed socket.
Ready(PreconnectedWebSocket),
}
/// A session-scoped client for model-provider API calls.
@@ -116,8 +208,9 @@ struct ModelClientState {
/// WebSocket fallback is session-scoped: once a turn activates the HTTP fallback, subsequent turns
/// will also use HTTP for the remainder of the session.
///
/// Turn-scoped settings (model selection, reasoning controls, telemetry context, and turn metadata)
/// are passed explicitly to the relevant methods to keep turn lifetime visible at the call site.
/// Turn-scoped settings (model selection, reasoning controls, telemetry context, and turn
/// metadata) are passed explicitly to the relevant methods to keep turn lifetime visible at the
/// call site.
///
/// This type is cheap to clone.
#[derive(Debug, Clone)]
@@ -127,14 +220,17 @@ pub struct ModelClient {
/// A turn-scoped streaming session created from a [`ModelClient`].
///
/// The session lazily establishes a Responses WebSocket connection (and reuses it across multiple
/// requests) and caches per-turn state:
/// The session establishes a Responses WebSocket connection lazily (or adopts a preconnected one)
/// and reuses it across multiple requests within the turn. It also caches per-turn state:
///
/// - The last request's input items, so subsequent calls can use `response.append` when the input
/// is an incremental extension of the previous request.
/// - The `x-codex-turn-state` sticky-routing token, which must be replayed for all requests within
/// the same turn.
///
/// When startup preconnect is still running, first use of this session awaits that in-flight task
/// before opening a new websocket so preconnect acts as the first connection attempt for the turn.
///
/// Create a fresh `ModelClientSession` for each Codex turn. Reusing it across turns would replay
/// the previous turn's sticky-routing token into the next turn, which violates the client/server
/// contract and can cause routing bugs.
@@ -184,14 +280,16 @@ impl ModelClient {
include_timing_metrics,
beta_features_header,
disable_websockets: AtomicBool::new(false),
preconnect: Mutex::new(PreconnectState::Idle),
}),
}
}
/// Creates a fresh turn-scoped streaming session.
///
/// This does not open any network connections; the WebSocket connection is established lazily
/// when the first WebSocket stream request is issued.
/// This constructor does not perform network I/O itself. The returned session either adopts a
/// previously preconnected websocket or opens a websocket lazily when the first stream request
/// is issued.
pub fn new_session(&self) -> ModelClientSession {
ModelClientSession {
client: self.clone(),
@@ -201,6 +299,79 @@ impl ModelClient {
}
}
/// Spawns a best-effort task that warms a websocket for the first turn.
///
/// This call performs only connection setup; it never sends prompt payloads.
///
/// A timeout when computing turn metadata is treated the same as "no metadata" so startup
/// cannot block indefinitely on optional preconnect context.
pub fn pre_establish_connection(&self, otel_manager: OtelManager, cwd: PathBuf) {
if !self.responses_websocket_enabled() || self.disable_websockets() {
return;
}
let model_client = self.clone();
let handle = tokio::spawn(async move {
let turn_metadata_header = resolve_turn_metadata_header_with_timeout(
build_turn_metadata_header(cwd.as_path()),
None,
)
.await;
let _ = model_client
.preconnect(&otel_manager, turn_metadata_header.as_deref())
.await;
});
self.store_preconnect_task(handle);
}
/// Opportunistically pre-establishes a Responses WebSocket connection for this session.
///
/// This method is best-effort: it returns `false` on any setup/connect failure and the caller
/// should continue normally. A successful preconnect reduces first-turn latency but never sends
/// an initial prompt; the first `response.create` is still sent only when a turn starts.
///
/// The preconnected slot is single-consumer and single-use: the next `ModelClientSession` may
/// adopt it once, after which later turns either keep using that same turn-local connection or
/// create a new one.
pub async fn preconnect(
&self,
otel_manager: &OtelManager,
turn_metadata_header: Option<&str>,
) -> bool {
if !self.responses_websocket_enabled() || self.disable_websockets() {
return false;
}
let client_setup = match self.current_client_setup().await {
Ok(client_setup) => client_setup,
Err(err) => {
warn!("failed to build websocket preconnect client setup: {err}");
return false;
}
};
let turn_state = Arc::new(OnceLock::new());
match self
.connect_websocket(
otel_manager,
client_setup.api_provider,
client_setup.api_auth,
Some(Arc::clone(&turn_state)),
turn_metadata_header,
)
.await
{
Ok(connection) => {
self.store_preconnected_websocket(connection, turn_state.get().cloned());
true
}
Err(err) => {
debug!("websocket preconnect failed: {err}");
false
}
}
}
/// Compacts the current conversation history using the Compact endpoint.
///
/// This is a unary call (no streaming) that returns a new list of
@@ -217,20 +388,12 @@ impl ModelClient {
if prompt.input.is_empty() {
return Ok(Vec::new());
}
let auth_manager = self.state.auth_manager.clone();
let auth = match auth_manager.as_ref() {
Some(manager) => manager.auth().await,
None => None,
};
let api_provider = self
.state
.provider
.to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?;
let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?;
let client_setup = self.current_client_setup().await?;
let transport = ReqwestTransport::new(build_reqwest_client());
let request_telemetry = Self::build_request_telemetry(otel_manager);
let client = ApiCompactClient::new(transport, api_provider, api_auth)
.with_telemetry(Some(request_telemetry));
let client =
ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth)
.with_telemetry(Some(request_telemetry));
let instructions = prompt.base_instructions.text.clone();
let payload = ApiCompactionInput {
@@ -263,20 +426,12 @@ impl ModelClient {
return Ok(Vec::new());
}
let auth_manager = self.state.auth_manager.clone();
let auth = match auth_manager.as_ref() {
Some(manager) => manager.auth().await,
None => None,
};
let api_provider = self
.state
.provider
.to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?;
let api_auth = auth_provider_from_auth(auth, &self.state.provider)?;
let client_setup = self.current_client_setup().await?;
let transport = ReqwestTransport::new(build_reqwest_client());
let request_telemetry = Self::build_request_telemetry(otel_manager);
let client = ApiMemoriesClient::new(transport, api_provider, api_auth)
.with_telemetry(Some(request_telemetry));
let client =
ApiMemoriesClient::new(transport, client_setup.api_provider, client_setup.api_auth)
.with_telemetry(Some(request_telemetry));
let payload = ApiMemoryTraceSummarizeInput {
model: model_info.slug.clone(),
@@ -315,13 +470,223 @@ impl ModelClient {
let request_telemetry: Arc<dyn RequestTelemetry> = telemetry;
request_telemetry
}
/// Returns whether this session is configured to use Responses-over-WebSocket.
///
/// This combines provider capability and feature gating; both must be true for websocket paths
/// to be eligible.
fn responses_websocket_enabled(&self) -> bool {
self.state.provider.supports_websockets && self.state.enable_responses_websockets
}
/// Returns whether websocket transport has been permanently disabled for this session.
///
/// Once set by fallback activation, subsequent turns must stay on HTTP transport.
fn disable_websockets(&self) -> bool {
self.state.disable_websockets.load(Ordering::Relaxed)
}
/// Returns auth + provider configuration resolved from the current session auth state.
///
/// This centralizes setup used by both preconnect and normal request paths so they stay in
/// lockstep when auth/provider resolution changes.
async fn current_client_setup(&self) -> Result<CurrentClientSetup> {
let auth = match self.state.auth_manager.as_ref() {
Some(manager) => manager.auth().await,
None => None,
};
let api_provider = self
.state
.provider
.to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?;
let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?;
Ok(CurrentClientSetup {
auth,
api_provider,
api_auth,
})
}
/// Opens a websocket connection using the same header and telemetry wiring as normal turns.
///
/// Both startup preconnect and in-turn `needs_new` reconnects call this path so handshake
/// behavior remains consistent across both flows.
async fn connect_websocket(
&self,
otel_manager: &OtelManager,
api_provider: codex_api::Provider,
api_auth: CoreAuthProvider,
turn_state: Option<Arc<OnceLock<String>>>,
turn_metadata_header: Option<&str>,
) -> std::result::Result<ApiWebSocketConnection, ApiError> {
let headers = self.build_websocket_headers(turn_state.as_ref(), turn_metadata_header);
let websocket_telemetry = ModelClientSession::build_websocket_telemetry(otel_manager);
ApiWebSocketResponsesClient::new(api_provider, api_auth)
.connect(headers, turn_state, Some(websocket_telemetry))
.await
}
/// Builds websocket handshake headers for both preconnect and turn-time reconnect.
///
/// Callers should pass the current turn-state lock when available so sticky-routing state is
/// replayed on reconnect within the same turn.
fn build_websocket_headers(
&self,
turn_state: Option<&Arc<OnceLock<String>>>,
turn_metadata_header: Option<&str>,
) -> ApiHeaderMap {
let turn_metadata_header = parse_turn_metadata_header(turn_metadata_header);
let mut headers = build_responses_headers(
self.state.beta_features_header.as_deref(),
turn_state,
turn_metadata_header.as_ref(),
);
headers.extend(build_conversation_headers(Some(
self.state.conversation_id.to_string(),
)));
headers.insert(
OPENAI_BETA_HEADER,
HeaderValue::from_static(OPENAI_BETA_RESPONSES_WEBSOCKETS),
);
if self.state.include_timing_metrics {
headers.insert(
X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER,
HeaderValue::from_static("true"),
);
}
headers
}
/// Consumes the warmed websocket slot.
fn take_preconnected_websocket(&self) -> Option<PreconnectedWebSocket> {
let mut state = self
.state
.preconnect
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let previous = std::mem::replace(&mut *state, PreconnectState::Idle);
match previous {
PreconnectState::Ready(preconnected) => Some(preconnected),
other => {
*state = other;
None
}
}
}
/// Stores a freshly preconnected websocket and optional captured turn-state token.
///
/// This overwrites any previously warmed socket because only one preconnect candidate is kept.
fn store_preconnected_websocket(
&self,
connection: ApiWebSocketConnection,
turn_state: Option<String>,
) {
let mut state = self
.state
.preconnect
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if self.disable_websockets() {
debug!("discarding startup websocket preconnect because websocket fallback is active");
*state = PreconnectState::Idle;
return;
}
*state = PreconnectState::Ready(PreconnectedWebSocket {
connection,
turn_state,
});
}
/// Stores the latest startup preconnect task handle.
///
/// If a previous task is still running, it is aborted so only one in-flight startup attempt
/// is tracked.
fn store_preconnect_task(&self, task: JoinHandle<()>) {
let mut task = Some(task);
let previous_in_flight = {
let mut state = self
.state
.preconnect
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
match &*state {
// A very fast startup preconnect can complete before this method stores the
// task handle; keep the warmed socket and drop the now-useless handle.
PreconnectState::Ready(_) => None,
_ => match task.take() {
Some(next_task) => {
match std::mem::replace(&mut *state, PreconnectState::InFlight(next_task)) {
PreconnectState::InFlight(previous) => Some(previous),
_ => None,
}
}
None => None,
},
}
};
if let Some(previous) = previous_in_flight {
previous.abort();
}
if let Some(task) = task {
task.abort();
}
}
/// Awaits the startup preconnect task once, if one is currently tracked.
///
/// This lets the first turn treat startup preconnect as the first websocket connection
/// attempt, avoiding a redundant second connect while the preconnect attempt is in flight.
///
/// This await intentionally has no separate timeout wrapper. WebSocket connect handshakes
/// already run without an app-level timeout, so waiting on the in-flight preconnect task does
/// not add a new unbounded wait class; it reuses the same first connection attempt.
async fn await_preconnect_task(&self) {
let task = {
let mut state = self
.state
.preconnect
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let previous = std::mem::replace(&mut *state, PreconnectState::Idle);
match previous {
PreconnectState::InFlight(task) => Some(task),
other => {
*state = other;
None
}
}
};
if let Some(task) = task {
let in_flight = !task.is_finished();
if in_flight {
debug!("awaiting startup websocket preconnect before opening a new websocket");
}
if let Err(err) = task.await {
debug!("startup websocket preconnect task failed: {err}");
}
}
}
/// Clears all startup preconnect state.
///
/// This aborts any in-flight startup preconnect task and drops any warmed socket.
fn clear_preconnect(&self) {
let previous = {
let mut state = self
.state
.preconnect
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
std::mem::replace(&mut *state, PreconnectState::Idle)
};
if let PreconnectState::InFlight(task) = previous {
task.abort();
}
}
}
impl ModelClientSession {
fn disable_websockets(&self) -> bool {
self.client.state.disable_websockets.load(Ordering::Relaxed)
}
fn activate_http_fallback(&self, websocket_enabled: bool) -> bool {
websocket_enabled
&& !self
@@ -331,11 +696,6 @@ impl ModelClientSession {
.swap(true, Ordering::Relaxed)
}
fn responses_websocket_enabled(&self) -> bool {
self.client.state.provider.supports_websockets
&& self.client.state.enable_responses_websockets
}
fn build_responses_request(prompt: &Prompt) -> Result<ApiPrompt> {
let instructions = prompt.base_instructions.text.clone();
let tools_json: Vec<Value> = create_tools_json_for_responses_api(&prompt.tools)?;
@@ -343,6 +703,10 @@ impl ModelClientSession {
}
#[allow(clippy::too_many_arguments)]
/// Builds shared Responses API request options for both HTTP and WebSocket streaming.
///
/// Keeping option construction in one place ensures request-scoped headers are consistent
/// regardless of transport choice.
fn build_responses_options(
&self,
prompt: &Prompt,
@@ -352,8 +716,7 @@ impl ModelClientSession {
turn_metadata_header: Option<&str>,
compression: Compression,
) -> ApiResponsesOptions {
let turn_metadata_header =
turn_metadata_header.and_then(|value| HeaderValue::from_str(value).ok());
let turn_metadata_header = parse_turn_metadata_header(turn_metadata_header);
let default_reasoning_effort = model_info.default_reasoning_level;
let reasoning = if model_info.supports_reasoning_summaries {
@@ -466,30 +829,54 @@ impl ModelClientSession {
ResponsesWsRequest::ResponseCreate(payload)
}
/// Returns a websocket connection for this turn, reusing preconnect when possible.
///
/// This method first tries to adopt the session-level preconnect slot, then falls back to a
/// fresh websocket handshake only when the turn has no live connection. If startup preconnect
/// is still running, it is awaited first so that task acts as the first connection attempt for
/// this turn instead of racing a second handshake. If that attempt fails, the normal connect
/// and stream retry flow continues unchanged.
async fn websocket_connection(
&mut self,
otel_manager: &OtelManager,
api_provider: codex_api::Provider,
api_auth: CoreAuthProvider,
turn_metadata_header: Option<&str>,
options: &ApiResponsesOptions,
) -> std::result::Result<&ApiWebSocketConnection, ApiError> {
// Prefer the session-level preconnect slot before creating a new websocket.
if self.connection.is_none() {
if let Some(preconnected) = self.try_use_preconnected_websocket() {
self.adopt_preconnected_websocket(preconnected);
} else {
self.client.await_preconnect_task().await;
if let Some(preconnected) = self.try_use_preconnected_websocket() {
self.adopt_preconnected_websocket(preconnected);
}
}
}
let needs_new = match self.connection.as_ref() {
Some(conn) => conn.is_closed().await,
None => true,
};
if needs_new {
let headers =
build_websocket_connect_headers(options, self.client.state.include_timing_metrics);
let websocket_telemetry = Self::build_websocket_telemetry(otel_manager);
let new_conn: ApiWebSocketConnection =
ApiWebSocketResponsesClient::new(api_provider, api_auth)
.connect(
headers,
options.turn_state.clone(),
Some(websocket_telemetry),
)
.await?;
self.client.clear_preconnect();
let turn_state = options
.turn_state
.clone()
.unwrap_or_else(|| Arc::clone(&self.turn_state));
let new_conn = self
.client
.connect_websocket(
otel_manager,
api_provider,
api_auth,
Some(turn_state),
turn_metadata_header,
)
.await?;
self.connection = Some(new_conn);
}
@@ -498,6 +885,33 @@ impl ModelClientSession {
))
}
/// Adopts the session-level preconnect slot for this turn.
///
/// If a turn-local connection already exists, this intentionally does nothing to avoid
/// replacing an active connection mid-turn.
fn try_use_preconnected_websocket(&mut self) -> Option<PreconnectedWebSocket> {
if self.connection.is_some() {
return None;
}
self.client.take_preconnected_websocket()
}
/// Moves a preconnected socket into the turn-local connection slot.
///
/// If the preconnect handshake captured sticky-routing turn state, this also seeds the
/// turn-local state lock so all later requests in the turn replay the same token.
fn adopt_preconnected_websocket(&mut self, preconnected: PreconnectedWebSocket) {
let PreconnectedWebSocket {
connection,
turn_state,
} = preconnected;
if let Some(turn_state) = turn_state {
let _ = self.turn_state.set(turn_state);
}
self.connection = Some(connection);
}
fn responses_request_compression(&self, auth: Option<&crate::auth::CodexAuth>) -> Compression {
if self.client.state.enable_request_compression
&& auth.is_some_and(CodexAuth::is_chatgpt_auth)
@@ -540,22 +954,17 @@ impl ModelClientSession {
.as_ref()
.map(super::auth::AuthManager::unauthorized_recovery);
loop {
let auth = match auth_manager.as_ref() {
Some(manager) => manager.auth().await,
None => None,
};
let api_provider = self
.client
.state
.provider
.to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?;
let api_auth = auth_provider_from_auth(auth.clone(), &self.client.state.provider)?;
let client_setup = self.client.current_client_setup().await?;
let transport = ReqwestTransport::new(build_reqwest_client());
let (request_telemetry, sse_telemetry) = Self::build_streaming_telemetry(otel_manager);
let compression = self.responses_request_compression(auth.as_ref());
let compression = self.responses_request_compression(client_setup.auth.as_ref());
let client = ApiResponsesClient::new(transport, api_provider, api_auth)
.with_telemetry(Some(request_telemetry), Some(sse_telemetry));
let client = ApiResponsesClient::new(
transport,
client_setup.api_provider,
client_setup.api_auth,
)
.with_telemetry(Some(request_telemetry), Some(sse_telemetry));
let options = self.build_responses_options(
prompt,
@@ -603,17 +1012,8 @@ impl ModelClientSession {
.as_ref()
.map(super::auth::AuthManager::unauthorized_recovery);
loop {
let auth = match auth_manager.as_ref() {
Some(manager) => manager.auth().await,
None => None,
};
let api_provider = self
.client
.state
.provider
.to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?;
let api_auth = auth_provider_from_auth(auth.clone(), &self.client.state.provider)?;
let compression = self.responses_request_compression(auth.as_ref());
let client_setup = self.client.current_client_setup().await?;
let compression = self.responses_request_compression(client_setup.auth.as_ref());
let options = self.build_responses_options(
prompt,
@@ -628,8 +1028,9 @@ impl ModelClientSession {
let connection = match self
.websocket_connection(
otel_manager,
api_provider.clone(),
api_auth.clone(),
client_setup.api_provider,
client_setup.api_auth,
turn_metadata_header,
&options,
)
.await
@@ -691,7 +1092,7 @@ impl ModelClientSession {
match wire_api {
WireApi::Responses => {
let websocket_enabled =
self.responses_websocket_enabled() && !self.disable_websockets();
self.client.responses_websocket_enabled() && !self.client.disable_websockets();
if websocket_enabled {
self.stream_responses_websocket(
@@ -721,10 +1122,16 @@ impl ModelClientSession {
/// Permanently disables WebSockets for this Codex session and resets WebSocket state.
///
/// This is used after exhausting the provider retry budget, to force subsequent requests onto
/// the HTTP transport. Returns `true` if this call activated fallback, or `false` if fallback
/// was already active.
/// the HTTP transport. It also clears any warmed websocket preconnect state so future turns
/// cannot accidentally adopt a stale socket after fallback has been activated.
///
/// Startup preconnect handshakes are intentionally not counted against `stream_max_retries`.
/// See [`crate::client`] module docs ("Retry-Budget Tradeoff") for rationale and future
/// alternatives.
///
/// Returns `true` if this call activated fallback, or `false` if fallback was already active.
pub(crate) fn try_switch_fallback_transport(&mut self, otel_manager: &OtelManager) -> bool {
let websocket_enabled = self.responses_websocket_enabled();
let websocket_enabled = self.client.responses_websocket_enabled();
let activated = self.activate_http_fallback(websocket_enabled);
if activated {
warn!("falling back to HTTP");
@@ -736,6 +1143,7 @@ impl ModelClientSession {
self.connection = None;
self.websocket_last_items.clear();
self.client.clear_preconnect();
}
activated
}
@@ -752,6 +1160,14 @@ fn build_api_prompt(prompt: &Prompt, instructions: String, tools_json: Vec<Value
}
}
/// Parses per-turn metadata into an HTTP header value.
///
/// Invalid values are treated as absent so callers can compare and propagate
/// metadata with the same sanitization path used when constructing headers.
fn parse_turn_metadata_header(turn_metadata_header: Option<&str>) -> Option<HeaderValue> {
turn_metadata_header.and_then(|value| HeaderValue::from_str(value).ok())
}
/// Builds the extra headers attached to Responses API requests.
///
/// These headers implement Codex-specific conventions:
@@ -783,25 +1199,6 @@ fn build_responses_headers(
headers
}
fn build_websocket_connect_headers(
options: &ApiResponsesOptions,
include_timing_metrics: bool,
) -> ApiHeaderMap {
let mut headers = options.extra_headers.clone();
headers.extend(build_conversation_headers(options.conversation_id.clone()));
headers.insert(
OPENAI_BETA_HEADER,
HeaderValue::from_static(OPENAI_BETA_RESPONSES_WEBSOCKETS),
);
if include_timing_metrics {
headers.insert(
X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER,
HeaderValue::from_static("true"),
);
}
headers
}
fn map_response_stream<S>(api_stream: S, otel_manager: OtelManager) -> ResponseStream
where
S: futures::Stream<Item = std::result::Result<ResponseEvent, ApiError>>
codex-rs/core/src/codex.rs
+19 -10
View File
@@ -38,6 +38,7 @@ use crate::stream_events_utils::last_assistant_message_from_item;
use crate::terminal;
use crate::truncate::TruncationPolicy;
use crate::turn_metadata::build_turn_metadata_header;
use crate::turn_metadata::resolve_turn_metadata_header_with_timeout;
use crate::util::error_or_panic;
use async_channel::Receiver;
use async_channel::Sender;
@@ -570,22 +571,22 @@ impl TurnContext {
.clone()
}
/// Resolves the per-turn metadata header under a shared timeout policy.
///
/// This uses the same timeout helper as websocket startup preconnect so both turn execution
/// and background preconnect observe identical "timeout means best-effort fallback" behavior.
pub async fn resolve_turn_metadata_header(&self) -> Option<String> {
const TURN_METADATA_HEADER_TIMEOUT_MS: u64 = 250;
match tokio::time::timeout(
std::time::Duration::from_millis(TURN_METADATA_HEADER_TIMEOUT_MS),
resolve_turn_metadata_header_with_timeout(
self.build_turn_metadata_header(),
self.turn_metadata_header.get().cloned().flatten(),
)
.await
{
Ok(header) => header,
Err(_) => {
warn!("timed out after 250ms while building turn metadata header");
self.turn_metadata_header.get().cloned().flatten()
}
}
}
/// Starts best-effort background computation of turn metadata.
///
/// This warms the cached value used by [`TurnContext::resolve_turn_metadata_header`] so turns
/// and websocket preconnect are less likely to pay metadata construction latency on demand.
pub fn spawn_turn_metadata_header_task(self: &Arc<Self>) {
let context = Arc::clone(self);
tokio::spawn(async move {
@@ -1106,6 +1107,14 @@ impl Session {
next_internal_sub_id: AtomicU64::new(0),
});
// Warm a websocket in the background so the first turn can reuse it.
// This performs only connection setup; user input is still sent later via response.create
// when submit_turn() runs.
sess.services.model_client.pre_establish_connection(
sess.services.otel_manager.clone(),
session_configuration.cwd.clone(),
);
// Dispatch the SessionConfiguredEvent first and then report any errors.
// If resuming, include converted initial messages in the payload so UIs can render them immediately.
let initial_messages = initial_history.get_event_msgs();
codex-rs/core/src/config/mod.rs
+5
View File
@@ -1,6 +1,7 @@
use crate::auth::AuthCredentialsStoreMode;
use crate::config::edit::ConfigEdit;
use crate::config::edit::ConfigEditsBuilder;
use crate::config::types::AppsConfigToml;
use crate::config::types::DEFAULT_OTEL_ENVIRONMENT;
use crate::config::types::History;
use crate::config::types::McpServerConfig;
@@ -1024,6 +1025,10 @@ pub struct ConfigToml {
/// Defaults to `true`.
pub feedback: Option<crate::config::types::FeedbackConfigToml>,
/// Settings for app-specific controls.
#[serde(default)]
pub apps: Option<AppsConfigToml>,
/// OTEL configuration.
pub otel: Option<crate::config::types::OtelConfigToml>,
codex-rs/core/src/config/service.rs
+59
View File
@@ -700,6 +700,9 @@ fn find_effective_layer(
mod tests {
use super::*;
use anyhow::Result;
use codex_app_server_protocol::AppConfig;
use codex_app_server_protocol::AppDisabledReason;
use codex_app_server_protocol::AppsConfig;
use codex_app_server_protocol::AskForApproval;
use codex_utils_absolute_path::AbsolutePathBuf;
use pretty_assertions::assert_eq;
@@ -800,6 +803,62 @@ remote_models = true
Ok(())
}
#[tokio::test]
async fn write_value_supports_nested_app_paths() -> Result<()> {
let tmp = tempdir().expect("tempdir");
std::fs::write(tmp.path().join(CONFIG_TOML_FILE), "")?;
let service = ConfigService::new_with_defaults(tmp.path().to_path_buf());
service
.write_value(ConfigValueWriteParams {
file_path: Some(tmp.path().join(CONFIG_TOML_FILE).display().to_string()),
key_path: "apps".to_string(),
value: serde_json::json!({
"app1": {
"enabled": false,
},
}),
merge_strategy: MergeStrategy::Replace,
expected_version: None,
})
.await
.expect("write apps succeeds");
service
.write_value(ConfigValueWriteParams {
file_path: Some(tmp.path().join(CONFIG_TOML_FILE).display().to_string()),
key_path: "apps.app1.disabled_reason".to_string(),
value: serde_json::json!("user"),
merge_strategy: MergeStrategy::Replace,
expected_version: None,
})
.await
.expect("write apps.app1.disabled_reason succeeds");
let read = service
.read(ConfigReadParams {
include_layers: false,
cwd: None,
})
.await
.expect("config read succeeds");
assert_eq!(
read.config.apps,
Some(AppsConfig {
apps: std::collections::HashMap::from([(
"app1".to_string(),
AppConfig {
enabled: false,
disabled_reason: Some(AppDisabledReason::User),
},
)]),
})
);
Ok(())
}
#[tokio::test]
async fn read_includes_origins_and_layers() {
let tmp = tempdir().expect("tempdir");
codex-rs/core/src/config/types.rs
+38
View File
@@ -340,6 +340,44 @@ pub struct FeedbackConfigToml {
pub enabled: Option<bool>,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum AppDisabledReason {
Unknown,
User,
}
impl fmt::Display for AppDisabledReason {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
AppDisabledReason::Unknown => write!(f, "unknown"),
AppDisabledReason::User => write!(f, "user"),
}
}
}
/// Config values for a single app/connector.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema)]
#[schemars(deny_unknown_fields)]
pub struct AppConfig {
/// When `false`, Codex does not surface this app.
#[serde(default = "default_enabled")]
pub enabled: bool,
/// Reason this app was disabled.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub disabled_reason: Option<AppDisabledReason>,
}
/// App/connector settings loaded from `config.toml`.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema)]
#[schemars(deny_unknown_fields)]
pub struct AppsConfigToml {
/// Per-app settings keyed by app ID (for example `[apps.google_drive]`).
#[serde(default, flatten)]
pub apps: HashMap<String, AppConfig>,
}
// ===== OTEL configuration =====
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema)]
codex-rs/core/src/turn_metadata.rs
+37
View File
@@ -1,12 +1,49 @@
//! Helpers for computing and resolving optional per-turn metadata headers.
//!
//! This module owns both metadata construction and the shared timeout policy used by
//! turn execution and startup websocket preconnect. Keeping timeout behavior centralized
//! ensures both call sites treat timeout as the same best-effort fallback condition.
use std::collections::BTreeMap;
use std::future::Future;
use std::path::Path;
use std::time::Duration;
use serde::Serialize;
use tracing::warn;
use crate::git_info::get_git_remote_urls_assume_git_repo;
use crate::git_info::get_git_repo_root;
use crate::git_info::get_head_commit_hash;
/// Timeout used when resolving the optional turn-metadata header.
pub(crate) const TURN_METADATA_HEADER_TIMEOUT: Duration = Duration::from_millis(250);
/// Resolves turn metadata with a shared timeout policy.
///
/// On timeout, this logs a warning and returns the provided fallback header.
///
/// Keeping this helper centralized avoids drift between turn-time metadata resolution and startup
/// websocket preconnect, both of which need identical timeout semantics.
pub(crate) async fn resolve_turn_metadata_header_with_timeout<F>(
build_header: F,
fallback_on_timeout: Option<String>,
) -> Option<String>
where
F: Future<Output = Option<String>>,
{
match tokio::time::timeout(TURN_METADATA_HEADER_TIMEOUT, build_header).await {
Ok(header) => header,
Err(_) => {
warn!(
"timed out after {}ms while building turn metadata header",
TURN_METADATA_HEADER_TIMEOUT.as_millis()
);
fallback_on_timeout
}
}
}
#[derive(Serialize)]
struct TurnMetadataWorkspace {
#[serde(skip_serializing_if = "Option::is_none")]
codex-rs/core/tests/common/responses.rs
+33
View File
@@ -268,6 +268,11 @@ impl WebSocketHandshake {
pub struct WebSocketConnectionConfig {
pub requests: Vec<Vec<Value>>,
pub response_headers: Vec<(String, String)>,
/// Optional delay inserted before accepting the websocket handshake.
///
/// Tests use this to force startup preconnect into an in-flight state so first-turn adoption
/// paths can be exercised deterministically.
pub accept_delay: Option<Duration>,
}
pub struct WebSocketTestServer {
@@ -299,6 +304,29 @@ impl WebSocketTestServer {
self.handshakes.lock().unwrap().clone()
}
/// Waits until at least `expected` websocket handshakes have been observed or timeout elapses.
///
/// Uses a short bounded polling interval so tests can deterministically wait for background
/// preconnect activity without busy-spinning.
pub async fn wait_for_handshakes(&self, expected: usize, timeout: Duration) -> bool {
if self.handshakes.lock().unwrap().len() >= expected {
return true;
}
let deadline = tokio::time::Instant::now() + timeout;
let poll_interval = Duration::from_millis(10);
loop {
if self.handshakes.lock().unwrap().len() >= expected {
return true;
}
let now = tokio::time::Instant::now();
if now >= deadline {
return false;
}
let sleep_for = std::cmp::min(poll_interval, deadline.saturating_duration_since(now));
tokio::time::sleep(sleep_for).await;
}
}
pub fn single_handshake(&self) -> WebSocketHandshake {
let handshakes = self.handshakes.lock().unwrap();
if handshakes.len() != 1 {
@@ -861,6 +889,7 @@ pub async fn start_websocket_server(connections: Vec<Vec<Vec<Value>>>) -> WebSoc
.map(|requests| WebSocketConnectionConfig {
requests,
response_headers: Vec::new(),
accept_delay: None,
})
.collect();
start_websocket_server_with_headers(connections).await
@@ -900,6 +929,10 @@ pub async fn start_websocket_server_with_headers(
continue;
};
if let Some(delay) = connection.accept_delay {
tokio::time::sleep(delay).await;
}
let response_headers = connection.response_headers.clone();
let handshake_log = Arc::clone(&handshakes);
let callback = move |req: &Request, mut response: Response| {
codex-rs/core/tests/suite/agent_websocket.rs
+53
View File
@@ -1,14 +1,17 @@
use anyhow::Result;
use core_test_support::responses::WebSocketConnectionConfig;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_done;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::ev_shell_command_call;
use core_test_support::responses::start_websocket_server;
use core_test_support::responses::start_websocket_server_with_headers;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::test_codex;
use pretty_assertions::assert_eq;
use serde_json::Value;
use std::time::Duration;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn websocket_test_codex_shell_chain() -> Result<()> {
@@ -67,3 +70,53 @@ async fn websocket_test_codex_shell_chain() -> Result<()> {
server.shutdown().await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn websocket_preconnect_happens_on_session_start() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_websocket_server(vec![vec![vec![
ev_response_created("resp-1"),
ev_completed("resp-1"),
]]])
.await;
let mut builder = test_codex();
let test = builder.build_with_websocket_server(&server).await?;
assert!(
server.wait_for_handshakes(1, Duration::from_secs(2)).await,
"expected websocket preconnect handshake during session startup"
);
test.submit_turn("hello").await?;
assert_eq!(server.handshakes().len(), 1);
assert_eq!(server.single_connection().len(), 1);
server.shutdown().await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn websocket_first_turn_waits_for_inflight_preconnect() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_websocket_server_with_headers(vec![WebSocketConnectionConfig {
requests: vec![vec![ev_response_created("resp-1"), ev_completed("resp-1")]],
response_headers: Vec::new(),
// Delay handshake so submit_turn() observes startup preconnect as in-flight.
accept_delay: Some(Duration::from_millis(150)),
}])
.await;
let mut builder = test_codex();
let test = builder.build_with_websocket_server(&server).await?;
test.submit_turn("hello").await?;
assert_eq!(server.handshakes().len(), 1);
assert_eq!(server.single_connection().len(), 1);
server.shutdown().await;
Ok(())
}
codex-rs/core/tests/suite/client_websockets.rs
+64
View File
@@ -85,6 +85,68 @@ async fn responses_websocket_streams_request() {
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_preconnect_reuses_connection() {
skip_if_no_network!();
let server = start_websocket_server(vec![vec![vec![
ev_response_created("resp-1"),
ev_completed("resp-1"),
]]])
.await;
let harness = websocket_harness(&server).await;
assert!(harness.client.preconnect(&harness.otel_manager, None).await);
let mut client_session = harness.client.new_session();
let prompt = prompt_with_input(vec![message_item("hello")]);
stream_until_complete(&mut client_session, &harness, &prompt).await;
assert_eq!(server.handshakes().len(), 1);
assert_eq!(server.single_connection().len(), 1);
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_preconnect_is_reused_even_with_header_changes() {
skip_if_no_network!();
let server = start_websocket_server(vec![vec![vec![
ev_response_created("resp-1"),
ev_completed("resp-1"),
]]])
.await;
let harness = websocket_harness(&server).await;
assert!(harness.client.preconnect(&harness.otel_manager, None).await);
let mut client_session = harness.client.new_session();
let prompt = prompt_with_input(vec![message_item("hello")]);
let mut stream = client_session
.stream(
&prompt,
&harness.model_info,
&harness.otel_manager,
harness.effort,
harness.summary,
None,
)
.await
.expect("websocket stream failed");
while let Some(event) = stream.next().await {
if matches!(event, Ok(ResponseEvent::Completed { .. })) {
break;
}
}
assert_eq!(server.handshakes().len(), 1);
assert_eq!(server.single_connection().len(), 1);
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[traced_test]
async fn responses_websocket_emits_websocket_telemetry_events() {
@@ -198,6 +260,7 @@ async fn responses_websocket_emits_reasoning_included_event() {
let server = start_websocket_server_with_headers(vec![WebSocketConnectionConfig {
requests: vec![vec![ev_response_created("resp-1"), ev_completed("resp-1")]],
response_headers: vec![("X-Reasoning-Included".to_string(), "true".to_string())],
accept_delay: None,
}])
.await;
@@ -268,6 +331,7 @@ async fn responses_websocket_emits_rate_limit_events() {
("X-Models-Etag".to_string(), "etag-123".to_string()),
("X-Reasoning-Included".to_string(), "true".to_string()),
],
accept_delay: None,
}])
.await;
codex-rs/core/tests/suite/turn_state.rs
+3
View File
@@ -83,6 +83,7 @@ async fn websocket_turn_state_persists_within_turn_and_resets_after() -> Result<
ev_done(),
]],
response_headers: vec![(TURN_STATE_HEADER.to_string(), "ts-1".to_string())],
accept_delay: None,
},
WebSocketConnectionConfig {
requests: vec![vec![
@@ -91,6 +92,7 @@ async fn websocket_turn_state_persists_within_turn_and_resets_after() -> Result<
ev_completed("resp-2"),
]],
response_headers: Vec::new(),
accept_delay: None,
},
WebSocketConnectionConfig {
requests: vec![vec![
@@ -99,6 +101,7 @@ async fn websocket_turn_state_persists_within_turn_and_resets_after() -> Result<
ev_completed("resp-3"),
]],
response_headers: Vec::new(),
accept_delay: None,
},
])
.await;
codex-rs/core/tests/suite/websocket_fallback.rs
+8 -2
View File
@@ -46,7 +46,10 @@ async fn websocket_fallback_switches_to_http_after_retries_exhausted() -> Result
.filter(|req| req.method == Method::POST && req.url.path().ends_with("/responses"))
.count();
assert_eq!(websocket_attempts, 1);
// One websocket attempt comes from startup preconnect and one from the first turn's stream
// attempt before fallback activates; after fallback, transport is HTTP. This matches the
// retry-budget tradeoff documented in [`codex_core::client`] module docs.
assert_eq!(websocket_attempts, 2);
assert_eq!(http_attempts, 1);
assert_eq!(response_mock.requests().len(), 1);
@@ -92,7 +95,10 @@ async fn websocket_fallback_is_sticky_across_turns() -> Result<()> {
.filter(|req| req.method == Method::POST && req.url.path().ends_with("/responses"))
.count();
assert_eq!(websocket_attempts, 1);
// The first turn issues exactly two websocket attempts (startup preconnect + first stream
// attempt). After fallback becomes sticky, subsequent turns stay on HTTP. This mirrors the
// retry-budget tradeoff documented in [`codex_core::client`] module docs.
assert_eq!(websocket_attempts, 2);
assert_eq!(http_attempts, 2);
assert_eq!(response_mock.requests().len(), 2);
codex-rs/linux-sandbox/src/landlock.rs
+3
View File
@@ -149,6 +149,9 @@ fn install_network_seccomp_filter_on_current_thread() -> std::result::Result<(),
deny_syscall(libc::SYS_getsockopt);
deny_syscall(libc::SYS_setsockopt);
deny_syscall(libc::SYS_ptrace);
deny_syscall(libc::SYS_io_uring_setup);
deny_syscall(libc::SYS_io_uring_enter);
deny_syscall(libc::SYS_io_uring_register);
// For `socket` we allow AF_UNIX (arg0 == AF_UNIX) and deny everything else.
let unix_only_rule = SeccompRule::new(vec![SeccompCondition::new(
codex-rs/network-proxy/src/http_proxy.rs
+102 -18
View File
@@ -1,6 +1,8 @@
use crate::config::NetworkMode;
use crate::network_policy::NetworkDecision;
use crate::network_policy::NetworkDecisionSource;
use crate::network_policy::NetworkPolicyDecider;
use crate::network_policy::NetworkPolicyDecision;
use crate::network_policy::NetworkPolicyRequest;
use crate::network_policy::NetworkPolicyRequestArgs;
use crate::network_policy::NetworkProtocol;
@@ -9,8 +11,12 @@ use crate::policy::normalize_host;
use crate::reasons::REASON_METHOD_NOT_ALLOWED;
use crate::reasons::REASON_NOT_ALLOWED;
use crate::reasons::REASON_PROXY_DISABLED;
use crate::responses::PolicyDecisionDetails;
use crate::responses::blocked_header_value;
use crate::responses::blocked_message_with_policy;
use crate::responses::blocked_text_response_with_policy;
use crate::responses::json_response;
use crate::responses::policy_decision_prefix;
use crate::runtime::unix_socket_permissions_supported;
use crate::state::BlockedRequest;
use crate::state::BlockedRequestArgs;
@@ -141,9 +147,10 @@ async fn http_connect_accept(
return Err(proxy_disabled_response(
&app_state,
host,
authority.port,
client_addr(&req),
Some("CONNECT".to_string()),
"http-connect",
NetworkProtocol::HttpsConnect,
)
.await);
}
@@ -159,7 +166,19 @@ async fn http_connect_accept(
});
match evaluate_host_policy(&app_state, policy_decider.as_ref(), &request).await {
Ok(NetworkDecision::Deny { reason }) => {
Ok(NetworkDecision::Deny {
reason,
source,
decision,
}) => {
let details = PolicyDecisionDetails {
decision,
reason: &reason,
source,
protocol: NetworkProtocol::HttpsConnect,
host: &host,
port: authority.port,
};
let _ = app_state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host: host.clone(),
@@ -172,7 +191,7 @@ async fn http_connect_accept(
.await;
let client = client.as_deref().unwrap_or_default();
warn!("CONNECT blocked (client={client}, host={host}, reason={reason})");
return Err(blocked_text(&reason));
return Err(blocked_text_with_details(&reason, &details));
}
Ok(NetworkDecision::Allow) => {
let client = client.as_deref().unwrap_or_default();
@@ -190,6 +209,14 @@ async fn http_connect_accept(
.map_err(|err| internal_error("failed to read network mode", err))?;
if mode == NetworkMode::Limited {
let details = PolicyDecisionDetails {
decision: NetworkPolicyDecision::Deny,
reason: REASON_METHOD_NOT_ALLOWED,
source: NetworkDecisionSource::ModeGuard,
protocol: NetworkProtocol::HttpsConnect,
host: &host,
port: authority.port,
};
let _ = app_state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host: host.clone(),
@@ -202,7 +229,10 @@ async fn http_connect_accept(
.await;
let client = client.as_deref().unwrap_or_default();
warn!("CONNECT blocked by method policy (client={client}, host={host}, mode=limited)");
return Err(blocked_text(REASON_METHOD_NOT_ALLOWED));
return Err(blocked_text_with_details(
REASON_METHOD_NOT_ALLOWED,
&details,
));
}
req.extensions_mut().insert(ProxyTarget(authority));
@@ -346,9 +376,10 @@ async fn http_plain_proxy(
return Ok(proxy_disabled_response(
&app_state,
socket_path,
0,
client_addr(&req),
Some(req.method().as_str().to_string()),
"unix-socket",
NetworkProtocol::Http,
)
.await);
}
@@ -358,7 +389,7 @@ async fn http_plain_proxy(
warn!(
"unix socket blocked by method policy (client={client}, method={method}, mode=limited, allowed_methods=GET, HEAD, OPTIONS)"
);
return Ok(json_blocked("unix-socket", REASON_METHOD_NOT_ALLOWED));
return Ok(json_blocked("unix-socket", REASON_METHOD_NOT_ALLOWED, None));
}
if !unix_socket_permissions_supported() {
@@ -387,7 +418,7 @@ async fn http_plain_proxy(
Ok(false) => {
let client = client.as_deref().unwrap_or_default();
warn!("unix socket blocked (client={client}, path={socket_path})");
Ok(json_blocked("unix-socket", REASON_NOT_ALLOWED))
Ok(json_blocked("unix-socket", REASON_NOT_ALLOWED, None))
}
Err(err) => {
warn!("unix socket check failed: {err}");
@@ -420,9 +451,10 @@ async fn http_plain_proxy(
return Ok(proxy_disabled_response(
&app_state,
host,
port,
client_addr(&req),
Some(req.method().as_str().to_string()),
"http",
NetworkProtocol::Http,
)
.await);
}
@@ -438,7 +470,19 @@ async fn http_plain_proxy(
});
match evaluate_host_policy(&app_state, policy_decider.as_ref(), &request).await {
Ok(NetworkDecision::Deny { reason }) => {
Ok(NetworkDecision::Deny {
reason,
source,
decision,
}) => {
let details = PolicyDecisionDetails {
decision,
reason: &reason,
source,
protocol: NetworkProtocol::Http,
host: &host,
port,
};
let _ = app_state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host: host.clone(),
@@ -451,7 +495,7 @@ async fn http_plain_proxy(
.await;
let client = client.as_deref().unwrap_or_default();
warn!("request blocked (client={client}, host={host}, reason={reason})");
return Ok(json_blocked(&host, &reason));
return Ok(json_blocked(&host, &reason, Some(&details)));
}
Ok(NetworkDecision::Allow) => {}
Err(err) => {
@@ -461,6 +505,14 @@ async fn http_plain_proxy(
}
if !method_allowed {
let details = PolicyDecisionDetails {
decision: NetworkPolicyDecision::Deny,
reason: REASON_METHOD_NOT_ALLOWED,
source: NetworkDecisionSource::ModeGuard,
protocol: NetworkProtocol::Http,
host: &host,
port,
};
let _ = app_state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host: host.clone(),
@@ -476,7 +528,11 @@ async fn http_plain_proxy(
warn!(
"request blocked by method policy (client={client}, host={host}, method={method}, mode=limited, allowed_methods=GET, HEAD, OPTIONS)"
);
return Ok(json_blocked(&host, REASON_METHOD_NOT_ALLOWED));
return Ok(json_blocked(
&host,
REASON_METHOD_NOT_ALLOWED,
Some(&details),
));
}
let client = client.as_deref().unwrap_or_default();
@@ -540,11 +596,21 @@ fn client_addr<T: ExtensionsRef>(input: &T) -> Option<String> {
.map(|info| info.peer_addr().to_string())
}
fn json_blocked(host: &str, reason: &str) -> Response {
fn json_blocked(host: &str, reason: &str, details: Option<&PolicyDecisionDetails<'_>>) -> Response {
let (policy_decision_prefix, message) = details
.map(|details| {
(
Some(policy_decision_prefix(details)),
Some(blocked_message_with_policy(reason, details)),
)
})
.unwrap_or((None, None));
let response = BlockedResponse {
status: "blocked",
host,
reason,
policy_decision_prefix,
message,
};
let mut resp = json_response(&response);
*resp.status_mut() = StatusCode::FORBIDDEN;
@@ -555,28 +621,42 @@ fn json_blocked(host: &str, reason: &str) -> Response {
resp
}
fn blocked_text(reason: &str) -> Response {
crate::responses::blocked_text_response(reason)
fn blocked_text_with_details(reason: &str, details: &PolicyDecisionDetails<'_>) -> Response {
blocked_text_response_with_policy(reason, details)
}
async fn proxy_disabled_response(
app_state: &NetworkProxyState,
host: String,
port: u16,
client: Option<String>,
method: Option<String>,
protocol: &str,
protocol: NetworkProtocol,
) -> Response {
let blocked_host = host.clone();
let _ = app_state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host,
host: blocked_host,
reason: REASON_PROXY_DISABLED.to_string(),
client,
method,
mode: None,
protocol: protocol.to_string(),
protocol: protocol.as_policy_protocol().to_string(),
}))
.await;
text_response(StatusCode::SERVICE_UNAVAILABLE, "proxy disabled")
let details = PolicyDecisionDetails {
decision: NetworkPolicyDecision::Deny,
reason: REASON_PROXY_DISABLED,
source: NetworkDecisionSource::ProxyState,
protocol,
host: &host,
port,
};
text_response(
StatusCode::SERVICE_UNAVAILABLE,
&blocked_message_with_policy(REASON_PROXY_DISABLED, &details),
)
}
fn internal_error(context: &str, err: impl std::fmt::Display) -> Response {
@@ -597,6 +677,10 @@ struct BlockedResponse<'a> {
status: &'static str,
host: &'a str,
reason: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
policy_decision_prefix: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
message: Option<String>,
}
#[cfg(test)]
codex-rs/network-proxy/src/network_policy.rs
+101 -7
View File
@@ -15,6 +15,51 @@ pub enum NetworkProtocol {
Socks5Udp,
}
impl NetworkProtocol {
pub const fn as_policy_protocol(self) -> &'static str {
match self {
Self::Http => "http",
Self::HttpsConnect => "https_connect",
Self::Socks5Tcp => "socks5_tcp",
Self::Socks5Udp => "socks5_udp",
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NetworkPolicyDecision {
Deny,
Ask,
}
impl NetworkPolicyDecision {
pub const fn as_str(self) -> &'static str {
match self {
Self::Deny => "deny",
Self::Ask => "ask",
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NetworkDecisionSource {
BaselinePolicy,
ModeGuard,
ProxyState,
Decider,
}
impl NetworkDecisionSource {
pub const fn as_str(self) -> &'static str {
match self {
Self::BaselinePolicy => "baseline_policy",
Self::ModeGuard => "mode_guard",
Self::ProxyState => "proxy_state",
Self::Decider => "decider",
}
}
}
#[derive(Clone, Debug)]
pub struct NetworkPolicyRequest {
pub protocol: NetworkProtocol,
@@ -62,18 +107,44 @@ impl NetworkPolicyRequest {
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum NetworkDecision {
Allow,
Deny { reason: String },
Deny {
reason: String,
source: NetworkDecisionSource,
decision: NetworkPolicyDecision,
},
}
impl NetworkDecision {
pub fn deny(reason: impl Into<String>) -> Self {
Self::deny_with_source(reason, NetworkDecisionSource::Decider)
}
pub fn deny_with_source(reason: impl Into<String>, source: NetworkDecisionSource) -> Self {
let reason = reason.into();
let reason = if reason.is_empty() {
REASON_POLICY_DENIED.to_string()
} else {
reason
};
Self::Deny { reason }
Self::Deny {
reason,
source,
decision: NetworkPolicyDecision::Deny,
}
}
pub fn ask_with_source(reason: impl Into<String>, source: NetworkDecisionSource) -> Self {
let reason = reason.into();
let reason = if reason.is_empty() {
REASON_POLICY_DENIED.to_string()
} else {
reason
};
Self::Deny {
reason,
source,
decision: NetworkPolicyDecision::Ask,
}
}
}
@@ -114,12 +185,31 @@ pub(crate) async fn evaluate_host_policy(
HostBlockDecision::Allowed => Ok(NetworkDecision::Allow),
HostBlockDecision::Blocked(HostBlockReason::NotAllowed) => {
if let Some(decider) = decider {
Ok(decider.decide(request.clone()).await)
Ok(map_decider_decision(decider.decide(request.clone()).await))
} else {
Ok(NetworkDecision::deny(HostBlockReason::NotAllowed.as_str()))
Ok(NetworkDecision::deny_with_source(
HostBlockReason::NotAllowed.as_str(),
NetworkDecisionSource::BaselinePolicy,
))
}
}
HostBlockDecision::Blocked(reason) => Ok(NetworkDecision::deny(reason.as_str())),
HostBlockDecision::Blocked(reason) => Ok(NetworkDecision::deny_with_source(
reason.as_str(),
NetworkDecisionSource::BaselinePolicy,
)),
}
}
fn map_decider_decision(decision: NetworkDecision) -> NetworkDecision {
match decision {
NetworkDecision::Allow => NetworkDecision::Allow,
NetworkDecision::Deny {
reason, decision, ..
} => NetworkDecision::Deny {
reason,
source: NetworkDecisionSource::Decider,
decision,
},
}
}
@@ -199,7 +289,9 @@ mod tests {
assert_eq!(
decision,
NetworkDecision::Deny {
reason: REASON_DENIED.to_string()
reason: REASON_DENIED.to_string(),
source: NetworkDecisionSource::BaselinePolicy,
decision: NetworkPolicyDecision::Deny,
}
);
assert_eq!(calls.load(Ordering::SeqCst), 0);
@@ -237,7 +329,9 @@ mod tests {
assert_eq!(
decision,
NetworkDecision::Deny {
reason: REASON_NOT_ALLOWED_LOCAL.to_string()
reason: REASON_NOT_ALLOWED_LOCAL.to_string(),
source: NetworkDecisionSource::BaselinePolicy,
decision: NetworkPolicyDecision::Deny,
}
);
assert_eq!(calls.load(Ordering::SeqCst), 0);
codex-rs/network-proxy/src/responses.rs
+101 -2
View File
@@ -1,3 +1,6 @@
use crate::network_policy::NetworkDecisionSource;
use crate::network_policy::NetworkPolicyDecision;
use crate::network_policy::NetworkProtocol;
use crate::reasons::REASON_DENIED;
use crate::reasons::REASON_METHOD_NOT_ALLOWED;
use crate::reasons::REASON_NOT_ALLOWED;
@@ -8,6 +11,28 @@ use rama_http::StatusCode;
use serde::Serialize;
use tracing::error;
const NETWORK_POLICY_DECISION_PREFIX: &str = "CODEX_NETWORK_POLICY_DECISION";
pub struct PolicyDecisionDetails<'a> {
pub decision: NetworkPolicyDecision,
pub reason: &'a str,
pub source: NetworkDecisionSource,
pub protocol: NetworkProtocol,
pub host: &'a str,
pub port: u16,
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct PolicyDecisionPayload<'a> {
decision: &'a str,
reason: &'a str,
source: &'a str,
protocol: &'a str,
host: &'a str,
port: u16,
}
pub fn text_response(status: StatusCode, body: &str) -> Response {
Response::builder()
.status(status)
@@ -57,11 +82,85 @@ pub fn blocked_message(reason: &str) -> &'static str {
}
}
pub fn blocked_text_response(reason: &str) -> Response {
pub fn policy_decision_prefix(details: &PolicyDecisionDetails<'_>) -> String {
let payload = PolicyDecisionPayload {
decision: details.decision.as_str(),
reason: details.reason,
source: details.source.as_str(),
protocol: details.protocol.as_policy_protocol(),
host: details.host,
port: details.port,
};
let payload_json = match serde_json::to_string(&payload) {
Ok(json) => json,
Err(err) => {
error!("failed to serialize policy decision payload: {err}");
"{}".to_string()
}
};
format!("{NETWORK_POLICY_DECISION_PREFIX} {payload_json}")
}
pub fn blocked_message_with_policy(reason: &str, details: &PolicyDecisionDetails<'_>) -> String {
format!(
"{}\n{}",
policy_decision_prefix(details),
blocked_message(reason)
)
}
pub fn blocked_text_response_with_policy(
reason: &str,
details: &PolicyDecisionDetails<'_>,
) -> Response {
Response::builder()
.status(StatusCode::FORBIDDEN)
.header("content-type", "text/plain")
.header("x-proxy-error", blocked_header_value(reason))
.body(Body::from(blocked_message(reason)))
.body(Body::from(blocked_message_with_policy(reason, details)))
.unwrap_or_else(|_| Response::new(Body::from("blocked")))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::reasons::REASON_NOT_ALLOWED;
use pretty_assertions::assert_eq;
#[test]
fn policy_decision_prefix_serializes_expected_payload() {
let details = PolicyDecisionDetails {
decision: NetworkPolicyDecision::Ask,
reason: REASON_NOT_ALLOWED,
source: NetworkDecisionSource::Decider,
protocol: NetworkProtocol::HttpsConnect,
host: "api.example.com",
port: 443,
};
let line = policy_decision_prefix(&details);
assert_eq!(
line,
r#"CODEX_NETWORK_POLICY_DECISION {"decision":"ask","reason":"not_allowed","source":"decider","protocol":"https_connect","host":"api.example.com","port":443}"#
);
}
#[test]
fn blocked_message_with_policy_includes_prefix_and_human_message() {
let details = PolicyDecisionDetails {
decision: NetworkPolicyDecision::Deny,
reason: REASON_NOT_ALLOWED,
source: NetworkDecisionSource::BaselinePolicy,
protocol: NetworkProtocol::Http,
host: "api.example.com",
port: 80,
};
let message = blocked_message_with_policy(REASON_NOT_ALLOWED, &details);
assert_eq!(
message,
r#"CODEX_NETWORK_POLICY_DECISION {"decision":"deny","reason":"not_allowed","source":"baseline_policy","protocol":"http","host":"api.example.com","port":80}
Codex blocked this request: domain not in allowlist."#
);
}
}
codex-rs/network-proxy/src/socks5.rs
+75 -17
View File
@@ -1,6 +1,8 @@
use crate::config::NetworkMode;
use crate::network_policy::NetworkDecision;
use crate::network_policy::NetworkDecisionSource;
use crate::network_policy::NetworkPolicyDecider;
use crate::network_policy::NetworkPolicyDecision;
use crate::network_policy::NetworkPolicyRequest;
use crate::network_policy::NetworkPolicyRequestArgs;
use crate::network_policy::NetworkProtocol;
@@ -8,6 +10,8 @@ use crate::network_policy::evaluate_host_policy;
use crate::policy::normalize_host;
use crate::reasons::REASON_METHOD_NOT_ALLOWED;
use crate::reasons::REASON_PROXY_DISABLED;
use crate::responses::PolicyDecisionDetails;
use crate::responses::blocked_message_with_policy;
use crate::state::BlockedRequest;
use crate::state::BlockedRequestArgs;
use crate::state::NetworkProxyState;
@@ -123,6 +127,14 @@ async fn handle_socks5_tcp(
match app_state.enabled().await {
Ok(true) => {}
Ok(false) => {
let details = PolicyDecisionDetails {
decision: NetworkPolicyDecision::Deny,
reason: REASON_PROXY_DISABLED,
source: NetworkDecisionSource::ProxyState,
protocol: NetworkProtocol::Socks5Tcp,
host: &host,
port,
};
let _ = app_state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host: host.clone(),
@@ -135,7 +147,7 @@ async fn handle_socks5_tcp(
.await;
let client = client.as_deref().unwrap_or_default();
warn!("SOCKS blocked; proxy disabled (client={client}, host={host})");
return Err(io::Error::new(io::ErrorKind::PermissionDenied, "proxy disabled").into());
return Err(policy_denied_error(REASON_PROXY_DISABLED, &details).into());
}
Err(err) => {
error!("failed to read enabled state: {err}");
@@ -145,6 +157,14 @@ async fn handle_socks5_tcp(
match app_state.network_mode().await {
Ok(NetworkMode::Limited) => {
let details = PolicyDecisionDetails {
decision: NetworkPolicyDecision::Deny,
reason: REASON_METHOD_NOT_ALLOWED,
source: NetworkDecisionSource::ModeGuard,
protocol: NetworkProtocol::Socks5Tcp,
host: &host,
port,
};
let _ = app_state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host: host.clone(),
@@ -159,7 +179,7 @@ async fn handle_socks5_tcp(
warn!(
"SOCKS blocked by method policy (client={client}, host={host}, mode=limited, allowed_methods=GET, HEAD, OPTIONS)"
);
return Err(io::Error::new(io::ErrorKind::PermissionDenied, "blocked").into());
return Err(policy_denied_error(REASON_METHOD_NOT_ALLOWED, &details).into());
}
Ok(NetworkMode::Full) => {}
Err(err) => {
@@ -179,7 +199,19 @@ async fn handle_socks5_tcp(
});
match evaluate_host_policy(&app_state, policy_decider.as_ref(), &request).await {
Ok(NetworkDecision::Deny { reason }) => {
Ok(NetworkDecision::Deny {
reason,
source,
decision,
}) => {
let details = PolicyDecisionDetails {
decision,
reason: &reason,
source,
protocol: NetworkProtocol::Socks5Tcp,
host: &host,
port,
};
let _ = app_state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host: host.clone(),
@@ -192,7 +224,7 @@ async fn handle_socks5_tcp(
.await;
let client = client.as_deref().unwrap_or_default();
warn!("SOCKS blocked (client={client}, host={host}, reason={reason})");
return Err(io::Error::new(io::ErrorKind::PermissionDenied, "blocked").into());
return Err(policy_denied_error(&reason, &details).into());
}
Ok(NetworkDecision::Allow) => {
let client = client.as_deref().unwrap_or_default();
@@ -232,6 +264,14 @@ async fn inspect_socks5_udp(
match state.enabled().await {
Ok(true) => {}
Ok(false) => {
let details = PolicyDecisionDetails {
decision: NetworkPolicyDecision::Deny,
reason: REASON_PROXY_DISABLED,
source: NetworkDecisionSource::ProxyState,
protocol: NetworkProtocol::Socks5Udp,
host: &host,
port,
};
let _ = state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host: host.clone(),
@@ -244,10 +284,7 @@ async fn inspect_socks5_udp(
.await;
let client = client.as_deref().unwrap_or_default();
warn!("SOCKS UDP blocked; proxy disabled (client={client}, host={host})");
return Ok(RelayResponse {
maybe_payload: None,
extensions,
});
return Err(policy_denied_error(REASON_PROXY_DISABLED, &details));
}
Err(err) => {
error!("failed to read enabled state: {err}");
@@ -257,6 +294,14 @@ async fn inspect_socks5_udp(
match state.network_mode().await {
Ok(NetworkMode::Limited) => {
let details = PolicyDecisionDetails {
decision: NetworkPolicyDecision::Deny,
reason: REASON_METHOD_NOT_ALLOWED,
source: NetworkDecisionSource::ModeGuard,
protocol: NetworkProtocol::Socks5Udp,
host: &host,
port,
};
let _ = state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host: host.clone(),
@@ -267,10 +312,7 @@ async fn inspect_socks5_udp(
protocol: "socks5-udp".to_string(),
}))
.await;
return Ok(RelayResponse {
maybe_payload: None,
extensions,
});
return Err(policy_denied_error(REASON_METHOD_NOT_ALLOWED, &details));
}
Ok(NetworkMode::Full) => {}
Err(err) => {
@@ -290,7 +332,19 @@ async fn inspect_socks5_udp(
});
match evaluate_host_policy(&state, policy_decider.as_ref(), &request).await {
Ok(NetworkDecision::Deny { reason }) => {
Ok(NetworkDecision::Deny {
reason,
source,
decision,
}) => {
let details = PolicyDecisionDetails {
decision,
reason: &reason,
source,
protocol: NetworkProtocol::Socks5Udp,
host: &host,
port,
};
let _ = state
.record_blocked(BlockedRequest::new(BlockedRequestArgs {
host: host.clone(),
@@ -303,10 +357,7 @@ async fn inspect_socks5_udp(
.await;
let client = client.as_deref().unwrap_or_default();
warn!("SOCKS UDP blocked (client={client}, host={host}, reason={reason})");
Ok(RelayResponse {
maybe_payload: None,
extensions,
})
Err(policy_denied_error(&reason, &details))
}
Ok(NetworkDecision::Allow) => Ok(RelayResponse {
maybe_payload: Some(payload),
@@ -318,3 +369,10 @@ async fn inspect_socks5_udp(
}
}
}
fn policy_denied_error(reason: &str, details: &PolicyDecisionDetails<'_>) -> io::Error {
io::Error::new(
io::ErrorKind::PermissionDenied,
blocked_message_with_policy(reason, details),
)
}
codex-rs/tui/src/chatwidget.rs
+40 -23
View File
@@ -1643,16 +1643,7 @@ impl ChatWidget {
}
if let Some(combined) = self.drain_queued_messages_for_restore() {
let combined_local_image_paths = combined
.local_images
.iter()
.map(|img| img.path.clone())
.collect();
self.bottom_pane.set_composer_text(
combined.text,
combined.text_elements,
combined_local_image_paths,
);
self.restore_user_message_to_composer(combined);
self.refresh_queued_user_messages();
}
@@ -1715,6 +1706,18 @@ impl ChatWidget {
Some(combined)
}
fn restore_user_message_to_composer(&mut self, user_message: UserMessage) {
let UserMessage {
text,
local_images,
text_elements,
mention_paths: _,
} = user_message;
let local_image_paths = local_images.into_iter().map(|img| img.path).collect();
self.bottom_pane
.set_composer_text(text, text_elements, local_image_paths);
}
fn on_plan_update(&mut self, update: UpdatePlanArgs) {
self.saw_plan_update_this_turn = true;
self.add_to_history(history_cell::new_plan_update(update));
@@ -3003,16 +3006,7 @@ impl ChatWidget {
} if !self.queued_user_messages.is_empty() => {
// Prefer the most recently queued item.
if let Some(user_message) = self.queued_user_messages.pop_back() {
let local_image_paths = user_message
.local_images
.iter()
.map(|img| img.path.clone())
.collect();
self.bottom_pane.set_composer_text(
user_message.text,
user_message.text_elements,
local_image_paths,
);
self.restore_user_message_to_composer(user_message);
self.refresh_queued_user_messages();
self.request_redraw();
}
@@ -3030,8 +3024,8 @@ impl ChatWidget {
text_elements,
mention_paths: self.bottom_pane.take_mention_paths(),
};
if self.is_session_configured() {
// Submitted is only emitted when steer is enabled (Enter sends immediately).
if self.is_session_configured() && !self.is_plan_streaming_in_tui() {
// Submitted is only emitted when steer is enabled.
// Reset any reasoning header only when we are actually submitting a turn.
self.reasoning_buffer.clear();
self.full_reasoning_buffer.clear();
@@ -6415,6 +6409,10 @@ impl ChatWidget {
self.bottom_pane.is_task_running() || self.is_review_mode
}
fn is_plan_streaming_in_tui(&self) -> bool {
self.plan_stream_controller.is_some()
}
pub(crate) fn composer_is_empty(&self) -> bool {
self.bottom_pane.composer_is_empty()
}
@@ -6424,8 +6422,27 @@ impl ChatWidget {
text: String,
collaboration_mode: CollaborationModeMask,
) {
if self.agent_turn_running
&& self.active_collaboration_mask.as_ref() != Some(&collaboration_mode)
{
self.add_error_message(
"Cannot switch collaboration mode while a turn is running.".to_string(),
);
return;
}
self.set_collaboration_mask(collaboration_mode);
self.submit_user_message(text.into());
let should_queue = self.is_plan_streaming_in_tui();
let user_message = UserMessage {
text,
local_images: Vec::new(),
text_elements: Vec::new(),
mention_paths: HashMap::new(),
};
if should_queue {
self.queue_user_message(user_message);
} else {
self.submit_user_message(user_message);
}
}
/// True when the UI is in the regular composer state with no running task,
codex-rs/tui/src/chatwidget/tests.rs
+188
View File
@@ -527,6 +527,54 @@ async fn interrupted_turn_restores_queued_messages_with_images_and_elements() {
);
}
#[tokio::test]
async fn interrupted_turn_restore_keeps_active_mode_for_resubmission() {
let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(Some("gpt-5")).await;
chat.thread_id = Some(ThreadId::new());
chat.set_feature_enabled(Feature::CollaborationModes, true);
let plan_mask = collaboration_modes::plan_mask(chat.models_manager.as_ref())
.expect("expected plan collaboration mode");
let expected_mode = plan_mask
.mode
.expect("expected mode kind on plan collaboration mode");
chat.set_collaboration_mask(plan_mask);
chat.on_task_started();
chat.queued_user_messages.push_back(UserMessage {
text: "Implement the plan.".to_string(),
local_images: Vec::new(),
text_elements: Vec::new(),
mention_paths: HashMap::new(),
});
chat.refresh_queued_user_messages();
chat.handle_codex_event(Event {
id: "interrupt".into(),
msg: EventMsg::TurnAborted(codex_core::protocol::TurnAbortedEvent {
reason: TurnAbortReason::Interrupted,
}),
});
assert_eq!(chat.bottom_pane.composer_text(), "Implement the plan.");
assert!(chat.queued_user_messages.is_empty());
assert_eq!(chat.active_collaboration_mode_kind(), expected_mode);
chat.handle_key_event(KeyEvent::from(KeyCode::Enter));
match next_submit_op(&mut op_rx) {
Op::UserTurn {
collaboration_mode: Some(CollaborationMode { mode, .. }),
personality: None,
..
} => assert_eq!(mode, expected_mode),
other => {
panic!("expected Op::UserTurn with active mode, got {other:?}")
}
}
assert_eq!(chat.active_collaboration_mode_kind(), expected_mode);
}
#[tokio::test]
async fn remap_placeholders_uses_attachment_labels() {
let placeholder_one = "[Image #1]";
@@ -1358,6 +1406,97 @@ async fn submit_user_message_with_mode_sets_coding_collaboration_mode() {
}
}
#[tokio::test]
async fn submit_user_message_with_mode_errors_when_mode_changes_during_running_turn() {
let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(Some("gpt-5")).await;
chat.thread_id = Some(ThreadId::new());
chat.set_feature_enabled(Feature::CollaborationModes, true);
let plan_mask =
collaboration_modes::mask_for_kind(chat.models_manager.as_ref(), ModeKind::Plan)
.expect("expected plan collaboration mask");
chat.set_collaboration_mask(plan_mask);
chat.on_task_started();
let default_mode = collaboration_modes::default_mask(chat.models_manager.as_ref())
.expect("expected default collaboration mode");
chat.submit_user_message_with_mode("Implement the plan.".to_string(), default_mode);
assert_eq!(chat.active_collaboration_mode_kind(), ModeKind::Plan);
assert!(chat.queued_user_messages.is_empty());
assert_matches!(op_rx.try_recv(), Err(TryRecvError::Empty));
let rendered = drain_insert_history(&mut rx)
.iter()
.map(|lines| lines_to_single_string(lines))
.collect::<Vec<_>>()
.join("\n");
assert!(
rendered.contains("Cannot switch collaboration mode while a turn is running."),
"expected running-turn error message, got: {rendered:?}"
);
}
#[tokio::test]
async fn submit_user_message_with_mode_allows_same_mode_during_running_turn() {
let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(Some("gpt-5")).await;
chat.thread_id = Some(ThreadId::new());
chat.set_feature_enabled(Feature::CollaborationModes, true);
let plan_mask =
collaboration_modes::mask_for_kind(chat.models_manager.as_ref(), ModeKind::Plan)
.expect("expected plan collaboration mask");
chat.set_collaboration_mask(plan_mask.clone());
chat.on_task_started();
chat.submit_user_message_with_mode("Continue planning.".to_string(), plan_mask);
assert_eq!(chat.active_collaboration_mode_kind(), ModeKind::Plan);
assert!(chat.queued_user_messages.is_empty());
match next_submit_op(&mut op_rx) {
Op::UserTurn {
collaboration_mode:
Some(CollaborationMode {
mode: ModeKind::Plan,
..
}),
personality: None,
..
} => {}
other => {
panic!("expected Op::UserTurn with plan collab mode, got {other:?}")
}
}
}
#[tokio::test]
async fn submit_user_message_with_mode_submits_when_plan_stream_is_not_active() {
let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(Some("gpt-5")).await;
chat.thread_id = Some(ThreadId::new());
chat.set_feature_enabled(Feature::CollaborationModes, true);
let plan_mask =
collaboration_modes::mask_for_kind(chat.models_manager.as_ref(), ModeKind::Plan)
.expect("expected plan collaboration mask");
chat.set_collaboration_mask(plan_mask);
let default_mode = collaboration_modes::default_mask(chat.models_manager.as_ref())
.expect("expected default collaboration mode");
let expected_mode = default_mode
.mode
.expect("expected default collaboration mode kind");
chat.submit_user_message_with_mode("Implement the plan.".to_string(), default_mode);
assert_eq!(chat.active_collaboration_mode_kind(), expected_mode);
assert!(chat.queued_user_messages.is_empty());
match next_submit_op(&mut op_rx) {
Op::UserTurn {
collaboration_mode: Some(CollaborationMode { mode, .. }),
personality: None,
..
} => assert_eq!(mode, expected_mode),
other => {
panic!("expected Op::UserTurn with default collab mode, got {other:?}")
}
}
}
#[tokio::test]
async fn plan_implementation_popup_skips_replayed_turn_complete() {
let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5")).await;
@@ -1986,6 +2125,55 @@ async fn unified_exec_begin_restores_working_status_snapshot() {
);
}
#[tokio::test]
async fn steer_enter_queues_while_plan_stream_is_active() {
let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(None).await;
chat.thread_id = Some(ThreadId::new());
chat.set_feature_enabled(Feature::CollaborationModes, true);
let plan_mask =
collaboration_modes::mask_for_kind(chat.models_manager.as_ref(), ModeKind::Plan)
.expect("expected plan collaboration mask");
chat.set_collaboration_mask(plan_mask);
chat.on_task_started();
chat.on_plan_delta("- Step 1".to_string());
chat.bottom_pane
.set_composer_text("queued submission".to_string(), Vec::new(), Vec::new());
chat.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE));
assert_eq!(chat.active_collaboration_mode_kind(), ModeKind::Plan);
assert_eq!(chat.queued_user_messages.len(), 1);
assert_eq!(
chat.queued_user_messages.front().unwrap().text,
"queued submission"
);
assert_matches!(op_rx.try_recv(), Err(TryRecvError::Empty));
}
#[tokio::test]
async fn steer_enter_submits_when_plan_stream_is_not_active() {
let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(None).await;
chat.thread_id = Some(ThreadId::new());
chat.set_feature_enabled(Feature::CollaborationModes, true);
let plan_mask =
collaboration_modes::mask_for_kind(chat.models_manager.as_ref(), ModeKind::Plan)
.expect("expected plan collaboration mask");
chat.set_collaboration_mask(plan_mask);
chat.on_task_started();
chat.bottom_pane
.set_composer_text("submitted immediately".to_string(), Vec::new(), Vec::new());
chat.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE));
assert!(chat.queued_user_messages.is_empty());
match next_submit_op(&mut op_rx) {
Op::UserTurn {
personality: None, ..
} => {}
other => panic!("expected Op::UserTurn, got {other:?}"),
}
}
#[tokio::test]
async fn ctrl_c_shutdown_works_with_caps_lock() {
let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await;