clone/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-04-23 22:24:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: inject code mode runtime into codex sessions

Browse Source 
Extract a CodeModeRuntime trait from codex-code-mode and make the core code-mode wrapper hold an injected runtime instead of constructing the native implementation inline. Thread an optional runtime override through Codex spawn/session startup and expose a ThreadManager entrypoint that can start a thread with a caller-provided runtime. Re-export the runtime trait from codex-core so downstream browser harness code can consume the seam without reaching into transitive dependencies.
This commit is contained in:
Jeremy Lewi
2026-04-13 09:40:19 -07:00
parent fa5d14a81f
commit 3ffd99cd3d
8 changed files with 99 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
codex-rs/code-mode/src/lib.rs
View File

@@ -22,9 +22,11 @@ pub use runtime::DEFAULT_WAIT_YIELD_TIME_MS;
pub use runtime::ExecuteRequest;
pub use runtime::RuntimeResponse;
pub use runtime::WaitRequest;
pub use service::CodeModeRuntime;
pub use service::CodeModeService;
pub use service::CodeModeTurnHost;
pub use service::CodeModeTurnWorker;
pub use service::CodeModeTurnWorkerHandle;
pub const PUBLIC_TOOL_NAME: &str = "exec";
pub const WAIT_TOOL_NAME: &str = "wait";

46
codex-rs/code-mode/src/service.rs
View File

