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

code-mode: add opt-in IPC consumer

Browse Source
This commit is contained in:
Channing Conger
2026-06-12 05:39:29 +00:00
parent c009b33335
commit b0cf57bdeb
26 changed files with 470 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+3
View File
@@ -2621,6 +2621,8 @@ dependencies = [
"codex-apply-patch",
"codex-async-utils",
"codex-code-mode",
"codex-code-mode-client",
"codex-code-mode-protocol",
"codex-config",
"codex-connectors",
"codex-context-fragments",
@@ -4624,6 +4626,7 @@ dependencies = [
"assert_cmd",
"base64 0.22.1",
"codex-arg0",
"codex-code-mode-client",
"codex-config",
"codex-core",
"codex-exec-server",
codex-rs/app-server/BUILD.bazel
+1
View File
@@ -16,6 +16,7 @@ codex_rust_crate(
},
extra_binaries = [
"//codex-rs/bwrap:bwrap",
"//codex-rs/code-mode-host:codex-code-mode-host",
],
test_tags = ["no-sandbox"],
)
codex-rs/app-server/tests/suite/v2/imagegen_extension.rs
+12 -2
View File
@@ -290,9 +290,19 @@ generatedImage(result);
ChatGptAuthFixture::new("access-chatgpt"),
AuthCredentialsStoreMode::File,
)?;
let code_mode_host = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?;
let code_mode_host = code_mode_host
.to_str()
.context("code-mode host path is not valid UTF-8")?;
let mut mcp =
TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?;
let mut mcp = TestAppServer::new_with_env(
codex_home.path(),
&[
("OPENAI_API_KEY", None),
("CODEX_CODE_MODE_HOST_PATH", Some(code_mode_host)),
],
)
.await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
start_image_generation_turn(&mut mcp).await?;
timeout(
codex-rs/core/BUILD.bazel
+1
View File
@@ -41,6 +41,7 @@ codex_rust_crate(
unit_test_timeout = "long",
extra_binaries = [
"//codex-rs/bwrap:bwrap",
"//codex-rs/code-mode-host:codex-code-mode-host",
"//codex-rs/linux-sandbox:codex-linux-sandbox",
"//codex-rs/rmcp-client:test_stdio_server",
"//codex-rs/rmcp-client:test_streamable_http_server",
codex-rs/core/Cargo.toml
+2
View File
@@ -29,6 +29,7 @@ codex-app-server-protocol = { workspace = true }
codex-apply-patch = { workspace = true }
codex-async-utils = { workspace = true }
codex-code-mode = { workspace = true }
codex-code-mode-protocol = { workspace = true }
codex-connectors = { workspace = true }
codex-context-fragments = { workspace = true }
codex-config = { workspace = true }
@@ -133,6 +134,7 @@ codex-shell-escalation = { workspace = true }
[dev-dependencies]
assert_cmd = { workspace = true }
assert_matches = { workspace = true }
codex-code-mode-client = { workspace = true }
codex-image-generation-extension = { workspace = true }
codex-home = { workspace = true }
codex-otel = { workspace = true }
codex-rs/core/src/codex_delegate.rs
+1
View File
@@ -113,6 +113,7 @@ pub(crate) async fn run_codex_thread_interactive(
analytics_events_client: Some(parent_session.services.analytics_events_client.clone()),
thread_store: Arc::clone(&parent_session.services.thread_store),
attestation_provider: parent_session.services.attestation_provider.clone(),
code_mode_session_provider: Arc::clone(&parent_session.services.code_mode_session_provider),
inherited_multi_agent_version: Some(MultiAgentVersion::Disabled),
}))
.or_cancel(&cancel_token)
codex-rs/core/src/session/mod.rs
+4
View File
@@ -432,6 +432,8 @@ pub(crate) struct CodexSpawnArgs {
pub(crate) analytics_events_client: Option<AnalyticsEventsClient>,
pub(crate) thread_store: Arc<dyn ThreadStore>,
pub(crate) attestation_provider: Option<Arc<dyn AttestationProvider>>,
pub(crate) code_mode_session_provider:
Arc<dyn codex_code_mode_protocol::CodeModeSessionProvider>,
pub(crate) inherited_multi_agent_version: Option<MultiAgentVersion>,
}
@@ -514,6 +516,7 @@ impl Codex {
analytics_events_client,
thread_store,
attestation_provider,
code_mode_session_provider,
inherited_multi_agent_version,
} = args;
let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
@@ -667,6 +670,7 @@ impl Codex {
thread_store,
parent_rollout_thread_trace,
attestation_provider,
code_mode_session_provider,
multi_agent_version,
))
.await
codex-rs/core/src/session/session.rs
+5 -1
View File
@@ -492,6 +492,7 @@ impl Session {
thread_store: Arc<dyn ThreadStore>,
parent_rollout_thread_trace: ThreadTraceContext,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
code_mode_session_provider: Arc<dyn codex_code_mode_protocol::CodeModeSessionProvider>,
multi_agent_version: Option<MultiAgentVersion>,
) -> anyhow::Result<Arc<Self>> {
debug!(
@@ -1031,7 +1032,10 @@ impl Session {
session_configuration.parent_thread_id,
),
),
code_mode_service: crate::tools::code_mode::CodeModeService::new(),
code_mode_service: crate::tools::code_mode::CodeModeService::new(
Arc::clone(&code_mode_session_provider),
),
code_mode_session_provider,
environment_manager,
};
let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) =
codex-rs/core/src/session/tests.rs
+15 -2
View File
@@ -4835,6 +4835,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() {
)),
codex_rollout_trace::ThreadTraceContext::disabled(),
/*attestation_provider*/ None,
Arc::new(codex_code_mode_client::IpcCodeModeSessionProvider::default()),
Some(config.multi_agent_version_from_features()),
)
.await;
@@ -5001,7 +5002,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
Session::build_model_client_beta_features_header(config.as_ref()),
/*attestation_provider*/ None,
),
code_mode_service: crate::tools::code_mode::CodeModeService::new(),
code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new(
codex_code_mode_client::IpcCodeModeSessionProvider::default(),
)),
code_mode_session_provider: Arc::new(
codex_code_mode_client::IpcCodeModeSessionProvider::default(),
),
environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
};
@@ -5173,6 +5179,7 @@ async fn make_session_with_config_and_rx(
)),
codex_rollout_trace::ThreadTraceContext::disabled(),
/*attestation_provider*/ None,
Arc::new(codex_code_mode_client::IpcCodeModeSessionProvider::default()),
Some(config.multi_agent_version_from_features()),
)
.await?;
@@ -5282,6 +5289,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx(
)),
codex_rollout_trace::ThreadTraceContext::disabled(),
/*attestation_provider*/ None,
Arc::new(codex_code_mode_client::IpcCodeModeSessionProvider::default()),
Some(config.multi_agent_version_from_features()),
)
.await?;
@@ -7000,7 +7008,12 @@ where
Session::build_model_client_beta_features_header(config.as_ref()),
/*attestation_provider*/ None,
),
code_mode_service: crate::tools::code_mode::CodeModeService::new(),
code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new(
codex_code_mode_client::IpcCodeModeSessionProvider::default(),
)),
code_mode_session_provider: Arc::new(
codex_code_mode_client::IpcCodeModeSessionProvider::default(),
),
environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
};
codex-rs/core/src/session/tests/guardian_tests.rs
+3
View File
@@ -736,6 +736,9 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
analytics_events_client: None,
thread_store,
attestation_provider: None,
code_mode_session_provider: Arc::new(
codex_code_mode_client::IpcCodeModeSessionProvider::default(),
),
inherited_multi_agent_version: None,
})
.await
codex-rs/core/src/state/service.rs
+2
View File
@@ -19,6 +19,7 @@ use anyhow::Result;
use arc_swap::ArcSwap;
use arc_swap::ArcSwapOption;
use codex_analytics::AnalyticsEventsClient;
use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_core_plugins::PluginsManager;
use codex_exec_server::EnvironmentManager;
use codex_extension_api::ExtensionData;
@@ -79,6 +80,7 @@ pub(crate) struct SessionServices {
/// Session-scoped model client shared across turns.
pub(crate) model_client: ModelClient,
pub(crate) code_mode_service: CodeModeService,
pub(crate) code_mode_session_provider: Arc<dyn CodeModeSessionProvider>,
/// Shared process-level environment registry. Sessions carry an `Arc` handle so they can pass
/// the same manager through child-thread spawn paths without reconstructing it.
pub(crate) environment_manager: Arc<EnvironmentManager>,
codex-rs/core/src/thread_manager.rs
+18
View File
@@ -19,6 +19,8 @@ use crate::tasks::interrupted_turn_history_marker;
use codex_analytics::AnalyticsEventsClient;
use codex_app_server_protocol::ThreadHistoryBuilder;
use codex_app_server_protocol::TurnStatus;
use codex_code_mode::InProcessCodeModeSessionProvider;
use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_core_plugins::PluginsManager;
use codex_exec_server::EnvironmentManager;
use codex_extension_api::ExtensionDataInit;
@@ -218,6 +220,7 @@ pub(crate) struct ThreadManagerState {
installation_id: String,
analytics_events_client: Option<AnalyticsEventsClient>,
state_db: Option<StateDbHandle>,
code_mode_session_provider: Arc<dyn CodeModeSessionProvider>,
// Captures submitted ops for testing purpose when test mode is enabled.
ops_log: Option<SharedCapturedOps>,
}
@@ -304,6 +307,7 @@ impl ThreadManager {
installation_id,
analytics_events_client,
state_db,
code_mode_session_provider: Arc::new(InProcessCodeModeSessionProvider),
ops_log: should_use_test_thread_manager_behavior()
.then(|| Arc::new(std::sync::Mutex::new(Vec::new()))),
}),
@@ -409,6 +413,7 @@ impl ThreadManager {
installation_id,
analytics_events_client: None,
state_db,
code_mode_session_provider: Arc::new(InProcessCodeModeSessionProvider),
ops_log: should_use_test_thread_manager_behavior()
.then(|| Arc::new(std::sync::Mutex::new(Vec::new()))),
}),
@@ -420,6 +425,18 @@ impl ThreadManager {
self.state.session_source.clone()
}
#[doc(hidden)]
pub fn with_code_mode_session_provider(
mut self,
provider: Arc<dyn CodeModeSessionProvider>,
) -> Self {
let Some(state) = Arc::get_mut(&mut self.state) else {
panic!("code-mode provider must be configured before sharing the thread manager");
};
state.code_mode_session_provider = provider;
self
}
pub fn auth_manager(&self) -> Arc<AuthManager> {
self.state.auth_manager.clone()
}
@@ -1430,6 +1447,7 @@ impl ThreadManagerState {
analytics_events_client: self.analytics_events_client.clone(),
thread_store: Arc::clone(&self.thread_store),
attestation_provider: self.attestation_provider.clone(),
code_mode_session_provider: Arc::clone(&self.code_mode_session_provider),
inherited_multi_agent_version: multi_agent_version,
}))
.await?;
codex-rs/core/src/tools/code_mode/delegate.rs
+5 -5
View File
@@ -2,11 +2,11 @@ use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use codex_code_mode::CellId;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeSessionDelegate;
use codex_code_mode::NotificationFuture;
use codex_code_mode::ToolInvocationFuture;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeNestedToolCall;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::NotificationFuture;
use codex_code_mode_protocol::ToolInvocationFuture;
use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseItem;
use serde_json::Value as JsonValue;
codex-rs/core/src/tools/code_mode/execute_handler.rs
+7 -4
View File
@@ -33,8 +33,8 @@ impl CodeModeExecuteHandler {
call_id: String,
code: String,
) -> Result<FunctionToolOutput, FunctionCallError> {
let args =
codex_code_mode::parse_exec_source(&code).map_err(FunctionCallError::RespondToModel)?;
let args = codex_code_mode_protocol::parse_exec_source(&code)
.map_err(FunctionCallError::RespondToModel)?;
let exec = ExecContext { session, turn };
let enabled_tools =
codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs);
@@ -43,7 +43,7 @@ impl CodeModeExecuteHandler {
.session
.services
.code_mode_service
.execute(codex_code_mode::ExecuteRequest {
.execute(codex_code_mode_protocol::ExecuteRequest {
tool_call_id: call_id.clone(),
enabled_tools,
source: args.code.clone(),
@@ -78,7 +78,10 @@ impl CodeModeExecuteHandler {
code_cell_trace.record_initial_response(&response);
// Yielded cells keep running, so terminal lifecycle is only emitted
// here when the first response also ended the runtime.
if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) {
if !matches!(
response,
codex_code_mode_protocol::RuntimeResponse::Yielded { .. }
) {
code_cell_trace.record_ended(&response);
exec.session
.services
codex-rs/core/src/tools/code_mode/execute_spec.rs
+8 -8
View File
@@ -1,4 +1,4 @@
use codex_code_mode::ToolDefinition as CodeModeToolDefinition;
use codex_code_mode_protocol::ToolDefinition as CodeModeToolDefinition;
use codex_tools::FreeformTool;
use codex_tools::FreeformToolFormat;
use codex_tools::ToolSpec;
@@ -6,7 +6,7 @@ use std::collections::BTreeMap;
pub(crate) fn create_code_mode_tool(
enabled_tools: &[CodeModeToolDefinition],
namespace_descriptions: &BTreeMap<String, codex_code_mode::ToolNamespaceDescription>,
namespace_descriptions: &BTreeMap<String, codex_code_mode_protocol::ToolNamespaceDescription>,
code_mode_only: bool,
deferred_tools_available: bool,
) -> ToolSpec {
@@ -21,8 +21,8 @@ SOURCE: /[\s\S]+/
"#;
ToolSpec::Freeform(FreeformTool {
name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(),
description: codex_code_mode::build_exec_tool_description(
name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(),
description: codex_code_mode_protocol::build_exec_tool_description(
enabled_tools,
namespace_descriptions,
code_mode_only,
@@ -44,11 +44,11 @@ mod tests {
#[test]
fn create_code_mode_tool_matches_expected_spec() {
let enabled_tools = vec![codex_code_mode::ToolDefinition {
let enabled_tools = vec![codex_code_mode_protocol::ToolDefinition {
name: "update_plan".to_string(),
tool_name: ToolName::plain("update_plan"),
description: "Update the plan".to_string(),
kind: codex_code_mode::CodeModeToolKind::Function,
kind: codex_code_mode_protocol::CodeModeToolKind::Function,
input_schema: None,
output_schema: None,
}];
@@ -61,8 +61,8 @@ mod tests {
/*deferred_tools_available*/ false,
),
ToolSpec::Freeform(FreeformTool {
name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(),
description: codex_code_mode::build_exec_tool_description(
name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(),
description: codex_code_mode_protocol::build_exec_tool_description(
&enabled_tools,
&BTreeMap::new(),
/*code_mode_only*/ true,
codex-rs/core/src/tools/code_mode/mod.rs
+227 -29
View File
@@ -8,13 +8,18 @@ pub(crate) mod wait_spec;
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode::CellId;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeSession;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::RuntimeResponse;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeNestedToolCall;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_code_mode_protocol::CodeModeToolKind;
use codex_code_mode_protocol::RuntimeResponse;
use codex_protocol::models::FunctionCallOutputContentItem;
use serde_json::Value as JsonValue;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use tokio::sync::Mutex;
use tokio::sync::Semaphore;
use tokio_util::sync::CancellationToken;
use crate::function_tool::FunctionCallError;
@@ -42,9 +47,10 @@ pub(crate) use execute_handler::CodeModeExecuteHandler;
use response_adapter::into_function_call_output_content_items;
pub(crate) use wait_handler::CodeModeWaitHandler;
pub(crate) const PUBLIC_TOOL_NAME: &str = codex_code_mode::PUBLIC_TOOL_NAME;
pub(crate) const WAIT_TOOL_NAME: &str = codex_code_mode::WAIT_TOOL_NAME;
pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 = codex_code_mode::DEFAULT_WAIT_YIELD_TIME_MS;
pub(crate) const PUBLIC_TOOL_NAME: &str = codex_code_mode_protocol::PUBLIC_TOOL_NAME;
pub(crate) const WAIT_TOOL_NAME: &str = codex_code_mode_protocol::WAIT_TOOL_NAME;
pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 =
codex_code_mode_protocol::DEFAULT_WAIT_YIELD_TIME_MS;
/// Returns true for the un-namespaced code-mode `exec` tool.
pub(crate) fn is_exec_tool_name(tool_name: &ToolName) -> bool {
@@ -58,50 +64,69 @@ pub(crate) struct ExecContext {
}
pub(crate) struct CodeModeService {
session: Option<Arc<dyn CodeModeSession>>,
session: Mutex<Option<Arc<dyn CodeModeSession>>>,
provider: Arc<dyn CodeModeSessionProvider>,
dispatch_broker: Arc<CodeModeDispatchBroker>,
session_init_permit: Semaphore,
shutting_down: AtomicBool,
}
impl CodeModeService {
pub(crate) fn new() -> Self {
pub(crate) fn new(provider: Arc<dyn CodeModeSessionProvider>) -> Self {
let dispatch_broker = Arc::new(CodeModeDispatchBroker::new());
Self {
session: Some(Arc::new(codex_code_mode::CodeModeService::with_delegate(
dispatch_broker.clone(),
))),
session: Mutex::new(None),
provider,
dispatch_broker,
session_init_permit: Semaphore::new(/*permits*/ 1),
shutting_down: AtomicBool::new(false),
}
}
pub(crate) async fn execute(
&self,
request: codex_code_mode::ExecuteRequest,
) -> Result<codex_code_mode::StartedCell, String> {
self.session()?.execute(request).await
request: codex_code_mode_protocol::ExecuteRequest,
) -> Result<codex_code_mode_protocol::StartedCell, String> {
self.session_for_execute().await?.execute(request).await
}
pub(crate) async fn wait(
&self,
request: codex_code_mode::WaitRequest,
) -> Result<codex_code_mode::WaitOutcome, String> {
self.session()?.wait(request).await
request: codex_code_mode_protocol::WaitRequest,
) -> Result<codex_code_mode_protocol::WaitOutcome, String> {
self.current_session()
.await
.ok_or_else(|| "code mode session is unavailable".to_string())?
.wait(request)
.await
}
pub(crate) async fn terminate(
&self,
cell_id: CellId,
) -> Result<codex_code_mode::WaitOutcome, String> {
self.session()?.terminate(cell_id).await
) -> Result<codex_code_mode_protocol::WaitOutcome, String> {
self.current_session()
.await
.ok_or_else(|| "code mode session is unavailable".to_string())?
.terminate(cell_id)
.await
}
pub(crate) async fn shutdown(&self) -> Result<(), String> {
match &self.session {
self.shutting_down.store(true, Ordering::Release);
let _permit = self
.session_init_permit
.acquire()
.await
.map_err(|_| "code mode session initializer closed".to_string())?;
let session = self.session.lock().await.clone();
match session {
Some(session) => session.shutdown().await,
None => Ok(()),
}
}
pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode::CellId) {
pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode_protocol::CellId) {
self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id);
}
@@ -116,9 +141,7 @@ impl CodeModeService {
router: Arc<ToolRouter>,
tracker: SharedTurnDiffTracker,
) -> Option<CodeModeDispatchWorker> {
if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly)
|| self.session.is_none()
{
if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) {
return None;
}
@@ -132,10 +155,43 @@ impl CodeModeService {
)
}
fn session(&self) -> Result<&Arc<dyn CodeModeSession>, String> {
async fn session_for_execute(&self) -> Result<Arc<dyn CodeModeSession>, String> {
if self.shutting_down.load(Ordering::Acquire) {
return Err("code mode session is shutting down".to_string());
}
if let Some(session) = self.current_session().await {
return Ok(session);
}
let _permit = self
.session_init_permit
.acquire()
.await
.map_err(|_| "code mode session initializer closed".to_string())?;
if self.shutting_down.load(Ordering::Acquire) {
return Err("code mode session is shutting down".to_string());
}
if let Some(session) = self.current_session().await {
return Ok(session);
}
let session = self
.provider
.create_session(self.dispatch_broker.clone())
.await?;
if self.shutting_down.load(Ordering::Acquire) {
let _ = session.shutdown().await;
return Err("code mode session is shutting down".to_string());
}
*self.session.lock().await = Some(Arc::clone(&session));
Ok(session)
}
async fn current_session(&self) -> Option<Arc<dyn CodeModeSession>> {
self.session
.lock()
.await
.as_ref()
.ok_or_else(|| "code mode is unavailable".to_string())
.filter(|session| session.is_alive())
.cloned()
}
}
@@ -320,11 +376,30 @@ fn build_freeform_tool_payload(
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use super::CodeModeService;
use super::build_nested_tool_payload;
use crate::tools::context::ToolPayload;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_code_mode_protocol::CodeModeSessionProviderFuture;
use codex_code_mode_protocol::CodeModeSessionResultFuture;
use codex_code_mode_protocol::CodeModeToolKind;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use codex_tools::ToolName;
use pretty_assertions::assert_eq;
use serde_json::json;
use tokio::sync::oneshot;
#[test]
fn build_nested_tool_payload_uses_function_kind() {
@@ -359,4 +434,127 @@ mod tests {
other => panic!("expected freeform payload, got {other:?}"),
}
}
struct RecoveringSessionProvider {
sessions_created: AtomicUsize,
}
impl CodeModeSessionProvider for RecoveringSessionProvider {
fn create_session<'a>(
&'a self,
_delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionProviderFuture<'a> {
let generation = self.sessions_created.fetch_add(1, Ordering::Relaxed) + 1;
Box::pin(async move {
let session: Arc<dyn CodeModeSession> = Arc::new(RecoveringSession {
generation,
alive: AtomicBool::new(true),
});
Ok(session)
})
}
}
struct RecoveringSession {
generation: usize,
alive: AtomicBool,
}
impl CodeModeSession for RecoveringSession {
fn is_alive(&self) -> bool {
self.alive.load(Ordering::Acquire)
}
fn execute<'a>(
&'a self,
_request: ExecuteRequest,
) -> CodeModeSessionResultFuture<'a, StartedCell> {
Box::pin(async move {
if self.generation == 1 {
self.alive.store(false, Ordering::Release);
return Err("host crashed".to_string());
}
let cell_id = CellId::new(format!("host{}_1", self.generation));
let (response_tx, response_rx) = oneshot::channel();
response_tx
.send(RuntimeResponse::Result {
cell_id: cell_id.clone(),
content_items: Vec::new(),
error_text: None,
})
.expect("test response receiver should be live");
Ok(StartedCell::new(cell_id, response_rx))
})
}
fn wait<'a>(
&'a self,
_request: WaitRequest,
) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(async { panic!("wait should not be sent to a failed session") })
}
fn terminate<'a>(
&'a self,
_cell_id: CellId,
) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(async { panic!("terminate should not be sent to a failed session") })
}
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
Box::pin(async { Ok(()) })
}
}
fn execute_request() -> ExecuteRequest {
ExecuteRequest {
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: "text('hello')".to_string(),
yield_time_ms: None,
max_output_tokens: None,
}
}
#[tokio::test]
async fn exec_replaces_failed_session_but_wait_does_not() {
let provider = Arc::new(RecoveringSessionProvider {
sessions_created: AtomicUsize::new(0),
});
let provider_trait: Arc<dyn CodeModeSessionProvider> = provider.clone();
let service = CodeModeService::new(provider_trait);
assert_eq!(
service.execute(execute_request()).await.err().as_deref(),
Some("host crashed")
);
assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1);
assert_eq!(
service
.wait(WaitRequest {
cell_id: CellId::new("host1_1".to_string()),
yield_time_ms: 1,
})
.await
.err()
.as_deref(),
Some("code mode session is unavailable")
);
assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1);
assert_eq!(
service
.terminate(CellId::new("host1_1".to_string()))
.await
.err()
.as_deref(),
Some("code mode session is unavailable")
);
assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1);
let started = service.execute(execute_request()).await.unwrap();
assert_eq!(started.cell_id, CellId::new("host2_1".to_string()));
assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 2);
}
}
codex-rs/core/src/tools/code_mode/response_adapter.rs
+13 -12
View File
@@ -1,4 +1,4 @@
use codex_code_mode::ImageDetail as CodeModeImageDetail;
use codex_code_mode_protocol::ImageDetail as CodeModeImageDetail;
use codex_protocol::models::DEFAULT_IMAGE_DETAIL;
use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::ImageDetail;
@@ -8,7 +8,7 @@ trait IntoProtocol<T> {
}
pub(super) fn into_function_call_output_content_items(
items: Vec<codex_code_mode::FunctionCallOutputContentItem>,
items: Vec<codex_code_mode_protocol::FunctionCallOutputContentItem>,
) -> Vec<FunctionCallOutputContentItem> {
items.into_iter().map(IntoProtocol::into_protocol).collect()
}
@@ -26,22 +26,23 @@ impl IntoProtocol<ImageDetail> for CodeModeImageDetail {
}
impl IntoProtocol<FunctionCallOutputContentItem>
for codex_code_mode::FunctionCallOutputContentItem
for codex_code_mode_protocol::FunctionCallOutputContentItem
{
fn into_protocol(self) -> FunctionCallOutputContentItem {
let value = self;
match value {
codex_code_mode::FunctionCallOutputContentItem::InputText { text } => {
codex_code_mode_protocol::FunctionCallOutputContentItem::InputText { text } => {
FunctionCallOutputContentItem::InputText { text }
}
codex_code_mode::FunctionCallOutputContentItem::InputImage { image_url, detail } => {
FunctionCallOutputContentItem::InputImage {
image_url,
detail: detail
.map(IntoProtocol::into_protocol)
.or(Some(DEFAULT_IMAGE_DETAIL)),
}
}
codex_code_mode_protocol::FunctionCallOutputContentItem::InputImage {
image_url,
detail,
} => FunctionCallOutputContentItem::InputImage {
image_url,
detail: detail
.map(IntoProtocol::into_protocol)
.or(Some(DEFAULT_IMAGE_DETAIL)),
},
}
}
}
codex-rs/core/src/tools/code_mode/wait_handler.rs
+14 -7
View File
@@ -78,7 +78,7 @@ impl CodeModeWaitHandler {
let args: ExecWaitArgs = parse_arguments(&arguments)?;
let exec = ExecContext { session, turn };
let started_at = std::time::Instant::now();
let cell_id = codex_code_mode::CellId::new(args.cell_id);
let cell_id = codex_code_mode_protocol::CellId::new(args.cell_id);
let wait_response = if args.terminate {
exec.session
.services
@@ -89,23 +89,30 @@ impl CodeModeWaitHandler {
exec.session
.services
.code_mode_service
.wait(codex_code_mode::WaitRequest {
.wait(codex_code_mode_protocol::WaitRequest {
cell_id,
yield_time_ms: args.yield_time_ms,
})
.await
}
.map_err(FunctionCallError::RespondToModel)?;
if let codex_code_mode::WaitOutcome::LiveCell(response) = &wait_response
&& !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. })
if let codex_code_mode_protocol::WaitOutcome::LiveCell(response) = &wait_response
&& !matches!(
response,
codex_code_mode_protocol::RuntimeResponse::Yielded { .. }
)
{
// Only a live-cell wait can close a CodeCell. A missing
// cell is still an ordinary `wait` tool result, but there
// is no runtime object for the reducer to complete.
let runtime_cell_id = match response {
codex_code_mode::RuntimeResponse::Yielded { cell_id, .. }
| codex_code_mode::RuntimeResponse::Terminated { cell_id, .. }
| codex_code_mode::RuntimeResponse::Result { cell_id, .. } => cell_id,
codex_code_mode_protocol::RuntimeResponse::Yielded { cell_id, .. }
| codex_code_mode_protocol::RuntimeResponse::Terminated {
cell_id, ..
}
| codex_code_mode_protocol::RuntimeResponse::Result { cell_id, .. } => {
cell_id
}
};
exec.session
.services
codex-rs/core/src/tools/code_mode/wait_spec.rs
+6 -6
View File
@@ -30,11 +30,11 @@ pub(crate) fn create_wait_tool() -> ToolSpec {
]);
ToolSpec::Function(ResponsesApiTool {
name: codex_code_mode::WAIT_TOOL_NAME.to_string(),
name: codex_code_mode_protocol::WAIT_TOOL_NAME.to_string(),
description: format!(
"Waits on a yielded `{}` cell and returns new output or completion.\n{}",
codex_code_mode::PUBLIC_TOOL_NAME,
codex_code_mode::build_wait_tool_description().trim()
codex_code_mode_protocol::PUBLIC_TOOL_NAME,
codex_code_mode_protocol::build_wait_tool_description().trim()
),
strict: false,
parameters: JsonSchema::object(
@@ -57,11 +57,11 @@ mod tests {
assert_eq!(
create_wait_tool(),
ToolSpec::Function(ResponsesApiTool {
name: codex_code_mode::WAIT_TOOL_NAME.to_string(),
name: codex_code_mode_protocol::WAIT_TOOL_NAME.to_string(),
description: format!(
"Waits on a yielded `{}` cell and returns new output or completion.\n{}",
codex_code_mode::PUBLIC_TOOL_NAME,
codex_code_mode::build_wait_tool_description().trim()
codex_code_mode_protocol::PUBLIC_TOOL_NAME,
codex_code_mode_protocol::build_wait_tool_description().trim()
),
strict: false,
defer_loading: None,
codex-rs/core/src/tools/spec_plan.rs
+16 -13
View File
@@ -242,7 +242,7 @@ fn spec_for_model_request(
ToolMode::CodeMode | ToolMode::CodeModeOnly
) && exposure != ToolExposure::DirectModelOnly
&& !is_excluded_from_code_mode(turn_context, tool_name)
&& codex_code_mode::is_code_mode_nested_tool(spec.name())
&& codex_code_mode_protocol::is_code_mode_nested_tool(spec.name())
{
codex_tools::augment_tool_spec_for_code_mode(spec)
} else {
@@ -408,9 +408,9 @@ fn is_hidden_by_code_mode_only(
) -> bool {
turn_context.tool_mode == ToolMode::CodeModeOnly
&& exposure != ToolExposure::DirectModelOnly
&& codex_code_mode::is_code_mode_nested_tool(&codex_tools::code_mode_name_for_tool_name(
tool_name,
))
&& codex_code_mode_protocol::is_code_mode_nested_tool(
&codex_tools::code_mode_name_for_tool_name(tool_name),
)
}
fn is_excluded_from_code_mode(turn_context: &TurnContext, tool_name: &ToolName) -> bool {
@@ -533,7 +533,7 @@ fn merge_into_namespaces(specs: Vec<ToolSpec>) -> Vec<ToolSpec> {
fn code_mode_namespace_descriptions(
specs: &[ToolSpec],
) -> BTreeMap<String, codex_code_mode::ToolNamespaceDescription> {
) -> BTreeMap<String, codex_code_mode_protocol::ToolNamespaceDescription> {
let mut namespace_descriptions = BTreeMap::new();
for spec in specs {
let ToolSpec::Namespace(namespace) = spec else {
@@ -542,7 +542,7 @@ fn code_mode_namespace_descriptions(
let entry = namespace_descriptions
.entry(namespace.name.clone())
.or_insert_with(|| codex_code_mode::ToolNamespaceDescription {
.or_insert_with(|| codex_code_mode_protocol::ToolNamespaceDescription {
name: namespace.name.clone(),
description: namespace.description.clone(),
});
@@ -883,8 +883,8 @@ fn append_extension_tool_executors(
turn_context.tool_mode,
ToolMode::CodeMode | ToolMode::CodeModeOnly
) {
reserved_tool_names.insert(ToolName::plain(codex_code_mode::PUBLIC_TOOL_NAME));
reserved_tool_names.insert(ToolName::plain(codex_code_mode::WAIT_TOOL_NAME));
reserved_tool_names.insert(ToolName::plain(codex_code_mode_protocol::PUBLIC_TOOL_NAME));
reserved_tool_names.insert(ToolName::plain(codex_code_mode_protocol::WAIT_TOOL_NAME));
}
if search_tool_enabled(turn_context)
&& namespace_tools_enabled(turn_context)
@@ -983,9 +983,9 @@ impl CoreToolRuntime for MultiAgentV2NamespaceOverride {
}
fn compare_code_mode_tools(
left: &codex_code_mode::ToolDefinition,
right: &codex_code_mode::ToolDefinition,
namespace_descriptions: &BTreeMap<String, codex_code_mode::ToolNamespaceDescription>,
left: &codex_code_mode_protocol::ToolDefinition,
right: &codex_code_mode_protocol::ToolDefinition,
namespace_descriptions: &BTreeMap<String, codex_code_mode_protocol::ToolNamespaceDescription>,
) -> std::cmp::Ordering {
let left_namespace = code_mode_namespace_name(left, namespace_descriptions);
let right_namespace = code_mode_namespace_name(right, namespace_descriptions);
@@ -997,8 +997,11 @@ fn compare_code_mode_tools(
}
fn code_mode_namespace_name<'a>(
tool: &codex_code_mode::ToolDefinition,
namespace_descriptions: &'a BTreeMap<String, codex_code_mode::ToolNamespaceDescription>,
tool: &codex_code_mode_protocol::ToolDefinition,
namespace_descriptions: &'a BTreeMap<
String,
codex_code_mode_protocol::ToolNamespaceDescription,
>,
) -> Option<&'a str> {
tool.tool_name
.namespace
codex-rs/core/src/tools/spec_plan_tests.rs
+10 -9
View File
@@ -876,8 +876,8 @@ async fn code_mode_only_exposes_code_executor_and_hides_nested_tools() {
&["lookup".to_string()]
);
plain.assert_visible_lacks(&[
codex_code_mode::PUBLIC_TOOL_NAME,
codex_code_mode::WAIT_TOOL_NAME,
codex_code_mode_protocol::PUBLIC_TOOL_NAME,
codex_code_mode_protocol::WAIT_TOOL_NAME,
]);
let code_mode_only = probe_with(
@@ -895,8 +895,8 @@ async fn code_mode_only_exposes_code_executor_and_hides_nested_tools() {
)
.await;
code_mode_only.assert_visible_contains(&[
codex_code_mode::PUBLIC_TOOL_NAME,
codex_code_mode::WAIT_TOOL_NAME,
codex_code_mode_protocol::PUBLIC_TOOL_NAME,
codex_code_mode_protocol::WAIT_TOOL_NAME,
]);
assert_eq!(
code_mode_only.namespace_function_names("codex_app"),
@@ -926,7 +926,8 @@ async fn excluded_deferred_namespaces_do_not_enable_nested_tool_guidance() {
)
.await;
let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode::PUBLIC_TOOL_NAME) else {
let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode_protocol::PUBLIC_TOOL_NAME)
else {
panic!("expected code mode exec tool");
};
assert!(
@@ -1071,8 +1072,8 @@ async fn tool_mode_selector_overrides_feature_flags() {
})
.await;
direct.assert_visible_lacks(&[
codex_code_mode::PUBLIC_TOOL_NAME,
codex_code_mode::WAIT_TOOL_NAME,
codex_code_mode_protocol::PUBLIC_TOOL_NAME,
codex_code_mode_protocol::WAIT_TOOL_NAME,
]);
}
@@ -1307,8 +1308,8 @@ async fn hosted_tools_follow_provider_auth_model_and_config_gates() {
code_mode_only.visible_names,
vec![
// Code-mode entrypoints.
codex_code_mode::PUBLIC_TOOL_NAME,
codex_code_mode::WAIT_TOOL_NAME,
codex_code_mode_protocol::PUBLIC_TOOL_NAME,
codex_code_mode_protocol::WAIT_TOOL_NAME,
// Multi-agent v2 tools.
"spawn_agent",
"send_message",
codex-rs/core/src/tools/tool_dispatch_trace_tests.rs
+76
View File
@@ -3,6 +3,17 @@ use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_code_mode_protocol::CodeModeSessionProviderFuture;
use codex_code_mode_protocol::CodeModeSessionResultFuture;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use codex_protocol::protocol::SessionSource;
use codex_rollout_trace::ExecutionStatus;
use codex_rollout_trace::ThreadStartedTraceMetadata;
@@ -58,6 +69,51 @@ impl ToolExecutor<ToolInvocation> for TestHandler {
impl CoreToolRuntime for TestHandler {}
struct MissingCellSessionProvider;
impl CodeModeSessionProvider for MissingCellSessionProvider {
fn create_session<'a>(
&'a self,
_delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionProviderFuture<'a> {
Box::pin(async {
let session: Arc<dyn CodeModeSession> = Arc::new(MissingCellSession);
Ok(session)
})
}
}
struct MissingCellSession;
impl CodeModeSession for MissingCellSession {
fn execute<'a>(
&'a self,
_request: ExecuteRequest,
) -> CodeModeSessionResultFuture<'a, StartedCell> {
Box::pin(async { Err("execute is not supported by this test session".to_string()) })
}
fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(async move { Ok(missing_cell_outcome(request.cell_id)) })
}
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(async move { Ok(missing_cell_outcome(cell_id)) })
}
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
Box::pin(async { Ok(()) })
}
}
fn missing_cell_outcome(cell_id: CellId) -> WaitOutcome {
WaitOutcome::MissingCell(RuntimeResponse::Result {
error_text: Some(format!("exec cell {cell_id} not found")),
cell_id,
content_items: Vec::new(),
})
}
#[tokio::test]
async fn dispatch_lifecycle_trace_records_direct_and_code_mode_requesters() -> anyhow::Result<()> {
let temp = TempDir::new()?;
@@ -213,6 +269,26 @@ async fn dispatch_lifecycle_trace_records_incompatible_payload_failures() -> any
async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Result<()> {
let temp = TempDir::new()?;
let (mut session, turn) = make_session_and_context().await;
let provider: Arc<dyn CodeModeSessionProvider> = Arc::new(MissingCellSessionProvider);
session.services.code_mode_service =
crate::tools::code_mode::CodeModeService::new(Arc::clone(&provider));
session.services.code_mode_session_provider = provider;
assert_eq!(
session
.services
.code_mode_service
.execute(ExecuteRequest {
tool_call_id: "setup-call".to_string(),
enabled_tools: Vec::new(),
source: String::new(),
yield_time_ms: None,
max_output_tokens: None,
})
.await
.err()
.as_deref(),
Some("execute is not supported by this test session")
);
attach_test_trace(&mut session, &turn, temp.path())?;
let registry = ToolRegistry::with_handler_for_test(Arc::new(CodeModeWaitHandler));
codex-rs/core/tests/common/Cargo.toml
+1
View File
@@ -17,6 +17,7 @@ assert_cmd = { workspace = true }
base64 = { workspace = true }
codex-arg0 = { workspace = true }
codex-config = { workspace = true }
codex-code-mode-client = { workspace = true }
codex-core = { workspace = true }
codex-extension-api = { workspace = true }
codex-exec-server = { workspace = true }
codex-rs/core/tests/common/test_codex.rs
+12
View File
@@ -13,6 +13,8 @@ use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
use anyhow::anyhow;
use codex_code_mode_client::CodeModeHostCommand;
use codex_code_mode_client::IpcCodeModeSessionProvider;
use codex_config::CloudConfigBundleLoader;
use codex_core::CodexThread;
use codex_core::ThreadManager;
@@ -567,6 +569,16 @@ impl TestCodexBuilder {
installation_id,
/*attestation_provider*/ None,
);
let code_mode_session_provider =
match codex_utils_cargo_bin::cargo_bin("codex-code-mode-host") {
Ok(program) => IpcCodeModeSessionProvider::new(CodeModeHostCommand {
program,
args: Vec::new(),
}),
Err(_) => IpcCodeModeSessionProvider::default(),
};
let thread_manager =
thread_manager.with_code_mode_session_provider(Arc::new(code_mode_session_provider));
let thread_manager = Arc::new(thread_manager);
let user_shell_override = self.user_shell_override.clone();
codex-rs/core/tests/suite/code_mode.rs
+4 -11
View File
@@ -1921,7 +1921,7 @@ text("after terminate");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_wait_returns_error_for_unknown_session() -> Result<()> {
async fn code_mode_wait_returns_error_without_a_session() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
@@ -1964,17 +1964,10 @@ async fn code_mode_wait_returns_error_for_unknown_session() -> Result<()> {
assert_ne!(success, Some(true));
let items = function_tool_output_items(&request, "call-1");
assert_eq!(items.len(), 2);
assert_regex_match(
concat!(
r"(?s)\A",
r"Script failed\nWall time \d+\.\d seconds\nOutput:\n\z"
),
text_item(&items, /*index*/ 0),
);
assert_eq!(items.len(), 1);
assert_eq!(
text_item(&items, /*index*/ 1),
"Script error:\nexec cell 999999 not found"
text_item(&items, /*index*/ 0),
"code mode session is unavailable"
);
Ok(())
codex-rs/core/tests/suite/model_runtime_selectors.rs
+4 -4
View File
@@ -158,8 +158,8 @@ async fn remote_tool_mode_selector_overrides_feature_flags() -> Result<()> {
assert!(
direct_tools
.iter()
.all(|name| name != codex_code_mode::PUBLIC_TOOL_NAME
&& name != codex_code_mode::WAIT_TOOL_NAME),
.all(|name| name != codex_code_mode_protocol::PUBLIC_TOOL_NAME
&& name != codex_code_mode_protocol::WAIT_TOOL_NAME),
"direct mode should override enabled code mode flags: {direct_tools:?}"
);
@@ -171,8 +171,8 @@ async fn remote_tool_mode_selector_overrides_feature_flags() -> Result<()> {
tool_names(&code_mode_only_body),
vec![
// Code-mode entrypoints.
codex_code_mode::PUBLIC_TOOL_NAME.to_string(),
codex_code_mode::WAIT_TOOL_NAME.to_string(),
codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(),
codex_code_mode_protocol::WAIT_TOOL_NAME.to_string(),
// Hosted Responses tools.
"web_search".to_string(),
"image_generation".to_string(),