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

code-mode: cut production over to stdio host

Browse Source
This commit is contained in:
Channing Conger
2026-06-21 08:35:01 +00:00
parent 60dad3fc7e
commit afef87eb07
40 changed files with 598 additions and 2866 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+5 -3
View File
@@ -2641,7 +2641,8 @@ dependencies = [
"codex-app-server-protocol",
"codex-apply-patch",
"codex-async-utils",
"codex-code-mode",
"codex-code-mode-client",
"codex-code-mode-protocol",
"codex-config",
"codex-connectors",
"codex-context-fragments",
@@ -3800,7 +3801,7 @@ name = "codex-rollout-trace"
version = "0.0.0"
dependencies = [
"anyhow",
"codex-code-mode",
"codex-code-mode-protocol",
"codex-protocol",
"http 1.4.0",
"pretty_assertions",
@@ -4016,7 +4017,7 @@ name = "codex-tools"
version = "0.0.0"
dependencies = [
"codex-app-server-protocol",
"codex-code-mode",
"codex-code-mode-protocol",
"codex-features",
"codex-file-system",
"codex-protocol",
@@ -4675,6 +4676,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
@@ -5,6 +5,7 @@ codex_rust_crate(
crate_name = "codex_app_server",
extra_binaries = [
"//codex-rs/bwrap:bwrap",
"//codex-rs/code-mode-host:codex-code-mode-host",
],
extra_binaries_non_windows = [
"//codex-rs/cli:codex",
codex-rs/app-server/tests/common/test_app_server.rs
+3
View File
@@ -224,6 +224,9 @@ impl TestAppServer {
cmd.current_dir(codex_home);
cmd.env("CODEX_HOME", codex_home);
cmd.env("RUST_LOG", "warn");
let code_mode_host = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")
.context("should find binary for codex-code-mode-host")?;
cmd.env("CODEX_CODE_MODE_HOST_PATH", code_mode_host);
// Keep integration tests isolated from host managed configuration.
cmd.env(
"CODEX_APP_SERVER_MANAGED_CONFIG_PATH",
codex-rs/code-mode-host/tests/provider_contract.rs
+21 -28
View File
@@ -3,23 +3,22 @@
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode::CellId;
use codex_code_mode::CellOutcome;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeSession;
use codex_code_mode::CodeModeSessionDelegate;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::CreateCellRequest;
use codex_code_mode::FunctionCallOutputContentItem;
use codex_code_mode::InProcessCodeModeSessionProvider;
use codex_code_mode::NotificationFuture;
use codex_code_mode::ObserveRequest;
use codex_code_mode::RuntimeResponse;
use codex_code_mode::ToolDefinition;
use codex_code_mode::ToolInvocationFuture;
use codex_code_mode_client::CodeModeHostCommand;
use codex_code_mode_client::StdioCodeModeSessionProvider;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CellOutcome;
use codex_code_mode_protocol::CodeModeNestedToolCall;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_code_mode_protocol::CodeModeToolKind;
use codex_code_mode_protocol::CreateCellRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
use codex_code_mode_protocol::NotificationFuture;
use codex_code_mode_protocol::ObserveRequest;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::ToolDefinition;
use codex_code_mode_protocol::ToolInvocationFuture;
use codex_protocol::ToolName;
use pretty_assertions::assert_eq;
use tokio::sync::mpsc;
@@ -33,19 +32,13 @@ struct ProviderCase {
fn provider_cases() -> Vec<ProviderCase> {
let host_program =
codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("resolve host binary");
vec![
ProviderCase {
name: "in-process",
provider: Arc::new(InProcessCodeModeSessionProvider),
},
ProviderCase {
name: "stdio",
provider: Arc::new(StdioCodeModeSessionProvider::new(CodeModeHostCommand {
program: host_program,
args: Vec::new(),
})),
},
]
vec![ProviderCase {
name: "stdio",
provider: Arc::new(StdioCodeModeSessionProvider::new(CodeModeHostCommand {
program: host_program,
args: Vec::new(),
})),
}]
}
#[derive(Debug, PartialEq)]
codex-rs/code-mode/src/cell_actor/tests.rs
+35 -1
View File
@@ -119,7 +119,6 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(host: Arc<H>) -> CellActorHar
source: "await new Promise(() => {});".to_string(),
},
runtime_event_tx,
PendingRuntimeMode::PauseUntilResumed,
)
.unwrap();
let (runtime_control_tx, runtime_control_rx) = std_mpsc::channel();
@@ -1102,3 +1101,38 @@ fn dropped_pending_observation_preserves_the_initial_yield_boundary() {
));
assert_eq!(response_rx.try_recv(), Ok(Ok(completion)));
}
#[tokio::test]
async fn repeated_termination_is_rejected_before_the_first_response() {
let harness = spawn_cell_actor_harness();
let initial_observation = harness
.handle
.observe(ObserveMode::YieldAfter(Duration::from_secs(60)));
assert_eq!(
harness
.handle
.observe(ObserveMode::YieldAfter(Duration::from_secs(60)))
.await,
Err(CellError::Busy)
);
let first_termination = harness.handle.terminate();
assert_eq!(
harness.handle.terminate().await,
Err(CellError::AlreadyTerminating)
);
harness
.event_tx
.send(RuntimeEvent::Result {
stored_value_writes: HashMap::new(),
error_text: None,
})
.unwrap();
let terminated = Ok(CellEvent::Terminated {
content_items: Vec::new(),
});
assert_eq!(first_termination.await, terminated.clone());
assert_eq!(initial_observation.await, terminated);
harness.task.await.unwrap();
}
codex-rs/code-mode/src/lib.rs
-5
View File
@@ -1,10 +1,5 @@
mod cell_actor;
mod runtime;
mod service;
pub mod session_runtime;
pub use codex_code_mode_protocol::*;
pub use service::CodeModeService;
pub use service::InProcessCodeModeSessionProvider;
pub use service::NoopCodeModeSessionDelegate;
pub use session_runtime::SessionRuntime;
codex-rs/code-mode/src/runtime/mod.rs
+1 -1
View File
@@ -318,7 +318,7 @@ mod tests {
use super::RuntimeControlCommand;
use super::RuntimeEvent;
use super::spawn_runtime;
use crate::FunctionCallOutputContentItem;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
fn execute_request(source: &str) -> CreateCellRequest {
CreateCellRequest {
codex-rs/code-mode/src/service.rs
-369
View File
@@ -1,369 +0,0 @@
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CellOutcome;
use codex_code_mode_protocol::CodeModeNestedToolCall;
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::CreateCellRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
use codex_code_mode_protocol::ImageDetail;
use codex_code_mode_protocol::NotificationFuture;
use codex_code_mode_protocol::ObserveRequest;
use codex_code_mode_protocol::ObserveToPendingOutcome;
use codex_code_mode_protocol::ObserveToPendingRequest;
use codex_code_mode_protocol::PendingOutcome;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::ToolInvocationFuture;
use serde_json::Value as JsonValue;
use tokio_util::sync::CancellationToken;
use crate::session_runtime as runtime;
use crate::session_runtime::SessionRuntime;
pub struct NoopCodeModeSessionDelegate;
impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async move {
cancellation_token.cancelled().await;
Err("code mode nested tools are unavailable".to_string())
})
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
#[derive(Default)]
pub struct InProcessCodeModeSessionProvider;
impl CodeModeSessionProvider for InProcessCodeModeSessionProvider {
fn create_session<'a>(
&'a self,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionProviderFuture<'a> {
Box::pin(async move {
let session: Arc<dyn CodeModeSession> =
Arc::new(CodeModeService::with_delegate(delegate));
Ok(session)
})
}
}
pub struct CodeModeService {
runtime: SessionRuntime<ProtocolDelegate>,
}
impl CodeModeService {
pub fn new() -> Self {
Self::with_delegate(Arc::new(NoopCodeModeSessionDelegate))
}
pub fn with_delegate(delegate: Arc<dyn CodeModeSessionDelegate>) -> Self {
#[cfg(not(test))]
let runtime = SessionRuntime::new(Arc::new(ProtocolDelegate { delegate }));
#[cfg(test)]
let runtime = SessionRuntime::new_for_test(Arc::new(ProtocolDelegate { delegate }));
Self { runtime }
}
pub async fn create_cell(&self, request: CreateCellRequest) -> Result<CellId, String> {
let runtime_cell_id = self
.runtime
.create_cell(runtime_request(request))
.await
.map_err(|error| error.to_string())?;
Ok(protocol_cell_id(&runtime_cell_id))
}
pub async fn observe(&self, request: ObserveRequest) -> Result<CellOutcome, String> {
self.begin_observe(request).await.await
}
async fn begin_observe(
&self,
request: ObserveRequest,
) -> CodeModeSessionResultFuture<'static, CellOutcome> {
let ObserveRequest {
cell_id,
yield_time_ms,
} = request;
let runtime_cell_id = runtime_cell_id(&cell_id);
match self
.runtime
.begin_observe(
&runtime_cell_id,
runtime::ObserveMode::YieldAfter(Duration::from_millis(yield_time_ms)),
)
.await
{
Ok(pending_event) => Box::pin(async move {
match pending_event.event().await {
Ok(event) => Ok(CellOutcome::LiveCell(runtime_response(&cell_id, event)?)),
Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => {
Ok(CellOutcome::MissingCell(missing_cell_response(cell_id)))
}
Err(error) => Err(error.to_string()),
}
}),
Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => {
missing_observation(cell_id)
}
Err(error) => Box::pin(async move { Err(error.to_string()) }),
}
}
pub async fn terminate(&self, cell_id: CellId) -> Result<CellOutcome, String> {
match self.runtime.terminate(&runtime_cell_id(&cell_id)).await {
Ok(event) => Ok(CellOutcome::LiveCell(runtime_response(&cell_id, event)?)),
Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => {
Ok(CellOutcome::MissingCell(missing_cell_response(cell_id)))
}
Err(error) => Err(error.to_string()),
}
}
pub async fn observe_to_pending(
&self,
request: ObserveToPendingRequest,
) -> Result<ObserveToPendingOutcome, String> {
let cell_id = request.cell_id;
match self
.runtime
.observe(
&runtime_cell_id(&cell_id),
runtime::ObserveMode::PendingFrontier,
)
.await
{
Ok(event) => Ok(ObserveToPendingOutcome::LiveCell(pending_outcome(
&cell_id, event,
)?)),
Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => Ok(
ObserveToPendingOutcome::MissingCell(missing_cell_response(cell_id)),
),
Err(error) => Err(error.to_string()),
}
}
pub async fn shutdown(&self) -> Result<(), String> {
self.runtime
.shutdown()
.await
.map_err(|error| error.to_string())
}
}
impl Default for CodeModeService {
fn default() -> Self {
Self::new()
}
}
impl CodeModeSession for CodeModeService {
fn is_alive(&self) -> bool {
self.runtime.is_alive()
}
fn create_cell<'a>(
&'a self,
request: CreateCellRequest,
) -> CodeModeSessionResultFuture<'a, CellId> {
Box::pin(CodeModeService::create_cell(self, request))
}
fn observe<'a>(
&'a self,
request: ObserveRequest,
) -> CodeModeSessionResultFuture<'a, CellOutcome> {
Box::pin(CodeModeService::observe(self, request))
}
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, CellOutcome> {
Box::pin(CodeModeService::terminate(self, cell_id))
}
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
Box::pin(CodeModeService::shutdown(self))
}
}
struct ProtocolDelegate {
delegate: Arc<dyn CodeModeSessionDelegate>,
}
impl runtime::SessionRuntimeDelegate for ProtocolDelegate {
async fn invoke_tool(
&self,
invocation: runtime::NestedToolCall,
cancellation_token: CancellationToken,
) -> Result<JsonValue, String> {
self.delegate
.invoke_tool(
CodeModeNestedToolCall {
cell_id: protocol_cell_id(&invocation.cell_id),
runtime_tool_call_id: invocation.runtime_tool_call_id,
tool_name: codex_protocol::ToolName {
name: invocation.tool_name.name,
namespace: invocation.tool_name.namespace,
},
tool_kind: match invocation.tool_kind {
runtime::ToolKind::Function => CodeModeToolKind::Function,
runtime::ToolKind::Freeform => CodeModeToolKind::Freeform,
},
input: invocation.input,
},
cancellation_token,
)
.await
}
async fn notify(
&self,
call_id: String,
cell_id: runtime::CellId,
text: String,
cancellation_token: CancellationToken,
) -> Result<(), String> {
self.delegate
.notify(
call_id,
protocol_cell_id(&cell_id),
text,
cancellation_token,
)
.await
}
fn cell_closed(&self, cell_id: &runtime::CellId) {
self.delegate.cell_closed(&protocol_cell_id(cell_id));
}
}
fn runtime_request(request: CreateCellRequest) -> runtime::CreateCellRequest {
runtime::CreateCellRequest {
tool_call_id: request.tool_call_id,
enabled_tools: request
.enabled_tools
.into_iter()
.map(|definition| runtime::ToolDefinition {
name: definition.name,
tool_name: runtime::ToolName {
name: definition.tool_name.name,
namespace: definition.tool_name.namespace,
},
description: definition.description,
kind: match definition.kind {
CodeModeToolKind::Function => runtime::ToolKind::Function,
CodeModeToolKind::Freeform => runtime::ToolKind::Freeform,
},
})
.collect(),
source: request.source,
}
}
fn runtime_cell_id(cell_id: &CellId) -> runtime::CellId {
runtime::CellId::new(cell_id.as_str())
}
fn protocol_cell_id(cell_id: &runtime::CellId) -> CellId {
CellId::new(cell_id.as_str().to_string())
}
fn pending_outcome(cell_id: &CellId, event: runtime::CellEvent) -> Result<PendingOutcome, String> {
match event {
runtime::CellEvent::Pending {
content_items,
pending_tool_call_ids,
} => Ok(PendingOutcome::Pending {
cell_id: cell_id.clone(),
content_items: content_items.into_iter().map(output_item).collect(),
pending_tool_call_ids,
}),
event => Ok(PendingOutcome::Completed(runtime_response(cell_id, event)?)),
}
}
fn runtime_response(
cell_id: &CellId,
event: runtime::CellEvent,
) -> Result<RuntimeResponse, String> {
match event {
runtime::CellEvent::Yielded { content_items } => Ok(RuntimeResponse::Yielded {
cell_id: cell_id.clone(),
content_items: content_items.into_iter().map(output_item).collect(),
}),
runtime::CellEvent::Completed {
content_items,
error_text,
} => Ok(RuntimeResponse::Result {
cell_id: cell_id.clone(),
content_items: content_items.into_iter().map(output_item).collect(),
error_text,
}),
runtime::CellEvent::Terminated { content_items } => Ok(RuntimeResponse::Terminated {
cell_id: cell_id.clone(),
content_items: content_items.into_iter().map(output_item).collect(),
}),
runtime::CellEvent::Pending { .. } => {
Err("cell returned a pending frontier unexpectedly".to_string())
}
}
}
fn output_item(item: runtime::OutputItem) -> FunctionCallOutputContentItem {
match item {
runtime::OutputItem::Text { text } => FunctionCallOutputContentItem::InputText { text },
runtime::OutputItem::Image { image_url, detail } => {
FunctionCallOutputContentItem::InputImage {
image_url,
detail: detail.map(|detail| match detail {
runtime::ImageDetail::Auto => ImageDetail::Auto,
runtime::ImageDetail::Low => ImageDetail::Low,
runtime::ImageDetail::High => ImageDetail::High,
runtime::ImageDetail::Original => ImageDetail::Original,
}),
}
}
}
}
fn missing_cell_response(cell_id: CellId) -> RuntimeResponse {
RuntimeResponse::Result {
error_text: Some(format!("exec cell {cell_id} not found")),
cell_id,
content_items: Vec::new(),
}
}
fn missing_observation(cell_id: CellId) -> CodeModeSessionResultFuture<'static, CellOutcome> {
Box::pin(async move { Ok(CellOutcome::MissingCell(missing_cell_response(cell_id))) })
}
#[cfg(test)]
#[path = "service_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "service_contract_tests.rs"]
mod contract_tests;
codex-rs/code-mode/src/service_contract_tests.rs
-830
View File
@@ -1,830 +0,0 @@
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_protocol::ToolName;
use pretty_assertions::assert_eq;
use tokio::sync::Notify;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use super::*;
use crate::CodeModeToolKind;
use crate::ToolDefinition;
#[derive(Debug, PartialEq)]
enum DelegateEvent {
NotificationStarted,
NotificationFinished,
ToolStarted,
CellClosed(CellId),
}
fn record_cell_closed(events_tx: &mpsc::UnboundedSender<DelegateEvent>, cell_id: &CellId) {
let _ = events_tx.send(DelegateEvent::CellClosed(cell_id.clone()));
}
struct BlockingDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
tool_future_dropped: AtomicBool,
tool_release: Notify,
}
struct DropFlag<'a>(&'a AtomicBool);
impl Drop for DropFlag<'_> {
fn drop(&mut self) {
self.0.store(true, Ordering::Release);
}
}
struct NeverResolvingNotificationDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
}
struct ReleasableNotificationDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
notification_release: Semaphore,
}
struct NeverResolvingToolDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
}
impl NeverResolvingNotificationDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
(Arc::new(Self { events_tx }), events_rx)
}
}
impl ReleasableNotificationDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
(
Arc::new(Self {
events_tx,
notification_release: Semaphore::new(/*permits*/ 0),
}),
events_rx,
)
}
fn release_notification(&self) {
self.notification_release.add_permits(/*permits*/ 1);
}
}
impl NeverResolvingToolDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
(Arc::new(Self { events_tx }), events_rx)
}
}
impl CodeModeSessionDelegate for NeverResolvingNotificationDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
_cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async { Err("unexpected tool call".to_string()) })
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
std::future::pending().await
})
}
fn cell_closed(&self, cell_id: &CellId) {
record_cell_closed(&self.events_tx, cell_id);
}
}
impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async move {
cancellation_token.cancelled().await;
Err("cancelled".to_string())
})
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
tokio::select! {
_ = self.notification_release.acquire() => {
let _ = self.events_tx.send(DelegateEvent::NotificationFinished);
Ok(())
}
_ = cancellation_token.cancelled() => {
Err("cancelled".to_string())
}
}
})
}
fn cell_closed(&self, cell_id: &CellId) {
record_cell_closed(&self.events_tx, cell_id);
}
}
impl CodeModeSessionDelegate for NeverResolvingToolDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
_cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::ToolStarted);
std::future::pending().await
})
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, cell_id: &CellId) {
record_cell_closed(&self.events_tx, cell_id);
}
}
impl BlockingDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
(
Arc::new(Self {
events_tx,
tool_future_dropped: AtomicBool::new(false),
tool_release: Notify::new(),
}),
events_rx,
)
}
fn release_tool(&self) {
self.tool_release.notify_one();
}
}
impl CodeModeSessionDelegate for BlockingDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
let drop_flag = DropFlag(&self.tool_future_dropped);
Box::pin(async move {
let _drop_flag = drop_flag;
let _ = self.events_tx.send(DelegateEvent::ToolStarted);
tokio::select! {
_ = self.tool_release.notified() => {
Ok(serde_json::Value::Null)
}
_ = cancellation_token.cancelled() => {
Err("cancelled".to_string())
}
}
})
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
cancellation_token.cancelled().await;
Err("cancelled".to_string())
})
}
fn cell_closed(&self, cell_id: &CellId) {
record_cell_closed(&self.events_tx, cell_id);
}
}
fn cell_id(value: &str) -> CellId {
CellId::new(value.to_string())
}
fn execute_request(source: &str) -> CreateCellRequest {
CreateCellRequest {
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
}
}
async fn execute(service: &CodeModeService, request: CreateCellRequest) -> RuntimeResponse {
execute_with_yield_time(service, request, /*yield_time_ms*/ 1).await
}
async fn execute_with_yield_time(
service: &CodeModeService,
request: CreateCellRequest,
yield_time_ms: u64,
) -> RuntimeResponse {
let cell_id = service.create_cell(request).await.unwrap();
service
.observe(ObserveRequest {
cell_id,
yield_time_ms,
})
.await
.unwrap()
.into()
}
async fn create_and_observe_to_pending(
service: &CodeModeService,
request: CreateCellRequest,
) -> Result<PendingOutcome, String> {
let cell_id = service.create_cell(request).await?;
match service
.observe_to_pending(ObserveToPendingRequest { cell_id })
.await?
{
ObserveToPendingOutcome::LiveCell(outcome) => Ok(outcome),
ObserveToPendingOutcome::MissingCell(response) => Ok(PendingOutcome::Completed(response)),
}
}
fn blocking_tool() -> ToolDefinition {
ToolDefinition {
name: "block".to_string(),
tool_name: ToolName::plain("block"),
description: String::new(),
kind: CodeModeToolKind::Function,
input_schema: None,
output_schema: None,
}
}
async fn next_event(events_rx: &mut mpsc::UnboundedReceiver<DelegateEvent>) -> DelegateEvent {
tokio::time::timeout(Duration::from_secs(2), events_rx.recv())
.await
.expect("delegate event timeout")
.expect("delegate event channel closed")
}
#[tokio::test]
async fn yields_and_resumes() {
let service = CodeModeService::new();
let cell = execute_with_yield_time(
&service,
CreateCellRequest {
source: r#"text("before"); yield_control(); text("after");"#.to_string(),
..execute_request("")
},
/*yield_time_ms*/ 60_000,
)
.await;
assert_eq!(
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "before".to_string(),
}],
}
);
assert_eq!(
service
.observe(ObserveRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
})
.await
.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Result {
cell_id: cell_id("1"),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "after".to_string(),
}],
error_text: None,
})
);
}
#[tokio::test]
async fn yield_before_first_observation_preserves_its_output_boundary() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let cell_id = service
.create_cell(CreateCellRequest {
enabled_tools: vec![blocking_tool()],
source: r#"
text("before");
yield_control();
text("after");
await tools.block({});
"#
.to_string(),
..execute_request("")
})
.await
.unwrap();
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
assert_eq!(
service
.observe(ObserveRequest {
cell_id: cell_id.clone(),
yield_time_ms: 60_000,
})
.await
.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Yielded {
cell_id: cell_id.clone(),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "before".to_string(),
}],
})
);
delegate.release_tool();
assert_eq!(
service
.observe(ObserveRequest {
cell_id: cell_id.clone(),
yield_time_ms: 60_000,
})
.await
.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Result {
cell_id,
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "after".to_string(),
}],
error_text: None,
})
);
}
#[tokio::test]
async fn background_completion_notifies_the_delegate_without_another_observation() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate);
let created_cell_id = service
.create_cell(execute_request(
r#"await new Promise(resolve => setTimeout(resolve, 100)); text("done");"#,
))
.await
.unwrap();
assert_eq!(
tokio::time::timeout(
Duration::from_secs(2),
service.observe(ObserveRequest {
cell_id: created_cell_id.clone(),
yield_time_ms: 1,
}),
)
.await
.expect("initial observation should yield while the cell is still running")
.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Yielded {
cell_id: created_cell_id.clone(),
content_items: Vec::new(),
})
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(created_cell_id)
);
}
#[tokio::test]
async fn returns_and_resumes_from_the_pending_frontier() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
assert_eq!(
create_and_observe_to_pending(
&service,
CreateCellRequest {
enabled_tools: vec![blocking_tool()],
source: r#"
await tools.block({});
text("after");
"#
.to_string(),
..execute_request("")
},
)
.await
.unwrap(),
PendingOutcome::Pending {
cell_id: cell_id("1"),
content_items: Vec::new(),
pending_tool_call_ids: vec!["tool-1".to_string()],
}
);
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
delegate.release_tool();
assert_eq!(
service
.observe_to_pending(ObserveToPendingRequest {
cell_id: cell_id("1"),
})
.await
.unwrap(),
ObserveToPendingOutcome::LiveCell(PendingOutcome::Completed(RuntimeResponse::Result {
cell_id: cell_id("1"),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "after".to_string(),
}],
error_text: None,
}))
);
}
#[tokio::test]
async fn observed_natural_completion_wins_over_termination() {
let service = CodeModeService::new();
let cell = execute(
&service,
execute_request(r#"yield_control(); store("finished", true); text("done");"#),
)
.await;
assert_eq!(
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
tokio::time::timeout(Duration::from_secs(10), async {
loop {
// Observe the synchronous probe through completion so scheduler load
// cannot turn its short yield deadline into a false lifecycle failure.
let response = create_and_observe_to_pending(
&service,
CreateCellRequest {
..execute_request(r#"text(String(load("finished")));"#)
},
)
.await
.unwrap();
let PendingOutcome::Completed(RuntimeResponse::Result { content_items, .. }) = response
else {
panic!("expected stored-value probe to complete");
};
if content_items
== vec![FunctionCallOutputContentItem::InputText {
text: "true".to_string(),
}]
{
break;
}
tokio::task::yield_now().await;
}
})
.await
.unwrap();
assert_eq!(
service.terminate(cell_id("1")).await.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Result {
cell_id: cell_id("1"),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "done".to_string(),
}],
error_text: None,
})
);
}
#[tokio::test]
async fn termination_discards_pending_callbacks_before_responding() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let cell = execute(
&service,
execute_request(r#"notify("pending"); await new Promise(() => {});"#),
)
.await;
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationStarted
);
assert_eq!(
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
assert_eq!(
service.terminate(cell_id("1")).await.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Terminated {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
}
#[tokio::test]
async fn termination_discards_stored_writes_before_the_next_cell_can_load_them() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let created_cell_id = service
.create_cell(CreateCellRequest {
enabled_tools: vec![blocking_tool()],
source: r#"
store("candidate", "leaked");
await tools.block({});
"#
.to_string(),
..execute_request("")
})
.await
.unwrap();
// Reaching the delegate proves that the store ran before execution became gated.
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
assert_eq!(
service.terminate(created_cell_id).await.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Terminated {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
delegate.release_tool();
assert_eq!(
execute_with_yield_time(
&service,
execute_request(r#"text(String(load("candidate")));"#),
/*yield_time_ms*/ 60_000,
)
.await,
RuntimeResponse::Result {
cell_id: cell_id("2"),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "undefined".to_string(),
}],
error_text: None,
}
);
}
#[tokio::test]
async fn shutdown_does_not_await_notifications_during_natural_completion() {
let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate));
service
.create_cell(execute_request(r#"notify("pending");"#))
.await
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationStarted
);
let shutdown_service = Arc::clone(&service);
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), async move {
shutdown_service.shutdown().await
})
.await
.expect("shutdown should not await a non-cooperative notification")
.unwrap();
}
#[tokio::test]
async fn shutdown_does_not_await_a_non_cooperative_nested_tool() {
let (delegate, mut events_rx) = NeverResolvingToolDelegate::new();
let service = CodeModeService::with_delegate(delegate);
let _cell_id = service
.create_cell(CreateCellRequest {
enabled_tools: vec![blocking_tool()],
source: r#"await tools.block({});"#.to_string(),
..execute_request("")
})
.await
.unwrap();
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), service.shutdown())
.await
.expect("shutdown should not await a non-cooperative nested tool")
.unwrap();
}
#[tokio::test]
async fn termination_does_not_await_a_non_cooperative_notification() {
let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new();
let service = CodeModeService::with_delegate(delegate);
let cell = execute(
&service,
execute_request(r#"notify("pending"); await new Promise(() => {});"#),
)
.await;
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationStarted
);
assert_eq!(
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
assert_eq!(
tokio::time::timeout(
Duration::from_millis(/*millis*/ 100),
service.terminate(cell_id("1")),
)
.await
.expect("termination should not await a non-cooperative notification")
.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Terminated {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
}
#[tokio::test]
async fn create_cell_returns_before_natural_completion() {
let (delegate, mut events_rx) = ReleasableNotificationDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let created_cell_id = service
.create_cell(execute_request(r#"notify("pending");"#))
.await
.unwrap();
assert_eq!(created_cell_id, cell_id("1"));
let mut observation = Box::pin(service.observe(ObserveRequest {
cell_id: created_cell_id,
yield_time_ms: 60_000,
}));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationStarted
);
std::future::poll_fn(|context| match observation.as_mut().poll(context) {
std::task::Poll::Pending => std::task::Poll::Ready(()),
std::task::Poll::Ready(result) => {
panic!("observation returned while the notification was blocked: {result:?}")
}
})
.await;
delegate.release_notification();
assert_eq!(
observation.await.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Result {
cell_id: cell_id("1"),
content_items: Vec::new(),
error_text: None,
})
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationFinished
);
}
#[tokio::test]
async fn created_cell_can_be_terminated_before_observation() {
let service = CodeModeService::new();
let created_cell_id = service
.create_cell(CreateCellRequest {
source: "await new Promise(() => {});".to_string(),
..execute_request("")
})
.await
.unwrap();
assert_eq!(created_cell_id, cell_id("1"));
assert_eq!(
service.terminate(created_cell_id.clone()).await.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Terminated {
cell_id: created_cell_id,
content_items: Vec::new(),
})
);
}
#[tokio::test]
async fn second_observer_is_rejected_without_displacing_the_first() {
let service = CodeModeService::new();
let cell = execute(&service, execute_request("await new Promise(() => {});")).await;
assert_eq!(
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
let first_observer = service
.begin_observe(ObserveRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
})
.await;
assert_eq!(
service
.observe(ObserveRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
})
.await
.unwrap_err(),
"exec cell 1 already has an active observer"
);
let terminated = RuntimeResponse::Terminated {
cell_id: cell_id("1"),
content_items: Vec::new(),
};
assert_eq!(
service.terminate(cell_id("1")).await.unwrap(),
CellOutcome::LiveCell(terminated.clone())
);
assert_eq!(
first_observer.await.unwrap(),
CellOutcome::LiveCell(terminated)
);
}
#[tokio::test]
async fn natural_completion_cleans_up_callbacks_before_responding() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let created_cell_id = service
.create_cell(CreateCellRequest {
enabled_tools: vec![blocking_tool()],
source: concat!(
"tools.block({});",
"await new Promise(resolve => setTimeout(resolve, 100));",
"text('done');",
)
.to_string(),
..execute_request("")
})
.await
.unwrap();
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
assert_eq!(
service
.observe(ObserveRequest {
cell_id: created_cell_id,
yield_time_ms: 60_000,
})
.await
.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Result {
cell_id: cell_id("1"),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "done".to_string(),
}],
error_text: None,
})
);
assert!(delegate.tool_future_dropped.load(Ordering::Acquire));
}
codex-rs/code-mode/src/service_tests.rs
-1408
View File
File diff suppressed because it is too large Load Diff
codex-rs/code-mode/src/session_runtime/mod.rs
-9
View File
@@ -58,8 +58,6 @@ struct Inner<D: SessionRuntimeDelegate> {
enum CellIdNamespace {
Runtime(uuid::Uuid),
#[cfg(test)]
Unscoped,
}
impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
@@ -67,11 +65,6 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
Self::with_cell_id_namespace(delegate, CellIdNamespace::Runtime(uuid::Uuid::new_v4()))
}
#[cfg(test)]
pub(crate) fn new_for_test(delegate: Arc<D>) -> Self {
Self::with_cell_id_namespace(delegate, CellIdNamespace::Unscoped)
}
fn with_cell_id_namespace(delegate: Arc<D>, cell_id_namespace: CellIdNamespace) -> Self {
Self {
inner: Arc::new(Inner {
@@ -167,8 +160,6 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
}
encoded
}
#[cfg(test)]
CellIdNamespace::Unscoped => sequence.to_string(),
};
CellId::new(value)
}
codex-rs/core/BUILD.bazel
+1
View File
@@ -14,6 +14,7 @@ codex_rust_crate(
crate_name = "codex_core",
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 -1
View File
@@ -28,7 +28,8 @@ codex-api = { workspace = true }
codex-app-server-protocol = { workspace = true }
codex-apply-patch = { workspace = true }
codex-async-utils = { workspace = true }
codex-code-mode = { workspace = true }
codex-code-mode-client = { workspace = true }
codex-code-mode-protocol = { workspace = true }
codex-connectors = { workspace = true }
codex-context-fragments = { workspace = true }
codex-config = { workspace = true }
codex-rs/core/src/codex_delegate.rs
+1
View File
@@ -121,6 +121,7 @@ pub(crate) async fn run_codex_thread_interactive(
thread_store: Arc::clone(&parent_session.services.thread_store),
attestation_provider: parent_session.services.attestation_provider.clone(),
external_time_provider: Some(Arc::clone(&parent_session.services.time_provider)),
code_mode_session_provider: Arc::clone(&parent_session.services.code_mode_session_provider),
inherited_multi_agent_version: Some(MultiAgentVersion::Disabled),
initial_multi_agent_mode: None,
}))
codex-rs/core/src/session/mod.rs
+4
View File
@@ -444,6 +444,8 @@ pub(crate) struct CodexSpawnArgs {
pub(crate) thread_store: Arc<dyn ThreadStore>,
pub(crate) attestation_provider: Option<Arc<dyn AttestationProvider>>,
pub(crate) external_time_provider: Option<Arc<dyn TimeProvider>>,
pub(crate) code_mode_session_provider:
Arc<dyn codex_code_mode_protocol::CodeModeSessionProvider>,
pub(crate) inherited_multi_agent_version: Option<MultiAgentVersion>,
pub(crate) initial_multi_agent_mode: Option<MultiAgentMode>,
}
@@ -529,6 +531,7 @@ impl Codex {
thread_store,
attestation_provider,
external_time_provider,
code_mode_session_provider,
inherited_multi_agent_version,
initial_multi_agent_mode,
} = args;
@@ -680,6 +683,7 @@ impl Codex {
parent_rollout_thread_trace,
attestation_provider,
external_time_provider,
code_mode_session_provider,
multi_agent_version,
))
.await
codex-rs/core/src/session/session.rs
+5 -1
View File
@@ -501,6 +501,7 @@ impl Session {
parent_rollout_thread_trace: ThreadTraceContext,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
external_time_provider: Option<Arc<dyn TimeProvider>>,
code_mode_session_provider: Arc<dyn codex_code_mode_protocol::CodeModeSessionProvider>,
multi_agent_version: Option<MultiAgentVersion>,
) -> anyhow::Result<Arc<Self>> {
debug!(
@@ -1055,7 +1056,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,
tool_search_handler_cache: Default::default(),
turn_environments: Arc::clone(&turn_environments),
};
codex-rs/core/src/session/tests.rs
+15 -2
View File
@@ -4924,6 +4924,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() {
codex_rollout_trace::ThreadTraceContext::disabled(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
Arc::new(codex_code_mode_client::StdioCodeModeSessionProvider::default()),
Some(config.multi_agent_version_from_features()),
)
.await;
@@ -5104,7 +5105,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
/*item_ids_enabled*/ config.features.enabled(Feature::ItemIds),
/*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::StdioCodeModeSessionProvider::default(),
)),
code_mode_session_provider: Arc::new(
codex_code_mode_client::StdioCodeModeSessionProvider::default(),
),
tool_search_handler_cache: Default::default(),
turn_environments: Arc::clone(&turn_environments),
};
@@ -5284,6 +5290,7 @@ async fn make_session_with_config_and_rx(
codex_rollout_trace::ThreadTraceContext::disabled(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
Arc::new(codex_code_mode_client::StdioCodeModeSessionProvider::default()),
Some(config.multi_agent_version_from_features()),
)
.await?;
@@ -5398,6 +5405,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx(
codex_rollout_trace::ThreadTraceContext::disabled(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
Arc::new(codex_code_mode_client::StdioCodeModeSessionProvider::default()),
Some(config.multi_agent_version_from_features()),
)
.await?;
@@ -7163,7 +7171,12 @@ where
/*item_ids_enabled*/ config.features.enabled(Feature::ItemIds),
/*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::StdioCodeModeSessionProvider::default(),
)),
code_mode_session_provider: Arc::new(
codex_code_mode_client::StdioCodeModeSessionProvider::default(),
),
tool_search_handler_cache: Default::default(),
turn_environments: Arc::clone(&turn_environments),
};
codex-rs/core/src/session/tests/guardian_tests.rs
+3
View File
@@ -737,6 +737,9 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
thread_store,
attestation_provider: None,
external_time_provider: None,
code_mode_session_provider: Arc::new(
codex_code_mode_client::StdioCodeModeSessionProvider::default(),
),
inherited_multi_agent_version: None,
initial_multi_agent_mode: None,
})
codex-rs/core/src/state/service.rs
+2
View File
@@ -23,6 +23,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_extension_api::ExtensionData;
use codex_extension_api::ExtensionDataInit;
@@ -84,6 +85,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>,
pub(crate) tool_search_handler_cache: ToolSearchHandlerCache,
pub(crate) turn_environments: Arc<ThreadEnvironments>,
}
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_client::StdioCodeModeSessionProvider;
use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_core_plugins::PluginsManager;
use codex_exec_server::EnvironmentManager;
use codex_extension_api::ExtensionDataInit;
@@ -223,6 +225,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>,
}
@@ -311,6 +314,7 @@ impl ThreadManager {
installation_id,
analytics_events_client,
state_db,
code_mode_session_provider: Arc::new(StdioCodeModeSessionProvider::default()),
ops_log: should_use_test_thread_manager_behavior()
.then(|| Arc::new(std::sync::Mutex::new(Vec::new()))),
}),
@@ -417,6 +421,7 @@ impl ThreadManager {
installation_id,
analytics_events_client: None,
state_db,
code_mode_session_provider: Arc::new(StdioCodeModeSessionProvider::default()),
ops_log: should_use_test_thread_manager_behavior()
.then(|| Arc::new(std::sync::Mutex::new(Vec::new()))),
}),
@@ -428,6 +433,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()
}
@@ -1499,6 +1516,7 @@ impl ThreadManagerState {
thread_store: Arc::clone(&self.thread_store),
attestation_provider: self.attestation_provider.clone(),
external_time_provider: self.external_time_provider.clone(),
code_mode_session_provider: Arc::clone(&self.code_mode_session_provider),
inherited_multi_agent_version: multi_agent_version,
initial_multi_agent_mode,
}))
codex-rs/core/src/tools/code_mode/delegate.rs
+68 -22
View File
@@ -1,12 +1,13 @@
use std::collections::HashMap;
use std::collections::HashSet;
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;
@@ -24,7 +25,13 @@ use crate::tools::parallel::ToolCallRuntime;
pub(super) struct CodeModeDispatchBroker {
dispatch_tx: async_channel::Sender<DispatchMessage>,
dispatch_rx: async_channel::Receiver<DispatchMessage>,
dispatch_gates: Arc<Mutex<HashMap<CellId, watch::Sender<bool>>>>,
dispatch_gates: Arc<Mutex<DispatchGateState>>,
}
#[derive(Default)]
struct DispatchGateState {
gates: HashMap<CellId, watch::Sender<bool>>,
terminal_before_ready: HashSet<CellId>,
}
impl CodeModeDispatchBroker {
@@ -33,18 +40,22 @@ impl CodeModeDispatchBroker {
Self {
dispatch_tx,
dispatch_rx,
dispatch_gates: Arc::new(Mutex::new(HashMap::new())),
dispatch_gates: Arc::new(Mutex::new(DispatchGateState::default())),
}
}
pub(super) fn mark_cell_ready_for_dispatch(&self, cell_id: &CellId) {
dispatch_gate(&self.dispatch_gates, cell_id).send_replace(true);
mark_dispatch_gate_ready(&self.dispatch_gates, cell_id);
}
pub(super) fn close_cell(&self, cell_id: &CellId) {
remove_dispatch_gate(&self.dispatch_gates, cell_id);
}
fn record_cell_closed(&self, cell_id: &CellId) {
record_terminal_dispatch_gate(&self.dispatch_gates, cell_id);
}
pub(super) fn start_turn_worker(
&self,
exec: ExecContext,
@@ -130,39 +141,74 @@ impl CodeModeDispatchBroker {
}
fn dispatch_gate(
dispatch_gates: &Mutex<HashMap<CellId, watch::Sender<bool>>>,
dispatch_gates: &Mutex<DispatchGateState>,
cell_id: &CellId,
) -> watch::Sender<bool> {
) -> Option<watch::Sender<bool>> {
let mut dispatch_gates = match dispatch_gates.lock() {
Ok(dispatch_gates) => dispatch_gates,
Err(poisoned) => poisoned.into_inner(),
};
dispatch_gates
.entry(cell_id.clone())
.or_insert_with(|| watch::channel(false).0)
.clone()
if dispatch_gates.terminal_before_ready.contains(cell_id) {
return None;
}
Some(
dispatch_gates
.gates
.entry(cell_id.clone())
.or_insert_with(|| watch::channel(false).0)
.clone(),
)
}
fn remove_dispatch_gate(
dispatch_gates: &Mutex<HashMap<CellId, watch::Sender<bool>>>,
cell_id: &CellId,
) {
fn mark_dispatch_gate_ready(dispatch_gates: &Mutex<DispatchGateState>, cell_id: &CellId) {
let mut dispatch_gates = match dispatch_gates.lock() {
Ok(dispatch_gates) => dispatch_gates,
Err(poisoned) => poisoned.into_inner(),
};
dispatch_gates.remove(cell_id);
if dispatch_gates.terminal_before_ready.remove(cell_id) {
return;
}
dispatch_gates
.gates
.entry(cell_id.clone())
.or_insert_with(|| watch::channel(false).0)
.send_replace(true);
}
fn remove_dispatch_gate(dispatch_gates: &Mutex<DispatchGateState>, cell_id: &CellId) {
let mut dispatch_gates = match dispatch_gates.lock() {
Ok(dispatch_gates) => dispatch_gates,
Err(poisoned) => poisoned.into_inner(),
};
dispatch_gates.gates.remove(cell_id);
}
fn record_terminal_dispatch_gate(dispatch_gates: &Mutex<DispatchGateState>, cell_id: &CellId) {
let mut dispatch_gates = match dispatch_gates.lock() {
Ok(dispatch_gates) => dispatch_gates,
Err(poisoned) => poisoned.into_inner(),
};
let ready = dispatch_gates
.gates
.remove(cell_id)
.is_some_and(|gate| *gate.borrow());
if !ready {
dispatch_gates.terminal_before_ready.insert(cell_id.clone());
}
}
async fn wait_until_cell_ready_for_dispatch(
dispatch_gates: &Mutex<HashMap<CellId, watch::Sender<bool>>>,
dispatch_gates: &Mutex<DispatchGateState>,
cell_id: &CellId,
cancellation_token: &CancellationToken,
) -> bool {
if cancellation_token.is_cancelled() {
return false;
}
let mut ready_rx = dispatch_gate(dispatch_gates, cell_id).subscribe();
let Some(dispatch_gate) = dispatch_gate(dispatch_gates, cell_id) else {
return false;
};
let mut ready_rx = dispatch_gate.subscribe();
loop {
if *ready_rx.borrow_and_update() {
return true;
@@ -240,7 +286,7 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker {
}
fn cell_closed(&self, cell_id: &CellId) {
self.close_cell(cell_id);
self.record_cell_closed(cell_id);
}
}
codex-rs/core/src/tools/code_mode/delegate_tests.rs
+28 -51
View File
@@ -1,66 +1,43 @@
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode::CellId;
use codex_code_mode::CellOutcome;
use codex_code_mode::CodeModeService;
use codex_code_mode::CodeModeSessionDelegate;
use codex_code_mode::CreateCellRequest;
use codex_code_mode::ObserveRequest;
use codex_code_mode::RuntimeResponse;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use super::CodeModeDispatchBroker;
#[test]
fn terminal_notification_removes_the_cell_dispatch_gate() {
fn terminal_notification_removes_a_ready_cell_dispatch_gate() {
let broker = CodeModeDispatchBroker::new();
let cell_id = CellId::new("cell-a7".to_string());
broker.mark_cell_ready_for_dispatch(&cell_id);
assert!(broker.dispatch_gates.lock().unwrap().contains_key(&cell_id));
assert!(
broker
.dispatch_gates
.lock()
.unwrap()
.gates
.contains_key(&cell_id)
);
CodeModeSessionDelegate::cell_closed(&broker, &cell_id);
assert!(!broker.dispatch_gates.lock().unwrap().contains_key(&cell_id));
assert!(
!broker
.dispatch_gates
.lock()
.unwrap()
.gates
.contains_key(&cell_id)
);
}
#[tokio::test]
async fn background_completion_removes_the_dispatch_gate_without_another_observation() {
let broker = Arc::new(CodeModeDispatchBroker::new());
let service = CodeModeService::with_delegate(broker.clone());
let cell_id = service
.create_cell(CreateCellRequest {
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: concat!(
"await new Promise(resolve => setTimeout(resolve, 100));",
"text('done');",
)
.to_string(),
})
.await
.unwrap();
#[test]
fn terminal_notification_before_readiness_does_not_leave_a_dispatch_gate() {
let broker = CodeModeDispatchBroker::new();
let cell_id = CellId::new("cell-b8".to_string());
CodeModeSessionDelegate::cell_closed(&broker, &cell_id);
broker.mark_cell_ready_for_dispatch(&cell_id);
assert_eq!(
service
.observe(ObserveRequest {
cell_id: cell_id.clone(),
yield_time_ms: 1,
})
.await
.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Yielded {
cell_id: cell_id.clone(),
content_items: Vec::new(),
})
);
tokio::time::timeout(Duration::from_secs(/*secs*/ 2), async {
while broker.dispatch_gates.lock().unwrap().contains_key(&cell_id) {
tokio::task::yield_now().await;
}
})
.await
.expect("terminal notification should remove the dispatch gate");
service.shutdown().await.unwrap();
let dispatch_gates = broker.dispatch_gates.lock().unwrap();
assert!(!dispatch_gates.gates.contains_key(&cell_id));
assert!(!dispatch_gates.terminal_before_ready.contains(&cell_id));
}
codex-rs/core/src/tools/code_mode/execute_handler.rs
+11 -7
View File
@@ -34,8 +34,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);
@@ -44,7 +44,7 @@ impl CodeModeExecuteHandler {
.session
.services
.code_mode_service
.create_cell(codex_code_mode::CreateCellRequest {
.create_cell(codex_code_mode_protocol::CreateCellRequest {
tool_call_id: call_id.clone(),
enabled_tools,
source: args.code.clone(),
@@ -67,6 +67,7 @@ impl CodeModeExecuteHandler {
cell_id.clone(),
code_cell_trace.clone(),
)
.await
.map_err(FunctionCallError::RespondToModel)?;
exec.session
.services
@@ -76,14 +77,14 @@ impl CodeModeExecuteHandler {
.session
.services
.code_mode_service
.observe(codex_code_mode::ObserveRequest {
.observe(codex_code_mode_protocol::ObserveRequest {
cell_id: cell_id.clone(),
yield_time_ms: args
.yield_time_ms
.unwrap_or(codex_code_mode::DEFAULT_EXEC_YIELD_TIME_MS),
.unwrap_or(codex_code_mode_protocol::DEFAULT_EXEC_YIELD_TIME_MS),
})
.await;
let response: codex_code_mode::RuntimeResponse = match observed {
let response: codex_code_mode_protocol::RuntimeResponse = match observed {
Ok(outcome) => outcome.into(),
Err(error) => {
initial_observation_guard.finish_with_error(&error);
@@ -96,7 +97,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
+257 -44
View File
@@ -6,15 +6,21 @@ mod wait_handler;
pub(crate) mod wait_spec;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
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::CellOutcome;
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 tokio::sync::Mutex;
use tokio::sync::Semaphore;
use tokio_util::sync::CancellationToken;
use crate::function_tool::FunctionCallError;
@@ -43,9 +49,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 {
@@ -59,8 +66,11 @@ 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,
}
struct InitialObservationGuard {
@@ -71,13 +81,19 @@ struct InitialObservationGuard {
}
impl InitialObservationGuard {
fn new(
async fn new(
service: &CodeModeService,
cell_id: CellId,
code_cell_trace: codex_rollout_trace::CodeCellTraceContext,
) -> Result<Self, String> {
let session = service
.session
.lock()
.await
.clone()
.ok_or_else(|| "code mode is unavailable".to_string())?;
Ok(Self {
session: Arc::clone(service.session()?),
session,
dispatch_broker: Arc::clone(&service.dispatch_broker),
cell_id: Some(cell_id),
code_cell_trace,
@@ -95,10 +111,10 @@ impl InitialObservationGuard {
};
self.code_cell_trace.record_initial_response(&response);
self.code_cell_trace.record_ended(&response);
self.dispatch_broker.close_cell(&cell_id);
let session = Arc::clone(&self.session);
let dispatch_broker = Arc::clone(&self.dispatch_broker);
let Ok(runtime) = tokio::runtime::Handle::try_current() else {
dispatch_broker.close_cell(&cell_id);
tracing::warn!("unable to terminate code-mode cell {cell_id}: no Tokio runtime");
return;
};
@@ -106,6 +122,7 @@ impl InitialObservationGuard {
if let Err(error) = session.terminate(cell_id.clone()).await {
tracing::warn!("failed to terminate code-mode cell {cell_id}: {error}");
}
dispatch_broker.close_cell(&cell_id);
});
}
@@ -121,45 +138,56 @@ impl Drop for InitialObservationGuard {
}
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 create_cell(
&self,
request: codex_code_mode::CreateCellRequest,
request: codex_code_mode_protocol::CreateCellRequest,
) -> Result<CellId, String> {
self.session()?.create_cell(request).await
self.session_for_create().await?.create_cell(request).await
}
pub(crate) async fn observe(
&self,
request: codex_code_mode::ObserveRequest,
) -> Result<codex_code_mode::CellOutcome, String> {
self.session()?.observe(request).await
request: codex_code_mode_protocol::ObserveRequest,
) -> Result<CellOutcome, String> {
match self.current_session().await {
Some(session) => session.observe(request).await,
None => Ok(missing_cell_outcome(request.cell_id)),
}
}
pub(crate) async fn terminate(
&self,
cell_id: CellId,
) -> Result<codex_code_mode::CellOutcome, String> {
self.session()?.terminate(cell_id).await
pub(crate) async fn terminate(&self, cell_id: CellId) -> Result<CellOutcome, String> {
match self.current_session().await {
Some(session) => session.terminate(cell_id).await,
None => Ok(missing_cell_outcome(cell_id)),
}
}
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: &CellId) {
self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id);
}
@@ -175,9 +203,7 @@ impl CodeModeService {
tracker: SharedTurnDiffTracker,
) -> Option<CodeModeDispatchWorker> {
let tool_mode = effective_tool_mode(turn);
if !matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly)
|| self.session.is_none()
{
if !matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) {
return None;
}
@@ -191,10 +217,43 @@ impl CodeModeService {
)
}
fn session(&self) -> Result<&Arc<dyn CodeModeSession>, String> {
async fn session_for_create(&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()
}
}
@@ -367,6 +426,14 @@ fn serialize_function_tool_arguments(
}
}
fn missing_cell_outcome(cell_id: CellId) -> CellOutcome {
CellOutcome::MissingCell(RuntimeResponse::Result {
error_text: Some(format!("exec cell {cell_id} not found")),
cell_id,
content_items: Vec::new(),
})
}
fn build_freeform_tool_payload(
tool_name: &ToolName,
input: Option<JsonValue>,
@@ -379,24 +446,33 @@ fn build_freeform_tool_payload(
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use super::CodeModeDispatchBroker;
use super::CodeModeService;
use super::InitialObservationGuard;
use super::build_nested_tool_payload;
use super::truncate_code_mode_result;
use crate::tools::context::ToolPayload;
use codex_code_mode::CellId;
use codex_code_mode::CellOutcome;
use codex_code_mode::CodeModeSession;
use codex_code_mode::CodeModeSessionResultFuture;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::CreateCellRequest;
use codex_code_mode::ObserveRequest;
use codex_code_mode::RuntimeResponse;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CellOutcome;
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::CreateCellRequest;
use codex_code_mode_protocol::ObserveRequest;
use codex_code_mode_protocol::RuntimeResponse;
use codex_protocol::models::FunctionCallOutputContentItem;
use codex_tools::ToolName;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::sync::Arc;
use std::sync::Mutex;
struct RecordingCodeModeSession {
terminated_cells: Mutex<Vec<CellId>>,
@@ -532,4 +608,141 @@ mod tests {
vec![cell_id]
);
}
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 create_cell<'a>(
&'a self,
_request: CreateCellRequest,
) -> CodeModeSessionResultFuture<'a, CellId> {
Box::pin(async move {
if self.generation == 1 {
self.alive.store(false, Ordering::Release);
return Err("host crashed".to_string());
}
Ok(CellId::new(format!("cell-{}", self.generation)))
})
}
fn observe<'a>(
&'a self,
request: ObserveRequest,
) -> CodeModeSessionResultFuture<'a, CellOutcome> {
Box::pin(async move {
Ok(CellOutcome::LiveCell(RuntimeResponse::Result {
cell_id: request.cell_id,
content_items: Vec::new(),
error_text: None,
}))
})
}
fn terminate<'a>(
&'a self,
_cell_id: CellId,
) -> CodeModeSessionResultFuture<'a, CellOutcome> {
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 create_cell_request() -> CreateCellRequest {
CreateCellRequest {
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: "text('hello')".to_string(),
}
}
#[tokio::test]
async fn create_replaces_failed_session_but_observe_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
.create_cell(create_cell_request())
.await
.err()
.as_deref(),
Some("host crashed")
);
assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1);
assert_eq!(
service
.observe(ObserveRequest {
cell_id: CellId::new("1".to_string()),
yield_time_ms: 1,
})
.await,
Ok(CellOutcome::MissingCell(RuntimeResponse::Result {
cell_id: CellId::new("1".to_string()),
content_items: Vec::new(),
error_text: Some("exec cell 1 not found".to_string()),
}))
);
assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1);
assert_eq!(
service.terminate(CellId::new("1".to_string())).await,
Ok(CellOutcome::MissingCell(RuntimeResponse::Result {
cell_id: CellId::new("1".to_string()),
content_items: Vec::new(),
error_text: Some("exec cell 1 not found".to_string()),
}))
);
assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1);
let cell_id = service.create_cell(create_cell_request()).await.unwrap();
assert_eq!(cell_id, CellId::new("cell-2".to_string()));
assert_eq!(
service
.observe(ObserveRequest {
cell_id: cell_id.clone(),
yield_time_ms: 1,
})
.await,
Ok(CellOutcome::LiveCell(RuntimeResponse::Result {
cell_id,
content_items: Vec::new(),
error_text: None,
}))
);
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
.observe(codex_code_mode::ObserveRequest {
.observe(codex_code_mode_protocol::ObserveRequest {
cell_id,
yield_time_ms: args.yield_time_ms,
})
.await
}
.map_err(FunctionCallError::RespondToModel)?;
if let codex_code_mode::CellOutcome::LiveCell(response) = &wait_response
&& !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. })
if let codex_code_mode_protocol::CellOutcome::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
@@ -279,7 +279,7 @@ fn spec_for_model_request(
if matches!(tool_mode, 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 {
@@ -457,9 +457,9 @@ fn is_hidden_by_code_mode_only(
let tool_mode = effective_tool_mode(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 {
@@ -581,7 +581,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 {
@@ -590,7 +590,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(),
});
@@ -986,8 +986,8 @@ fn append_extension_tool_executors(
.collect::<HashSet<_>>();
let tool_mode = effective_tool_mode(turn_context);
if matches!(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)
@@ -1086,9 +1086,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);
@@ -1100,8 +1100,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
+18 -15
View File
@@ -453,8 +453,8 @@ async fn request_user_input_stays_direct_in_code_mode_only() {
plan.assert_visible_contains(&[
"request_user_input",
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,
]);
plan.assert_registered_contains(&["request_user_input"]);
assert_eq!(
@@ -462,7 +462,8 @@ async fn request_user_input_stays_direct_in_code_mode_only() {
ToolExposure::DirectModelOnly
);
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!(!exec.description.contains("request_user_input"));
@@ -983,8 +984,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(
@@ -1002,8 +1003,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"),
@@ -1033,8 +1034,8 @@ async fn code_mode_only_exposes_configured_dynamic_namespace_directly() {
.await;
plan.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,
"direct_only",
]);
plan.assert_visible_lacks(&["tool_search"]);
@@ -1047,7 +1048,8 @@ async fn code_mode_only_exposes_configured_dynamic_namespace_directly() {
};
let ResponsesApiNamespaceTool::Function(tool) = &namespace.tools[0];
assert_eq!(tool.defer_loading, None);
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!(!exec.description.contains("direct_only_lookup(args:"));
@@ -1075,7 +1077,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!(
@@ -1222,8 +1225,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,
]);
}
@@ -1460,8 +1463,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,
"request_user_input",
// Multi-agent v2 tools.
"spawn_agent",
codex-rs/core/tests/common/Cargo.toml
+1
View File
@@ -16,6 +16,7 @@ anyhow = { workspace = true }
assert_cmd = { workspace = true }
base64 = { workspace = true }
codex-arg0 = { workspace = true }
codex-code-mode-client = { workspace = true }
codex-config = { workspace = true }
codex-core = { workspace = true }
codex-extension-api = { workspace = true }
codex-rs/core/tests/common/test_codex.rs
+10
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::StdioCodeModeSessionProvider;
use codex_config::CloudConfigBundleLoader;
use codex_core::CodexThread;
use codex_core::StartThreadOptions;
@@ -591,6 +593,14 @@ impl TestCodexBuilder {
/*attestation_provider*/ None,
/*external_time_provider*/ self.external_time_provider.clone(),
);
let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")
.context("resolve code-mode host for stdio core tests")?;
let thread_manager = thread_manager.with_code_mode_session_provider(Arc::new(
StdioCodeModeSessionProvider::new(CodeModeHostCommand {
program,
args: Vec::new(),
}),
));
let thread_manager = Arc::new(thread_manager);
let user_shell_override = self.user_shell_override.clone();
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(),
"request_user_input".to_string(),
// Hosted Responses tools.
"web_search".to_string(),
codex-rs/rollout-trace/Cargo.toml
+1 -1
View File
@@ -14,7 +14,7 @@ workspace = true
[dependencies]
anyhow = { workspace = true }
codex-code-mode = { workspace = true }
codex-code-mode-protocol = { workspace = true }
codex-protocol = { workspace = true }
http = { workspace = true }
serde = { workspace = true, features = ["derive"] }
codex-rs/rollout-trace/src/code_cell.rs
+1 -1
View File
@@ -7,7 +7,7 @@
use std::sync::Arc;
use codex_code_mode::RuntimeResponse;
use codex_code_mode_protocol::RuntimeResponse;
use serde::Serialize;
use tracing::warn;
codex-rs/rollout-trace/src/tool_dispatch.rs
+3 -3
View File
@@ -194,7 +194,7 @@ impl ToolDispatchTraceContext {
fn suppresses_tool_dispatch_trace(invocation: &ToolDispatchInvocation) -> bool {
matches!(invocation.payload, ToolDispatchPayload::Custom { .. })
&& invocation.tool_namespace.is_none()
&& invocation.tool_name == codex_code_mode::PUBLIC_TOOL_NAME
&& invocation.tool_name == codex_code_mode_protocol::PUBLIC_TOOL_NAME
}
fn record_started(context: &EnabledToolDispatchTraceContext, invocation: ToolDispatchInvocation) {
@@ -394,7 +394,7 @@ mod tests {
#[test]
fn suppresses_only_noncanonical_dispatch_boundaries() {
assert!(suppresses_tool_dispatch_trace(&invocation(
codex_code_mode::PUBLIC_TOOL_NAME,
codex_code_mode_protocol::PUBLIC_TOOL_NAME,
/*tool_namespace*/ None,
ToolDispatchRequester::Model {
model_visible_call_id: "call-exec".to_string(),
@@ -414,7 +414,7 @@ mod tests {
},
)));
assert!(!suppresses_tool_dispatch_trace(&invocation(
codex_code_mode::PUBLIC_TOOL_NAME,
codex_code_mode_protocol::PUBLIC_TOOL_NAME,
Some("mcp__server".to_string()),
ToolDispatchRequester::Model {
model_visible_call_id: "call-namespaced".to_string(),
codex-rs/tools/Cargo.toml
+1 -1
View File
@@ -9,7 +9,7 @@ workspace = true
[dependencies]
codex-app-server-protocol = { workspace = true }
codex-code-mode = { workspace = true }
codex-code-mode-protocol = { workspace = true }
codex-features = { workspace = true }
codex-file-system = { workspace = true }
codex-protocol = { workspace = true }
codex-rs/tools/src/code_mode.rs
+10 -9
View File
@@ -1,8 +1,8 @@
use crate::ResponsesApiNamespaceTool;
use crate::ToolName;
use crate::ToolSpec;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::ToolDefinition as CodeModeToolDefinition;
use codex_code_mode_protocol::CodeModeToolKind;
use codex_code_mode_protocol::ToolDefinition as CodeModeToolDefinition;
/// Augment tool descriptions with code-mode-specific exec samples.
pub fn augment_tool_spec_for_code_mode(spec: ToolSpec) -> ToolSpec {
@@ -40,7 +40,8 @@ pub fn augment_tool_spec_for_code_mode(spec: ToolSpec) -> ToolSpec {
output_schema: tool.output_schema.clone(),
};
tool.description =
codex_code_mode::augment_tool_definition(definition).description;
codex_code_mode_protocol::augment_tool_definition(definition)
.description;
}
}
}
@@ -54,8 +55,8 @@ pub fn augment_tool_spec_for_code_mode(spec: ToolSpec) -> ToolSpec {
/// including the code-mode-specific description sample.
pub fn tool_spec_to_code_mode_tool_definition(spec: &ToolSpec) -> Option<CodeModeToolDefinition> {
let definition = code_mode_tool_definition_for_spec(spec)?;
codex_code_mode::is_code_mode_nested_tool(&definition.name)
.then(|| codex_code_mode::augment_tool_definition(definition))
codex_code_mode_protocol::is_code_mode_nested_tool(&definition.name)
.then(|| codex_code_mode_protocol::augment_tool_definition(definition))
}
pub fn collect_code_mode_tool_definitions<'a>(
@@ -76,8 +77,8 @@ pub fn collect_code_mode_tool_definitions<'a>(
}
definitions
})
.filter(|definition| codex_code_mode::is_code_mode_nested_tool(&definition.name))
.map(codex_code_mode::augment_tool_definition)
.filter(|definition| codex_code_mode_protocol::is_code_mode_nested_tool(&definition.name))
.map(codex_code_mode_protocol::augment_tool_definition)
.collect::<Vec<_>>();
tool_definitions.sort_by(|left, right| left.name.cmp(&right.name));
tool_definitions.dedup_by(|left, right| left.name == right.name);
@@ -90,7 +91,7 @@ pub fn collect_code_mode_exec_prompt_tool_definitions<'a>(
let mut tool_definitions = specs
.into_iter()
.flat_map(code_mode_tool_definitions_for_spec)
.filter(|definition| codex_code_mode::is_code_mode_nested_tool(&definition.name))
.filter(|definition| codex_code_mode_protocol::is_code_mode_nested_tool(&definition.name))
.collect::<Vec<_>>();
tool_definitions.sort_by(|left, right| left.name.cmp(&right.name));
tool_definitions.dedup_by(|left, right| left.name == right.name);
@@ -99,7 +100,7 @@ pub fn collect_code_mode_exec_prompt_tool_definitions<'a>(
fn augmented_description_for_spec(spec: &ToolSpec) -> Option<String> {
code_mode_tool_definition_for_spec(spec)
.map(codex_code_mode::augment_tool_definition)
.map(codex_code_mode_protocol::augment_tool_definition)
.map(|definition| definition.description)
}
codex-rs/tools/src/code_mode_tests.rs
+4 -4
View File
@@ -69,7 +69,7 @@ declare const tools: { lookup_order(args: { order_id: string; }): Promise<{ ok:
fn augment_tool_spec_for_code_mode_preserves_exec_tool_description() {
assert_eq!(
augment_tool_spec_for_code_mode(ToolSpec::Freeform(FreeformTool {
name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(),
name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(),
description: "Run code".to_string(),
format: FreeformToolFormat {
r#type: "grammar".to_string(),
@@ -78,7 +78,7 @@ fn augment_tool_spec_for_code_mode_preserves_exec_tool_description() {
},
})),
ToolSpec::Freeform(FreeformTool {
name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(),
name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(),
description: "Run code".to_string(),
format: FreeformToolFormat {
r#type: "grammar".to_string(),
@@ -103,7 +103,7 @@ fn tool_spec_to_code_mode_tool_definition_returns_augmented_nested_tools() {
assert_eq!(
tool_spec_to_code_mode_tool_definition(&spec),
Some(codex_code_mode::ToolDefinition {
Some(codex_code_mode_protocol::ToolDefinition {
name: "apply_patch".to_string(),
tool_name: ToolName::plain("apply_patch"),
description: r#"Apply a patch
@@ -113,7 +113,7 @@ exec tool declaration:
declare const tools: { apply_patch(input: string): Promise<unknown>; };
```"#
.to_string(),
kind: codex_code_mode::CodeModeToolKind::Freeform,
kind: codex_code_mode_protocol::CodeModeToolKind::Freeform,
input_schema: None,
output_schema: None,
})
scripts/run-nextest.py
+7
View File
@@ -54,6 +54,13 @@ def main() -> int:
check=True,
env=env,
)
target_dir = Path(env.get("CARGO_TARGET_DIR", CODEX_RS_ROOT / "target"))
if not target_dir.is_absolute():
target_dir = CODEX_RS_ROOT / target_dir
host_name = (
"codex-code-mode-host.exe" if os.name == "nt" else "codex-code-mode-host"
)
env["CODEX_CODE_MODE_HOST_PATH"] = str(target_dir / "debug" / host_name)
return subprocess.run(
[cargo, "nextest", "run", "--no-fail-fast", *sys.argv[1:]],
cwd=CODEX_RS_ROOT,