diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 968adba7a6..5d17c0f603 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2816,6 +2816,7 @@ dependencies = [ "cc", "clap", "codex-core", + "codex-process-hardening", "codex-protocol", "codex-sandboxing", "codex-utils-absolute-path", diff --git a/codex-rs/agent-identity/src/lib.rs b/codex-rs/agent-identity/src/lib.rs index cdfb258f4d..7aad81a34f 100644 --- a/codex-rs/agent-identity/src/lib.rs +++ b/codex-rs/agent-identity/src/lib.rs @@ -8,7 +8,7 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use chrono::SecondsFormat; use chrono::Utc; -use codex_protocol::account::PlanType as AccountPlanType; +use codex_protocol::auth::PlanType as AuthPlanType; use codex_protocol::protocol::SessionSource; use crypto_box::SecretKey as Curve25519SecretKey; use ed25519_dalek::Signer as _; @@ -73,7 +73,7 @@ pub struct AgentIdentityJwtClaims { pub account_id: String, pub chatgpt_user_id: String, pub email: String, - pub plan_type: AccountPlanType, + pub plan_type: AuthPlanType, pub chatgpt_account_is_fedramp: bool, } @@ -408,6 +408,8 @@ mod tests { use jsonwebtoken::Header; use pretty_assertions::assert_eq; + use codex_protocol::auth::KnownPlan; + use super::*; #[test] @@ -517,12 +519,33 @@ mod tests { account_id: "account-id".to_string(), chatgpt_user_id: "user-id".to_string(), email: "user@example.com".to_string(), - plan_type: AccountPlanType::Pro, + plan_type: AuthPlanType::Known(KnownPlan::Pro), chatgpt_account_is_fedramp: false, } ); } + #[test] + fn decode_agent_identity_jwt_maps_raw_plan_aliases() { + let jwt = jwt_with_payload(serde_json::json!({ + "iss": AGENT_IDENTITY_JWT_ISSUER, + "aud": AGENT_IDENTITY_JWT_AUDIENCE, + "iat": 1_700_000_000usize, + "exp": 4_000_000_000usize, + "agent_runtime_id": "agent-runtime-id", + "agent_private_key": "private-key", + "account_id": "account-id", + "chatgpt_user_id": "user-id", + "email": "user@example.com", + "plan_type": "hc", + "chatgpt_account_is_fedramp": false, + })); + + let claims = decode_agent_identity_jwt(&jwt, /*jwks*/ None).expect("JWT should decode"); + + assert_eq!(claims.plan_type, AuthPlanType::Known(KnownPlan::Enterprise)); + } + #[test] fn decode_agent_identity_jwt_verifies_when_jwks_is_present() { let jwks = test_jwks("test-key"); @@ -536,7 +559,7 @@ mod tests { account_id: "account-id".to_string(), chatgpt_user_id: "user-id".to_string(), email: "user@example.com".to_string(), - plan_type: AccountPlanType::Pro, + plan_type: AuthPlanType::Known(KnownPlan::Pro), chatgpt_account_is_fedramp: false, }; let jwt = jsonwebtoken::encode( @@ -568,7 +591,7 @@ mod tests { account_id: "account-id".to_string(), chatgpt_user_id: "user-id".to_string(), email: "user@example.com".to_string(), - plan_type: AccountPlanType::Pro, + plan_type: AuthPlanType::Known(KnownPlan::Pro), chatgpt_account_is_fedramp: false, }; assert_eq!( diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 167435ac88..99f4d7d3e1 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -1,4 +1,5 @@ use std::path::Path; +use std::path::PathBuf; use crate::JSONRPCNotification; use crate::JSONRPCRequest; @@ -73,6 +74,76 @@ macro_rules! experimental_type_entry { }; } +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ClientRequestSerializationScope { + Global(&'static str), + Thread { thread_id: String }, + ThreadPath { path: PathBuf }, + CommandExecProcess { process_id: String }, + FuzzyFileSearchSession { session_id: String }, + FsWatch { watch_id: String }, + McpOauth { server_name: String }, +} + +macro_rules! serialization_scope_expr { + ($actual_params:ident, None) => { + None + }; + ($actual_params:ident, global($key:literal)) => { + Some(ClientRequestSerializationScope::Global($key)) + }; + ($actual_params:ident, thread_id($params:ident . $field:ident)) => { + Some(ClientRequestSerializationScope::Thread { + thread_id: $actual_params.$field.clone(), + }) + }; + ($actual_params:ident, optional_thread_id($params:ident . $field:ident)) => { + $actual_params + .$field + .clone() + .map(|thread_id| ClientRequestSerializationScope::Thread { thread_id }) + }; + ($actual_params:ident, thread_or_path($params:ident . $thread_field:ident, $params2:ident . $path_field:ident)) => { + if !$actual_params.$thread_field.is_empty() { + Some(ClientRequestSerializationScope::Thread { + thread_id: $actual_params.$thread_field.clone(), + }) + } else if let Some(path) = $actual_params.$path_field.clone() { + Some(ClientRequestSerializationScope::ThreadPath { path }) + } else { + Some(ClientRequestSerializationScope::Thread { + thread_id: $actual_params.$thread_field.clone(), + }) + } + }; + ($actual_params:ident, optional_command_process_id($params:ident . $field:ident)) => { + $actual_params + .$field + .clone() + .map(|process_id| ClientRequestSerializationScope::CommandExecProcess { process_id }) + }; + ($actual_params:ident, command_process_id($params:ident . $field:ident)) => { + Some(ClientRequestSerializationScope::CommandExecProcess { + process_id: $actual_params.$field.clone(), + }) + }; + ($actual_params:ident, fuzzy_session_id($params:ident . $field:ident)) => { + Some(ClientRequestSerializationScope::FuzzyFileSearchSession { + session_id: $actual_params.$field.clone(), + }) + }; + ($actual_params:ident, fs_watch_id($params:ident . $field:ident)) => { + Some(ClientRequestSerializationScope::FsWatch { + watch_id: $actual_params.$field.clone(), + }) + }; + ($actual_params:ident, mcp_oauth_server($params:ident . $field:ident)) => { + Some(ClientRequestSerializationScope::McpOauth { + server_name: $actual_params.$field.clone(), + }) + }; +} + /// Generates an `enum ClientRequest` where each variant is a request that the /// client can send to the server. Each variant has associated `params` and /// `response` types. Also generates a `export_client_responses()` function to @@ -85,6 +156,7 @@ macro_rules! client_request_definitions { $variant:ident $(=> $wire:literal)? { params: $(#[$params_meta:meta])* $params:ty, $(inspect_params: $inspect_params:tt,)? + serialization: $serialization:ident $( ( $($serialization_args:tt)* ) )?, response: $response:ty, } ),* $(,)? @@ -123,6 +195,19 @@ macro_rules! client_request_definitions { }) .unwrap_or_else(|| "".to_string()) } + + pub fn serialization_scope(&self) -> Option { + match self { + $( + Self::$variant { params, .. } => { + let _ = params; + serialization_scope_expr!( + params, $serialization $( ( $($serialization_args)* ) )? + ) + } + )* + } + } } /// Typed response from the server to the client. @@ -235,6 +320,7 @@ macro_rules! client_request_definitions { client_request_definitions! { Initialize { params: v1::InitializeParams, + serialization: None, response: v1::InitializeResponse, }, @@ -244,24 +330,29 @@ client_request_definitions! { ThreadStart => "thread/start" { params: v2::ThreadStartParams, inspect_params: true, + serialization: None, response: v2::ThreadStartResponse, }, ThreadResume => "thread/resume" { params: v2::ThreadResumeParams, inspect_params: true, + serialization: thread_or_path(params.thread_id, params.path), response: v2::ThreadResumeResponse, }, ThreadFork => "thread/fork" { params: v2::ThreadForkParams, inspect_params: true, + serialization: thread_or_path(params.thread_id, params.path), response: v2::ThreadForkResponse, }, ThreadArchive => "thread/archive" { params: v2::ThreadArchiveParams, + serialization: thread_id(params.thread_id), response: v2::ThreadArchiveResponse, }, ThreadUnsubscribe => "thread/unsubscribe" { params: v2::ThreadUnsubscribeParams, + serialization: thread_id(params.thread_id), response: v2::ThreadUnsubscribeResponse, }, #[experimental("thread/increment_elicitation")] @@ -271,6 +362,7 @@ client_request_definitions! { /// approval or other elicitation is pending outside the app-server request flow. ThreadIncrementElicitation => "thread/increment_elicitation" { params: v2::ThreadIncrementElicitationParams, + serialization: thread_id(params.thread_id), response: v2::ThreadIncrementElicitationResponse, }, #[experimental("thread/decrement_elicitation")] @@ -279,302 +371,372 @@ client_request_definitions! { /// When the count reaches zero, timeout accounting resumes for the thread. ThreadDecrementElicitation => "thread/decrement_elicitation" { params: v2::ThreadDecrementElicitationParams, + serialization: thread_id(params.thread_id), response: v2::ThreadDecrementElicitationResponse, }, ThreadSetName => "thread/name/set" { params: v2::ThreadSetNameParams, + serialization: thread_id(params.thread_id), response: v2::ThreadSetNameResponse, }, #[experimental("thread/goal/set")] ThreadGoalSet => "thread/goal/set" { params: v2::ThreadGoalSetParams, + serialization: thread_id(params.thread_id), response: v2::ThreadGoalSetResponse, }, #[experimental("thread/goal/get")] ThreadGoalGet => "thread/goal/get" { params: v2::ThreadGoalGetParams, + serialization: thread_id(params.thread_id), response: v2::ThreadGoalGetResponse, }, #[experimental("thread/goal/clear")] ThreadGoalClear => "thread/goal/clear" { params: v2::ThreadGoalClearParams, + serialization: thread_id(params.thread_id), response: v2::ThreadGoalClearResponse, }, ThreadMetadataUpdate => "thread/metadata/update" { params: v2::ThreadMetadataUpdateParams, + serialization: thread_id(params.thread_id), response: v2::ThreadMetadataUpdateResponse, }, #[experimental("thread/memoryMode/set")] ThreadMemoryModeSet => "thread/memoryMode/set" { params: v2::ThreadMemoryModeSetParams, + serialization: thread_id(params.thread_id), response: v2::ThreadMemoryModeSetResponse, }, #[experimental("memory/reset")] MemoryReset => "memory/reset" { params: #[ts(type = "undefined")] #[serde(skip_serializing_if = "Option::is_none")] Option<()>, + serialization: global("memory"), response: v2::MemoryResetResponse, }, ThreadUnarchive => "thread/unarchive" { params: v2::ThreadUnarchiveParams, + serialization: thread_id(params.thread_id), response: v2::ThreadUnarchiveResponse, }, ThreadCompactStart => "thread/compact/start" { params: v2::ThreadCompactStartParams, + serialization: thread_id(params.thread_id), response: v2::ThreadCompactStartResponse, }, ThreadShellCommand => "thread/shellCommand" { params: v2::ThreadShellCommandParams, + serialization: thread_id(params.thread_id), response: v2::ThreadShellCommandResponse, }, ThreadApproveGuardianDeniedAction => "thread/approveGuardianDeniedAction" { params: v2::ThreadApproveGuardianDeniedActionParams, + serialization: thread_id(params.thread_id), response: v2::ThreadApproveGuardianDeniedActionResponse, }, #[experimental("thread/backgroundTerminals/clean")] ThreadBackgroundTerminalsClean => "thread/backgroundTerminals/clean" { params: v2::ThreadBackgroundTerminalsCleanParams, + serialization: thread_id(params.thread_id), response: v2::ThreadBackgroundTerminalsCleanResponse, }, ThreadRollback => "thread/rollback" { params: v2::ThreadRollbackParams, + serialization: thread_id(params.thread_id), response: v2::ThreadRollbackResponse, }, ThreadList => "thread/list" { params: v2::ThreadListParams, + serialization: None, response: v2::ThreadListResponse, }, ThreadLoadedList => "thread/loaded/list" { params: v2::ThreadLoadedListParams, + serialization: None, response: v2::ThreadLoadedListResponse, }, ThreadRead => "thread/read" { params: v2::ThreadReadParams, + serialization: thread_id(params.thread_id), response: v2::ThreadReadResponse, }, ThreadTurnsList => "thread/turns/list" { params: v2::ThreadTurnsListParams, + // Explicitly concurrent: this primarily reads append-only rollout storage. + serialization: None, response: v2::ThreadTurnsListResponse, }, /// Append raw Responses API items to the thread history without starting a user turn. ThreadInjectItems => "thread/inject_items" { params: v2::ThreadInjectItemsParams, + serialization: thread_id(params.thread_id), response: v2::ThreadInjectItemsResponse, }, SkillsList => "skills/list" { params: v2::SkillsListParams, + serialization: global("config"), response: v2::SkillsListResponse, }, MarketplaceAdd => "marketplace/add" { params: v2::MarketplaceAddParams, + serialization: global("config"), response: v2::MarketplaceAddResponse, }, MarketplaceRemove => "marketplace/remove" { params: v2::MarketplaceRemoveParams, + serialization: global("config"), response: v2::MarketplaceRemoveResponse, }, MarketplaceUpgrade => "marketplace/upgrade" { params: v2::MarketplaceUpgradeParams, + serialization: global("config"), response: v2::MarketplaceUpgradeResponse, }, PluginList => "plugin/list" { params: v2::PluginListParams, + serialization: global("config"), response: v2::PluginListResponse, }, PluginRead => "plugin/read" { params: v2::PluginReadParams, + serialization: global("config"), response: v2::PluginReadResponse, }, AppsList => "app/list" { params: v2::AppsListParams, + serialization: None, response: v2::AppsListResponse, }, DeviceKeyCreate => "device/key/create" { params: v2::DeviceKeyCreateParams, + serialization: global("device-key"), response: v2::DeviceKeyCreateResponse, }, DeviceKeyPublic => "device/key/public" { params: v2::DeviceKeyPublicParams, + serialization: global("device-key"), response: v2::DeviceKeyPublicResponse, }, DeviceKeySign => "device/key/sign" { params: v2::DeviceKeySignParams, + serialization: global("device-key"), response: v2::DeviceKeySignResponse, }, + // File system requests are intentionally concurrent. Desktop already treats local + // file system operations as concurrent, and app-server remote fs mirrors that model. FsReadFile => "fs/readFile" { params: v2::FsReadFileParams, + serialization: None, response: v2::FsReadFileResponse, }, FsWriteFile => "fs/writeFile" { params: v2::FsWriteFileParams, + serialization: None, response: v2::FsWriteFileResponse, }, FsCreateDirectory => "fs/createDirectory" { params: v2::FsCreateDirectoryParams, + serialization: None, response: v2::FsCreateDirectoryResponse, }, FsGetMetadata => "fs/getMetadata" { params: v2::FsGetMetadataParams, + serialization: None, response: v2::FsGetMetadataResponse, }, FsReadDirectory => "fs/readDirectory" { params: v2::FsReadDirectoryParams, + serialization: None, response: v2::FsReadDirectoryResponse, }, FsRemove => "fs/remove" { params: v2::FsRemoveParams, + serialization: None, response: v2::FsRemoveResponse, }, FsCopy => "fs/copy" { params: v2::FsCopyParams, + serialization: None, response: v2::FsCopyResponse, }, FsWatch => "fs/watch" { params: v2::FsWatchParams, + serialization: fs_watch_id(params.watch_id), response: v2::FsWatchResponse, }, FsUnwatch => "fs/unwatch" { params: v2::FsUnwatchParams, + serialization: fs_watch_id(params.watch_id), response: v2::FsUnwatchResponse, }, SkillsConfigWrite => "skills/config/write" { params: v2::SkillsConfigWriteParams, + serialization: global("config"), response: v2::SkillsConfigWriteResponse, }, PluginInstall => "plugin/install" { params: v2::PluginInstallParams, + serialization: global("config"), response: v2::PluginInstallResponse, }, PluginUninstall => "plugin/uninstall" { params: v2::PluginUninstallParams, + serialization: global("config"), response: v2::PluginUninstallResponse, }, TurnStart => "turn/start" { params: v2::TurnStartParams, inspect_params: true, + serialization: thread_id(params.thread_id), response: v2::TurnStartResponse, }, TurnSteer => "turn/steer" { params: v2::TurnSteerParams, inspect_params: true, + serialization: thread_id(params.thread_id), response: v2::TurnSteerResponse, }, TurnInterrupt => "turn/interrupt" { params: v2::TurnInterruptParams, + serialization: thread_id(params.thread_id), response: v2::TurnInterruptResponse, }, #[experimental("thread/realtime/start")] ThreadRealtimeStart => "thread/realtime/start" { params: v2::ThreadRealtimeStartParams, + serialization: thread_id(params.thread_id), response: v2::ThreadRealtimeStartResponse, }, #[experimental("thread/realtime/appendAudio")] ThreadRealtimeAppendAudio => "thread/realtime/appendAudio" { params: v2::ThreadRealtimeAppendAudioParams, + serialization: thread_id(params.thread_id), response: v2::ThreadRealtimeAppendAudioResponse, }, #[experimental("thread/realtime/appendText")] ThreadRealtimeAppendText => "thread/realtime/appendText" { params: v2::ThreadRealtimeAppendTextParams, + serialization: thread_id(params.thread_id), response: v2::ThreadRealtimeAppendTextResponse, }, #[experimental("thread/realtime/stop")] ThreadRealtimeStop => "thread/realtime/stop" { params: v2::ThreadRealtimeStopParams, + serialization: thread_id(params.thread_id), response: v2::ThreadRealtimeStopResponse, }, #[experimental("thread/realtime/listVoices")] ThreadRealtimeListVoices => "thread/realtime/listVoices" { params: v2::ThreadRealtimeListVoicesParams, + serialization: None, response: v2::ThreadRealtimeListVoicesResponse, }, ReviewStart => "review/start" { params: v2::ReviewStartParams, + serialization: thread_id(params.thread_id), response: v2::ReviewStartResponse, }, ModelList => "model/list" { params: v2::ModelListParams, + serialization: None, response: v2::ModelListResponse, }, ExperimentalFeatureList => "experimentalFeature/list" { params: v2::ExperimentalFeatureListParams, + serialization: global("config"), response: v2::ExperimentalFeatureListResponse, }, ExperimentalFeatureEnablementSet => "experimentalFeature/enablement/set" { params: v2::ExperimentalFeatureEnablementSetParams, + serialization: global("config"), response: v2::ExperimentalFeatureEnablementSetResponse, }, #[experimental("collaborationMode/list")] /// Lists collaboration mode presets. CollaborationModeList => "collaborationMode/list" { params: v2::CollaborationModeListParams, + serialization: None, response: v2::CollaborationModeListResponse, }, #[experimental("mock/experimentalMethod")] /// Test-only method used to validate experimental gating. MockExperimentalMethod => "mock/experimentalMethod" { params: v2::MockExperimentalMethodParams, + serialization: None, response: v2::MockExperimentalMethodResponse, }, McpServerOauthLogin => "mcpServer/oauth/login" { params: v2::McpServerOauthLoginParams, + serialization: mcp_oauth_server(params.name), response: v2::McpServerOauthLoginResponse, }, McpServerRefresh => "config/mcpServer/reload" { params: #[ts(type = "undefined")] #[serde(skip_serializing_if = "Option::is_none")] Option<()>, + serialization: global("mcp-registry"), response: v2::McpServerRefreshResponse, }, McpServerStatusList => "mcpServerStatus/list" { params: v2::ListMcpServerStatusParams, + serialization: global("mcp-registry"), response: v2::ListMcpServerStatusResponse, }, McpResourceRead => "mcpServer/resource/read" { params: v2::McpResourceReadParams, + serialization: optional_thread_id(params.thread_id), response: v2::McpResourceReadResponse, }, McpServerToolCall => "mcpServer/tool/call" { params: v2::McpServerToolCallParams, + serialization: thread_id(params.thread_id), response: v2::McpServerToolCallResponse, }, WindowsSandboxSetupStart => "windowsSandbox/setupStart" { params: v2::WindowsSandboxSetupStartParams, + serialization: global("windows-sandbox-setup"), response: v2::WindowsSandboxSetupStartResponse, }, LoginAccount => "account/login/start" { params: v2::LoginAccountParams, inspect_params: true, + serialization: global("account-auth"), response: v2::LoginAccountResponse, }, CancelLoginAccount => "account/login/cancel" { params: v2::CancelLoginAccountParams, + serialization: global("account-auth"), response: v2::CancelLoginAccountResponse, }, LogoutAccount => "account/logout" { params: #[ts(type = "undefined")] #[serde(skip_serializing_if = "Option::is_none")] Option<()>, + serialization: global("account-auth"), response: v2::LogoutAccountResponse, }, GetAccountRateLimits => "account/rateLimits/read" { params: #[ts(type = "undefined")] #[serde(skip_serializing_if = "Option::is_none")] Option<()>, + serialization: None, response: v2::GetAccountRateLimitsResponse, }, SendAddCreditsNudgeEmail => "account/sendAddCreditsNudgeEmail" { params: v2::SendAddCreditsNudgeEmailParams, + serialization: global("account-auth"), response: v2::SendAddCreditsNudgeEmailResponse, }, FeedbackUpload => "feedback/upload" { params: v2::FeedbackUploadParams, + serialization: None, response: v2::FeedbackUploadResponse, }, @@ -582,86 +744,106 @@ client_request_definitions! { OneOffCommandExec => "command/exec" { params: v2::CommandExecParams, inspect_params: true, + serialization: optional_command_process_id(params.process_id), response: v2::CommandExecResponse, }, /// Write stdin bytes to a running `command/exec` session or close stdin. CommandExecWrite => "command/exec/write" { params: v2::CommandExecWriteParams, + serialization: command_process_id(params.process_id), response: v2::CommandExecWriteResponse, }, /// Terminate a running `command/exec` session by client-supplied `processId`. CommandExecTerminate => "command/exec/terminate" { params: v2::CommandExecTerminateParams, + serialization: command_process_id(params.process_id), response: v2::CommandExecTerminateResponse, }, /// Resize a running PTY-backed `command/exec` session by client-supplied `processId`. CommandExecResize => "command/exec/resize" { params: v2::CommandExecResizeParams, + serialization: command_process_id(params.process_id), response: v2::CommandExecResizeResponse, }, ConfigRead => "config/read" { params: v2::ConfigReadParams, + serialization: global("config"), response: v2::ConfigReadResponse, }, ExternalAgentConfigDetect => "externalAgentConfig/detect" { params: v2::ExternalAgentConfigDetectParams, + serialization: global("config"), response: v2::ExternalAgentConfigDetectResponse, }, ExternalAgentConfigImport => "externalAgentConfig/import" { params: v2::ExternalAgentConfigImportParams, + serialization: global("config"), response: v2::ExternalAgentConfigImportResponse, }, ConfigValueWrite => "config/value/write" { params: v2::ConfigValueWriteParams, + serialization: global("config"), response: v2::ConfigWriteResponse, }, ConfigBatchWrite => "config/batchWrite" { params: v2::ConfigBatchWriteParams, + serialization: global("config"), response: v2::ConfigWriteResponse, }, ConfigRequirementsRead => "configRequirements/read" { params: #[ts(type = "undefined")] #[serde(skip_serializing_if = "Option::is_none")] Option<()>, + serialization: global("config"), response: v2::ConfigRequirementsReadResponse, }, GetAccount => "account/read" { params: v2::GetAccountParams, + serialization: global("account-auth"), response: v2::GetAccountResponse, }, /// DEPRECATED APIs below GetConversationSummary { params: v1::GetConversationSummaryParams, + serialization: None, response: v1::GetConversationSummaryResponse, }, GitDiffToRemote { params: v1::GitDiffToRemoteParams, + serialization: None, response: v1::GitDiffToRemoteResponse, }, /// DEPRECATED in favor of GetAccount GetAuthStatus { params: v1::GetAuthStatusParams, + serialization: global("account-auth"), response: v1::GetAuthStatusResponse, }, + // Legacy fuzzy search cancellation is intentionally concurrent: clients reuse a + // cancellation token so a newer request can cancel an older in-flight search. FuzzyFileSearch { params: FuzzyFileSearchParams, + serialization: None, response: FuzzyFileSearchResponse, }, #[experimental("fuzzyFileSearch/sessionStart")] FuzzyFileSearchSessionStart => "fuzzyFileSearch/sessionStart" { params: FuzzyFileSearchSessionStartParams, + serialization: fuzzy_session_id(params.session_id), response: FuzzyFileSearchSessionStartResponse, }, #[experimental("fuzzyFileSearch/sessionUpdate")] FuzzyFileSearchSessionUpdate => "fuzzyFileSearch/sessionUpdate" { params: FuzzyFileSearchSessionUpdateParams, + serialization: fuzzy_session_id(params.session_id), response: FuzzyFileSearchSessionUpdateResponse, }, #[experimental("fuzzyFileSearch/sessionStop")] FuzzyFileSearchSessionStop => "fuzzyFileSearch/sessionStop" { params: FuzzyFileSearchSessionStopParams, + serialization: fuzzy_session_id(params.session_id), response: FuzzyFileSearchSessionStopResponse, }, } @@ -1150,6 +1332,325 @@ mod tests { test_path_buf(&path).abs() } + fn request_id() -> RequestId { + const REQUEST_ID: i64 = 1; + RequestId::Integer(REQUEST_ID) + } + + #[test] + fn client_request_serialization_scope_covers_keyed_families() { + let thread_id = "thread-1".to_string(); + let thread_resume = ClientRequest::ThreadResume { + request_id: request_id(), + params: v2::ThreadResumeParams { + thread_id: thread_id.clone(), + ..Default::default() + }, + }; + assert_eq!( + thread_resume.serialization_scope(), + Some(ClientRequestSerializationScope::Thread { + thread_id: thread_id.clone() + }) + ); + + let thread_resume_with_path = ClientRequest::ThreadResume { + request_id: request_id(), + params: v2::ThreadResumeParams { + thread_id: thread_id.clone(), + path: Some(PathBuf::from("/tmp/resume-thread.jsonl")), + ..Default::default() + }, + }; + assert_eq!( + thread_resume_with_path.serialization_scope(), + Some(ClientRequestSerializationScope::Thread { + thread_id: thread_id.clone() + }) + ); + + let thread_fork = ClientRequest::ThreadFork { + request_id: request_id(), + params: v2::ThreadForkParams { + thread_id: thread_id.clone(), + path: Some(PathBuf::from("/tmp/source-thread.jsonl")), + ..Default::default() + }, + }; + assert_eq!( + thread_fork.serialization_scope(), + Some(ClientRequestSerializationScope::Thread { thread_id }) + ); + + let command_exec = ClientRequest::OneOffCommandExec { + request_id: request_id(), + params: v2::CommandExecParams { + command: vec!["sleep".to_string(), "10".to_string()], + process_id: Some("proc-1".to_string()), + tty: false, + stream_stdin: false, + stream_stdout_stderr: false, + output_bytes_cap: None, + disable_output_cap: false, + disable_timeout: false, + timeout_ms: None, + cwd: None, + env: None, + size: None, + sandbox_policy: None, + permission_profile: None, + }, + }; + assert_eq!( + command_exec.serialization_scope(), + Some(ClientRequestSerializationScope::CommandExecProcess { + process_id: "proc-1".to_string() + }) + ); + + let fuzzy_update = ClientRequest::FuzzyFileSearchSessionUpdate { + request_id: request_id(), + params: FuzzyFileSearchSessionUpdateParams { + session_id: "search-1".to_string(), + query: "lib".to_string(), + }, + }; + assert_eq!( + fuzzy_update.serialization_scope(), + Some(ClientRequestSerializationScope::FuzzyFileSearchSession { + session_id: "search-1".to_string() + }) + ); + + let fs_watch = ClientRequest::FsWatch { + request_id: request_id(), + params: v2::FsWatchParams { + watch_id: "watch-1".to_string(), + path: absolute_path("/tmp/repo"), + }, + }; + assert_eq!( + fs_watch.serialization_scope(), + Some(ClientRequestSerializationScope::FsWatch { + watch_id: "watch-1".to_string() + }) + ); + + let plugin_install = ClientRequest::PluginInstall { + request_id: request_id(), + params: v2::PluginInstallParams { + marketplace_path: Some(absolute_path("/tmp/marketplace")), + remote_marketplace_name: None, + plugin_name: "plugin-a".to_string(), + }, + }; + assert_eq!( + plugin_install.serialization_scope(), + Some(ClientRequestSerializationScope::Global("config")) + ); + + let plugin_uninstall = ClientRequest::PluginUninstall { + request_id: request_id(), + params: v2::PluginUninstallParams { + plugin_id: "plugin-a".to_string(), + }, + }; + assert_eq!( + plugin_uninstall.serialization_scope(), + Some(ClientRequestSerializationScope::Global("config")) + ); + + let mcp_oauth = ClientRequest::McpServerOauthLogin { + request_id: request_id(), + params: v2::McpServerOauthLoginParams { + name: "server-a".to_string(), + scopes: None, + timeout_secs: None, + }, + }; + assert_eq!( + mcp_oauth.serialization_scope(), + Some(ClientRequestSerializationScope::McpOauth { + server_name: "server-a".to_string() + }) + ); + + let mcp_resource_read = ClientRequest::McpResourceRead { + request_id: request_id(), + params: v2::McpResourceReadParams { + thread_id: Some("thread-1".to_string()), + server: "server-a".to_string(), + uri: "file:///tmp/resource".to_string(), + }, + }; + assert_eq!( + mcp_resource_read.serialization_scope(), + Some(ClientRequestSerializationScope::Thread { + thread_id: "thread-1".to_string() + }) + ); + + let config_read = ClientRequest::ConfigRead { + request_id: request_id(), + params: v2::ConfigReadParams { + include_layers: false, + cwd: None, + }, + }; + assert_eq!( + config_read.serialization_scope(), + Some(ClientRequestSerializationScope::Global("config")) + ); + + let account_read = ClientRequest::GetAccount { + request_id: request_id(), + params: v2::GetAccountParams { + refresh_token: false, + }, + }; + assert_eq!( + account_read.serialization_scope(), + Some(ClientRequestSerializationScope::Global("account-auth")) + ); + + let thread_goal_set = ClientRequest::ThreadGoalSet { + request_id: request_id(), + params: v2::ThreadGoalSetParams { + thread_id: "goal-thread".to_string(), + objective: Some("ship it".to_string()), + status: None, + token_budget: None, + }, + }; + assert_eq!( + thread_goal_set.serialization_scope(), + Some(ClientRequestSerializationScope::Thread { + thread_id: "goal-thread".to_string() + }) + ); + + let guardian_approval = ClientRequest::ThreadApproveGuardianDeniedAction { + request_id: request_id(), + params: v2::ThreadApproveGuardianDeniedActionParams { + thread_id: "guardian-thread".to_string(), + event: json!({ "type": "guardian" }), + }, + }; + assert_eq!( + guardian_approval.serialization_scope(), + Some(ClientRequestSerializationScope::Thread { + thread_id: "guardian-thread".to_string() + }) + ); + + let marketplace_remove = ClientRequest::MarketplaceRemove { + request_id: request_id(), + params: v2::MarketplaceRemoveParams { + marketplace_name: "marketplace".to_string(), + }, + }; + assert_eq!( + marketplace_remove.serialization_scope(), + Some(ClientRequestSerializationScope::Global("config")) + ); + + let device_key_create = ClientRequest::DeviceKeyCreate { + request_id: request_id(), + params: v2::DeviceKeyCreateParams { + protection_policy: None, + account_user_id: "user".to_string(), + client_id: "client".to_string(), + }, + }; + assert_eq!( + device_key_create.serialization_scope(), + Some(ClientRequestSerializationScope::Global("device-key")) + ); + + let add_credits_nudge = ClientRequest::SendAddCreditsNudgeEmail { + request_id: request_id(), + params: v2::SendAddCreditsNudgeEmailParams { + credit_type: v2::AddCreditsNudgeCreditType::Credits, + }, + }; + assert_eq!( + add_credits_nudge.serialization_scope(), + Some(ClientRequestSerializationScope::Global("account-auth")) + ); + } + + #[test] + fn client_request_serialization_scope_covers_unkeyed_representatives() { + let initialize = ClientRequest::Initialize { + request_id: request_id(), + params: v1::InitializeParams { + client_info: v1::ClientInfo { + name: "test".to_string(), + title: None, + version: "0.1.0".to_string(), + }, + capabilities: None, + }, + }; + assert_eq!(initialize.serialization_scope(), None); + + let thread_start = ClientRequest::ThreadStart { + request_id: request_id(), + params: v2::ThreadStartParams::default(), + }; + assert_eq!(thread_start.serialization_scope(), None); + + let command_exec = ClientRequest::OneOffCommandExec { + request_id: request_id(), + params: v2::CommandExecParams { + command: vec!["true".to_string()], + process_id: None, + tty: false, + stream_stdin: false, + stream_stdout_stderr: false, + output_bytes_cap: None, + disable_output_cap: false, + disable_timeout: false, + timeout_ms: None, + cwd: None, + env: None, + size: None, + sandbox_policy: None, + permission_profile: None, + }, + }; + assert_eq!(command_exec.serialization_scope(), None); + + let fs_read = ClientRequest::FsReadFile { + request_id: request_id(), + params: v2::FsReadFileParams { + path: absolute_path("/tmp/file.txt"), + }, + }; + assert_eq!(fs_read.serialization_scope(), None); + + let thread_turns_list = ClientRequest::ThreadTurnsList { + request_id: request_id(), + params: v2::ThreadTurnsListParams { + thread_id: "thread-1".to_string(), + cursor: None, + limit: None, + sort_direction: None, + }, + }; + assert_eq!(thread_turns_list.serialization_scope(), None); + + let mcp_resource_read = ClientRequest::McpResourceRead { + request_id: request_id(), + params: v2::McpResourceReadParams { + thread_id: None, + server: "server-a".to_string(), + uri: "file:///tmp/resource".to_string(), + }, + }; + assert_eq!(mcp_resource_read.serialization_scope(), None); + } + #[test] fn serialize_get_conversation_summary() -> Result<()> { let request = ClientRequest::GetConversationSummary { diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 2f968789e4..248f5d4158 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -192,7 +192,7 @@ Example with notification opt-out: - `fs/changed` — notification emitted when watched paths change, including the `watchId` and `changedPaths`. - `model/list` — list available models (set `includeHidden: true` to include entries with `hidden: true`), with reasoning effort options, `additionalSpeedTiers`, optional legacy `upgrade` model ids, optional `upgradeInfo` metadata (`model`, `upgradeCopy`, `modelLink`, `migrationMarkdown`), and optional `availabilityNux` metadata. - `experimentalFeature/list` — list feature flags with stage metadata (`beta`, `underDevelopment`, `stable`, etc.), enabled/default-enabled state, and cursor pagination. For non-beta flags, `displayName`/`description`/`announcement` are `null`. -- `experimentalFeature/enablement/set` — patch the in-memory process-wide runtime feature enablement for the currently supported feature keys (`apps`, `plugins`). For each feature, precedence is: cloud requirements > --enable > config.toml > experimentalFeature/enablement/set (new) > code default. +- `experimentalFeature/enablement/set` — patch the in-memory process-wide runtime feature enablement for the currently supported feature keys (`apps`, `memories`, `plugins`, `remote_control`, `tool_search`, `tool_suggest`, `tool_call_mcp_elicitation`). For each feature, precedence is: cloud requirements > --enable > config.toml > experimentalFeature/enablement/set (new) > code default. - `collaborationMode/list` — list available collaboration mode presets (experimental, no pagination). This response omits built-in developer instructions; clients should either pass `settings.developer_instructions: null` when setting a mode to use Codex's built-in instructions, or provide their own instructions explicitly. - `skills/list` — list skills for one or more `cwd` values (optional `forceReload`). - `marketplace/add` — add a remote plugin marketplace from an HTTP(S) Git URL, SSH Git URL, or GitHub `owner/repo` shorthand, then persist it into the user marketplace config. Returns the installed root path plus whether the marketplace was already present. diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 04de06f469..d67efed164 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -178,6 +178,7 @@ pub(crate) async fn apply_bespoke_event_handling( outgoing: ThreadScopedOutgoingMessageSender, thread_state: Arc>, thread_watch_manager: ThreadWatchManager, + thread_list_state_permit: Arc, api_version: ApiVersion, fallback_model_provider: String, codex_home: &Path, @@ -1874,6 +1875,20 @@ pub(crate) async fn apply_bespoke_event_handling( }; if let Some(request_id) = pending { + let _thread_list_state_permit = match thread_list_state_permit.acquire().await { + Ok(permit) => permit, + Err(err) => { + outgoing + .send_error( + request_id, + internal_error(format!( + "failed to acquire thread list state permit: {err}" + )), + ) + .await; + return; + } + }; let Some(rollout_path) = conversation.rollout_path() else { outgoing .send_error( @@ -3271,6 +3286,7 @@ mod tests { self.outgoing.clone(), self.thread_state.clone(), self.thread_watch_manager.clone(), + Arc::new(tokio::sync::Semaphore::new(/*permits*/ 1)), ApiVersion::V2, "test-provider".to_string(), &self.codex_home, diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 1ab9c0c789..362f8692b6 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -393,6 +393,8 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; use tokio::sync::Mutex; +use tokio::sync::Semaphore; +use tokio::sync::SemaphorePermit; use tokio::sync::broadcast; use tokio::sync::oneshot; use tokio::sync::watch; @@ -524,6 +526,10 @@ pub(crate) struct CodexMessageProcessor { pending_thread_unloads: Arc>>, thread_state_manager: ThreadStateManager, thread_watch_manager: ThreadWatchManager, + /// Serializes mutations of list membership or fields rendered from list + /// results. `thread/list` is intentionally not serialized so it can run + /// concurrently against mostly append-only storage. + thread_list_state_permit: Arc, command_exec_manager: CommandExecManager, workspace_settings_cache: Arc, pending_fuzzy_searches: Arc>>>, @@ -549,6 +555,7 @@ struct ListenerTaskContext { pending_thread_unloads: Arc>>, analytics_events_client: AnalyticsEventsClient, thread_watch_manager: ThreadWatchManager, + thread_list_state_permit: Arc, fallback_model_provider: String, codex_home: PathBuf, } @@ -790,6 +797,7 @@ impl CodexMessageProcessor { pending_thread_unloads: Arc::new(Mutex::new(HashSet::new())), thread_state_manager: ThreadStateManager::new(), thread_watch_manager: ThreadWatchManager::new_with_outgoing(outgoing), + thread_list_state_permit: Arc::new(Semaphore::new(/*permits*/ 1)), command_exec_manager: CommandExecManager::default(), workspace_settings_cache: Arc::new( workspace_settings::WorkspaceSettingsCache::default(), @@ -1326,6 +1334,17 @@ impl CodexMessageProcessor { } } + async fn acquire_thread_list_state_permit( + &self, + ) -> Result, JSONRPCErrorError> { + self.thread_list_state_permit + .acquire() + .await + .map_err(|err| { + internal_error(format!("failed to acquire thread list state permit: {err}")) + }) + } + async fn login_api_key_common( &self, params: &LoginApiKeyParams, @@ -2421,6 +2440,7 @@ impl CodexMessageProcessor { pending_thread_unloads: Arc::clone(&self.pending_thread_unloads), analytics_events_client: self.analytics_events_client.clone(), thread_watch_manager: self.thread_watch_manager.clone(), + thread_list_state_permit: self.thread_list_state_permit.clone(), fallback_model_provider: self.config.model_provider_id.clone(), codex_home: self.config.codex_home.to_path_buf(), }; @@ -2871,6 +2891,13 @@ impl CodexMessageProcessor { } async fn thread_archive(&self, request_id: ConnectionRequestId, params: ThreadArchiveParams) { + let _thread_list_state_permit = match self.acquire_thread_list_state_permit().await { + Ok(permit) => permit, + Err(error) => { + self.outgoing.send_error(request_id, error).await; + return; + } + }; let result = self.thread_archive_response(params).await; let archived_thread_ids = result .as_ref() @@ -3077,6 +3104,7 @@ impl CodexMessageProcessor { return Err(invalid_request("thread name must not be empty")); }; + let _thread_list_state_permit = self.acquire_thread_list_state_permit().await?; if let Ok(thread) = self.thread_manager.get_thread(thread_id).await { self.submit_core_op(request_id, thread.as_ref(), Op::SetThreadName { name }) .await @@ -3218,6 +3246,7 @@ impl CodexMessageProcessor { return Err(invalid_request("gitInfo must include at least one field")); } + let _thread_list_state_permit = self.acquire_thread_list_state_permit().await?; let loaded_thread = self.thread_manager.get_thread(thread_uuid).await.ok(); let mut state_db_ctx = loaded_thread.as_ref().and_then(|thread| thread.state_db()); if state_db_ctx.is_none() { @@ -3414,6 +3443,13 @@ impl CodexMessageProcessor { request_id: ConnectionRequestId, params: ThreadUnarchiveParams, ) { + let _thread_list_state_permit = match self.acquire_thread_list_state_permit().await { + Ok(permit) => permit, + Err(error) => { + self.outgoing.send_error(request_id, error).await; + return; + } + }; let result = self.thread_unarchive_response(params).await; let notification = result @@ -4143,6 +4179,13 @@ impl CodexMessageProcessor { return; } + let _thread_list_state_permit = match self.acquire_thread_list_state_permit().await { + Ok(permit) => permit, + Err(error) => { + self.outgoing.send_error(request_id, error).await; + return; + } + }; match self.resume_running_thread(&request_id, ¶ms).await { Ok(true) => return, Ok(false) => {} @@ -7237,6 +7280,7 @@ impl CodexMessageProcessor { pending_thread_unloads: Arc::clone(&self.pending_thread_unloads), analytics_events_client: self.analytics_events_client.clone(), thread_watch_manager: self.thread_watch_manager.clone(), + thread_list_state_permit: self.thread_list_state_permit.clone(), fallback_model_provider: self.config.model_provider_id.clone(), codex_home: self.config.codex_home.to_path_buf(), }, @@ -7354,6 +7398,7 @@ impl CodexMessageProcessor { pending_thread_unloads: Arc::clone(&self.pending_thread_unloads), analytics_events_client: self.analytics_events_client.clone(), thread_watch_manager: self.thread_watch_manager.clone(), + thread_list_state_permit: self.thread_list_state_permit.clone(), fallback_model_provider: self.config.model_provider_id.clone(), codex_home: self.config.codex_home.to_path_buf(), }, @@ -7402,6 +7447,7 @@ impl CodexMessageProcessor { pending_thread_unloads, analytics_events_client: _, thread_watch_manager, + thread_list_state_permit, fallback_model_provider, codex_home, } = listener_task_context; @@ -7480,6 +7526,7 @@ impl CodexMessageProcessor { thread_outgoing, thread_state.clone(), thread_watch_manager.clone(), + thread_list_state_permit.clone(), api_version, fallback_model_provider.clone(), codex_home.as_path(), diff --git a/codex-rs/app-server/src/config_api.rs b/codex-rs/app-server/src/config_api.rs index e8bb82777c..ac7c1510c8 100644 --- a/codex-rs/app-server/src/config_api.rs +++ b/codex-rs/app-server/src/config_api.rs @@ -48,6 +48,7 @@ const SUPPORTED_EXPERIMENTAL_FEATURE_ENABLEMENT: &[&str] = &[ "apps", "memories", "plugins", + "remote_control", "tool_search", "tool_suggest", "tool_call_mcp_elicitation", diff --git a/codex-rs/app-server/src/connection_rpc_gate.rs b/codex-rs/app-server/src/connection_rpc_gate.rs new file mode 100644 index 0000000000..12fed79b36 --- /dev/null +++ b/codex-rs/app-server/src/connection_rpc_gate.rs @@ -0,0 +1,209 @@ +use std::future::Future; + +use tokio::sync::Mutex; +use tokio_util::task::TaskTracker; + +/// Per-connection gate for initialized RPC handler execution. +/// +/// Closing the gate prevents queued handlers from starting while allowing +/// handlers that already acquired a token to finish. +#[derive(Debug)] +pub(crate) struct ConnectionRpcGate { + accepting: Mutex, + tasks: TaskTracker, +} + +impl ConnectionRpcGate { + pub(crate) fn new() -> Self { + let accepting = true; + Self { + accepting: Mutex::new(accepting), + tasks: TaskTracker::new(), + } + } + + pub(crate) async fn run(&self, future: F) + where + F: Future, + { + let token = { + let accepting = self.accepting.lock().await; + if !*accepting { + return; + } + self.tasks.token() + }; + + future.await; + drop(token); + } + + pub(crate) async fn shutdown(&self) { + { + let mut accepting = self.accepting.lock().await; + *accepting = false; + self.tasks.close(); + } + self.tasks.wait().await; + } + + #[cfg(test)] + async fn is_accepting(&self) -> bool { + *self.accepting.lock().await + } + + #[cfg(test)] + fn inflight_count(&self) -> usize { + self.tasks.len() + } +} + +impl Default for ConnectionRpcGate { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; + use tokio::sync::oneshot; + use tokio::time::Duration; + use tokio::time::timeout; + + #[tokio::test] + async fn run_executes_while_open() { + let gate = ConnectionRpcGate::new(); + let ran = Arc::new(AtomicBool::new(/*v*/ false)); + let ran_clone = Arc::clone(&ran); + + gate.run(async move { + ran_clone.store(/*val*/ true, Ordering::Release); + }) + .await; + + assert!(ran.load(Ordering::Acquire)); + } + + #[tokio::test] + async fn run_drops_future_without_polling_after_shutdown() { + let gate = ConnectionRpcGate::new(); + gate.shutdown().await; + let polled = Arc::new(AtomicBool::new(/*v*/ false)); + let polled_clone = Arc::clone(&polled); + + gate.run(async move { + polled_clone.store(/*val*/ true, Ordering::Release); + }) + .await; + + assert!(!polled.load(Ordering::Acquire)); + assert!(!gate.is_accepting().await); + } + + #[tokio::test] + async fn shutdown_waits_for_started_run_to_finish() { + let gate = Arc::new(ConnectionRpcGate::new()); + let (started_tx, started_rx) = oneshot::channel(); + let (finish_tx, finish_rx) = oneshot::channel(); + let gate_for_run = Arc::clone(&gate); + let run_task = tokio::spawn(async move { + gate_for_run + .run(async move { + started_tx.send(()).expect("receiver should be open"); + let _ = finish_rx.await; + }) + .await; + }); + + started_rx.await.expect("run should start"); + assert_eq!(gate.inflight_count(), 1); + + let gate_for_shutdown = Arc::clone(&gate); + let shutdown_task = tokio::spawn(async move { + gate_for_shutdown.shutdown().await; + }); + + timeout(Duration::from_millis(/*millis*/ 50), shutdown_task) + .await + .expect_err("shutdown should wait for the running future"); + + finish_tx + .send(()) + .expect("running future should be waiting"); + run_task.await.expect("run task should complete"); + gate.shutdown().await; + assert_eq!(gate.inflight_count(), 0); + } + + #[tokio::test] + async fn shutdown_drops_late_runs_while_waiting_for_inflight_work() { + let gate = Arc::new(ConnectionRpcGate::new()); + let (started_tx, started_rx) = oneshot::channel(); + let (finish_tx, finish_rx) = oneshot::channel(); + let gate_for_run = Arc::clone(&gate); + let run_task = tokio::spawn(async move { + gate_for_run + .run(async move { + started_tx.send(()).expect("receiver should be open"); + let _ = finish_rx.await; + }) + .await; + }); + + started_rx.await.expect("run should start"); + let gate_for_shutdown = Arc::clone(&gate); + let shutdown_task = tokio::spawn(async move { + gate_for_shutdown.shutdown().await; + }); + + timeout(Duration::from_millis(/*millis*/ 50), shutdown_task) + .await + .expect_err("shutdown should wait for the running future"); + + let late_polled = Arc::new(AtomicBool::new(/*v*/ false)); + let late_polled_clone = Arc::clone(&late_polled); + gate.run(async move { + late_polled_clone.store(/*val*/ true, Ordering::Release); + }) + .await; + + assert!(!late_polled.load(Ordering::Acquire)); + + finish_tx + .send(()) + .expect("running future should still be waiting"); + run_task.await.expect("run task should complete"); + gate.shutdown().await; + assert_eq!(gate.inflight_count(), 0); + } + + #[tokio::test] + async fn run_is_counted_before_handler_body_continues() { + let gate = Arc::new(ConnectionRpcGate::new()); + let (entered_tx, entered_rx) = oneshot::channel(); + let (continue_tx, continue_rx) = oneshot::channel(); + let gate_for_run = Arc::clone(&gate); + let run_task = tokio::spawn(async move { + gate_for_run + .run(async move { + entered_tx.send(()).expect("receiver should be open"); + let _ = continue_rx.await; + }) + .await; + }); + + entered_rx.await.expect("handler body should be entered"); + assert_eq!(gate.inflight_count(), 1); + + continue_tx + .send(()) + .expect("handler body should still be waiting"); + run_task.await.expect("run task should complete"); + assert_eq!(gate.inflight_count(), 0); + } +} diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 73332394f4..5ee044b14a 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -490,7 +490,9 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { processor.clear_runtime_references(); processor.cancel_active_login().await; - processor.connection_closed(IN_PROCESS_CONNECTION_ID).await; + processor + .connection_closed(IN_PROCESS_CONNECTION_ID, &session) + .await; processor.clear_all_thread_listeners().await; processor.drain_background_tasks().await; processor.shutdown_threads().await; diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index de89608101..4369bbee31 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -75,6 +75,7 @@ mod config; mod config_api; mod config_manager; mod config_manager_service; +mod connection_rpc_gate; mod device_key_api; mod dynamic_tools; mod error_code; @@ -87,6 +88,7 @@ pub mod in_process; mod message_processor; mod models; mod outgoing_message; +mod request_serialization; mod server_request_error; mod thread_state; mod thread_status; @@ -809,9 +811,9 @@ pub async fn run_main_with_transport_options( ); } TransportEvent::ConnectionClosed { connection_id } => { - if connections.remove(&connection_id).is_none() { + let Some(connection_state) = connections.remove(&connection_id) else { continue; - } + }; if outbound_control_tx .send(OutboundControlEvent::Closed { connection_id }) .await @@ -819,7 +821,7 @@ pub async fn run_main_with_transport_options( { break; } - processor.connection_closed(connection_id).await; + processor.connection_closed(connection_id, &connection_state.session).await; if shutdown_when_no_connections && connections.is_empty() { break; } @@ -933,6 +935,12 @@ pub async fn run_main_with_transport_options( } if !shutdown_state.forced() { + futures::future::join_all( + connections + .values() + .map(|connection_state| connection_state.session.rpc_gate.shutdown()), + ) + .await; processor.drain_background_tasks().await; processor.shutdown_threads().await; } diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 5cb1c1da36..df77e8ca2b 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -9,6 +9,7 @@ use crate::codex_message_processor::CodexMessageProcessor; use crate::codex_message_processor::CodexMessageProcessorArgs; use crate::config_api::ConfigApi; use crate::config_manager::ConfigManager; +use crate::connection_rpc_gate::ConnectionRpcGate; use crate::device_key_api::DeviceKeyApi; use crate::error_code::invalid_request; use crate::external_agent_config_api::ExternalAgentConfigApi; @@ -18,6 +19,9 @@ use crate::outgoing_message::ConnectionId; use crate::outgoing_message::ConnectionRequestId; use crate::outgoing_message::OutgoingMessageSender; use crate::outgoing_message::RequestContext; +use crate::request_serialization::QueuedInitializedRequest; +use crate::request_serialization::RequestSerializationQueueKey; +use crate::request_serialization::RequestSerializationQueues; use crate::transport::AppServerTransport; use crate::transport::ConnectionOrigin; use crate::transport::RemoteControlHandle; @@ -167,11 +171,13 @@ pub(crate) struct MessageProcessor { config_warnings: Arc>, rpc_transport: AppServerRpcTransport, remote_control_handle: Option, + request_serialization_queues: RequestSerializationQueues, } #[derive(Debug)] pub(crate) struct ConnectionSessionState { origin: ConnectionOrigin, + pub(crate) rpc_gate: Arc, initialized: OnceLock, } @@ -193,6 +199,7 @@ impl ConnectionSessionState { pub(crate) fn new(origin: ConnectionOrigin) -> Self { Self { origin, + rpc_gate: Arc::new(ConnectionRpcGate::new()), initialized: OnceLock::new(), } } @@ -344,6 +351,7 @@ impl MessageProcessor { config_warnings: Arc::new(config_warnings), rpc_transport, remote_control_handle, + request_serialization_queues: RequestSerializationQueues::default(), } } @@ -540,7 +548,12 @@ impl MessageProcessor { self.codex_message_processor.shutdown_threads().await; } - pub(crate) async fn connection_closed(&self, connection_id: ConnectionId) { + pub(crate) async fn connection_closed( + &self, + connection_id: ConnectionId, + session_state: &ConnectionSessionState, + ) { + session_state.rpc_gate.shutdown().await; self.outgoing.connection_closed(connection_id).await; self.fs_watch_manager.connection_closed(connection_id).await; self.codex_message_processor @@ -724,19 +737,46 @@ impl MessageProcessor { ); } + let serialization_scope = codex_request.serialization_scope(); let app_server_client_name = session.app_server_client_name().map(str::to_string); let client_version = session.client_version().map(str::to_string); let device_key_requests_allowed = session.allows_device_key_requests(); - Arc::clone(self) - .handle_initialized_client_request( - connection_request_id, - codex_request, - request_context, - app_server_client_name, - client_version, - device_key_requests_allowed, - ) - .await + let error_request_id = connection_request_id.clone(); + let rpc_gate = Arc::clone(&session.rpc_gate); + let processor = Arc::clone(self); + let span = request_context.span(); + let request = QueuedInitializedRequest::new( + rpc_gate, + async move { + let processor_for_request = Arc::clone(&processor); + let result = processor_for_request + .handle_initialized_client_request( + connection_request_id, + codex_request, + request_context, + app_server_client_name, + client_version, + device_key_requests_allowed, + ) + .await; + if let Err(error) = result { + processor.outgoing.send_error(error_request_id, error).await; + } + } + .instrument(span), + ); + + if let Some(scope) = serialization_scope { + let key = RequestSerializationQueueKey::from_scope(connection_id, scope); + self.request_serialization_queues + .enqueue(key, request) + .await; + } else { + tokio::spawn(async move { + request.run().await; + }); + } + Ok(()) } async fn handle_initialized_client_request( diff --git a/codex-rs/app-server/src/request_serialization.rs b/codex-rs/app-server/src/request_serialization.rs new file mode 100644 index 0000000000..c3e21d134e --- /dev/null +++ b/codex-rs/app-server/src/request_serialization.rs @@ -0,0 +1,380 @@ +use std::collections::HashMap; +use std::collections::VecDeque; +use std::future::Future; +use std::path::PathBuf; +use std::pin::Pin; +use std::sync::Arc; + +use codex_app_server_protocol::ClientRequestSerializationScope; +use tokio::sync::Mutex; +use tracing::Instrument; + +use crate::connection_rpc_gate::ConnectionRpcGate; +use crate::outgoing_message::ConnectionId; + +type BoxFutureUnit = Pin + Send + 'static>>; + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(crate) enum RequestSerializationQueueKey { + Global(&'static str), + Thread { + thread_id: String, + }, + ThreadPath { + path: PathBuf, + }, + CommandExecProcess { + connection_id: ConnectionId, + process_id: String, + }, + FuzzyFileSearchSession { + session_id: String, + }, + FsWatch { + connection_id: ConnectionId, + watch_id: String, + }, + McpOauth { + server_name: String, + }, +} + +impl RequestSerializationQueueKey { + pub(crate) fn from_scope( + connection_id: ConnectionId, + scope: ClientRequestSerializationScope, + ) -> Self { + match scope { + ClientRequestSerializationScope::Global(name) => Self::Global(name), + ClientRequestSerializationScope::Thread { thread_id } => Self::Thread { thread_id }, + ClientRequestSerializationScope::ThreadPath { path } => Self::ThreadPath { path }, + ClientRequestSerializationScope::CommandExecProcess { process_id } => { + Self::CommandExecProcess { + connection_id, + process_id, + } + } + ClientRequestSerializationScope::FuzzyFileSearchSession { session_id } => { + Self::FuzzyFileSearchSession { session_id } + } + ClientRequestSerializationScope::FsWatch { watch_id } => Self::FsWatch { + connection_id, + watch_id, + }, + ClientRequestSerializationScope::McpOauth { server_name } => { + Self::McpOauth { server_name } + } + } + } +} + +pub(crate) struct QueuedInitializedRequest { + gate: Arc, + future: BoxFutureUnit, +} + +impl QueuedInitializedRequest { + pub(crate) fn new( + gate: Arc, + future: impl Future + Send + 'static, + ) -> Self { + Self { + gate, + future: Box::pin(future), + } + } + + pub(crate) async fn run(self) { + let Self { gate, future } = self; + gate.run(future).await; + } +} + +#[derive(Clone, Default)] +pub(crate) struct RequestSerializationQueues { + inner: Arc>>>, +} + +impl RequestSerializationQueues { + pub(crate) async fn enqueue( + &self, + key: RequestSerializationQueueKey, + request: QueuedInitializedRequest, + ) { + let should_spawn = { + let mut queues = self.inner.lock().await; + match queues.get_mut(&key) { + Some(queue) => { + queue.push_back(request); + false + } + None => { + let mut queue = VecDeque::new(); + queue.push_back(request); + queues.insert(key.clone(), queue); + true + } + } + }; + + if should_spawn { + let queues = self.clone(); + let span = tracing::debug_span!("app_server.serialized_request_queue", ?key); + tokio::spawn(async move { queues.drain(key).await }.instrument(span)); + } + } + + async fn drain(self, key: RequestSerializationQueueKey) { + loop { + let request = { + let mut queues = self.inner.lock().await; + let Some(queue) = queues.get_mut(&key) else { + return; + }; + match queue.pop_front() { + Some(request) => request, + None => { + queues.remove(&key); + return; + } + } + }; + + request.run().await; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use std::sync::Arc; + use tokio::sync::mpsc; + use tokio::sync::oneshot; + use tokio::time::Duration; + use tokio::time::timeout; + + const FIRST_REQUEST_VALUE: i32 = 1; + const SECOND_REQUEST_VALUE: i32 = 2; + const THIRD_REQUEST_VALUE: i32 = 3; + + fn gate() -> Arc { + Arc::new(ConnectionRpcGate::new()) + } + + fn queue_drain_timeout() -> Duration { + Duration::from_secs(/*secs*/ 1) + } + + fn shutdown_wait_timeout() -> Duration { + Duration::from_millis(/*millis*/ 50) + } + + #[tokio::test] + async fn same_key_requests_run_fifo() { + let queues = RequestSerializationQueues::default(); + let key = RequestSerializationQueueKey::Global("test"); + let gate = gate(); + let (tx, mut rx) = mpsc::unbounded_channel(); + + for value in [ + FIRST_REQUEST_VALUE, + SECOND_REQUEST_VALUE, + THIRD_REQUEST_VALUE, + ] { + let tx = tx.clone(); + queues + .enqueue( + key.clone(), + QueuedInitializedRequest::new(Arc::clone(&gate), async move { + tx.send(value).expect("receiver should be open"); + }), + ) + .await; + } + drop(tx); + + let mut values = Vec::new(); + while let Some(value) = timeout(queue_drain_timeout(), rx.recv()) + .await + .expect("timed out waiting for queued request") + { + values.push(value); + } + + assert_eq!( + values, + vec![ + FIRST_REQUEST_VALUE, + SECOND_REQUEST_VALUE, + THIRD_REQUEST_VALUE + ] + ); + } + + #[tokio::test] + async fn different_keys_run_concurrently() { + let queues = RequestSerializationQueues::default(); + let (blocked_tx, blocked_rx) = oneshot::channel::<()>(); + let (ran_tx, ran_rx) = oneshot::channel::<()>(); + + queues + .enqueue( + RequestSerializationQueueKey::Global("blocked"), + QueuedInitializedRequest::new(gate(), async move { + let _ = blocked_rx.await; + }), + ) + .await; + queues + .enqueue( + RequestSerializationQueueKey::Global("other"), + QueuedInitializedRequest::new(gate(), async move { + ran_tx.send(()).expect("receiver should be open"); + }), + ) + .await; + + timeout(queue_drain_timeout(), ran_rx) + .await + .expect("other key should not be blocked") + .expect("sender should be open"); + blocked_tx + .send(()) + .expect("blocked request should be waiting"); + } + + #[tokio::test] + async fn closed_gate_request_is_skipped_and_following_requests_continue() { + let queues = RequestSerializationQueues::default(); + let key = RequestSerializationQueueKey::Global("test"); + let live_gate = gate(); + let closed_gate = gate(); + closed_gate.shutdown().await; + let (tx, mut rx) = mpsc::unbounded_channel(); + let (blocked_tx, blocked_rx) = oneshot::channel::<()>(); + + { + let tx = tx.clone(); + queues + .enqueue( + key.clone(), + QueuedInitializedRequest::new(Arc::clone(&live_gate), async move { + tx.send(FIRST_REQUEST_VALUE) + .expect("receiver should be open"); + let _ = blocked_rx.await; + }), + ) + .await; + } + { + let tx = tx.clone(); + queues + .enqueue( + key.clone(), + QueuedInitializedRequest::new(closed_gate, async move { + tx.send(SECOND_REQUEST_VALUE) + .expect("receiver should be open"); + }), + ) + .await; + } + { + let tx = tx.clone(); + queues + .enqueue( + key, + QueuedInitializedRequest::new(live_gate, async move { + tx.send(THIRD_REQUEST_VALUE) + .expect("receiver should be open"); + }), + ) + .await; + } + drop(tx); + + assert_eq!( + timeout(queue_drain_timeout(), rx.recv()) + .await + .expect("timed out waiting for first request"), + Some(FIRST_REQUEST_VALUE) + ); + blocked_tx + .send(()) + .expect("blocked request should be waiting"); + + let mut values = Vec::new(); + while let Some(value) = timeout(queue_drain_timeout(), rx.recv()) + .await + .expect("timed out waiting for queue to drain") + { + values.push(value); + } + + assert_eq!(values, vec![THIRD_REQUEST_VALUE]); + } + + #[tokio::test] + async fn shutdown_of_live_gate_skips_already_queued_requests() { + let queues = RequestSerializationQueues::default(); + let key = RequestSerializationQueueKey::Global("test"); + let live_gate = gate(); + let (tx, mut rx) = mpsc::unbounded_channel(); + let (blocked_tx, blocked_rx) = oneshot::channel::<()>(); + + { + let tx = tx.clone(); + queues + .enqueue( + key.clone(), + QueuedInitializedRequest::new(Arc::clone(&live_gate), async move { + tx.send(FIRST_REQUEST_VALUE) + .expect("receiver should be open"); + let _ = blocked_rx.await; + }), + ) + .await; + } + { + let tx = tx.clone(); + queues + .enqueue( + key, + QueuedInitializedRequest::new(live_gate.clone(), async move { + tx.send(SECOND_REQUEST_VALUE) + .expect("receiver should be open"); + }), + ) + .await; + } + drop(tx); + + assert_eq!( + timeout(queue_drain_timeout(), rx.recv()) + .await + .expect("timed out waiting for first request"), + Some(FIRST_REQUEST_VALUE) + ); + + let gate_for_shutdown = Arc::clone(&live_gate); + let shutdown_task = tokio::spawn(async move { + gate_for_shutdown.shutdown().await; + }); + + timeout(shutdown_wait_timeout(), shutdown_task) + .await + .expect_err("shutdown should wait for the running request"); + + blocked_tx + .send(()) + .expect("blocked request should still be waiting"); + + assert_eq!( + timeout(queue_drain_timeout(), rx.recv()) + .await + .expect("timed out waiting for queue to drain"), + None + ); + } +} diff --git a/codex-rs/app-server/tests/suite/v2/experimental_feature_list.rs b/codex-rs/app-server/tests/suite/v2/experimental_feature_list.rs index 57520a2d6c..54b1f319e9 100644 --- a/codex-rs/app-server/tests/suite/v2/experimental_feature_list.rs +++ b/codex-rs/app-server/tests/suite/v2/experimental_feature_list.rs @@ -306,6 +306,24 @@ async fn experimental_feature_enablement_set_only_updates_named_features() -> Re Ok(()) } +#[tokio::test] +async fn experimental_feature_enablement_set_allows_remote_control() -> Result<()> { + let codex_home = TempDir::new()?; + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + let remote_control_enabled = false; + let enablement = BTreeMap::from([("remote_control".to_string(), remote_control_enabled)]); + + let actual = set_experimental_feature_enablement(&mut mcp, enablement.clone()).await?; + + assert_eq!( + actual, + ExperimentalFeatureEnablementSetResponse { enablement } + ); + + Ok(()) +} + #[tokio::test] async fn experimental_feature_enablement_set_empty_map_is_no_op() -> Result<()> { let codex_home = TempDir::new()?; @@ -363,7 +381,7 @@ async fn experimental_feature_enablement_set_rejects_non_allowlisted_feature() - ); assert!( error.message.contains( - "apps, memories, plugins, tool_search, tool_suggest, tool_call_mcp_elicitation" + "apps, memories, plugins, remote_control, tool_search, tool_suggest, tool_call_mcp_elicitation" ), "{}", error.message diff --git a/codex-rs/cloud-requirements/src/lib.rs b/codex-rs/cloud-requirements/src/lib.rs index d38fbb0846..c548179646 100644 --- a/codex-rs/cloud-requirements/src/lib.rs +++ b/codex-rs/cloud-requirements/src/lib.rs @@ -179,6 +179,14 @@ fn auth_identity(auth: &CodexAuth) -> (Option, Option) { (auth.get_chatgpt_user_id(), auth.get_account_id()) } +fn cloud_requirements_eligible_auth(auth: &CodexAuth) -> bool { + let Some(plan_type) = auth.account_plan_type() else { + return false; + }; + auth.uses_codex_backend() + && (plan_type.is_business_like() || matches!(plan_type, PlanType::Enterprise)) +} + fn cache_payload_bytes(payload: &CloudRequirementsCacheSignedPayload) -> Option> { serde_json::to_vec(&payload).ok() } @@ -329,17 +337,7 @@ impl CloudRequirementsService { let Some(auth) = self.auth_manager.auth().await else { return Ok(None); }; - if matches!(auth, CodexAuth::AgentIdentity(_)) { - // AgentIdentity does not carry a human bearer token, and identity-edge - // only allowlists task-scoped AgentAssertion calls for the Codex runtime. - return Ok(None); - } - let Some(plan_type) = auth.account_plan_type() else { - return Ok(None); - }; - if !auth.uses_codex_backend() - || !(plan_type.is_business_like() || matches!(plan_type, PlanType::Enterprise)) - { + if !cloud_requirements_eligible_auth(&auth) { return Ok(None); } let (chatgpt_user_id, account_id) = auth_identity(&auth); @@ -554,12 +552,7 @@ impl CloudRequirementsService { let Some(auth) = self.auth_manager.auth().await else { return false; }; - let Some(plan_type) = auth.account_plan_type() else { - return false; - }; - if !auth.uses_codex_backend() - || !(plan_type.is_business_like() || matches!(plan_type, PlanType::Enterprise)) - { + if !cloud_requirements_eligible_auth(&auth) { return false; } @@ -837,18 +830,41 @@ mod tests { use base64::Engine; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use codex_config::types::AuthCredentialsStoreMode; + use codex_login::auth::AgentIdentityAuth; + use codex_login::auth::AgentIdentityAuthRecord; use codex_protocol::protocol::AskForApproval; use pretty_assertions::assert_eq; use serde_json::json; use std::collections::BTreeMap; use std::collections::VecDeque; + use std::ffi::OsString; use std::future::pending; + use std::io::Read; + use std::io::Write; + use std::net::TcpListener; use std::path::Path; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use std::thread; use tempfile::TempDir; use tempfile::tempdir; + struct EnvVarGuard { + key: &'static str, + original: Option, + } + + impl Drop for EnvVarGuard { + fn drop(&mut self) { + unsafe { + match &self.original { + Some(value) => std::env::set_var(self.key, value), + None => std::env::remove_var(self.key), + } + } + } + } + fn write_auth_json(codex_home: &Path, value: serde_json::Value) -> std::io::Result<()> { std::fs::write(codex_home.join("auth.json"), serde_json::to_string(&value)?)?; Ok(()) @@ -1200,6 +1216,55 @@ mod tests { ); } + #[tokio::test] + async fn cloud_requirements_eligible_auth_allows_agent_identity_business_plan() { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind task registration server"); + let addr = listener + .local_addr() + .expect("task registration server addr"); + let server = thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept task registration request"); + let mut request = [0; 4096]; + let _ = stream + .read(&mut request) + .expect("read task registration request"); + let body = r#"{"task_id":"task-123"}"#; + write!( + stream, + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ) + .expect("write task registration response"); + }); + let record = AgentIdentityAuthRecord { + agent_runtime_id: "agent-runtime-123".to_string(), + agent_private_key: "MC4CAQAwBQYDK2VwBCIEIDQg14jybCLydjHQwXeBzsDM7oB6BSAenodx6oCovQ/D" + .to_string(), + account_id: "account-12345".to_string(), + chatgpt_user_id: "user-12345".to_string(), + email: "user@example.com".to_string(), + plan_type: PlanType::Business, + chatgpt_account_is_fedramp: false, + }; + let authapi_base_url = format!("http://{addr}/backend-api"); + let original_authapi_base_url = std::env::var_os("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL"); + unsafe { + std::env::set_var("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &authapi_base_url); + } + let _authapi_guard = EnvVarGuard { + key: "CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", + original: original_authapi_base_url, + }; + let auth = AgentIdentityAuth::load(record) + .await + .map(CodexAuth::AgentIdentity) + .expect("agent identity auth"); + server.join().expect("task registration server joined"); + + assert!(cloud_requirements_eligible_auth(&auth)); + } + #[tokio::test] async fn fetch_cloud_requirements_allows_business_like_usage_based_plan() { let codex_home = tempdir().expect("tempdir"); diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index f0e7c23c3b..1f8a93da4e 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1336,6 +1336,12 @@ "minimum": 1.0, "type": "integer" }, + "min_wait_timeout_ms": { + "format": "int64", + "maximum": 3600000.0, + "minimum": 1.0, + "type": "integer" + }, "root_agent_usage_hint_text": { "type": "string" }, diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index c9ddaa1bab..f075a4b7fc 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -7804,6 +7804,7 @@ async fn multi_agent_v2_config_from_feature_table() -> std::io::Result<()> { r#"[features.multi_agent_v2] enabled = true max_concurrent_threads_per_session = 5 +min_wait_timeout_ms = 2500 usage_hint_enabled = false usage_hint_text = "Custom delegation guidance." root_agent_usage_hint_text = "Root guidance." @@ -7820,6 +7821,7 @@ hide_spawn_agent_metadata = true assert!(config.features.enabled(Feature::MultiAgentV2)); assert_eq!(config.multi_agent_v2.max_concurrent_threads_per_session, 5); + assert_eq!(config.multi_agent_v2.min_wait_timeout_ms, 2500); assert_eq!(config.agent_max_threads, Some(4)); assert!(!config.multi_agent_v2.usage_hint_enabled); assert_eq!( @@ -7848,6 +7850,7 @@ async fn profile_multi_agent_v2_config_overrides_base() -> std::io::Result<()> { [features.multi_agent_v2] max_concurrent_threads_per_session = 4 +min_wait_timeout_ms = 3000 usage_hint_enabled = true usage_hint_text = "base hint" root_agent_usage_hint_text = "base root hint" @@ -7856,6 +7859,7 @@ hide_spawn_agent_metadata = true [profiles.no_hint.features.multi_agent_v2] max_concurrent_threads_per_session = 6 +min_wait_timeout_ms = 1500 usage_hint_enabled = false usage_hint_text = "profile hint" root_agent_usage_hint_text = "profile root hint" @@ -7871,6 +7875,7 @@ hide_spawn_agent_metadata = false .await?; assert_eq!(config.multi_agent_v2.max_concurrent_threads_per_session, 6); + assert_eq!(config.multi_agent_v2.min_wait_timeout_ms, 1500); assert!(!config.multi_agent_v2.usage_hint_enabled); assert_eq!( config.multi_agent_v2.usage_hint_text.as_deref(), @@ -7906,6 +7911,7 @@ enabled = true .await?; assert_eq!(config.multi_agent_v2.max_concurrent_threads_per_session, 4); + assert_eq!(config.multi_agent_v2.min_wait_timeout_ms, 10_000); assert_eq!(config.agent_max_threads, Some(3)); Ok(()) @@ -7940,6 +7946,54 @@ max_threads = 3 Ok(()) } +#[tokio::test] +async fn multi_agent_v2_rejects_invalid_min_wait_timeout() -> std::io::Result<()> { + let codex_home = TempDir::new()?; + std::fs::write( + codex_home.path().join(CONFIG_TOML_FILE), + r#"[features.multi_agent_v2] +enabled = true +min_wait_timeout_ms = 0 +"#, + )?; + + let err = ConfigBuilder::without_managed_config_for_tests() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(codex_home.path().to_path_buf())) + .build() + .await + .expect_err("zero min_wait_timeout_ms should be rejected"); + + assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); + assert_eq!( + err.to_string(), + "features.multi_agent_v2.min_wait_timeout_ms must be at least 1" + ); + + std::fs::write( + codex_home.path().join(CONFIG_TOML_FILE), + r#"[features.multi_agent_v2] +enabled = true +min_wait_timeout_ms = 3600001 +"#, + )?; + + let err = ConfigBuilder::without_managed_config_for_tests() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(codex_home.path().to_path_buf())) + .build() + .await + .expect_err("too large min_wait_timeout_ms should be rejected"); + + assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); + assert_eq!( + err.to_string(), + "features.multi_agent_v2.min_wait_timeout_ms must be at most 3600000" + ); + + Ok(()) +} + #[tokio::test] async fn multi_agent_v2_session_thread_cap_one_disallows_subagents() -> std::io::Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index dcf443b84e..6c8258c519 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -156,6 +156,8 @@ impl Default for GhostSnapshotConfig { pub(crate) const AGENTS_MD_MAX_BYTES: usize = 32 * 1024; // 32 KiB pub(crate) const DEFAULT_AGENT_MAX_THREADS: Option = Some(6); pub(crate) const DEFAULT_MULTI_AGENT_V2_MAX_CONCURRENT_THREADS_PER_SESSION: usize = 4; +pub(crate) const DEFAULT_MULTI_AGENT_V2_MIN_WAIT_TIMEOUT_MS: i64 = 10_000; +pub(crate) const MAX_MULTI_AGENT_V2_WAIT_TIMEOUT_MS: i64 = 3600 * 1000; pub(crate) const DEFAULT_AGENT_MAX_DEPTH: i32 = 1; pub(crate) const DEFAULT_AGENT_JOB_MAX_RUNTIME_SECONDS: Option = None; const LOCAL_DEV_BUILD_VERSION: &str = "0.0.0"; @@ -753,6 +755,7 @@ pub struct Config { #[derive(Debug, Clone, PartialEq, Eq)] pub struct MultiAgentV2Config { pub max_concurrent_threads_per_session: usize, + pub min_wait_timeout_ms: i64, pub usage_hint_enabled: bool, pub usage_hint_text: Option, pub root_agent_usage_hint_text: Option, @@ -765,6 +768,7 @@ impl Default for MultiAgentV2Config { Self { max_concurrent_threads_per_session: DEFAULT_MULTI_AGENT_V2_MAX_CONCURRENT_THREADS_PER_SESSION, + min_wait_timeout_ms: DEFAULT_MULTI_AGENT_V2_MIN_WAIT_TIMEOUT_MS, usage_hint_enabled: true, usage_hint_text: None, root_agent_usage_hint_text: None, @@ -1638,6 +1642,10 @@ fn resolve_multi_agent_v2_config( .and_then(|config| config.max_concurrent_threads_per_session) .or_else(|| base.and_then(|config| config.max_concurrent_threads_per_session)) .unwrap_or(default.max_concurrent_threads_per_session); + let min_wait_timeout_ms = profile + .and_then(|config| config.min_wait_timeout_ms) + .or_else(|| base.and_then(|config| config.min_wait_timeout_ms)) + .unwrap_or(default.min_wait_timeout_ms); let usage_hint_enabled = profile .and_then(|config| config.usage_hint_enabled) .or_else(|| base.and_then(|config| config.usage_hint_enabled)) @@ -1664,6 +1672,7 @@ fn resolve_multi_agent_v2_config( MultiAgentV2Config { max_concurrent_threads_per_session, + min_wait_timeout_ms, usage_hint_enabled, usage_hint_text, root_agent_usage_hint_text, @@ -2186,6 +2195,20 @@ impl Config { "features.multi_agent_v2.max_concurrent_threads_per_session must be at least 1", )); } + if multi_agent_v2.min_wait_timeout_ms <= 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "features.multi_agent_v2.min_wait_timeout_ms must be at least 1", + )); + } + if multi_agent_v2.min_wait_timeout_ms > MAX_MULTI_AGENT_V2_WAIT_TIMEOUT_MS { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "features.multi_agent_v2.min_wait_timeout_ms must be at most {MAX_MULTI_AGENT_V2_WAIT_TIMEOUT_MS}" + ), + )); + } let agent_max_threads_from_config = cfg.agents.as_ref().and_then(|agents| agents.max_threads); let agent_max_threads = if features.enabled(Feature::MultiAgentV2) { if agent_max_threads_from_config.is_some() { diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 9401c2d0ba..c22a9faaf1 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -54,6 +54,11 @@ pub(super) async fn spawn_review_thread( .with_hide_spawn_agent_metadata(config.multi_agent_v2.hide_spawn_agent_metadata) .with_goal_tools_allowed(goal_tools_supported) .with_max_concurrent_threads_per_session(config.agent_max_threads) + .with_wait_agent_min_timeout_ms( + review_features + .enabled(Feature::MultiAgentV2) + .then_some(config.multi_agent_v2.min_wait_timeout_ms), + ) .with_agent_type_description(crate::agent::role::spawn_tool_spec::build( &config.agent_roles, )); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index ec5705adfc..823dcbfc3d 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1864,6 +1864,11 @@ async fn try_run_sampling_request( otel.name = field::Empty, tool_name = field::Empty, from = field::Empty, + gen_ai.usage.input_tokens = field::Empty, + gen_ai.usage.cache_read.input_tokens = field::Empty, + gen_ai.usage.output_tokens = field::Empty, + codex.usage.reasoning_output_tokens = field::Empty, + codex.usage.total_tokens = field::Empty, ); let event = match stream diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index a8244303d6..35b517b720 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -201,6 +201,12 @@ impl TurnContext { .enabled(Feature::MultiAgentV2) .then_some(config.multi_agent_v2.max_concurrent_threads_per_session), ) + .with_wait_agent_min_timeout_ms( + config + .features + .enabled(Feature::MultiAgentV2) + .then_some(config.multi_agent_v2.min_wait_timeout_ms), + ) .with_agent_type_description(crate::agent::role::spawn_tool_spec::build( &config.agent_roles, )); @@ -475,6 +481,12 @@ impl Session { .max_concurrent_threads_per_session, ), ) + .with_wait_agent_min_timeout_ms( + per_turn_config + .features + .enabled(Feature::MultiAgentV2) + .then_some(per_turn_config.multi_agent_v2.min_wait_timeout_ms), + ) .with_agent_type_description(crate::agent::role::spawn_tool_spec::build( &per_turn_config.agent_roles, )); diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index 1439e28bb3..0dd0c09379 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -79,17 +79,25 @@ pub(crate) struct RunningTask { pub(crate) _timer: Option, } +pub(crate) struct RemovedTask { + pub(crate) records_turn_token_usage_on_span: bool, + pub(crate) active_turn_is_empty: bool, +} + impl ActiveTurn { pub(crate) fn add_task(&mut self, task: RunningTask) { let sub_id = task.turn_context.sub_id.clone(); self.tasks.insert(sub_id, task); } - pub(crate) fn remove_task(&mut self, sub_id: &str) -> bool { - if let Some(task) = self.tasks.swap_remove(sub_id) { - task.handle.detach(); - } - self.tasks.is_empty() + pub(crate) fn remove_task(&mut self, sub_id: &str) -> Option { + let task = self.tasks.swap_remove(sub_id)?; + let records_turn_token_usage_on_span = task.task.records_turn_token_usage_on_span(); + task.handle.detach(); + Some(RemovedTask { + records_turn_token_usage_on_span, + active_turn_is_empty: self.tasks.is_empty(), + }) } pub(crate) fn drain_tasks(&mut self) -> Vec { diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index e0286f3090..91078c50ce 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -14,6 +14,8 @@ use tokio::sync::Notify; use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; use tracing::Instrument; +use tracing::Span; +use tracing::field; use tracing::info_span; use tracing::trace; use tracing::warn; @@ -188,6 +190,11 @@ pub(crate) trait SessionTask: Send + Sync + 'static { /// Returns the tracing name for a spawned task span. fn span_name(&self) -> &'static str; + /// Returns whether turn token usage should be recorded on this task's turn span. + fn records_turn_token_usage_on_span(&self) -> bool { + false + } + /// Executes the task until completion or cancellation. /// /// Implementations typically stream protocol events using `session` and @@ -225,6 +232,8 @@ pub(crate) trait AnySessionTask: Send + Sync + 'static { fn span_name(&self) -> &'static str; + fn records_turn_token_usage_on_span(&self) -> bool; + fn run( self: Arc, session: Arc, @@ -252,6 +261,10 @@ where SessionTask::span_name(self) } + fn records_turn_token_usage_on_span(&self) -> bool { + SessionTask::records_turn_token_usage_on_span(self) + } + fn run( self: Arc, session: Arc, @@ -361,6 +374,12 @@ impl Session { thread.id = %self.conversation_id, turn.id = %turn_context.sub_id, model = %turn_context.model_info.slug, + codex.turn.token_usage.input_tokens = field::Empty, + codex.turn.token_usage.cached_input_tokens = field::Empty, + codex.turn.token_usage.non_cached_input_tokens = field::Empty, + codex.turn.token_usage.output_tokens = field::Empty, + codex.turn.token_usage.reasoning_output_tokens = field::Empty, + codex.turn.token_usage.total_tokens = field::Empty, ); let handle = tokio::spawn( async move { @@ -548,14 +567,20 @@ impl Session { let mut token_usage_at_turn_start = None; let mut turn_had_memory_citation = false; let mut turn_tool_calls = 0_u64; + let mut records_turn_token_usage_on_span = false; let turn_state = { let mut active = self.active_turn.lock().await; if let Some(at) = active.as_mut() - && at.remove_task(&turn_context.sub_id) + && let Some(removed_task) = at.remove_task(&turn_context.sub_id) { - should_clear_active_turn = true; - let turn_state = Arc::clone(&at.turn_state); - Some(turn_state) + records_turn_token_usage_on_span = removed_task.records_turn_token_usage_on_span; + if removed_task.active_turn_is_empty { + should_clear_active_turn = true; + let turn_state = Arc::clone(&at.turn_state); + Some(turn_state) + } else { + None + } } else { None } @@ -634,6 +659,33 @@ impl Session { - token_usage_at_turn_start.total_tokens) .max(0), }; + if records_turn_token_usage_on_span { + let current_span = Span::current(); + current_span.record( + "codex.turn.token_usage.input_tokens", + turn_token_usage.input_tokens, + ); + current_span.record( + "codex.turn.token_usage.cached_input_tokens", + turn_token_usage.cached_input(), + ); + current_span.record( + "codex.turn.token_usage.non_cached_input_tokens", + turn_token_usage.non_cached_input(), + ); + current_span.record( + "codex.turn.token_usage.output_tokens", + turn_token_usage.output_tokens, + ); + current_span.record( + "codex.turn.token_usage.reasoning_output_tokens", + turn_token_usage.reasoning_output_tokens, + ); + current_span.record( + "codex.turn.token_usage.total_tokens", + turn_token_usage.total_tokens, + ); + } self.services .analytics_events_client .track_turn_token_usage(TurnTokenUsageFact { diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index d64b83f01b..08c4933488 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -33,6 +33,10 @@ impl SessionTask for RegularTask { "session_task.turn" } + fn records_turn_token_usage_on_span(&self) -> bool { + true + } + async fn run( self: Arc, session: Arc, diff --git a/codex-rs/core/src/tools/handlers/multi_agents_common.rs b/codex-rs/core/src/tools/handlers/multi_agents_common.rs index c722ddb8d3..c01755cb2b 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_common.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_common.rs @@ -1,5 +1,7 @@ use crate::agent::AgentStatus; use crate::config::Config; +use crate::config::DEFAULT_MULTI_AGENT_V2_MIN_WAIT_TIMEOUT_MS; +use crate::config::MAX_MULTI_AGENT_V2_WAIT_TIMEOUT_MS; use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::turn_context::TurnContext; @@ -26,9 +28,9 @@ use serde_json::Value as JsonValue; use std::collections::HashMap; /// Minimum wait timeout to prevent tight polling loops from burning CPU. -pub(crate) const MIN_WAIT_TIMEOUT_MS: i64 = 10_000; +pub(crate) const MIN_WAIT_TIMEOUT_MS: i64 = DEFAULT_MULTI_AGENT_V2_MIN_WAIT_TIMEOUT_MS; pub(crate) const DEFAULT_WAIT_TIMEOUT_MS: i64 = 30_000; -pub(crate) const MAX_WAIT_TIMEOUT_MS: i64 = 3600 * 1000; +pub(crate) const MAX_WAIT_TIMEOUT_MS: i64 = MAX_MULTI_AGENT_V2_WAIT_TIMEOUT_MS; pub(crate) fn function_arguments(payload: ToolPayload) -> Result { match payload { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index fc861be61c..2fecb77800 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -2742,6 +2742,59 @@ async fn multi_agent_v2_wait_agent_accepts_timeout_only_argument() { assert_eq!(success, None); } +#[tokio::test] +async fn multi_agent_v2_wait_agent_uses_configured_min_timeout() { + let (session, mut turn) = make_session_and_context().await; + let mut config = (*turn.config).clone(); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + config.multi_agent_v2.min_wait_timeout_ms = 50; + turn.config = Arc::new(config); + let session = Arc::new(session); + let turn = Arc::new(turn); + + let early = timeout( + Duration::from_millis(/*millis*/ 20), + WaitAgentHandlerV2.handle(invocation( + session.clone(), + turn.clone(), + "wait_agent", + function_payload(json!({"timeout_ms": 1})), + )), + ) + .await; + assert!( + early.is_err(), + "wait_agent should not return before the configured minimum timeout" + ); + + let output = timeout( + Duration::from_secs(/*secs*/ 1), + WaitAgentHandlerV2.handle(invocation( + session, + turn, + "wait_agent", + function_payload(json!({"timeout_ms": 1})), + )), + ) + .await + .expect("configured minimum should be shorter than the test timeout") + .expect("wait_agent should succeed"); + let (content, success) = expect_text_output(output); + let result: crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult = + serde_json::from_str(&content).expect("wait_agent result should be json"); + assert_eq!( + result, + crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult { + message: "Wait timed out.".to_string(), + timed_out: true, + } + ); + assert_eq!(success, None); +} + #[tokio::test] async fn wait_agent_returns_not_found_for_missing_agents() { let (mut session, turn) = make_session_and_context().await; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs index e50c6cab23..778c57be21 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs @@ -28,13 +28,18 @@ impl ToolHandler for Handler { let arguments = function_arguments(payload)?; let args: WaitArgs = parse_arguments(&arguments)?; let timeout_ms = args.timeout_ms.unwrap_or(DEFAULT_WAIT_TIMEOUT_MS); + let min_timeout_ms = turn + .config + .multi_agent_v2 + .min_wait_timeout_ms + .clamp(1, MAX_WAIT_TIMEOUT_MS); let timeout_ms = match timeout_ms { ms if ms <= 0 => { return Err(FunctionCallError::RespondToModel( "timeout_ms must be greater than zero".to_owned(), )); } - ms => ms.clamp(MIN_WAIT_TIMEOUT_MS, MAX_WAIT_TIMEOUT_MS), + ms => ms.clamp(min_timeout_ms, MAX_WAIT_TIMEOUT_MS), }; let mut mailbox_seq_rx = session.subscribe_mailbox_seq(); diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index ebbc38b8be..91d38f25b7 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -124,6 +124,16 @@ pub(crate) fn build_specs_with_discoverable_tools( }); let default_agent_type_description = crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); + let min_wait_timeout_ms = if config.multi_agent_v2 { + config + .wait_agent_min_timeout_ms + .unwrap_or(MIN_WAIT_TIMEOUT_MS) + .clamp(1, MAX_WAIT_TIMEOUT_MS) + } else { + MIN_WAIT_TIMEOUT_MS + }; + let default_wait_timeout_ms = + DEFAULT_WAIT_TIMEOUT_MS.clamp(min_wait_timeout_ms, MAX_WAIT_TIMEOUT_MS); let plan = build_tool_registry_plan( config, ToolRegistryPlanParams { @@ -138,8 +148,8 @@ pub(crate) fn build_specs_with_discoverable_tools( dynamic_tools, default_agent_type_description: &default_agent_type_description, wait_agent_timeouts: WaitAgentTimeoutOptions { - default_timeout_ms: DEFAULT_WAIT_TIMEOUT_MS, - min_timeout_ms: MIN_WAIT_TIMEOUT_MS, + default_timeout_ms: default_wait_timeout_ms, + min_timeout_ms: min_wait_timeout_ms, max_timeout_ms: MAX_WAIT_TIMEOUT_MS, }, }, diff --git a/codex-rs/core/src/tools/spec_tests.rs b/codex-rs/core/src/tools/spec_tests.rs index f88bb04f3a..5942bf8061 100644 --- a/codex-rs/core/src/tools/spec_tests.rs +++ b/codex-rs/core/src/tools/spec_tests.rs @@ -784,6 +784,35 @@ async fn spawn_agent_description_uses_configured_usage_hint_text() { ); } +#[tokio::test] +async fn multi_agent_v2_wait_agent_schema_uses_configured_min_timeout() { + let wait_agent_min_timeout_ms = Some(60_000); + let tools_config = multi_agent_v2_tools_config() + .await + .with_wait_agent_min_timeout_ms(wait_agent_min_timeout_ms); + let (tools, _) = build_specs( + &tools_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ) + .build(); + let wait_agent = find_tool(&tools, "wait_agent"); + let ToolSpec::Function(ResponsesApiTool { parameters, .. }) = &wait_agent.spec else { + panic!("wait_agent should be a function tool"); + }; + let timeout_description = parameters + .properties + .as_ref() + .and_then(|properties| properties.get("timeout_ms")) + .and_then(|schema| schema.description.as_deref()); + + assert_eq!( + timeout_description, + Some("Optional timeout in milliseconds. Defaults to 60000, min 60000, max 3600000.") + ); +} + #[tokio::test] async fn tool_suggest_requires_apps_and_plugins_features() { let model_info = search_capable_model_info().await; diff --git a/codex-rs/core/tests/suite/otel.rs b/codex-rs/core/tests/suite/otel.rs index 5bf41bd547..deeffcd855 100644 --- a/codex-rs/core/tests/suite/otel.rs +++ b/codex-rs/core/tests/suite/otel.rs @@ -563,6 +563,91 @@ async fn process_sse_emits_completed_telemetry() { }); } +#[tokio::test(flavor = "current_thread")] +async fn turn_and_completed_response_spans_record_token_usage() { + let buffer: &'static Mutex> = Box::leak(Box::new(Mutex::new(Vec::new()))); + let subscriber = tracing_subscriber::fmt() + .with_level(true) + .with_ansi(false) + .with_max_level(Level::TRACE) + .with_span_events(FmtSpan::FULL) + .with_writer(MockWriter::new(buffer)) + .finish(); + let _guard = tracing::subscriber::set_default(subscriber); + + let server = start_mock_server().await; + + mount_sse_once( + &server, + sse(vec![serde_json::json!({ + "type": "response.completed", + "response": { + "id": "resp1", + "usage": { + "input_tokens": 3, + "input_tokens_details": { "cached_tokens": 1 }, + "output_tokens": 5, + "output_tokens_details": { "reasoning_tokens": 2 }, + "total_tokens": 9 + } + } + })]), + ) + .await; + + let TestCodex { codex, .. } = test_codex() + .with_config(|config| { + config + .features + .disable(Feature::GhostCommit) + .expect("test config should allow feature update"); + }) + .build(&server) + .await + .unwrap(); + + codex + .submit(Op::UserInput { + environments: None, + items: vec![UserInput::Text { + text: "hello".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + }) + .await + .unwrap(); + + wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + let logs = String::from_utf8(buffer.lock().unwrap().clone()).unwrap(); + + assert!( + logs.lines().any(|line| { + line.contains("handle_responses{otel.name=\"completed\"") + && line.contains("gen_ai.usage.input_tokens=3") + && line.contains("gen_ai.usage.cache_read.input_tokens=1") + && line.contains("gen_ai.usage.output_tokens=5") + && line.contains("codex.usage.reasoning_output_tokens=2") + && line.contains("codex.usage.total_tokens=9") + }), + "missing completed response span token usage\nlogs:\n{logs}" + ); + assert!( + logs.lines().any(|line| { + line.contains("turn{otel.name=\"session_task.turn\"") + && line.contains("codex.turn.token_usage.input_tokens=3") + && line.contains("codex.turn.token_usage.cached_input_tokens=1") + && line.contains("codex.turn.token_usage.non_cached_input_tokens=2") + && line.contains("codex.turn.token_usage.output_tokens=5") + && line.contains("codex.turn.token_usage.reasoning_output_tokens=2") + && line.contains("codex.turn.token_usage.total_tokens=9") + }), + "missing regular turn span token usage\nlogs:\n{logs}" + ); +} + #[tokio::test] async fn handle_responses_span_records_response_kind_and_tool_name() { let buffer: &'static Mutex> = Box::leak(Box::new(Mutex::new(Vec::new()))); diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index 3906bc2e83..fd489db505 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -12,6 +12,9 @@ pub struct MultiAgentV2ConfigToml { #[schemars(range(min = 1))] pub max_concurrent_threads_per_session: Option, #[serde(skip_serializing_if = "Option::is_none")] + #[schemars(range(min = 1, max = 3600000))] + pub min_wait_timeout_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub usage_hint_enabled: Option, #[serde(skip_serializing_if = "Option::is_none")] pub usage_hint_text: Option, diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index 4e1f9a0951..0ed1d7ecc0 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -413,6 +413,7 @@ fn multi_agent_v2_feature_config_deserializes_table() { [multi_agent_v2] enabled = true max_concurrent_threads_per_session = 4 +min_wait_timeout_ms = 2500 usage_hint_enabled = false usage_hint_text = "Custom delegation guidance." root_agent_usage_hint_text = "Root guidance." @@ -431,6 +432,7 @@ hide_spawn_agent_metadata = true Some(crate::FeatureToml::Config(crate::MultiAgentV2ConfigToml { enabled: Some(true), max_concurrent_threads_per_session: Some(4), + min_wait_timeout_ms: Some(2500), usage_hint_enabled: Some(false), usage_hint_text: Some("Custom delegation guidance.".to_string()), root_agent_usage_hint_text: Some("Root guidance.".to_string()), @@ -465,6 +467,7 @@ usage_hint_enabled = false Some(crate::FeatureToml::Config(crate::MultiAgentV2ConfigToml { enabled: None, max_concurrent_threads_per_session: None, + min_wait_timeout_ms: None, usage_hint_enabled: Some(false), usage_hint_text: None, root_agent_usage_hint_text: None, diff --git a/codex-rs/linux-sandbox/Cargo.toml b/codex-rs/linux-sandbox/Cargo.toml index 519ae5138e..05967661e2 100644 --- a/codex-rs/linux-sandbox/Cargo.toml +++ b/codex-rs/linux-sandbox/Cargo.toml @@ -17,6 +17,7 @@ workspace = true [target.'cfg(target_os = "linux")'.dependencies] clap = { workspace = true, features = ["derive"] } +codex-process-hardening = { workspace = true } codex-protocol = { workspace = true } codex-sandboxing = { workspace = true } codex-utils-absolute-path = { workspace = true } diff --git a/codex-rs/linux-sandbox/src/landlock.rs b/codex-rs/linux-sandbox/src/landlock.rs index 5794982530..50ea87bd1d 100644 --- a/codex-rs/linux-sandbox/src/landlock.rs +++ b/codex-rs/linux-sandbox/src/landlock.rs @@ -176,6 +176,8 @@ fn install_network_seccomp_filter_on_current_thread( let mut rules: BTreeMap> = BTreeMap::new(); deny_syscall(&mut rules, libc::SYS_ptrace); + deny_syscall(&mut rules, libc::SYS_process_vm_readv); + deny_syscall(&mut rules, libc::SYS_process_vm_writev); deny_syscall(&mut rules, libc::SYS_io_uring_setup); deny_syscall(&mut rules, libc::SYS_io_uring_enter); deny_syscall(&mut rules, libc::SYS_io_uring_register); diff --git a/codex-rs/linux-sandbox/src/proxy_routing.rs b/codex-rs/linux-sandbox/src/proxy_routing.rs index 07e1893ee3..d28b864666 100644 --- a/codex-rs/linux-sandbox/src/proxy_routing.rs +++ b/codex-rs/linux-sandbox/src/proxy_routing.rs @@ -450,7 +450,7 @@ fn spawn_host_bridge(endpoint: SocketAddr, uds_path: &Path) -> io::Result io::Result<()> { - set_parent_death_signal()?; + harden_bridge_process()?; if uds_path.exists() { std::fs::remove_file(uds_path)?; } @@ -501,7 +501,7 @@ fn spawn_local_bridge(uds_path: &Path) -> io::Result { } fn run_local_bridge(uds_path: &Path, ready_fd: libc::c_int) -> io::Result<()> { - set_parent_death_signal()?; + harden_bridge_process()?; let listener = bind_local_loopback_listener()?; let port = listener.local_addr()?.port(); @@ -614,6 +614,11 @@ fn set_parent_death_signal() -> io::Result<()> { } } +fn harden_bridge_process() -> io::Result<()> { + set_parent_death_signal()?; + codex_process_hardening::disable_process_dumping() +} + fn proxy_bidirectional(mut tcp_stream: TcpStream, mut unix_stream: UnixStream) -> io::Result<()> { let mut tcp_reader = tcp_stream.try_clone()?; let mut unix_writer = unix_stream.try_clone()?; diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 90748f7860..80ec9d07a4 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -88,7 +88,8 @@ async fn login_with_agent_identity_writes_only_token() { let dir = tempdir().unwrap(); let auth_path = dir.path().join("auth.json"); let record = agent_identity_record("account-123"); - let agent_identity = signed_agent_identity_jwt(&record).expect("signed agent identity"); + let agent_identity = + signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity"); let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/backend-api/wham/agent-identities/jwks")) @@ -709,7 +710,8 @@ async fn load_auth_reads_agent_identity_from_env() { let codex_home = tempdir().unwrap(); let expected_record = agent_identity_record("account-123"); let agent_identity = - signed_agent_identity_jwt(&expected_record).expect("signed agent identity"); + signed_agent_identity_jwt(&expected_record, json!(expected_record.plan_type)) + .expect("signed agent identity"); let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/backend-api/wham/agent-identities/jwks")) @@ -925,6 +927,13 @@ fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord { } fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result { + fake_agent_identity_jwt_with_plan_type(record, serde_json::to_value(record.plan_type)?) +} + +fn fake_agent_identity_jwt_with_plan_type( + record: &AgentIdentityAuthRecord, + plan_type: serde_json::Value, +) -> std::io::Result { let encode = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes); let header_b64 = encode(br#"{"alg":"EdDSA","typ":"JWT"}"#); let payload = json!({ @@ -937,7 +946,7 @@ fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result< "account_id": record.account_id, "chatgpt_user_id": record.chatgpt_user_id, "email": record.email, - "plan_type": record.plan_type, + "plan_type": plan_type, "chatgpt_account_is_fedramp": record.chatgpt_account_is_fedramp, }); let payload_b64 = encode(&serde_json::to_vec(&payload)?); @@ -947,6 +956,7 @@ fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result< fn signed_agent_identity_jwt( record: &AgentIdentityAuthRecord, + plan_type: serde_json::Value, ) -> jsonwebtoken::errors::Result { let mut header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256); header.kid = Some("test-key".to_string()); @@ -962,7 +972,7 @@ fn signed_agent_identity_jwt( "account_id": record.account_id, "chatgpt_user_id": record.chatgpt_user_id, "email": record.email, - "plan_type": record.plan_type, + "plan_type": plan_type, "chatgpt_account_is_fedramp": record.chatgpt_account_is_fedramp, }), &jsonwebtoken::EncodingKey::from_rsa_pem(TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM)?, @@ -1011,6 +1021,50 @@ J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN 8U4M2TSWCKUY/A6sT4W8+mT9 -----END PRIVATE KEY-----"#; +#[tokio::test] +#[serial(codex_auth_env)] +async fn agent_identity_plan_type_maps_raw_enterprise_alias() { + assert_agent_identity_plan_alias(json!("hc"), AccountPlanType::Enterprise).await; +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn agent_identity_plan_type_maps_raw_education_alias() { + assert_agent_identity_plan_alias(json!("education"), AccountPlanType::Edu).await; +} + +async fn assert_agent_identity_plan_alias( + plan_type: serde_json::Value, + expected_plan_type: AccountPlanType, +) { + let record = agent_identity_record("account-id"); + let jwt = signed_agent_identity_jwt(&record, plan_type).expect("agent identity jwt"); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/backend-api/wham/agent-identities/jwks")) + .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body())) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/backend-api/v1/agent/agent-runtime-id/task/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "task_id": "task-123", + }))) + .expect(1) + .mount(&server) + .await; + let chatgpt_base_url = format!("{}/backend-api", server.uri()); + let _authapi_guard = + EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url); + let auth = CodexAuth::from_agent_identity_jwt(&jwt, Some(&chatgpt_base_url)) + .await + .expect("agent identity auth"); + + pretty_assertions::assert_eq!(auth.account_plan_type(), Some(expected_plan_type)); + server.verify().await; +} + #[tokio::test] #[serial(codex_auth_env)] async fn plan_type_maps_known_plan() { diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 9f140c2768..29897db7be 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -39,7 +39,6 @@ use crate::token_data::parse_jwt_expiration; use codex_client::CodexHttpClient; use codex_config::types::AuthCredentialsStoreMode; use codex_protocol::account::PlanType as AccountPlanType; -use codex_protocol::auth::KnownPlan as InternalKnownPlan; use codex_protocol::auth::PlanType as InternalPlanType; use codex_protocol::auth::RefreshTokenFailedError; use codex_protocol::auth::RefreshTokenFailedReason; @@ -380,29 +379,10 @@ impl CodexAuth { return Some(auth.plan_type()); } - let map_known = |kp: &InternalKnownPlan| match kp { - InternalKnownPlan::Free => AccountPlanType::Free, - InternalKnownPlan::Go => AccountPlanType::Go, - InternalKnownPlan::Plus => AccountPlanType::Plus, - InternalKnownPlan::Pro => AccountPlanType::Pro, - InternalKnownPlan::ProLite => AccountPlanType::ProLite, - InternalKnownPlan::Team => AccountPlanType::Team, - InternalKnownPlan::SelfServeBusinessUsageBased => { - AccountPlanType::SelfServeBusinessUsageBased - } - InternalKnownPlan::Business => AccountPlanType::Business, - InternalKnownPlan::EnterpriseCbpUsageBased => AccountPlanType::EnterpriseCbpUsageBased, - InternalKnownPlan::Enterprise => AccountPlanType::Enterprise, - InternalKnownPlan::Edu => AccountPlanType::Edu, - }; - self.get_current_token_data().map(|t| { t.id_token .chatgpt_plan_type - .map(|pt| match pt { - InternalPlanType::Known(k) => map_known(&k), - InternalPlanType::Unknown(_) => AccountPlanType::Unknown, - }) + .map(AccountPlanType::from) .unwrap_or(AccountPlanType::Unknown) }) } diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index b2e8946fa9..3a1c8ae6aa 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -75,7 +75,7 @@ impl From for AgentIdentityAuthRecord { account_id: claims.account_id, chatgpt_user_id: claims.chatgpt_user_id, email: claims.email, - plan_type: claims.plan_type, + plan_type: claims.plan_type.into(), chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp, } } diff --git a/codex-rs/network-proxy/src/connect_policy.rs b/codex-rs/network-proxy/src/connect_policy.rs new file mode 100644 index 0000000000..b9425db799 --- /dev/null +++ b/codex-rs/network-proxy/src/connect_policy.rs @@ -0,0 +1,157 @@ +use crate::policy::is_non_public_ip; +use crate::state::NetworkProxyState; +use rama_core::Service; +use rama_core::error::BoxError; +use rama_core::error::ErrorExt as _; +use rama_core::error::OpaqueError; +use rama_core::extensions::ExtensionsMut; +use rama_net::address::ProxyAddress; +use rama_net::client::EstablishedClientConnection; +use rama_net::transport::TryRefIntoTransportContext; +use rama_tcp::TcpStream; +use rama_tcp::client::TcpStreamConnector; +use rama_tcp::client::service::TcpConnector; +use std::io; +use std::net::SocketAddr; +use std::sync::Arc; + +#[derive(Clone)] +pub(crate) struct TargetCheckedTcpConnector { + policy: TargetPolicy, +} + +impl TargetCheckedTcpConnector { + pub(crate) fn new(state: Arc) -> Self { + Self { + policy: TargetPolicy::State(state), + } + } + + pub(crate) fn from_allow_local_binding(allow_local_binding: bool) -> Self { + Self { + policy: TargetPolicy::Config { + allow_local_binding, + }, + } + } +} + +impl Service for TargetCheckedTcpConnector +where + Input: TryRefIntoTransportContext + Send + ExtensionsMut + 'static, + Input::Error: Into + Send + Sync + 'static, +{ + type Output = EstablishedClientConnection; + type Error = BoxError; + + async fn serve(&self, input: Input) -> Result { + if input.extensions().get::().is_some() { + return TcpConnector::new().serve(input).await; + } + + TcpConnector::new() + .with_connector(TargetCheckedStreamConnector { + policy: self.policy.clone(), + }) + .serve(input) + .await + } +} + +#[derive(Clone)] +struct TargetCheckedStreamConnector { + policy: TargetPolicy, +} + +impl TcpStreamConnector for TargetCheckedStreamConnector { + type Error = BoxError; + + async fn connect(&self, addr: SocketAddr) -> Result { + if !self.policy.allow_local_binding().await? && is_non_public_ip(addr.ip()) { + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "network target rejected by policy", + ) + .into()); + } + + tokio::net::TcpStream::connect(addr) + .await + .map(TcpStream::from) + .map_err(Into::into) + } +} + +#[derive(Clone)] +enum TargetPolicy { + Config { allow_local_binding: bool }, + State(Arc), +} + +impl TargetPolicy { + async fn allow_local_binding(&self) -> Result { + match self { + Self::Config { + allow_local_binding, + } => Ok(*allow_local_binding), + Self::State(state) => state.allow_local_binding().await.map_err(|err| { + let err: BoxError = err.into(); + OpaqueError::from_boxed(err) + .context("read network proxy config") + .into_boxed() + }), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::NetworkProxySettings; + use crate::state::network_proxy_state_for_policy; + use rama_net::address::HostWithPort; + use std::net::Ipv4Addr; + use tokio::net::TcpListener; + + #[tokio::test(flavor = "current_thread")] + async fn direct_connector_rejects_non_public_target_when_local_binding_disabled() { + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) + .await + .expect("bind local listener"); + let target = listener.local_addr().expect("local addr"); + let connector = TargetCheckedTcpConnector::new(Arc::new(network_proxy_state_for_policy( + NetworkProxySettings::default(), + ))); + + let request: rama_tcp::client::Request = + rama_tcp::client::Request::new(HostWithPort::from(target)); + let err = Service::serve(&connector, request) + .await + .expect_err("local target should be rejected"); + + assert!( + format!("{err:?}").contains("network target rejected by policy"), + "unexpected error: {err:?}" + ); + } + + #[tokio::test(flavor = "current_thread")] + async fn direct_connector_allows_non_public_target_when_local_binding_enabled() { + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) + .await + .expect("bind local listener"); + let target = listener.local_addr().expect("local addr"); + let connector = TargetCheckedTcpConnector::new(Arc::new(network_proxy_state_for_policy( + NetworkProxySettings { + allow_local_binding: true, + ..NetworkProxySettings::default() + }, + ))); + + let request: rama_tcp::client::Request = + rama_tcp::client::Request::new(HostWithPort::from(target)); + let result = Service::serve(&connector, request).await; + + assert!(result.is_ok(), "local target should be allowed: {result:?}"); + } +} diff --git a/codex-rs/network-proxy/src/http_proxy.rs b/codex-rs/network-proxy/src/http_proxy.rs index a29af08b1c..fa3616844e 100644 --- a/codex-rs/network-proxy/src/http_proxy.rs +++ b/codex-rs/network-proxy/src/http_proxy.rs @@ -1,4 +1,5 @@ use crate::config::NetworkMode; +use crate::connect_policy::TargetCheckedTcpConnector; use crate::mitm; use crate::network_policy::BlockDecisionAuditEventArgs; use crate::network_policy::NetworkDecision; @@ -66,7 +67,6 @@ use rama_net::proxy::ProxyTarget; use rama_net::proxy::StreamForwardService; use rama_net::stream::SocketInfo; use rama_tcp::client::Request as TcpRequest; -use rama_tcp::client::service::TcpConnector; use rama_tcp::server::TcpListener; use rama_tls_rustls::client::TlsConnectorDataBuilder; use rama_tls_rustls::client::TlsConnectorLayer; @@ -345,20 +345,22 @@ async fn http_connect_proxy(upgraded: Upgraded) -> Result<(), Infallible> { return Ok(()); } - let allow_upstream_proxy = match upgraded + let app_state = match upgraded .extensions() .get::>() .cloned() { - Some(state) => match state.allow_upstream_proxy().await { - Ok(allowed) => allowed, - Err(err) => { - error!("failed to read upstream proxy setting: {err}"); - false - } - }, + Some(state) => state, None => { error!("missing app state"); + return Ok(()); + } + }; + + let allow_upstream_proxy = match app_state.allow_upstream_proxy().await { + Ok(allowed) => allowed, + Err(err) => { + error!("failed to read upstream proxy setting: {err}"); false } }; @@ -369,7 +371,7 @@ async fn http_connect_proxy(upgraded: Upgraded) -> Result<(), Infallible> { None }; - if let Err(err) = forward_connect_tunnel(upgraded, proxy).await { + if let Err(err) = forward_connect_tunnel(upgraded, proxy, app_state).await { warn!("tunnel error: {err}"); } Ok(()) @@ -378,6 +380,7 @@ async fn http_connect_proxy(upgraded: Upgraded) -> Result<(), Infallible> { async fn forward_connect_tunnel( upgraded: Upgraded, proxy: Option, + app_state: Arc, ) -> Result<(), BoxError> { let authority = upgraded .extensions() @@ -392,7 +395,7 @@ async fn forward_connect_tunnel( let req = TcpRequest::new_with_extensions(authority.clone(), extensions) .with_protocol(Protocol::HTTPS); - let proxy_connector = HttpProxyConnector::optional(TcpConnector::new()); + let proxy_connector = HttpProxyConnector::optional(TargetCheckedTcpConnector::new(app_state)); let tls_config = TlsConnectorDataBuilder::new() .with_alpn_protocols_http_auto() .build(); @@ -730,9 +733,9 @@ async fn http_plain_proxy( Err(resp) => return Ok(resp), }; let client = if allow_upstream_proxy { - UpstreamClient::from_env_proxy() + UpstreamClient::from_env_proxy(app_state.clone()) } else { - UpstreamClient::direct() + UpstreamClient::direct(app_state.clone()) }; // Strip hop-by-hop headers only after extracting metadata used for policy correlation. diff --git a/codex-rs/network-proxy/src/lib.rs b/codex-rs/network-proxy/src/lib.rs index 760b09e21e..416ebb0a41 100644 --- a/codex-rs/network-proxy/src/lib.rs +++ b/codex-rs/network-proxy/src/lib.rs @@ -2,6 +2,7 @@ mod certs; mod config; +mod connect_policy; mod http_proxy; mod mitm; mod network_policy; diff --git a/codex-rs/network-proxy/src/mitm.rs b/codex-rs/network-proxy/src/mitm.rs index 8f20afaa3d..7be700b1df 100644 --- a/codex-rs/network-proxy/src/mitm.rs +++ b/codex-rs/network-proxy/src/mitm.rs @@ -52,6 +52,11 @@ pub struct MitmState { max_body_bytes: usize, } +pub(crate) struct MitmUpstreamConfig { + pub(crate) allow_upstream_proxy: bool, + pub(crate) allow_local_binding: bool, +} + #[derive(Clone)] struct MitmPolicyContext { target_host: String, @@ -80,16 +85,16 @@ impl std::fmt::Debug for MitmState { } impl MitmState { - pub(crate) fn new(allow_upstream_proxy: bool) -> Result { + pub(crate) fn new(config: MitmUpstreamConfig) -> Result { // MITM exists to make limited-mode HTTPS enforceable: once CONNECT is established, plain // proxying would lose visibility into the inner HTTP request. We generate/load a local CA // and issue per-host leaf certs so we can terminate TLS and apply policy. let ca = ManagedMitmCa::load_or_create()?; - let upstream = if allow_upstream_proxy { - UpstreamClient::from_env_proxy() + let upstream = if config.allow_upstream_proxy { + UpstreamClient::from_env_proxy_with_allow_local_binding(config.allow_local_binding) } else { - UpstreamClient::direct() + UpstreamClient::direct_with_allow_local_binding(config.allow_local_binding) }; Ok(Self { diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index daafeec737..bddf814415 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -532,6 +532,12 @@ impl NetworkProxyState { Ok(guard.config.network.allow_upstream_proxy) } + pub async fn allow_local_binding(&self) -> Result { + self.reload_if_needed().await?; + let guard = self.state.read().await; + Ok(guard.config.network.allow_local_binding) + } + pub async fn network_mode(&self) -> Result { self.reload_if_needed().await?; let guard = self.state.read().await; diff --git a/codex-rs/network-proxy/src/socks5.rs b/codex-rs/network-proxy/src/socks5.rs index b360d6853a..2d4c05f95c 100644 --- a/codex-rs/network-proxy/src/socks5.rs +++ b/codex-rs/network-proxy/src/socks5.rs @@ -1,4 +1,5 @@ use crate::config::NetworkMode; +use crate::connect_policy::TargetCheckedTcpConnector; use crate::network_policy::BlockDecisionAuditEventArgs; use crate::network_policy::NetworkDecision; use crate::network_policy::NetworkDecisionSource; @@ -34,7 +35,6 @@ use rama_socks5::server::udp::RelayRequest; use rama_socks5::server::udp::RelayResponse; use rama_tcp::TcpStream; use rama_tcp::client::Request as TcpRequest; -use rama_tcp::client::service::TcpConnector; use rama_tcp::server::TcpListener; use std::io; use std::net::SocketAddr; @@ -94,7 +94,7 @@ async fn run_socks5_with_listener( } } - let tcp_connector = TcpConnector::default(); + let tcp_connector = TargetCheckedTcpConnector::new(state.clone()); let policy_tcp_connector = service_fn({ let policy_decider = policy_decider.clone(); move |req: TcpRequest| { @@ -131,7 +131,7 @@ async fn run_socks5_with_listener( async fn handle_socks5_tcp( req: TcpRequest, - tcp_connector: TcpConnector, + tcp_connector: TargetCheckedTcpConnector, policy_decider: Option>, ) -> Result, BoxError> { let app_state = req @@ -548,7 +548,7 @@ mod tests { let (result, events) = capture_events(|| async { handle_socks5_tcp( request, - TcpConnector::default(), + TargetCheckedTcpConnector::new(state.clone()), /*policy_decider*/ None, ) .await diff --git a/codex-rs/network-proxy/src/state.rs b/codex-rs/network-proxy/src/state.rs index 01e4966b5a..b14968e9a8 100644 --- a/codex-rs/network-proxy/src/state.rs +++ b/codex-rs/network-proxy/src/state.rs @@ -3,6 +3,7 @@ use crate::config::NetworkMode; use crate::config::NetworkProxyConfig; use crate::config::NetworkUnixSocketPermissions; use crate::mitm::MitmState; +use crate::mitm::MitmUpstreamConfig; use crate::policy::DomainPattern; use crate::policy::compile_allowlist_globset; use crate::policy::compile_denylist_globset; @@ -66,9 +67,10 @@ pub fn build_config_state( let deny_set = compile_denylist_globset(&denied_domains)?; let allow_set = compile_allowlist_globset(&allowed_domains)?; let mitm = if config.network.mitm { - Some(Arc::new(MitmState::new( - config.network.allow_upstream_proxy, - )?)) + Some(Arc::new(MitmState::new(MitmUpstreamConfig { + allow_upstream_proxy: config.network.allow_upstream_proxy, + allow_local_binding: config.network.allow_local_binding, + })?)) } else { None }; diff --git a/codex-rs/network-proxy/src/upstream.rs b/codex-rs/network-proxy/src/upstream.rs index 97e78f303b..c7b67cc18a 100644 --- a/codex-rs/network-proxy/src/upstream.rs +++ b/codex-rs/network-proxy/src/upstream.rs @@ -1,3 +1,5 @@ +use crate::connect_policy::TargetCheckedTcpConnector; +use crate::state::NetworkProxyState; use rama_core::Layer; use rama_core::Service; use rama_core::error::BoxError; @@ -16,9 +18,9 @@ use rama_http_backend::client::proxy::layer::HttpProxyConnectorLayer; use rama_net::address::ProxyAddress; use rama_net::client::EstablishedClientConnection; use rama_net::http::RequestContext; -use rama_tcp::client::service::TcpConnector; use rama_tls_rustls::client::TlsConnectorDataBuilder; use rama_tls_rustls::client::TlsConnectorLayer; +use std::sync::Arc; use tracing::warn; #[cfg(target_os = "macos")] @@ -102,12 +104,32 @@ pub(crate) struct UpstreamClient { } impl UpstreamClient { - pub(crate) fn direct() -> Self { - Self::new(ProxyConfig::default()) + pub(crate) fn direct(state: Arc) -> Self { + Self::new( + ProxyConfig::default(), + TargetCheckedTcpConnector::new(state), + ) } - pub(crate) fn from_env_proxy() -> Self { - Self::new(ProxyConfig::from_env()) + pub(crate) fn from_env_proxy(state: Arc) -> Self { + Self::new( + ProxyConfig::from_env(), + TargetCheckedTcpConnector::new(state), + ) + } + + pub(crate) fn direct_with_allow_local_binding(allow_local_binding: bool) -> Self { + Self::new( + ProxyConfig::default(), + TargetCheckedTcpConnector::from_allow_local_binding(allow_local_binding), + ) + } + + pub(crate) fn from_env_proxy_with_allow_local_binding(allow_local_binding: bool) -> Self { + Self::new( + ProxyConfig::from_env(), + TargetCheckedTcpConnector::from_allow_local_binding(allow_local_binding), + ) } #[cfg(target_os = "macos")] @@ -119,8 +141,8 @@ impl UpstreamClient { } } - fn new(proxy_config: ProxyConfig) -> Self { - let connector = build_http_connector(); + fn new(proxy_config: ProxyConfig, transport: TargetCheckedTcpConnector) -> Self { + let connector = build_http_connector(transport); Self { connector, proxy_config, @@ -158,12 +180,13 @@ impl Service> for UpstreamClient { } } -fn build_http_connector() -> BoxService< +fn build_http_connector( + transport: TargetCheckedTcpConnector, +) -> BoxService< Request, EstablishedClientConnection, Request>, BoxError, > { - let transport = TcpConnector::default(); let proxy = HttpProxyConnectorLayer::optional().into_layer(transport); let tls_config = TlsConnectorDataBuilder::new() .with_alpn_protocols_http_auto() diff --git a/codex-rs/otel/src/events/session_telemetry.rs b/codex-rs/otel/src/events/session_telemetry.rs index 38d361e7a7..bb64bb2bba 100644 --- a/codex-rs/otel/src/events/session_telemetry.rs +++ b/codex-rs/otel/src/events/session_telemetry.rs @@ -305,6 +305,23 @@ impl SessionTelemetry { handle_responses_span.record("tool_name", name.as_str()); } } + ResponseEvent::Completed { + token_usage: Some(token_usage), + .. + } => { + handle_responses_span.record("gen_ai.usage.input_tokens", token_usage.input_tokens); + handle_responses_span.record( + "gen_ai.usage.cache_read.input_tokens", + token_usage.cached_input(), + ); + handle_responses_span + .record("gen_ai.usage.output_tokens", token_usage.output_tokens); + handle_responses_span.record( + "codex.usage.reasoning_output_tokens", + token_usage.reasoning_output_tokens, + ); + handle_responses_span.record("codex.usage.total_tokens", token_usage.total_tokens); + } _ => {} } } diff --git a/codex-rs/process-hardening/src/lib.rs b/codex-rs/process-hardening/src/lib.rs index 1c206aae87..f9695fcbdc 100644 --- a/codex-rs/process-hardening/src/lib.rs +++ b/codex-rs/process-hardening/src/lib.rs @@ -61,6 +61,17 @@ pub(crate) fn pre_main_hardening_linux() { remove_env_vars_with_prefix(b"LD_"); } +/// Mark the current Linux process non-dumpable so same-user processes cannot attach with ptrace. +#[cfg(target_os = "linux")] +pub fn disable_process_dumping() -> std::io::Result<()> { + let ret_code = unsafe { libc::prctl(libc::PR_SET_DUMPABLE, 0, 0, 0, 0) }; + if ret_code == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()) + } +} + #[cfg(any(target_os = "freebsd", target_os = "openbsd"))] pub(crate) fn pre_main_hardening_bsd() { // FreeBSD/OpenBSD: set RLIMIT_CORE to 0 and clear LD_* env vars diff --git a/codex-rs/protocol/src/account.rs b/codex-rs/protocol/src/account.rs index aea9ad843a..bfbccbeb24 100644 --- a/codex-rs/protocol/src/account.rs +++ b/codex-rs/protocol/src/account.rs @@ -3,6 +3,9 @@ use serde::Deserialize; use serde::Serialize; use ts_rs::TS; +use crate::auth::KnownPlan; +use crate::auth::PlanType as AuthPlanType; + #[derive(Serialize, Deserialize, Copy, Clone, Debug, PartialEq, Eq, JsonSchema, TS, Default)] #[serde(rename_all = "lowercase")] #[ts(rename_all = "lowercase")] @@ -57,9 +60,38 @@ impl PlanType { } } +impl From for PlanType { + fn from(plan_type: AuthPlanType) -> Self { + match plan_type { + AuthPlanType::Known(plan) => plan.into(), + AuthPlanType::Unknown(_) => Self::Unknown, + } + } +} + +impl From for PlanType { + fn from(plan: KnownPlan) -> Self { + match plan { + KnownPlan::Free => Self::Free, + KnownPlan::Go => Self::Go, + KnownPlan::Plus => Self::Plus, + KnownPlan::Pro => Self::Pro, + KnownPlan::ProLite => Self::ProLite, + KnownPlan::Team => Self::Team, + KnownPlan::SelfServeBusinessUsageBased => Self::SelfServeBusinessUsageBased, + KnownPlan::Business => Self::Business, + KnownPlan::EnterpriseCbpUsageBased => Self::EnterpriseCbpUsageBased, + KnownPlan::Enterprise => Self::Enterprise, + KnownPlan::Edu => Self::Edu, + } + } +} + #[cfg(test)] mod tests { use super::PlanType; + use crate::auth::KnownPlan; + use crate::auth::PlanType as AuthPlanType; use pretty_assertions::assert_eq; #[test] @@ -121,4 +153,20 @@ mod tests { assert_eq!(PlanType::Edu.is_workspace_account(), true); assert_eq!(PlanType::Pro.is_workspace_account(), false); } + + #[test] + fn auth_plan_type_converts_to_account_plan_type() { + assert_eq!( + PlanType::from(AuthPlanType::Known(KnownPlan::EnterpriseCbpUsageBased)), + PlanType::EnterpriseCbpUsageBased + ); + assert_eq!( + PlanType::from(AuthPlanType::Known(KnownPlan::Enterprise)), + PlanType::Enterprise + ); + assert_eq!( + PlanType::from(AuthPlanType::Unknown("mystery-tier".to_string())), + PlanType::Unknown + ); + } } diff --git a/codex-rs/protocol/src/auth.rs b/codex-rs/protocol/src/auth.rs index 99e067bf24..495f138f42 100644 --- a/codex-rs/protocol/src/auth.rs +++ b/codex-rs/protocol/src/auth.rs @@ -46,6 +46,7 @@ pub enum KnownPlan { EnterpriseCbpUsageBased, #[serde(alias = "hc")] Enterprise, + #[serde(alias = "education")] Edu, } @@ -118,3 +119,23 @@ pub enum RefreshTokenFailedReason { Revoked, Other, } + +#[cfg(test)] +mod tests { + use super::KnownPlan; + use super::PlanType; + use pretty_assertions::assert_eq; + + #[test] + fn plan_type_deserializes_raw_aliases() { + assert_eq!( + serde_json::from_str::("\"hc\"").expect("hc should deserialize"), + PlanType::Known(KnownPlan::Enterprise) + ); + assert_eq!( + serde_json::from_str::("\"education\"") + .expect("education should deserialize"), + PlanType::Known(KnownPlan::Edu) + ); + } +} diff --git a/codex-rs/state/src/log_db.rs b/codex-rs/state/src/log_db.rs index a9da475329..6024d6bdaa 100644 --- a/codex-rs/state/src/log_db.rs +++ b/codex-rs/state/src/log_db.rs @@ -367,6 +367,8 @@ async fn run_inserter( ) { let mut buffer = Vec::with_capacity(config.batch_size); let mut ticker = tokio::time::interval(config.flush_interval); + // Consume the immediate startup tick so entries flush after the interval. + ticker.tick().await; loop { tokio::select! { maybe_command = receiver.recv() => { @@ -645,7 +647,6 @@ mod tests { flush_interval: std::time::Duration::from_secs(60), }, ); - tokio::time::sleep(std::time::Duration::from_millis(10)).await; let guard = tracing_subscriber::registry() .with( diff --git a/codex-rs/tools/src/tool_config.rs b/codex-rs/tools/src/tool_config.rs index 3ddcf481a9..208efead9c 100644 --- a/codex-rs/tools/src/tool_config.rs +++ b/codex-rs/tools/src/tool_config.rs @@ -107,6 +107,7 @@ pub struct ToolsConfig { pub spawn_agent_usage_hint: bool, pub spawn_agent_usage_hint_text: Option, pub max_concurrent_threads_per_session: Option, + pub wait_agent_min_timeout_ms: Option, pub default_mode_request_user_input: bool, pub experimental_supported_tools: Vec, pub agent_jobs_tools: bool, @@ -226,6 +227,7 @@ impl ToolsConfig { spawn_agent_usage_hint: true, spawn_agent_usage_hint_text: None, max_concurrent_threads_per_session: None, + wait_agent_min_timeout_ms: None, default_mode_request_user_input: include_default_mode_request_user_input, experimental_supported_tools: model_info.experimental_supported_tools.clone(), agent_jobs_tools: include_agent_jobs, @@ -270,6 +272,14 @@ impl ToolsConfig { self } + pub fn with_wait_agent_min_timeout_ms( + mut self, + wait_agent_min_timeout_ms: Option, + ) -> Self { + self.wait_agent_min_timeout_ms = wait_agent_min_timeout_ms; + self + } + pub fn with_allow_login_shell(mut self, allow_login_shell: bool) -> Self { self.allow_login_shell = allow_login_shell; self