openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into fail-agent-loop-on-auto-compaction-failure

Browse Source
This commit is contained in:
Charles Cunningham
2026-02-06 13:11:02 -08:00
parent 14bb780e86 1751116ec6
commit 15c838bf27
29 changed files with 1912 additions and 374 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AGENTS.md
+10 -5
View File
@@ -128,15 +128,11 @@ These guidelines apply to app-server protocol work in `codex-rs`, especially:
`*Params` for request payloads, `*Response` for responses, and `*Notification` for notifications.
- Expose RPC methods as `<resource>/<method>` and keep `<resource>` singular (for example, `thread/read`, `app/list`).
- Always expose fields as camelCase on the wire with `#[serde(rename_all = "camelCase")]` unless a tagged union or explicit compatibility requirement needs a targeted rename.
- Exception: config RPC payloads are expected to use snake_case to mirror config.toml keys (see the config read/write/list APIs in `app-server-protocol/src/protocol/v2.rs`).
- Always set `#[ts(export_to = "v2/")]` on v2 request/response/notification types so generated TypeScript lands in the correct namespace.
- Never use `#[serde(skip_serializing_if = "Option::is_none")]` for v2 API payload fields.
Exception: client->server requests that intentionally have no params may use:
`params: #[ts(type = "undefined")] #[serde(skip_serializing_if = "Option::is_none")] Option<()>`.
- For client->server JSON-RPC request payloads (`*Params`) only, every optional field must be annotated with `#[ts(optional = nullable)]`. Do not use `#[ts(optional = nullable)]` outside client->server request payloads (`*Params`).
- For client->server JSON-RPC request payloads only, and you want to express a boolean field where omission means `false`, use `#[serde(default, skip_serializing_if = "std::ops::Not::not")] pub field: bool` over `Option<bool>`.
- For new list methods, implement cursor pagination by default:
request fields `pub cursor: Option<String>` and `pub limit: Option<u32>`,
response fields `pub data: Vec<...>` and `pub next_cursor: Option<String>`.
- Keep Rust and TS wire renames aligned. If a field or variant uses `#[serde(rename = "...")]`, add matching `#[ts(rename = "...")]`.
- For discriminated unions, use explicit tagging in both serializers:
`#[serde(tag = "type", ...)]` and `#[ts(tag = "type", ...)]`.
@@ -145,6 +141,15 @@ These guidelines apply to app-server protocol work in `codex-rs`, especially:
- For experimental API surface area:
use `#[experimental("method/or/field")]`, derive `ExperimentalApi` when field-level gating is needed, and use `inspect_params: true` in `common.rs` when only some fields of a method are experimental.
### Client->server request payloads (`*Params`)
- Every optional field must be annotated with `#[ts(optional = nullable)]`. Do not use `#[ts(optional = nullable)]` outside client->server request payloads (`*Params`).
- Optional collection fields (for example `Vec`, `HashMap`) must use `Option<...>` + `#[ts(optional = nullable)]`. Do not use `#[serde(default)]` to model optional collections, and do not use `skip_serializing_if` on v2 payload fields.
- When you want omission to mean `false` for boolean fields, use `#[serde(default, skip_serializing_if = "std::ops::Not::not")] pub field: bool` over `Option<bool>`.
- For new list methods, implement cursor pagination by default:
request fields `pub cursor: Option<String>` and `pub limit: Option<u32>`,
response fields `pub data: Vec<...>` and `pub next_cursor: Option<String>`.
### Development Workflow
- Update docs/examples when API behavior changes (at minimum `app-server/README.md`).
codex-rs/app-server-protocol/schema/json/ClientRequest.json
-40
View File
@@ -2323,13 +2323,6 @@
"null"
]
},
"path": {
"description": "[UNSTABLE] Specify the rollout path to fork from. If specified, the thread_id param will be ignored.",
"type": [
"string",
"null"
]
},
"sandbox": {
"anyOf": [
{
@@ -2486,16 +2479,6 @@
"null"
]
},
"history": {
"description": "[UNSTABLE] FOR CODEX CLOUD - DO NOT USE. If specified, the thread will be resumed with the provided history instead of loaded from disk.",
"items": {
"$ref": "#/definitions/ResponseItem"
},
"type": [
"array",
"null"
]
},
"model": {
"description": "Configuration overrides for the resumed thread, if any.",
"type": [
@@ -2509,13 +2492,6 @@
"null"
]
},
"path": {
"description": "[UNSTABLE] Specify the rollout path to resume from. If specified, the thread_id param will be ignored.",
"type": [
"string",
"null"
]
},
"personality": {
"anyOf": [
{
@@ -2643,11 +2619,6 @@
"null"
]
},
"experimentalRawEvents": {
"default": false,
"description": "If true, opt into emitting raw response items on the event stream.\n\nThis is for internal use only (e.g. Codex Cloud). (TODO): Figure out a better way to categorize internal / experimental events & protocols.",
"type": "boolean"
},
"model": {
"type": [
"string",
@@ -2722,17 +2693,6 @@
],
"description": "Override the approval policy for this turn and subsequent turns."
},
"collaborationMode": {
"anyOf": [
{
"$ref": "#/definitions/CollaborationMode"
},
{
"type": "null"
}
],
"description": "EXPERIMENTAL - set a pre-set collaboration mode. Takes precedence over model, reasoning_effort, and developer instructions if set."
},
"cwd": {
"description": "Override the working directory for this turn and subsequent turns.",
"type": [
codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json
-40
View File
@@ -14230,13 +14230,6 @@
"null"
]
},
"path": {
"description": "[UNSTABLE] Specify the rollout path to fork from. If specified, the thread_id param will be ignored.",
"type": [
"string",
"null"
]
},
"sandbox": {
"anyOf": [
{
@@ -14995,16 +14988,6 @@
"null"
]
},
"history": {
"description": "[UNSTABLE] FOR CODEX CLOUD - DO NOT USE. If specified, the thread will be resumed with the provided history instead of loaded from disk.",
"items": {
"$ref": "#/definitions/v2/ResponseItem"
},
"type": [
"array",
"null"
]
},
"model": {
"description": "Configuration overrides for the resumed thread, if any.",
"type": [
@@ -15018,13 +15001,6 @@
"null"
]
},
"path": {
"description": "[UNSTABLE] Specify the rollout path to resume from. If specified, the thread_id param will be ignored.",
"type": [
"string",
"null"
]
},
"personality": {
"anyOf": [
{
@@ -15224,11 +15200,6 @@
"null"
]
},
"experimentalRawEvents": {
"default": false,
"description": "If true, opt into emitting raw response items on the event stream.\n\nThis is for internal use only (e.g. Codex Cloud). (TODO): Figure out a better way to categorize internal / experimental events & protocols.",
"type": "boolean"
},
"model": {
"type": [
"string",
@@ -15665,17 +15636,6 @@
],
"description": "Override the approval policy for this turn and subsequent turns."
},
"collaborationMode": {
"anyOf": [
{
"$ref": "#/definitions/v2/CollaborationMode"
},
{
"type": "null"
}
],
"description": "EXPERIMENTAL - set a pre-set collaboration mode. Takes precedence over model, reasoning_effort, and developer instructions if set."
},
"cwd": {
"description": "Override the working directory for this turn and subsequent turns.",
"type": [
codex-rs/app-server-protocol/schema/json/v2/ThreadForkParams.json
-7
View File
@@ -69,13 +69,6 @@
"null"
]
},
"path": {
"description": "[UNSTABLE] Specify the rollout path to fork from. If specified, the thread_id param will be ignored.",
"type": [
"string",
"null"
]
},
"sandbox": {
"anyOf": [
{
codex-rs/app-server-protocol/schema/json/v2/ThreadResumeParams.json
-17
View File
@@ -832,16 +832,6 @@
"null"
]
},
"history": {
"description": "[UNSTABLE] FOR CODEX CLOUD - DO NOT USE. If specified, the thread will be resumed with the provided history instead of loaded from disk.",
"items": {
"$ref": "#/definitions/ResponseItem"
},
"type": [
"array",
"null"
]
},
"model": {
"description": "Configuration overrides for the resumed thread, if any.",
"type": [
@@ -855,13 +845,6 @@
"null"
]
},
"path": {
"description": "[UNSTABLE] Specify the rollout path to resume from. If specified, the thread_id param will be ignored.",
"type": [
"string",
"null"
]
},
"personality": {
"anyOf": [
{
codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json
-5
View File
@@ -86,11 +86,6 @@
"null"
]
},
"experimentalRawEvents": {
"default": false,
"description": "If true, opt into emitting raw response items on the event stream.\n\nThis is for internal use only (e.g. Codex Cloud). (TODO): Figure out a better way to categorize internal / experimental events & protocols.",
"type": "boolean"
},
"model": {
"type": [
"string",
codex-rs/app-server-protocol/schema/json/v2/TurnStartParams.json
-11
View File
@@ -383,17 +383,6 @@
],
"description": "Override the approval policy for this turn and subsequent turns."
},
"collaborationMode": {
"anyOf": [
{
"$ref": "#/definitions/CollaborationMode"
},
{
"type": "null"
}
],
"description": "EXPERIMENTAL - set a pre-set collaboration mode. Takes precedence over model, reasoning_effort, and developer instructions if set."
},
"cwd": {
"description": "Override the working directory for this turn and subsequent turns.",
"type": [
codex-rs/app-server-protocol/schema/typescript/v2/ThreadForkParams.ts
+3 -5
View File
@@ -14,13 +14,11 @@ import type { SandboxMode } from "./SandboxMode";
*
* Prefer using thread_id whenever possible.
*/
export type ThreadForkParams = { threadId: string,
/**
export type ThreadForkParams = {threadId: string, /**
* [UNSTABLE] Specify the rollout path to fork from.
* If specified, the thread_id param will be ignored.
*/
path?: string | null,
/**
path?: string | null, /**
* Configuration overrides for the forked thread, if any.
*/
model?: string | null, modelProvider?: string | null, cwd?: string | null, approvalPolicy?: AskForApproval | null, sandbox?: SandboxMode | null, config?: { [key in string]?: JsonValue } | null, baseInstructions?: string | null, developerInstructions?: string | null, };
model?: string | null, modelProvider?: string | null, cwd?: string | null, approvalPolicy?: AskForApproval | null, sandbox?: SandboxMode | null, config?: { [key in string]?: JsonValue } | null, baseInstructions?: string | null, developerInstructions?: string | null};
codex-rs/app-server-protocol/schema/typescript/v2/ThreadResumeParams.ts
+4 -7
View File
@@ -18,19 +18,16 @@ import type { SandboxMode } from "./SandboxMode";
*
* Prefer using thread_id whenever possible.
*/
export type ThreadResumeParams = { threadId: string,
/**
export type ThreadResumeParams = {threadId: string, /**
* [UNSTABLE] FOR CODEX CLOUD - DO NOT USE.
* If specified, the thread will be resumed with the provided history
* instead of loaded from disk.
*/
history?: Array<ResponseItem> | null,
/**
history?: Array<ResponseItem> | null, /**
* [UNSTABLE] Specify the rollout path to resume from.
* If specified, the thread_id param will be ignored.
*/
path?: string | null,
/**
path?: string | null, /**
* Configuration overrides for the resumed thread, if any.
*/
model?: string | null, modelProvider?: string | null, cwd?: string | null, approvalPolicy?: AskForApproval | null, sandbox?: SandboxMode | null, config?: { [key in string]?: JsonValue } | null, baseInstructions?: string | null, developerInstructions?: string | null, personality?: Personality | null, };
model?: string | null, modelProvider?: string | null, cwd?: string | null, approvalPolicy?: AskForApproval | null, sandbox?: SandboxMode | null, config?: { [key in string]?: JsonValue } | null, baseInstructions?: string | null, developerInstructions?: string | null, personality?: Personality | null};
codex-rs/app-server-protocol/schema/typescript/v2/ThreadStartParams.ts
+1 -3
View File
@@ -7,9 +7,7 @@ import type { AskForApproval } from "./AskForApproval";
import type { SandboxMode } from "./SandboxMode";
export type ThreadStartParams = {model?: string | null, modelProvider?: string | null, cwd?: string | null, approvalPolicy?: AskForApproval | null, sandbox?: SandboxMode | null, config?: { [key in string]?: JsonValue } | null, baseInstructions?: string | null, developerInstructions?: string | null, personality?: Personality | null, ephemeral?: boolean | null, /**
* If true, opt into emitting raw response items on the event stream.
*
* If true, opt into emitting raw Responses API items on the event stream.
* This is for internal use only (e.g. Codex Cloud).
* (TODO): Figure out a better way to categorize internal / experimental events & protocols.
*/
experimentalRawEvents: boolean};
codex-rs/app-server-protocol/schema/typescript/v2/TurnStartParams.ts
+11 -20
View File
@@ -10,41 +10,32 @@ import type { AskForApproval } from "./AskForApproval";
import type { SandboxPolicy } from "./SandboxPolicy";
import type { UserInput } from "./UserInput";
export type TurnStartParams = { threadId: string, input: Array<UserInput>,
/**
export type TurnStartParams = {threadId: string, input: Array<UserInput>, /**
* Override the working directory for this turn and subsequent turns.
*/
cwd?: string | null,
/**
cwd?: string | null, /**
* Override the approval policy for this turn and subsequent turns.
*/
approvalPolicy?: AskForApproval | null,
/**
approvalPolicy?: AskForApproval | null, /**
* Override the sandbox policy for this turn and subsequent turns.
*/
sandboxPolicy?: SandboxPolicy | null,
/**
sandboxPolicy?: SandboxPolicy | null, /**
* Override the model for this turn and subsequent turns.
*/
model?: string | null,
/**
model?: string | null, /**
* Override the reasoning effort for this turn and subsequent turns.
*/
effort?: ReasoningEffort | null,
/**
effort?: ReasoningEffort | null, /**
* Override the reasoning summary for this turn and subsequent turns.
*/
summary?: ReasoningSummary | null,
/**
summary?: ReasoningSummary | null, /**
* Override the personality for this turn and subsequent turns.
*/
personality?: Personality | null,
/**
personality?: Personality | null, /**
* Optional JSON Schema used to constrain the final assistant message for this turn.
*/
outputSchema?: JsonValue | null,
/**
* EXPERIMENTAL - set a pre-set collaboration mode.
outputSchema?: JsonValue | null, /**
* EXPERIMENTAL - Set a pre-set collaboration mode.
* Takes precedence over model, reasoning_effort, and developer instructions if set.
*/
collaborationMode?: CollaborationMode | null, };
collaborationMode?: CollaborationMode | null};
codex-rs/app-server-protocol/src/protocol/common.rs
+4
View File
@@ -190,10 +190,12 @@ client_request_definitions! {
},
ThreadResume => "thread/resume" {
params: v2::ThreadResumeParams,
inspect_params: true,
response: v2::ThreadResumeResponse,
},
ThreadFork => "thread/fork" {
params: v2::ThreadForkParams,
inspect_params: true,
response: v2::ThreadForkResponse,
},
ThreadArchive => "thread/archive" {
@@ -250,6 +252,7 @@ client_request_definitions! {
},
TurnStart => "turn/start" {
params: v2::TurnStartParams,
inspect_params: true,
response: v2::TurnStartResponse,
},
TurnSteer => "turn/steer" {
@@ -303,6 +306,7 @@ client_request_definitions! {
LoginAccount => "account/login/start" {
params: v2::LoginAccountParams,
inspect_params: true,
response: v2::LoginAccountResponse,
},
codex-rs/app-server-protocol/src/protocol/v2.rs
+18 -8
View File
@@ -867,7 +867,7 @@ pub enum Account {
Chatgpt { email: String, plan_type: PlanType },
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS, ExperimentalApi)]
#[serde(tag = "type")]
#[ts(tag = "type")]
#[ts(export_to = "v2/")]
@@ -884,6 +884,7 @@ pub enum LoginAccountParams {
Chatgpt,
/// [UNSTABLE] FOR OPENAI INTERNAL USE ONLY - DO NOT USE.
/// The access token must contain the same scopes that Codex-managed ChatGPT auth tokens have.
#[experimental("account/login/start.chatgptAuthTokens")]
#[serde(rename = "chatgptAuthTokens")]
#[ts(rename = "chatgptAuthTokens")]
ChatgptAuthTokens {
@@ -1314,10 +1315,9 @@ pub struct ThreadStartParams {
#[experimental("thread/start.mockExperimentalField")]
#[ts(optional = nullable)]
pub mock_experimental_field: Option<String>,
/// If true, opt into emitting raw response items on the event stream.
///
/// If true, opt into emitting raw Responses API items on the event stream.
/// This is for internal use only (e.g. Codex Cloud).
/// (TODO): Figure out a better way to categorize internal / experimental events & protocols.
#[experimental("thread/start.experimentalRawEvents")]
#[serde(default)]
pub experimental_raw_events: bool,
}
@@ -1352,7 +1352,9 @@ pub struct ThreadStartResponse {
pub reasoning_effort: Option<ReasoningEffort>,
}
#[derive(Serialize, Deserialize, Debug, Default, Clone, PartialEq, JsonSchema, TS)]
#[derive(
Serialize, Deserialize, Debug, Default, Clone, PartialEq, JsonSchema, TS, ExperimentalApi,
)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
/// There are three ways to resume a thread:
@@ -1370,11 +1372,13 @@ pub struct ThreadResumeParams {
/// [UNSTABLE] FOR CODEX CLOUD - DO NOT USE.
/// If specified, the thread will be resumed with the provided history
/// instead of loaded from disk.
#[experimental("thread/resume.history")]
#[ts(optional = nullable)]
pub history: Option<Vec<ResponseItem>>,
/// [UNSTABLE] Specify the rollout path to resume from.
/// If specified, the thread_id param will be ignored.
#[experimental("thread/resume.path")]
#[ts(optional = nullable)]
pub path: Option<PathBuf>,
@@ -1412,7 +1416,9 @@ pub struct ThreadResumeResponse {
pub reasoning_effort: Option<ReasoningEffort>,
}
#[derive(Serialize, Deserialize, Debug, Default, Clone, PartialEq, JsonSchema, TS)]
#[derive(
Serialize, Deserialize, Debug, Default, Clone, PartialEq, JsonSchema, TS, ExperimentalApi,
)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
/// There are two ways to fork a thread:
@@ -1427,6 +1433,7 @@ pub struct ThreadForkParams {
/// [UNSTABLE] Specify the rollout path to fork from.
/// If specified, the thread_id param will be ignored.
#[experimental("thread/fork.path")]
#[ts(optional = nullable)]
pub path: Option<PathBuf>,
@@ -2027,7 +2034,9 @@ pub enum TurnStatus {
}
// Turn APIs
#[derive(Serialize, Deserialize, Debug, Default, Clone, PartialEq, JsonSchema, TS)]
#[derive(
Serialize, Deserialize, Debug, Default, Clone, PartialEq, JsonSchema, TS, ExperimentalApi,
)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
pub struct TurnStartParams {
@@ -2058,8 +2067,9 @@ pub struct TurnStartParams {
#[ts(optional = nullable)]
pub output_schema: Option<JsonValue>,
/// EXPERIMENTAL - set a pre-set collaboration mode.
/// EXPERIMENTAL - Set a pre-set collaboration mode.
/// Takes precedence over model, reasoning_effort, and developer instructions if set.
#[experimental("turn/start.collaborationMode")]
#[ts(optional = nullable)]
pub collaboration_mode: Option<CollaborationMode>,
}
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
+64 -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;
@@ -563,22 +564,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 {
@@ -1078,6 +1079,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();
@@ -2016,6 +2025,15 @@ impl Session {
history.raw_items().to_vec()
}
pub(crate) async fn process_compacted_history(
&self,
turn_context: &TurnContext,
compacted_history: Vec<ResponseItem>,
) -> Vec<ResponseItem> {
let initial_context = self.build_initial_context(turn_context).await;
compact::process_compacted_history(compacted_history, &initial_context)
}
/// Append ResponseItems to the in-memory conversation history only.
pub(crate) async fn record_into_history(
&self,
@@ -5084,6 +5102,42 @@ mod tests {
assert_eq!(expected, reconstructed);
}
#[tokio::test]
async fn reconstruct_history_uses_replacement_history_verbatim() {
let (session, turn_context) = make_session_and_context().await;
let summary_item = ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
};
let replacement_history = vec![
summary_item.clone(),
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "stale developer instructions".to_string(),
}],
end_turn: None,
phase: None,
},
];
let rollout_items = vec![RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(replacement_history.clone()),
})];
let reconstructed = session
.reconstruct_history_from_rollout(&turn_context, &rollout_items)
.await;
assert_eq!(reconstructed, replacement_history);
}
#[tokio::test]
async fn record_initial_history_reconstructs_resumed_transcript() {
let (session, turn_context) = make_session_and_context().await;
codex-rs/core/src/compact.rs
+400 -43
View File
@@ -14,7 +14,6 @@ use crate::protocol::EventMsg;
use crate::protocol::TurnContextItem;
use crate::protocol::TurnStartedEvent;
use crate::protocol::WarningEvent;
use crate::session_prefix::TURN_ABORTED_OPEN_TAG;
use crate::truncate::TruncationPolicy;
use crate::truncate::approx_token_count;
use crate::truncate::truncate_text;
@@ -243,35 +242,64 @@ pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec<String> {
Some(user.message())
}
}
_ => collect_turn_aborted_marker(item),
_ => None,
})
.collect()
}
fn collect_turn_aborted_marker(item: &ResponseItem) -> Option<String> {
let ResponseItem::Message { role, content, .. } = item else {
return None;
};
if role != "user" {
return None;
}
let text = content_items_to_text(content)?;
if text
.trim_start()
.to_ascii_lowercase()
.starts_with(TURN_ABORTED_OPEN_TAG)
{
Some(text)
} else {
None
}
}
pub(crate) fn is_summary_message(message: &str) -> bool {
message.starts_with(format!("{SUMMARY_PREFIX}\n").as_str())
}
pub(crate) fn process_compacted_history(
mut compacted_history: Vec<ResponseItem>,
initial_context: &[ResponseItem],
) -> Vec<ResponseItem> {
compacted_history.retain(should_keep_compacted_history_item);
let initial_context = initial_context.to_vec();
// Re-inject canonical context from the current session since we stripped it
// from the pre-compaction history. Keep it right before the last user
// message so older user messages remain earlier in the transcript.
if let Some(last_user_index) = compacted_history.iter().rposition(|item| {
matches!(
crate::event_mapping::parse_turn_item(item),
Some(TurnItem::UserMessage(_))
)
}) {
compacted_history.splice(last_user_index..last_user_index, initial_context);
} else {
compacted_history.extend(initial_context);
}
compacted_history
}
/// Returns whether an item from remote compaction output should be preserved.
///
/// Called while processing the model-provided compacted transcript, before we
/// append fresh canonical context from the current session.
///
/// We drop:
/// - `developer` messages because remote output can include stale/duplicated
/// instruction content.
/// - non-user-content `user` messages (session prefix/instruction wrappers),
/// keeping only real user messages as parsed by `parse_turn_item`.
///
/// This intentionally keeps `user`-role warnings and compaction-generated
/// summary messages because they parse as `TurnItem::UserMessage`.
fn should_keep_compacted_history_item(item: &ResponseItem) -> bool {
match item {
ResponseItem::Message { role, .. } if role == "developer" => false,
ResponseItem::Message { role, .. } if role == "user" => matches!(
crate::event_mapping::parse_turn_item(item),
Some(TurnItem::UserMessage(_))
),
_ => true,
}
}
pub(crate) fn build_compacted_history(
initial_context: Vec<ResponseItem>,
user_messages: &[String],
@@ -391,7 +419,6 @@ async fn drain_to_completed(
mod tests {
use super::*;
use crate::session_prefix::TURN_ABORTED_OPEN_TAG;
use pretty_assertions::assert_eq;
#[test]
@@ -556,16 +583,13 @@ mod tests {
}
#[test]
fn build_compacted_history_preserves_turn_aborted_markers() {
let marker = format!(
"{TURN_ABORTED_OPEN_TAG}\n <turn_id>turn-1</turn_id>\n <reason>interrupted</reason>\n</turn_aborted>"
);
let items = vec![
fn process_compacted_history_replaces_developer_messages() {
let compacted_history = vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: marker.clone(),
text: "stale permissions".to_string(),
}],
end_turn: None,
phase: None,
@@ -574,25 +598,358 @@ mod tests {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "real user message".to_string(),
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "stale personality".to_string(),
}],
end_turn: None,
phase: None,
},
];
let initial_context = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<environment_context>cwd=/tmp</environment_context>".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh personality".to_string(),
}],
end_turn: None,
phase: None,
},
];
let user_messages = collect_user_messages(&items);
let history = build_compacted_history(Vec::new(), &user_messages, "SUMMARY");
let refreshed = process_compacted_history(compacted_history, &initial_context);
let expected = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<environment_context>cwd=/tmp</environment_context>".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh personality".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
},
];
assert_eq!(refreshed, expected);
}
let found_marker = history.iter().any(|item| match item {
ResponseItem::Message { role, content, .. } if role == "user" => {
content_items_to_text(content).is_some_and(|text| text == marker)
}
_ => false,
});
assert!(
found_marker,
"expected compacted history to retain <turn_aborted> marker"
);
#[test]
fn process_compacted_history_reinjects_full_initial_context() {
let compacted_history = vec![ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
}];
let initial_context = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "# AGENTS.md instructions for /repo\n\n<INSTRUCTIONS>\nkeep me updated\n</INSTRUCTIONS>".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<environment_context>\n <cwd>/repo</cwd>\n <shell>zsh</shell>\n</environment_context>".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<turn_aborted>\n <turn_id>turn-1</turn_id>\n <reason>interrupted</reason>\n</turn_aborted>".to_string(),
}],
end_turn: None,
phase: None,
},
];
let refreshed = process_compacted_history(compacted_history, &initial_context);
let expected = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "# AGENTS.md instructions for /repo\n\n<INSTRUCTIONS>\nkeep me updated\n</INSTRUCTIONS>".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<environment_context>\n <cwd>/repo</cwd>\n <shell>zsh</shell>\n</environment_context>".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<turn_aborted>\n <turn_id>turn-1</turn_id>\n <reason>interrupted</reason>\n</turn_aborted>"
.to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
},
];
assert_eq!(refreshed, expected);
}
#[test]
fn process_compacted_history_drops_non_user_content_messages() {
let compacted_history = vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "# AGENTS.md instructions for /repo\n\n<INSTRUCTIONS>\nkeep me updated\n</INSTRUCTIONS>".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<environment_context>\n <cwd>/repo</cwd>\n <shell>zsh</shell>\n</environment_context>".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "<turn_aborted>\n <turn_id>turn-1</turn_id>\n <reason>interrupted</reason>\n</turn_aborted>".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "stale developer instructions".to_string(),
}],
end_turn: None,
phase: None,
},
];
let initial_context = vec![ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh developer instructions".to_string(),
}],
end_turn: None,
phase: None,
}];
let refreshed = process_compacted_history(compacted_history, &initial_context);
let expected = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh developer instructions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
},
];
assert_eq!(refreshed, expected);
}
#[test]
fn process_compacted_history_inserts_context_before_last_real_user_message_only() {
let compacted_history = vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "older user".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: format!("{SUMMARY_PREFIX}\nsummary text"),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "latest user".to_string(),
}],
end_turn: None,
phase: None,
},
];
let initial_context = vec![ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
}];
let refreshed = process_compacted_history(compacted_history, &initial_context);
let expected = vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "older user".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: format!("{SUMMARY_PREFIX}\nsummary text"),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "latest user".to_string(),
}],
end_turn: None,
phase: None,
},
];
assert_eq!(refreshed, expected);
}
}
codex-rs/core/src/compact_remote.rs
+3
View File
@@ -84,6 +84,9 @@ pub(crate) async fn run_remote_compaction(
&turn_context.otel_manager,
)
.await?;
new_history = sess
.process_compacted_history(turn_context, new_history)
.await;
if !ghost_snapshots.is_empty() {
new_history.extend(ghost_snapshots);
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 {
@@ -874,6 +902,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
@@ -913,6 +942,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/compact_remote.rs
+313 -3
View File
@@ -926,10 +926,58 @@ async fn remote_compact_persists_replacement_history_in_rollout() -> Result<()>
};
if let RolloutItem::Compacted(compacted) = entry.item
&& compacted.message.is_empty()
&& compacted.replacement_history.as_ref() == Some(&compacted_history)
&& let Some(replacement_history) = compacted.replacement_history.as_ref()
{
saw_compacted_history = true;
break;
let has_compacted_user_summary = replacement_history.iter().any(|item| {
matches!(
item,
ResponseItem::Message { role, content, .. }
if role == "user"
&& content.iter().any(|part| matches!(
part,
ContentItem::InputText { text } if text == "COMPACTED_USER_SUMMARY"
))
)
});
let has_compaction_item = replacement_history.iter().any(|item| {
matches!(
item,
ResponseItem::Compaction { encrypted_content }
if encrypted_content == "ENCRYPTED_COMPACTION_SUMMARY"
)
});
let has_compacted_assistant_note = replacement_history.iter().any(|item| {
matches!(
item,
ResponseItem::Message { role, content, .. }
if role == "assistant"
&& content.iter().any(|part| matches!(
part,
ContentItem::OutputText { text } if text == "COMPACTED_ASSISTANT_NOTE"
))
)
});
let has_permissions_developer_message = replacement_history.iter().any(|item| {
matches!(
item,
ResponseItem::Message { role, content, .. }
if role == "developer"
&& content.iter().any(|part| matches!(
part,
ContentItem::InputText { text }
if text.contains("<permissions instructions>")
))
)
});
if has_compacted_user_summary
&& has_compaction_item
&& has_compacted_assistant_note
&& has_permissions_developer_message
{
saw_compacted_history = true;
break;
}
}
}
@@ -940,3 +988,265 @@ async fn remote_compact_persists_replacement_history_in_rollout() -> Result<()>
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn remote_compact_and_resume_refresh_stale_developer_instructions() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = wiremock::MockServer::start().await;
let stale_developer_message = "STALE_DEVELOPER_INSTRUCTIONS_SHOULD_BE_REMOVED";
let mut start_builder = test_codex()
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
.with_config(|config| {
config.features.enable(Feature::RemoteCompaction);
});
let initial = start_builder.build(&server).await?;
let home = initial.home.clone();
let rollout_path = initial
.session_configured
.rollout_path
.clone()
.expect("rollout path");
let responses_mock = responses::mount_sse_sequence(
&server,
vec![
responses::sse(vec![
responses::ev_assistant_message("m1", "BASELINE_REPLY"),
responses::ev_completed("resp-1"),
]),
responses::sse(vec![
responses::ev_assistant_message("m2", "AFTER_COMPACT_REPLY"),
responses::ev_completed("resp-2"),
]),
responses::sse(vec![
responses::ev_assistant_message("m3", "AFTER_RESUME_REPLY"),
responses::ev_completed("resp-3"),
]),
],
)
.await;
let compacted_history = vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "REMOTE_COMPACTED_SUMMARY".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: stale_developer_message.to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Compaction {
encrypted_content: "ENCRYPTED_COMPACTION_SUMMARY".to_string(),
},
];
let compact_mock = responses::mount_compact_json_once(
&server,
serde_json::json!({ "output": compacted_history }),
)
.await;
initial
.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "start remote compact flow".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
})
.await?;
wait_for_event(&initial.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
initial.codex.submit(Op::Compact).await?;
wait_for_event(&initial.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
initial
.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "after compact in same session".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
})
.await?;
wait_for_event(&initial.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
initial.codex.submit(Op::Shutdown).await?;
wait_for_event(&initial.codex, |ev| {
matches!(ev, EventMsg::ShutdownComplete)
})
.await;
let mut resume_builder = test_codex()
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
.with_config(|config| {
config.features.enable(Feature::RemoteCompaction);
});
let resumed = resume_builder.resume(&server, home, rollout_path).await?;
resumed
.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "after resume".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
})
.await?;
wait_for_event(&resumed.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
assert_eq!(compact_mock.requests().len(), 1);
let requests = responses_mock.requests();
assert_eq!(requests.len(), 3, "expected three model requests");
let after_compact_request = &requests[1];
let after_resume_request = &requests[2];
let after_compact_body = after_compact_request.body_json().to_string();
assert!(
!after_compact_body.contains(stale_developer_message),
"stale developer instructions should be removed immediately after compaction"
);
assert!(
after_compact_body.contains("<permissions instructions>"),
"fresh developer instructions should be present after compaction"
);
assert!(
after_compact_body.contains("REMOTE_COMPACTED_SUMMARY"),
"compacted summary should be present after compaction"
);
let after_resume_body = after_resume_request.body_json().to_string();
assert!(
!after_resume_body.contains(stale_developer_message),
"stale developer instructions should be removed after resume"
);
assert!(
after_resume_body.contains("<permissions instructions>"),
"fresh developer instructions should be present after resume"
);
assert!(
after_resume_body.contains("REMOTE_COMPACTED_SUMMARY"),
"compacted summary should persist after resume"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn remote_compact_refreshes_stale_developer_instructions_without_resume() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = wiremock::MockServer::start().await;
let stale_developer_message = "STALE_DEVELOPER_INSTRUCTIONS_SHOULD_BE_REMOVED";
let mut builder = test_codex()
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
.with_config(|config| {
config.features.enable(Feature::RemoteCompaction);
});
let test = builder.build(&server).await?;
let responses_mock = responses::mount_sse_sequence(
&server,
vec![
responses::sse(vec![
responses::ev_assistant_message("m1", "BASELINE_REPLY"),
responses::ev_completed("resp-1"),
]),
responses::sse(vec![
responses::ev_assistant_message("m2", "AFTER_COMPACT_REPLY"),
responses::ev_completed("resp-2"),
]),
],
)
.await;
let compacted_history = vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "REMOTE_COMPACTED_SUMMARY".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: stale_developer_message.to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Compaction {
encrypted_content: "ENCRYPTED_COMPACTION_SUMMARY".to_string(),
},
];
let compact_mock = responses::mount_compact_json_once(
&server,
serde_json::json!({ "output": compacted_history }),
)
.await;
test.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "start remote compact flow".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
})
.await?;
wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
test.codex.submit(Op::Compact).await?;
wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
test.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "after compact in same session".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
})
.await?;
wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
assert_eq!(compact_mock.requests().len(), 1);
let requests = responses_mock.requests();
assert_eq!(requests.len(), 2, "expected two model requests");
let after_compact_body = requests[1].body_json().to_string();
assert!(
!after_compact_body.contains(stale_developer_message),
"stale developer instructions should be removed immediately after compaction"
);
assert!(
after_compact_body.contains("<permissions instructions>"),
"fresh developer instructions should be present after compaction"
);
assert!(
after_compact_body.contains("REMOTE_COMPACTED_SUMMARY"),
"compacted summary should be present after compaction"
);
Ok(())
}
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),
)
}