@@ -34,6 +34,22 @@ pub trait CodeModeTurnHost: Send + Sync {
async fn notify(&self, call_id: String, cell_id: String, text: String) -> Result<(), String>;
}
#[async_trait]
pub trait CodeModeRuntime: Send + Sync {
async fn stored_values(&self) -> HashMap<String, JsonValue>;
async fn replace_stored_values(&self, values: HashMap<String, JsonValue>);
async fn execute(&self, request: ExecuteRequest) -> Result<RuntimeResponse, String>;
async fn wait(&self, request: WaitRequest) -> Result<RuntimeResponse, String>;
fn start_turn_worker(
&self,
host: Arc<dyn CodeModeTurnHost>,
) -> Box<dyn CodeModeTurnWorkerHandle>;
}
#[derive(Clone)]
struct SessionHandle {
control_tx: mpsc::UnboundedSender<SessionControlCommand>,
@@ -219,6 +235,10 @@ pub struct CodeModeTurnWorker {
shutdown_tx: Option<oneshot::Sender<()>>,
}
pub trait CodeModeTurnWorkerHandle: Send {}
impl CodeModeTurnWorkerHandle for CodeModeTurnWorker {}
impl Drop for CodeModeTurnWorker {
fn drop(&mut self) {
if let Some(shutdown_tx) = self.shutdown_tx.take() {
@@ -227,6 +247,32 @@ impl Drop for CodeModeTurnWorker {
}
}
#[async_trait]
impl CodeModeRuntime for CodeModeService {
async fn stored_values(&self) -> HashMap<String, JsonValue> {
CodeModeService::stored_values(self).await
}
async fn replace_stored_values(&self, values: HashMap<String, JsonValue>) {
CodeModeService::replace_stored_values(self, values).await;
}
async fn execute(&self, request: ExecuteRequest) -> Result<RuntimeResponse, String> {
CodeModeService::execute(self, request).await
}
async fn wait(&self, request: WaitRequest) -> Result<RuntimeResponse, String> {
CodeModeService::wait(self, request).await
}
fn start_turn_worker(
&self,
host: Arc<dyn CodeModeTurnHost>,
) -> Box<dyn CodeModeTurnWorkerHandle> {
Box::new(CodeModeService::start_turn_worker(self, host))
}
}
enum SessionControlCommand {
Poll {
yield_time_ms: u64,

9
codex-rs/core/src/codex.rs
View File

@@ -422,6 +422,7 @@ pub(crate) struct CodexSpawnArgs {
pub(crate) inherited_shell_snapshot: Option<Arc<ShellSnapshot>>,
pub(crate) inherited_exec_policy: Option<Arc<ExecPolicyManager>>,
pub(crate) user_shell_override: Option<shell::Shell>,
pub(crate) code_mode_runtime: Option<Arc<dyn codex_code_mode::CodeModeRuntime>>,
pub(crate) parent_trace: Option<W3cTraceContext>,
}
@@ -476,6 +477,7 @@ impl Codex {
inherited_shell_snapshot,
user_shell_override,
inherited_exec_policy,
code_mode_runtime,
parent_trace: _,
} = args;
let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
@@ -660,6 +662,7 @@ impl Codex {
mcp_manager.clone(),
skills_watcher,
agent_control,
code_mode_runtime,
)
.await
.map_err(|e| {
@@ -1478,6 +1481,7 @@ impl Session {
mcp_manager: Arc<McpManager>,
skills_watcher: Arc<SkillsWatcher>,
agent_control: AgentControl,
code_mode_runtime: Option<Arc<dyn codex_code_mode::CodeModeRuntime>>,
) -> anyhow::Result<Arc<Self>> {
debug!(
"Configuring session: model={}; provider={:?}",
@@ -1903,8 +1907,9 @@ impl Session {
config.features.enabled(Feature::RuntimeMetrics),
Self::build_model_client_beta_features_header(config.as_ref()),
),
code_mode_service: crate::tools::code_mode::CodeModeService::new(
config.js_repl_node_path.clone(),
code_mode_service: code_mode_runtime.map_or_else(
|| crate::tools::code_mode::CodeModeService::new(config.js_repl_node_path.clone()),
crate::tools::code_mode::CodeModeService::from_runtime,
),
environment: environment_manager.current().await?,
};

1
codex-rs/core/src/codex_delegate.rs
View File

@@ -93,6 +93,7 @@ pub(crate) async fn run_codex_thread_interactive(
inherited_shell_snapshot: None,
user_shell_override: None,
inherited_exec_policy: Some(Arc::clone(&parent_session.services.exec_policy)),
code_mode_runtime: None,
parent_trace: None,
})
.await?;

1
codex-rs/core/src/codex_tests_guardian.rs
View File

@@ -456,6 +456,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
inherited_shell_snapshot: None,
inherited_exec_policy: Some(Arc::new(parent_exec_policy)),
user_shell_override: None,
code_mode_runtime: None,
parent_trace: None,
})
.await

1
codex-rs/core/src/lib.rs
View File

@@ -201,6 +201,7 @@ pub use client_common::Prompt;
pub use client_common::REVIEW_PROMPT;
pub use client_common::ResponseEvent;
pub use client_common::ResponseStream;
pub use codex_code_mode::CodeModeRuntime;
pub use codex_sandboxing::get_platform_sandbox;
pub use codex_tools::parse_tool_input_schema;
pub use compact::content_items_to_text;

32
codex-rs/core/src/thread_manager.rs
View File

@@ -447,6 +447,27 @@ impl ThreadManager {
persist_extended_history,
metrics_service_name,
parent_trace,
/*code_mode_runtime*/ None,
/*user_shell_override*/ None,
))
.await
}
pub async fn start_thread_with_code_mode_runtime(
&self,
config: Config,
code_mode_runtime: Arc<dyn crate::CodeModeRuntime>,
) -> CodexResult<NewThread> {
Box::pin(self.state.spawn_thread(
config,
InitialHistory::New,
Arc::clone(&self.state.auth_manager),
self.agent_control(),
Vec::new(),
/*persist_extended_history*/ false,
/*metrics_service_name*/ None,
/*parent_trace*/ None,
Some(code_mode_runtime),
/*user_shell_override*/ None,
))
.await
@@ -487,6 +508,7 @@ impl ThreadManager {
persist_extended_history,
/*metrics_service_name*/ None,
parent_trace,
/*code_mode_runtime*/ None,
/*user_shell_override*/ None,
))
.await
@@ -506,6 +528,7 @@ impl ThreadManager {
/*persist_extended_history*/ false,
/*metrics_service_name*/ None,
/*parent_trace*/ None,
/*code_mode_runtime*/ None,
/*user_shell_override*/ Some(user_shell_override),
))
.await
@@ -528,6 +551,7 @@ impl ThreadManager {
/*persist_extended_history*/ false,
/*metrics_service_name*/ None,
/*parent_trace*/ None,
/*code_mode_runtime*/ None,
/*user_shell_override*/ Some(user_shell_override),
))
.await
@@ -635,6 +659,7 @@ impl ThreadManager {
persist_extended_history,
/*metrics_service_name*/ None,
parent_trace,
/*code_mode_runtime*/ None,
/*user_shell_override*/ None,
))
.await
@@ -736,6 +761,7 @@ impl ThreadManagerState {
inherited_shell_snapshot,
inherited_exec_policy,
/*parent_trace*/ None,
/*code_mode_runtime*/ None,
/*user_shell_override*/ None,
))
.await
@@ -763,6 +789,7 @@ impl ThreadManagerState {
inherited_shell_snapshot,
inherited_exec_policy,
/*parent_trace*/ None,
/*code_mode_runtime*/ None,
/*user_shell_override*/ None,
))
.await
@@ -791,6 +818,7 @@ impl ThreadManagerState {
inherited_shell_snapshot,
inherited_exec_policy,
/*parent_trace*/ None,
/*code_mode_runtime*/ None,
/*user_shell_override*/ None,
))
.await
@@ -808,6 +836,7 @@ impl ThreadManagerState {
persist_extended_history: bool,
metrics_service_name: Option<String>,
parent_trace: Option<W3cTraceContext>,
code_mode_runtime: Option<Arc<dyn codex_code_mode::CodeModeRuntime>>,
user_shell_override: Option<crate::shell::Shell>,
) -> CodexResult<NewThread> {
Box::pin(self.spawn_thread_with_source(
@@ -822,6 +851,7 @@ impl ThreadManagerState {
/*inherited_shell_snapshot*/ None,
/*inherited_exec_policy*/ None,
parent_trace,
code_mode_runtime,
user_shell_override,
))
.await
@@ -841,6 +871,7 @@ impl ThreadManagerState {
inherited_shell_snapshot: Option<Arc<ShellSnapshot>>,
inherited_exec_policy: Option<Arc<crate::exec_policy::ExecPolicyManager>>,
parent_trace: Option<W3cTraceContext>,
code_mode_runtime: Option<Arc<dyn codex_code_mode::CodeModeRuntime>>,
user_shell_override: Option<crate::shell::Shell>,
) -> CodexResult<NewThread> {
let watch_registration = self.skills_watcher.register_config(
@@ -868,6 +899,7 @@ impl ThreadManagerState {
inherited_shell_snapshot,
inherited_exec_policy,
user_shell_override,
code_mode_runtime,
parent_trace,
})
.await?;

14
codex-rs/core/src/tools/code_mode/mod.rs
View File

@@ -6,7 +6,9 @@ use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode::CodeModeRuntime;
use codex_code_mode::CodeModeTurnHost;
use codex_code_mode::CodeModeTurnWorkerHandle;
use codex_code_mode::RuntimeResponse;
use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::FunctionCallOutputPayload;
@@ -48,14 +50,16 @@ pub(crate) struct ExecContext {
}
pub(crate) struct CodeModeService {
inner: codex_code_mode::CodeModeService,
inner: Arc<dyn CodeModeRuntime>,
}
impl CodeModeService {
pub(crate) fn new(_js_repl_node_path: Option<PathBuf>) -> Self {
Self {
inner: codex_code_mode::CodeModeService::new(),
}
Self::from_runtime(Arc::new(codex_code_mode::CodeModeService::new()))
}
pub(crate) fn from_runtime(inner: Arc<dyn CodeModeRuntime>) -> Self {
Self { inner }
}
pub(crate) async fn stored_values(&self) -> std::collections::HashMap<String, JsonValue> {
@@ -89,7 +93,7 @@ impl CodeModeService {
turn: &Arc<TurnContext>,
router: Arc<ToolRouter>,
tracker: SharedTurnDiffTracker,
) -> Option<codex_code_mode::CodeModeTurnWorker> {
) -> Option<Box<dyn CodeModeTurnWorkerHandle>> {
if !turn.features.enabled(Feature::CodeMode) {
return None;
}