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

code-mode: move cell lifecycle ownership to runtime

Browse Source
This commit is contained in:
Channing Conger
2026-06-19 18:11:46 +00:00
parent f1d0826a54
commit dfc1727e0d
20 changed files with 703 additions and 645 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-protocol/src/lib.rs
-1
View File
@@ -38,7 +38,6 @@ pub use session::CodeModeSessionProvider;
pub use session::CodeModeSessionProviderFuture;
pub use session::CodeModeSessionResultFuture;
pub use session::NotificationFuture;
pub use session::StartedCell;
pub use session::ToolInvocationFuture;
pub const PUBLIC_TOOL_NAME: &str = "exec";
codex-rs/code-mode-protocol/src/runtime.rs
+2
View File
@@ -14,6 +14,8 @@ pub const DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL: usize = 10_000;
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct ExecuteRequest {
/// A client-generated nonce, unique within the owning code-mode session.
pub cell_id: CellId,
pub tool_call_id: String,
pub enabled_tools: Vec<ToolDefinition>,
pub source: String,
codex-rs/code-mode-protocol/src/session.rs
+1 -45
View File
@@ -6,7 +6,6 @@ use std::sync::Arc;
use serde::Deserialize;
use serde::Serialize;
use serde_json::Value as JsonValue;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use crate::CodeModeNestedToolCall;
@@ -48,42 +47,6 @@ impl fmt::Display for CellId {
}
}
pub struct StartedCell {
pub cell_id: CellId,
initial_response: CodeModeSessionResultFuture<'static, RuntimeResponse>,
}
impl StartedCell {
pub fn new(cell_id: CellId, initial_response_rx: oneshot::Receiver<RuntimeResponse>) -> Self {
Self {
cell_id,
initial_response: Box::pin(async move {
initial_response_rx
.await
.map_err(|_| "exec runtime ended unexpectedly".to_string())
}),
}
}
pub fn from_result_receiver(
cell_id: CellId,
initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,
) -> Self {
Self {
cell_id,
initial_response: Box::pin(async move {
initial_response_rx
.await
.map_err(|_| "exec runtime ended unexpectedly".to_string())?
}),
}
}
pub async fn initial_response(self) -> Result<RuntimeResponse, String> {
self.initial_response.await
}
}
/// Host callbacks used by a code-mode session while cells are executing.
pub trait CodeModeSessionDelegate: Send + Sync {
fn invoke_tool<'a>(
@@ -99,9 +62,6 @@ pub trait CodeModeSessionDelegate: Send + Sync {
text: String,
cancellation_token: CancellationToken,
) -> NotificationFuture<'a>;
/// Releases delegate state associated with a cell after it reaches a terminal state.
fn cell_closed(&self, cell_id: &CellId);
}
/// A durable code-mode session owned by one Codex thread.
@@ -119,7 +79,7 @@ pub trait CodeModeSession: Send + Sync {
fn execute<'a>(
&'a self,
request: ExecuteRequest,
) -> CodeModeSessionResultFuture<'a, StartedCell>;
) -> CodeModeSessionResultFuture<'a, RuntimeResponse>;
fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
@@ -138,7 +98,3 @@ pub trait CodeModeSessionProvider: Send + Sync {
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionProviderFuture<'a>;
}
#[cfg(test)]
#[path = "session_tests.rs"]
mod tests;
codex-rs/code-mode-protocol/src/session_tests.rs
-19
View File
@@ -1,19 +0,0 @@
use pretty_assertions::assert_eq;
use tokio::sync::oneshot;
use super::CellId;
use super::StartedCell;
#[tokio::test]
async fn started_cell_preserves_remote_initial_response_errors() {
let (response_tx, response_rx) = oneshot::channel();
response_tx
.send(Err("remote runtime failed".to_string()))
.expect("initial response receiver should be open");
let started = StartedCell::from_result_receiver(CellId::new("1".to_string()), response_rx);
assert_eq!(
started.initial_response().await,
Err("remote runtime failed".to_string())
);
}
codex-rs/code-mode/src/cell_actor/conversions.rs
+1
View File
@@ -12,6 +12,7 @@ use crate::session_runtime::OutputItem as CellOutputItem;
pub(super) fn runtime_request(request: CellRequest) -> ExecuteRequest {
ExecuteRequest {
cell_id: codex_code_mode_protocol::CellId::new(request.cell_id.as_str().to_string()),
tool_call_id: request.tool_call_id,
enabled_tools: request
.enabled_tools
codex-rs/code-mode/src/cell_actor/mod.rs
+120 -6
View File
@@ -5,6 +5,8 @@ mod types;
use std::collections::HashMap;
use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use serde_json::Value as JsonValue;
use tokio::sync::mpsc;
@@ -64,7 +66,14 @@ impl CellActor {
PendingRuntimeMode::PauseUntilResumed,
)?;
let cancellation_token = session_shutdown_token.child_token();
let handle = CellHandle::new(command_tx);
let termination_token = CancellationToken::new();
let accepting_requests = Arc::new(AtomicBool::new(true));
let handle = CellHandle::new(
command_tx,
cancellation_token.clone(),
termination_token.clone(),
Arc::clone(&accepting_requests),
);
let task = run_cell(
host,
CellContext {
@@ -72,6 +81,8 @@ impl CellActor {
runtime_control_tx,
runtime_terminate_handle,
cancellation_token,
termination_token,
accepting_requests,
session_shutdown_token,
},
event_rx,
@@ -92,6 +103,8 @@ struct CellContext {
runtime_control_tx: std::sync::mpsc::Sender<RuntimeControlCommand>,
runtime_terminate_handle: v8::IsolateHandle,
cancellation_token: CancellationToken,
termination_token: CancellationToken,
accepting_requests: Arc<AtomicBool>,
session_shutdown_token: CancellationToken,
}
@@ -116,6 +129,8 @@ async fn run_cell<H: CellHost>(
runtime_control_tx,
runtime_terminate_handle,
cancellation_token,
termination_token,
accepting_requests,
session_shutdown_token,
} = context;
let mut content_items = Vec::new();
@@ -174,6 +189,7 @@ async fn run_cell<H: CellHost>(
match command {
CellCommand::Observe { mode, response_tx } => {
if let Some(event) = completed_event.take() {
accepting_requests.store(false, Ordering::Release);
let _ = response_tx.send(Ok(event));
break;
}
@@ -199,6 +215,7 @@ async fn run_cell<H: CellHost>(
}
CellCommand::Terminate { response_tx } => {
if let Some(event) = completed_event.take() {
accepting_requests.store(false, Ordering::Release);
if let Some(response_tx) = response_tx {
let _ = response_tx.send(Ok(event));
}
@@ -226,6 +243,7 @@ async fn run_cell<H: CellHost>(
CallbackCompletion::Cancel,
).await;
send_termination_events(
&accepting_requests,
observer.take(),
termination.take(),
CellEvent::Terminated {
@@ -261,7 +279,7 @@ async fn run_cell<H: CellHost>(
}, if !yield_deadline_elapsed => {
let Some(event) = maybe_event else {
runtime_closed = true;
if termination.is_some() {
if termination.is_some() || session_shutdown_token.is_cancelled() {
finish_callbacks(
&cancellation_token,
&mut notification_tasks,
@@ -269,6 +287,7 @@ async fn run_cell<H: CellHost>(
CallbackCompletion::Cancel,
).await;
send_termination_events(
&accepting_requests,
observer.take(),
termination.take(),
CellEvent::Terminated {
@@ -283,7 +302,32 @@ async fn run_cell<H: CellHost>(
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::DrainNotifications,
).await;
)
.await;
if termination_token.is_cancelled() {
let termination =
receive_queued_termination(&mut command_rx).await;
send_termination_events(
&accepting_requests,
observer.take(),
termination,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
if session_shutdown_token.is_cancelled() {
send_termination_events(
&accepting_requests,
observer.take(),
None,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
let event = CellEvent::Completed {
content_items: std::mem::take(&mut content_items),
error_text: Some("exec runtime ended unexpectedly".to_string()),
@@ -292,6 +336,7 @@ async fn run_cell<H: CellHost>(
event,
&mut observer,
&mut completed_event,
&accepting_requests,
) {
break;
}
@@ -369,7 +414,7 @@ async fn run_cell<H: CellHost>(
RuntimeEvent::Result { stored_value_writes, error_text } => {
runtime_closed = true;
yield_timer = None;
if termination.is_some() {
if termination.is_some() || session_shutdown_token.is_cancelled() {
finish_callbacks(
&cancellation_token,
&mut notification_tasks,
@@ -377,6 +422,7 @@ async fn run_cell<H: CellHost>(
CallbackCompletion::Cancel,
).await;
send_termination_events(
&accepting_requests,
observer.take(),
termination.take(),
CellEvent::Terminated {
@@ -390,8 +436,54 @@ async fn run_cell<H: CellHost>(
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::DrainNotifications,
).await;
host.commit_stored_values(stored_value_writes).await;
)
.await;
if termination_token.is_cancelled() {
let termination =
receive_queued_termination(&mut command_rx).await;
send_termination_events(
&accepting_requests,
observer.take(),
termination,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
if session_shutdown_token.is_cancelled() {
send_termination_events(
&accepting_requests,
observer.take(),
None,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
let committed = tokio::select! {
biased;
_ = termination_token.cancelled() => false,
_ = session_shutdown_token.cancelled() => false,
_ = host.commit_stored_values(stored_value_writes) => true,
};
if !committed || session_shutdown_token.is_cancelled() {
let termination = if termination_token.is_cancelled() {
receive_queued_termination(&mut command_rx).await
} else {
None
};
send_termination_events(
&accepting_requests,
observer.take(),
termination,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
let event = CellEvent::Completed {
content_items: std::mem::take(&mut content_items),
error_text,
@@ -400,6 +492,7 @@ async fn run_cell<H: CellHost>(
event,
&mut observer,
&mut completed_event,
&accepting_requests,
) {
break;
}
@@ -415,6 +508,7 @@ async fn run_cell<H: CellHost>(
}
}
// Reject requests that arrive while asynchronous terminal cleanup runs.
accepting_requests.store(false, Ordering::Release);
drop(command_rx.take());
begin_termination(
&runtime_tx,
@@ -432,12 +526,30 @@ async fn run_cell<H: CellHost>(
host.closed().await;
}
async fn receive_queued_termination(
command_rx: &mut Option<mpsc::UnboundedReceiver<CellCommand>>,
) -> Option<Termination> {
loop {
let command = command_rx.as_mut()?.recv().await?;
match command {
CellCommand::Observe { response_tx, .. } => {
let _ = response_tx.send(Err(CellError::Busy));
}
CellCommand::Terminate { response_tx } => {
return Some(Termination { response_tx });
}
}
}
}
fn send_or_buffer_completion(
event: CellEvent,
observer: &mut Option<Observer>,
completed_event: &mut Option<CellEvent>,
accepting_requests: &AtomicBool,
) -> bool {
if observer.is_some() {
accepting_requests.store(false, Ordering::Release);
send_observer_event(observer.take(), event);
true
} else {
@@ -453,10 +565,12 @@ fn send_observer_event(observer: Option<Observer>, event: CellEvent) {
}
fn send_termination_events(
accepting_requests: &AtomicBool,
observer: Option<Observer>,
termination: Option<Termination>,
event: CellEvent,
) {
accepting_requests.store(false, Ordering::Release);
send_observer_event(observer, event.clone());
if let Some(response_tx) = termination.and_then(|termination| termination.response_tx) {
let _ = response_tx.send(Ok(event));
codex-rs/code-mode/src/cell_actor/tests.rs
+248 -1
View File
@@ -25,6 +25,12 @@ struct RecordingHost {
closed: AtomicBool,
}
#[derive(Default)]
struct PendingNotificationHost {
committed: AtomicBool,
closed: AtomicBool,
}
struct BlockingCommitHost {
commit_started_tx: mpsc::UnboundedSender<()>,
commit_release: Semaphore,
@@ -80,6 +86,33 @@ impl CellHost for RecordingHost {
}
}
impl CellHost for PendingNotificationHost {
async fn invoke_tool(
&self,
_invocation: CellToolCall,
_cancellation_token: CancellationToken,
) -> Result<JsonValue, String> {
Err("unexpected tool call".to_string())
}
async fn notify(
&self,
_call_id: String,
_text: String,
_cancellation_token: CancellationToken,
) -> Result<(), String> {
std::future::pending().await
}
async fn commit_stored_values(&self, _stored_value_writes: HashMap<String, JsonValue>) {
self.committed.store(true, Ordering::Release);
}
async fn closed(&self) {
self.closed.store(true, Ordering::Release);
}
}
impl BlockingCommitHost {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<()>) {
let (commit_started_tx, commit_started_rx) = mpsc::unbounded_channel();
@@ -150,6 +183,7 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = spawn_runtime(
HashMap::new(),
ExecuteRequest {
cell_id: codex_code_mode_protocol::CellId::new("cell-1".to_string()),
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: "await new Promise(() => {});".to_string(),
@@ -162,7 +196,14 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
.unwrap();
let session_shutdown_token = CancellationToken::new();
let cancellation_token = session_shutdown_token.child_token();
let handle = CellHandle::new(command_tx.clone());
let termination_token = CancellationToken::new();
let accepting_requests = Arc::new(AtomicBool::new(true));
let handle = CellHandle::new(
command_tx.clone(),
cancellation_token.clone(),
termination_token.clone(),
Arc::clone(&accepting_requests),
);
let task = tokio::spawn(run_cell(
host,
CellContext {
@@ -170,6 +211,8 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
runtime_control_tx,
runtime_terminate_handle,
cancellation_token,
termination_token,
accepting_requests,
session_shutdown_token: session_shutdown_token.clone(),
},
event_rx,
@@ -270,6 +313,210 @@ async fn termination_preempts_result_without_committing_stored_values() {
assert!(host.closed.load(Ordering::Acquire));
}
#[tokio::test]
async fn session_shutdown_during_result_callback_drain_prevents_stored_value_commit() {
let host = Arc::new(PendingNotificationHost::default());
let (event_tx, event_rx) = mpsc::unbounded_channel();
let (_command_tx, command_rx) = mpsc::unbounded_channel();
let (initial_event_tx, initial_event_rx) = oneshot::channel();
let (runtime_event_tx, _runtime_event_rx) = mpsc::unbounded_channel();
let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = spawn_runtime(
HashMap::new(),
ExecuteRequest {
cell_id: codex_code_mode_protocol::CellId::new("cell-1".to_string()),
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: "await new Promise(() => {});".to_string(),
yield_time_ms: None,
max_output_tokens: None,
},
runtime_event_tx,
PendingRuntimeMode::PauseUntilResumed,
)
.unwrap();
let session_shutdown_token = CancellationToken::new();
let termination_token = CancellationToken::new();
let accepting_requests = Arc::new(AtomicBool::new(true));
let mut task = Box::pin(run_cell(
Arc::clone(&host),
CellContext {
runtime_tx,
runtime_control_tx,
runtime_terminate_handle,
cancellation_token: session_shutdown_token.child_token(),
termination_token,
accepting_requests,
session_shutdown_token: session_shutdown_token.clone(),
},
event_rx,
command_rx,
Observer {
mode: ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
response_tx: initial_event_tx,
},
));
event_tx
.send(RuntimeEvent::Notify {
call_id: "call-1".to_string(),
text: "pending".to_string(),
})
.unwrap();
std::future::poll_fn(|context| match task.as_mut().poll(context) {
std::task::Poll::Pending => std::task::Poll::Ready(()),
std::task::Poll::Ready(()) => panic!("cell actor ended after notification"),
})
.await;
event_tx
.send(RuntimeEvent::Result {
stored_value_writes: HashMap::from([("key".to_string(), JsonValue::Bool(true))]),
error_text: None,
})
.unwrap();
std::future::poll_fn(|context| match task.as_mut().poll(context) {
std::task::Poll::Pending => std::task::Poll::Ready(()),
std::task::Poll::Ready(()) => panic!("cell actor ended before callback cleanup"),
})
.await;
session_shutdown_token.cancel();
drop(event_tx);
task.await;
assert_eq!(
initial_event_rx.await.unwrap(),
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
assert!(!host.committed.load(Ordering::Acquire));
assert!(host.closed.load(Ordering::Acquire));
}
#[tokio::test]
async fn termination_during_result_callback_drain_does_not_wait_for_the_notification() {
let host = Arc::new(PendingNotificationHost::default());
let harness = spawn_cell_actor_harness_with_host(
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
Arc::clone(&host),
);
harness
.event_tx
.send(RuntimeEvent::Notify {
call_id: "call-1".to_string(),
text: "pending".to_string(),
})
.unwrap();
tokio::task::yield_now().await;
harness
.event_tx
.send(RuntimeEvent::Result {
stored_value_writes: HashMap::from([("key".to_string(), JsonValue::Bool(true))]),
error_text: None,
})
.unwrap();
tokio::task::yield_now().await;
let termination = harness.handle.terminate();
let terminated = tokio::time::timeout(Duration::from_millis(/*millis*/ 100), termination)
.await
.expect("termination should not await a non-cooperative notification");
assert_eq!(
terminated,
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
assert_eq!(
harness.initial_event_rx.await.unwrap(),
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
drop(harness.event_tx);
harness.task.await.unwrap();
assert!(!host.committed.load(Ordering::Acquire));
assert!(host.closed.load(Ordering::Acquire));
}
#[tokio::test]
async fn session_shutdown_during_stored_value_commit_does_not_wait_for_the_host() {
let (host, mut commit_started_rx) = BlockingCommitHost::new();
let mut harness = spawn_cell_actor_harness_with_host(
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
host,
);
harness
.event_tx
.send(RuntimeEvent::Result {
stored_value_writes: HashMap::from([("key".to_string(), JsonValue::Bool(true))]),
error_text: None,
})
.unwrap();
assert_eq!(commit_started_rx.recv().await, Some(()));
harness.session_shutdown_token.cancel();
drop(harness.event_tx);
let event = tokio::time::timeout(
Duration::from_millis(/*millis*/ 100),
&mut harness.initial_event_rx,
)
.await
.expect("session shutdown should not await stored value commit")
.unwrap();
assert_eq!(
event,
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut harness.task)
.await
.expect("cell actor should exit after session shutdown")
.unwrap();
}
#[tokio::test]
async fn termination_during_stored_value_commit_does_not_wait_for_the_host() {
let (host, mut commit_started_rx) = BlockingCommitHost::new();
let mut harness = spawn_cell_actor_harness_with_host(
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
host,
);
harness
.event_tx
.send(RuntimeEvent::Result {
stored_value_writes: HashMap::from([("key".to_string(), JsonValue::Bool(true))]),
error_text: None,
})
.unwrap();
assert_eq!(commit_started_rx.recv().await, Some(()));
let termination = harness.handle.terminate();
assert_eq!(
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), termination)
.await
.expect("termination should not await stored value commit"),
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
assert_eq!(
(&mut harness.initial_event_rx).await.unwrap(),
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
drop(harness.event_tx);
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut harness.task)
.await
.expect("cell actor should exit after termination")
.unwrap();
}
#[tokio::test]
async fn observer_receives_a_buffered_completion_after_commit_finishes() {
let (host, mut commit_started_rx) = BlockingCommitHost::new();
codex-rs/code-mode/src/cell_actor/types.rs
+21 -1
View File
@@ -63,18 +63,32 @@ pub(crate) trait CellHost: Send + Sync + 'static {
#[derive(Clone)]
pub(crate) struct CellHandle {
command_tx: mpsc::UnboundedSender<CellCommand>,
cancellation_token: CancellationToken,
termination_token: CancellationToken,
termination_requested: Arc<AtomicBool>,
accepting_requests: Arc<AtomicBool>,
}
impl CellHandle {
pub(super) fn new(command_tx: mpsc::UnboundedSender<CellCommand>) -> Self {
pub(super) fn new(
command_tx: mpsc::UnboundedSender<CellCommand>,
cancellation_token: CancellationToken,
termination_token: CancellationToken,
accepting_requests: Arc<AtomicBool>,
) -> Self {
Self {
command_tx,
cancellation_token,
termination_token,
termination_requested: Arc::new(AtomicBool::new(false)),
accepting_requests,
}
}
pub(crate) fn observe(&self, mode: ObserveMode) -> CellEventFuture {
if !self.accepting_requests.load(Ordering::Acquire) {
return closed_event();
}
let (response_tx, response_rx) = oneshot::channel();
if self
.command_tx
@@ -87,6 +101,9 @@ impl CellHandle {
}
pub(crate) fn terminate(&self) -> CellEventFuture {
if !self.accepting_requests.load(Ordering::Acquire) {
return closed_event();
}
if self
.termination_requested
.compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
@@ -105,6 +122,9 @@ impl CellHandle {
self.termination_requested.store(false, Ordering::Relaxed);
return closed_event();
}
// The command carries the terminal response; these tokens interrupt in-flight cleanup.
self.termination_token.cancel();
self.cancellation_token.cancel();
response_event(response_rx)
}
}
codex-rs/code-mode/src/runtime/mod.rs
+2
View File
@@ -336,10 +336,12 @@ mod tests {
use super::RuntimeControlCommand;
use super::RuntimeEvent;
use super::spawn_runtime;
use crate::CellId;
use crate::FunctionCallOutputContentItem;
fn execute_request(source: &str) -> ExecuteRequest {
ExecuteRequest {
cell_id: CellId::new("cell-1".to_string()),
tool_call_id: "call_1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
codex-rs/code-mode/src/service.rs
+8 -30
View File
@@ -16,14 +16,12 @@ use codex_code_mode_protocol::FunctionCallOutputContentItem;
use codex_code_mode_protocol::ImageDetail;
use codex_code_mode_protocol::NotificationFuture;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::ToolInvocationFuture;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::WaitToPendingOutcome;
use codex_code_mode_protocol::WaitToPendingRequest;
use serde_json::Value as JsonValue;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use crate::SessionRuntime;
@@ -52,8 +50,6 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
#[derive(Default)]
@@ -87,9 +83,10 @@ impl CodeModeService {
}
}
pub async fn execute(&self, request: ExecuteRequest) -> Result<StartedCell, String> {
pub async fn execute(&self, request: ExecuteRequest) -> Result<RuntimeResponse, String> {
let yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS);
let started = self
let cell_id = request.cell_id.clone();
let event = self
.runtime
.execute(
runtime_request(request),
@@ -97,25 +94,15 @@ impl CodeModeService {
)
.await
.map_err(|error| error.to_string())?;
let cell_id = protocol_cell_id(&started.cell_id);
let response_cell_id = cell_id.clone();
let (response_tx, response_rx) = oneshot::channel();
tokio::spawn(async move {
let response = started
.initial_event()
.await
.map_err(|error| error.to_string())
.and_then(|event| runtime_response(&response_cell_id, event));
let _ = response_tx.send(response);
});
Ok(StartedCell::from_result_receiver(cell_id, response_rx))
runtime_response(&cell_id, event)
}
pub async fn execute_to_pending(
&self,
request: ExecuteRequest,
) -> Result<ExecuteToPendingOutcome, String> {
let started = self
let cell_id = request.cell_id.clone();
let event = self
.runtime
.execute(
runtime_request(request),
@@ -123,11 +110,6 @@ impl CodeModeService {
)
.await
.map_err(|error| error.to_string())?;
let cell_id = protocol_cell_id(&started.cell_id);
let event = started
.initial_event()
.await
.map_err(|error| error.to_string())?;
pending_outcome(&cell_id, event)
}
@@ -221,7 +203,7 @@ impl CodeModeSession for CodeModeService {
fn execute<'a>(
&'a self,
request: ExecuteRequest,
) -> CodeModeSessionResultFuture<'a, StartedCell> {
) -> CodeModeSessionResultFuture<'a, RuntimeResponse> {
Box::pin(CodeModeService::execute(self, request))
}
@@ -284,15 +266,11 @@ impl runtime::SessionRuntimeDelegate for ProtocolDelegate {
)
.await
}
async fn cell_closed(&self, cell_id: &runtime::CellId) -> Result<(), String> {
self.delegate.cell_closed(&protocol_cell_id(cell_id));
Ok(())
}
}
fn runtime_request(request: ExecuteRequest) -> runtime::ExecuteRequest {
runtime::ExecuteRequest {
cell_id: runtime_cell_id(&request.cell_id),
tool_call_id: request.tool_call_id,
enabled_tools: request
.enabled_tools
codex-rs/code-mode/src/service_contract_tests.rs
+96 -68
View File
@@ -4,6 +4,7 @@ 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;
@@ -16,7 +17,6 @@ enum DelegateEvent {
NotificationStarted,
NotificationFinished,
ToolStarted,
CellClosed(CellId),
}
struct BlockingDelegate {
@@ -30,7 +30,7 @@ struct NeverResolvingNotificationDelegate {
struct ReleasableNotificationDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
notification_release: Notify,
notification_release: Semaphore,
}
struct NeverResolvingToolDelegate {
@@ -50,14 +50,14 @@ impl ReleasableNotificationDelegate {
(
Arc::new(Self {
events_tx,
notification_release: Notify::new(),
notification_release: Semaphore::new(/* permits */ 0),
}),
events_rx,
)
}
fn release_notification(&self) {
self.notification_release.notify_one();
self.notification_release.add_permits(/* permits */ 1);
}
}
@@ -89,12 +89,6 @@ impl CodeModeSessionDelegate for NeverResolvingNotificationDelegate {
std::future::pending().await
})
}
fn cell_closed(&self, cell_id: &CellId) {
let _ = self
.events_tx
.send(DelegateEvent::CellClosed(cell_id.clone()));
}
}
impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
@@ -119,7 +113,7 @@ impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
tokio::select! {
_ = self.notification_release.notified() => {
_ = self.notification_release.acquire() => {
let _ = self.events_tx.send(DelegateEvent::NotificationFinished);
Ok(())
}
@@ -129,12 +123,6 @@ impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
}
})
}
fn cell_closed(&self, cell_id: &CellId) {
let _ = self
.events_tx
.send(DelegateEvent::CellClosed(cell_id.clone()));
}
}
impl CodeModeSessionDelegate for NeverResolvingToolDelegate {
@@ -158,12 +146,6 @@ impl CodeModeSessionDelegate for NeverResolvingToolDelegate {
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, cell_id: &CellId) {
let _ = self
.events_tx
.send(DelegateEvent::CellClosed(cell_id.clone()));
}
}
impl BlockingDelegate {
@@ -215,12 +197,6 @@ impl CodeModeSessionDelegate for BlockingDelegate {
Err("cancelled".to_string())
})
}
fn cell_closed(&self, cell_id: &CellId) {
let _ = self
.events_tx
.send(DelegateEvent::CellClosed(cell_id.clone()));
}
}
fn cell_id(value: &str) -> CellId {
@@ -229,6 +205,7 @@ fn cell_id(value: &str) -> CellId {
fn execute_request(source: &str) -> ExecuteRequest {
ExecuteRequest {
cell_id: cell_id("1"),
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
@@ -260,7 +237,13 @@ async fn yields_and_resumes() {
let service = CodeModeService::new();
let cell = service
.execute(ExecuteRequest {
source: r#"text("before"); yield_control(); text("after");"#.to_string(),
source: r#"
text("before");
yield_control();
store("finished", true);
text("after");
"#
.to_string(),
yield_time_ms: Some(60_000),
..execute_request("")
})
@@ -268,7 +251,7 @@ async fn yields_and_resumes() {
.unwrap();
assert_eq!(
cell.initial_response().await.unwrap(),
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: vec![FunctionCallOutputContentItem::InputText {
@@ -276,6 +259,30 @@ async fn yields_and_resumes() {
}],
}
);
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), async {
for cell_number in 2.. {
let response = service
.execute(ExecuteRequest {
cell_id: cell_id(&cell_number.to_string()),
..execute_request(r#"text(String(load("finished")));"#)
})
.await
.unwrap();
let 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
.wait(WaitRequest {
@@ -292,6 +299,20 @@ async fn yields_and_resumes() {
error_text: None,
})
);
assert_eq!(
service
.wait(WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
})
.await
.unwrap(),
WaitOutcome::MissingCell(RuntimeResponse::Result {
cell_id: cell_id("1"),
content_items: Vec::new(),
error_text: Some("exec cell 1 not found".to_string()),
})
);
}
#[tokio::test]
@@ -353,19 +374,19 @@ async fn observed_natural_completion_wins_over_termination() {
.unwrap();
assert_eq!(
cell.initial_response().await.unwrap(),
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
tokio::time::timeout(Duration::from_secs(1), async {
loop {
for cell_number in 2.. {
let response = service
.execute(execute_request(r#"text(String(load("finished")));"#))
.await
.unwrap()
.initial_response()
.execute(ExecuteRequest {
cell_id: cell_id(&cell_number.to_string()),
..execute_request(r#"text(String(load("finished")));"#)
})
.await
.unwrap();
let RuntimeResponse::Result { content_items, .. } = response else {
@@ -411,7 +432,7 @@ async fn termination_discards_pending_callbacks_before_responding() {
DelegateEvent::NotificationStarted
);
assert_eq!(
cell.initial_response().await.unwrap(),
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
@@ -424,10 +445,6 @@ async fn termination_discards_pending_callbacks_before_responding() {
content_items: Vec::new(),
})
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}
#[tokio::test]
@@ -435,7 +452,7 @@ 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
.execute(execute_request(r#"notify("pending");"#))
.execute(execute_request(r#"yield_control(); notify("pending");"#))
.await
.unwrap();
@@ -451,9 +468,33 @@ async fn shutdown_does_not_await_notifications_during_natural_completion() {
.await
.expect("shutdown should not await a non-cooperative notification")
.unwrap();
}
#[tokio::test]
async fn termination_does_not_await_notifications_during_natural_completion() {
let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new();
let service = CodeModeService::with_delegate(delegate);
service
.execute(execute_request(r#"yield_control(); notify("pending");"#))
.await
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
DelegateEvent::NotificationStarted
);
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(),
WaitOutcome::LiveCell(RuntimeResponse::Terminated {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
}
@@ -461,11 +502,14 @@ async fn shutdown_does_not_await_notifications_during_natural_completion() {
async fn natural_completion_waits_for_notifications_before_responding() {
let (delegate, mut events_rx) = ReleasableNotificationDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let cell = service
.execute(execute_request(r#"notify("pending");"#))
.await
.unwrap();
let mut initial_response = Box::pin(cell.initial_response());
let mut initial_response = Box::pin(service.execute(execute_request(r#"notify("pending");"#)));
std::future::poll_fn(|context| match initial_response.as_mut().poll(context) {
std::task::Poll::Pending => std::task::Poll::Ready(()),
std::task::Poll::Ready(result) => {
panic!("execute returned before the notification was released: {result:?}")
}
})
.await;
assert_eq!(
next_event(&mut events_rx).await,
@@ -491,10 +535,6 @@ async fn natural_completion_waits_for_notifications_before_responding() {
next_event(&mut events_rx).await,
DelegateEvent::NotificationFinished
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}
#[tokio::test]
@@ -512,7 +552,7 @@ async fn termination_discards_pending_tools_before_responding() {
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
assert_eq!(
cell.initial_response().await.unwrap(),
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
@@ -526,10 +566,6 @@ async fn termination_discards_pending_tools_before_responding() {
content_items: Vec::new(),
})
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}
#[tokio::test]
@@ -549,10 +585,6 @@ async fn shutdown_discards_pending_tools_before_returning() {
let shutdown_service = Arc::clone(&service);
assert_eq!(shutdown_service.shutdown().await, Ok(()));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}
#[tokio::test]
@@ -577,10 +609,6 @@ async fn shutdown_completes_when_a_nested_tool_ignores_cancellation() {
.await
.expect("shutdown should not await a non-cooperative nested tool")
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}
#[tokio::test]
@@ -592,7 +620,7 @@ async fn second_observer_is_rejected_without_displacing_the_first() {
.unwrap();
assert_eq!(
cell.initial_response().await.unwrap(),
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
@@ -641,7 +669,7 @@ async fn dropping_a_wait_allows_a_later_wait_to_observe_the_cell() {
.unwrap();
assert_eq!(
cell.initial_response().await.unwrap(),
cell,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
codex-rs/code-mode/src/service_tests.rs
+5 -11
View File
@@ -57,12 +57,11 @@ impl CodeModeSessionDelegate for ReleasableToolDelegate {
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
fn execute_request(source: &str) -> ExecuteRequest {
ExecuteRequest {
cell_id: cell_id("1"),
tool_call_id: "call_1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
@@ -87,13 +86,7 @@ fn echo_tool() -> ToolDefinition {
}
async fn execute(service: &CodeModeService, request: ExecuteRequest) -> RuntimeResponse {
service
.execute(request)
.await
.unwrap()
.initial_response()
.await
.unwrap()
service.execute(request).await.unwrap()
}
#[tokio::test]
@@ -140,6 +133,7 @@ async fn stored_values_are_shared_between_cells_but_not_sessions() {
let same_session = execute(
&first_session,
ExecuteRequest {
cell_id: cell_id("2"),
source: r#"text(String(load("key")));"#.to_string(),
yield_time_ms: None,
..execute_request("")
@@ -190,7 +184,7 @@ async fn stored_values_are_shared_between_cells_but_not_sessions() {
async fn shutdown_interrupts_cpu_bound_cells() {
let service = CodeModeService::new();
let cell = service
let response = service
.execute(ExecuteRequest {
source: "while (true) {}".to_string(),
..execute_request("")
@@ -198,7 +192,7 @@ async fn shutdown_interrupts_cpu_bound_cells() {
.await
.unwrap();
assert_eq!(
cell.initial_response().await.unwrap(),
response,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
codex-rs/code-mode/src/session_runtime/mod.rs
+7 -56
View File
@@ -4,14 +4,11 @@ use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use serde_json::Value as JsonValue;
use tokio::sync::Mutex;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tracing::warn;
pub use self::types::CellEvent;
pub use self::types::CellId;
@@ -41,16 +38,10 @@ pub struct SessionRuntime<D: SessionRuntimeDelegate> {
struct Inner<D: SessionRuntimeDelegate> {
stored_values: Mutex<HashMap<String, JsonValue>>,
cells: Mutex<HashMap<CellId, CellState>>,
cells: Mutex<HashMap<CellId, CellHandle>>,
cell_count_tx: watch::Sender<usize>,
shutdown_token: CancellationToken,
delegate: Arc<D>,
next_cell_id: AtomicU64,
}
enum CellState {
Live(CellHandle),
Closing,
}
impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
@@ -63,7 +54,6 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
cell_count_tx,
shutdown_token: CancellationToken::new(),
delegate,
next_cell_id: AtomicU64::new(1),
}),
}
}
@@ -76,18 +66,11 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
&self,
request: ExecuteRequest,
initial_observe_mode: ObserveMode,
) -> Result<StartedCell, Error> {
) -> Result<CellEvent, Error> {
if self.inner.shutdown_token.is_cancelled() {
return Err(Error::ShuttingDown);
}
let cell_id = self.allocate_cell_id();
let initial_event = self
.start_cell(cell_id.clone(), request, initial_observe_mode)
.await?;
Ok(StartedCell {
cell_id,
initial_event,
})
self.start_cell(request, initial_observe_mode).await?.await
}
pub async fn observe(&self, cell_id: &CellId, mode: ObserveMode) -> Result<CellEvent, Error> {
@@ -101,7 +84,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
) -> Result<PendingEvent, Error> {
let handle = {
let cells = self.inner.cells.lock().await;
let Some(CellState::Live(handle)) = cells.get(cell_id) else {
let Some(handle) = cells.get(cell_id) else {
return Err(Error::MissingCell(cell_id.clone()));
};
handle.clone()
@@ -114,7 +97,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
pub async fn terminate(&self, cell_id: &CellId) -> Result<CellEvent, Error> {
let handle = {
let cells = self.inner.cells.lock().await;
let Some(CellState::Live(handle)) = cells.get(cell_id) else {
let Some(handle) = cells.get(cell_id) else {
return Err(Error::MissingCell(cell_id.clone()));
};
handle.clone()
@@ -139,21 +122,12 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
Ok(())
}
fn allocate_cell_id(&self) -> CellId {
CellId::new(
self.inner
.next_cell_id
.fetch_add(1, Ordering::Relaxed)
.to_string(),
)
}
async fn start_cell(
&self,
cell_id: CellId,
request: ExecuteRequest,
initial_observe_mode: ObserveMode,
) -> Result<RuntimeEventFuture, Error> {
let cell_id = request.cell_id.clone();
let stored_values = self.inner.stored_values.lock().await.clone();
let host = Arc::new(RuntimeCellHost {
cell_id: cell_id.clone(),
@@ -174,7 +148,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
self.inner.shutdown_token.clone(),
)
.map_err(Error::Runtime)?;
cells.insert(cell_id.clone(), CellState::Live(handle));
cells.insert(cell_id.clone(), handle);
self.inner.cell_count_tx.send_replace(cells.len());
drop(cells);
tokio::spawn(task);
@@ -192,12 +166,6 @@ impl<D: SessionRuntimeDelegate> Drop for SessionRuntime<D> {
}
}
/// A cell admitted by [`SessionRuntime::execute`].
pub struct StartedCell {
pub cell_id: CellId,
initial_event: RuntimeEventFuture,
}
/// An admitted observation that has not reached its requested frontier yet.
pub struct PendingEvent {
event: RuntimeEventFuture,
@@ -209,12 +177,6 @@ impl PendingEvent {
}
}
impl StartedCell {
pub async fn initial_event(self) -> Result<CellEvent, Error> {
self.initial_event.await
}
}
struct RuntimeCellHost<D: SessionRuntimeDelegate> {
cell_id: CellId,
inner: Arc<Inner<D>>,
@@ -262,17 +224,6 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
}
async fn closed(&self) {
self.inner
.cells
.lock()
.await
.insert(self.cell_id.clone(), CellState::Closing);
if let Err(err) = self.inner.delegate.cell_closed(&self.cell_id).await {
warn!(
"failed to close code mode delegate state for cell {}: {err}",
self.cell_id
);
}
let mut cells = self.inner.cells.lock().await;
if cells.remove(&self.cell_id).is_some() {
self.inner.cell_count_tx.send_replace(cells.len());
codex-rs/code-mode/src/session_runtime/tests.rs
+79 -278
View File
@@ -1,63 +1,33 @@
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::task::Poll;
use std::time::Duration;
use pretty_assertions::assert_eq;
use serde_json::Value as JsonValue;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use super::*;
enum FirstCloseOutcome {
Succeeds,
Fails,
}
#[derive(Debug, PartialEq)]
enum NotificationEvent {
Started,
Closed(CellId),
}
struct BlockingCloseDelegate {
close_started_tx: mpsc::UnboundedSender<CellId>,
close_release: Semaphore,
fail_first_close: AtomicBool,
}
struct RecordingDelegate;
struct BlockingNotificationDelegate {
events_tx: mpsc::UnboundedSender<NotificationEvent>,
notification_started_tx: mpsc::UnboundedSender<()>,
}
impl BlockingCloseDelegate {
fn new(first_close_outcome: FirstCloseOutcome) -> (Arc<Self>, mpsc::UnboundedReceiver<CellId>) {
let (close_started_tx, close_started_rx) = mpsc::unbounded_channel();
impl BlockingNotificationDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<()>) {
let (notification_started_tx, notification_started_rx) = mpsc::unbounded_channel();
(
Arc::new(Self {
close_started_tx,
close_release: Semaphore::new(/*permits*/ 0),
fail_first_close: AtomicBool::new(matches!(
first_close_outcome,
FirstCloseOutcome::Fails
)),
notification_started_tx,
}),
close_started_rx,
notification_started_rx,
)
}
}
impl BlockingNotificationDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<NotificationEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
(Arc::new(Self { events_tx }), events_rx)
}
}
impl SessionRuntimeDelegate for BlockingCloseDelegate {
impl SessionRuntimeDelegate for RecordingDelegate {
async fn invoke_tool(
&self,
_invocation: NestedToolCall,
@@ -75,21 +45,6 @@ impl SessionRuntimeDelegate for BlockingCloseDelegate {
) -> Result<(), String> {
Ok(())
}
async fn cell_closed(&self, cell_id: &CellId) -> Result<(), String> {
self.close_started_tx
.send(cell_id.clone())
.map_err(|_| "test did not receive cell close".to_string())?;
self.close_release
.acquire()
.await
.map_err(|_| "test did not release cell close".to_string())?
.forget();
if self.fail_first_close.swap(false, Ordering::AcqRel) {
return Err("test close failure".to_string());
}
Ok(())
}
}
impl SessionRuntimeDelegate for BlockingNotificationDelegate {
@@ -108,22 +63,21 @@ impl SessionRuntimeDelegate for BlockingNotificationDelegate {
_text: String,
cancellation_token: CancellationToken,
) -> Result<(), String> {
self.events_tx
.send(NotificationEvent::Started)
self.notification_started_tx
.send(())
.map_err(|_| "test did not receive notification start".to_string())?;
cancellation_token.cancelled().await;
Ok(())
}
async fn cell_closed(&self, cell_id: &CellId) -> Result<(), String> {
self.events_tx
.send(NotificationEvent::Closed(cell_id.clone()))
.map_err(|_| "test did not receive cell close".to_string())
}
}
fn execute_request(source: &str) -> ExecuteRequest {
execute_request_with_id("1", source)
}
fn execute_request_with_id(cell_id: &str, source: &str) -> ExecuteRequest {
ExecuteRequest {
cell_id: CellId::new(cell_id),
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
@@ -131,20 +85,19 @@ fn execute_request(source: &str) -> ExecuteRequest {
}
#[tokio::test]
async fn closing_cells_reject_requests_while_delegate_cleanup_runs() {
let (delegate, mut close_started_rx) = BlockingCloseDelegate::new(FirstCloseOutcome::Succeeds);
let runtime = Arc::new(SessionRuntime::new(Arc::clone(&delegate)));
let started = runtime
async fn terminal_cells_are_unrouted_before_they_are_unregistered() {
let runtime = Arc::new(SessionRuntime::new(Arc::new(RecordingDelegate)));
let cell_id = CellId::new("1");
let event = runtime
.execute(
execute_request(r#"text("done");"#),
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
)
.await
.unwrap();
let cell_id = started.cell_id.clone();
assert_eq!(
started.initial_event().await.unwrap(),
event,
CellEvent::Completed {
content_items: vec![OutputItem::Text {
text: "done".to_string(),
@@ -152,8 +105,6 @@ async fn closing_cells_reject_requests_while_delegate_cleanup_runs() {
error_text: None,
}
);
assert_eq!(close_started_rx.recv().await, Some(cell_id.clone()));
assert_eq!(
runtime
.observe(&cell_id, ObserveMode::PendingFrontier)
@@ -164,181 +115,12 @@ async fn closing_cells_reject_requests_while_delegate_cleanup_runs() {
runtime.terminate(&cell_id).await,
Err(Error::MissingCell(cell_id.clone()))
);
let shutdown_runtime = Arc::clone(&runtime);
let mut shutdown = tokio::spawn(async move { shutdown_runtime.shutdown().await });
assert!(
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut shutdown)
.await
.is_err()
);
delegate.close_release.add_permits(/*n*/ 1);
assert_eq!(shutdown.await.unwrap(), Ok(()));
}
#[tokio::test]
async fn shutdown_waits_for_each_cell_to_finish_delegate_cleanup() {
let (delegate, mut close_started_rx) = BlockingCloseDelegate::new(FirstCloseOutcome::Succeeds);
let runtime = Arc::new(SessionRuntime::new(Arc::clone(&delegate)));
let completed = runtime
.execute(
execute_request(r#"text("done");"#),
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
)
.await
.unwrap();
assert_eq!(
completed.initial_event().await.unwrap(),
CellEvent::Completed {
content_items: vec![OutputItem::Text {
text: "done".to_string(),
}],
error_text: None,
}
);
let closing_cell_id = close_started_rx.recv().await.unwrap();
let live = runtime
.execute(
execute_request("while (true) {}"),
ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
)
.await
.unwrap();
assert_eq!(
live.initial_event().await.unwrap(),
CellEvent::Yielded {
content_items: Vec::new(),
}
);
let shutdown_runtime = Arc::clone(&runtime);
let mut shutdown = tokio::spawn(async move { shutdown_runtime.shutdown().await });
let terminated_cell_id = close_started_rx.recv().await.unwrap();
assert_ne!(closing_cell_id, terminated_cell_id);
assert!(
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut shutdown)
.await
.is_err()
);
delegate.close_release.add_permits(/*n*/ 1);
assert!(
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut shutdown)
.await
.is_err()
);
delegate.close_release.add_permits(/*n*/ 1);
assert_eq!(shutdown.await.unwrap(), Ok(()));
}
#[tokio::test]
async fn shutdown_waits_for_remaining_cells_when_delegate_cleanup_fails() {
let (delegate, mut close_started_rx) = BlockingCloseDelegate::new(FirstCloseOutcome::Fails);
let runtime = Arc::new(SessionRuntime::new(Arc::clone(&delegate)));
let completed = runtime
.execute(
execute_request(r#"text("done");"#),
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
)
.await
.unwrap();
assert_eq!(
completed.initial_event().await.unwrap(),
CellEvent::Completed {
content_items: vec![OutputItem::Text {
text: "done".to_string(),
}],
error_text: None,
}
);
let closing_cell_id = close_started_rx.recv().await.unwrap();
let live = runtime
.execute(
execute_request("while (true) {}"),
ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
)
.await
.unwrap();
assert_eq!(
live.initial_event().await.unwrap(),
CellEvent::Yielded {
content_items: Vec::new(),
}
);
let shutdown_runtime = Arc::clone(&runtime);
let mut shutdown = tokio::spawn(async move { shutdown_runtime.shutdown().await });
let terminated_cell_id = close_started_rx.recv().await.unwrap();
assert_ne!(closing_cell_id, terminated_cell_id);
delegate.close_release.add_permits(/*n*/ 1);
assert!(
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut shutdown)
.await
.is_err()
);
delegate.close_release.add_permits(/*n*/ 1);
assert_eq!(shutdown.await.unwrap(), Ok(()));
}
#[tokio::test]
async fn concurrent_shutdowns_wait_for_the_same_cell_cleanup() {
let (delegate, mut close_started_rx) = BlockingCloseDelegate::new(FirstCloseOutcome::Succeeds);
let runtime = Arc::new(SessionRuntime::new(Arc::clone(&delegate)));
let completed = runtime
.execute(
execute_request(r#"text("done");"#),
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
)
.await
.unwrap();
let cell_id = completed.cell_id.clone();
assert_eq!(
completed.initial_event().await.unwrap(),
CellEvent::Completed {
content_items: vec![OutputItem::Text {
text: "done".to_string(),
}],
error_text: None,
}
);
assert_eq!(close_started_rx.recv().await, Some(cell_id));
let first_runtime = Arc::clone(&runtime);
let mut first_shutdown = tokio::spawn(async move { first_runtime.shutdown().await });
let second_runtime = Arc::clone(&runtime);
let mut second_shutdown = tokio::spawn(async move { second_runtime.shutdown().await });
assert!(
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut first_shutdown)
.await
.is_err()
);
assert!(
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut second_shutdown)
.await
.is_err()
);
delegate.close_release.add_permits(/*n*/ 1);
assert_eq!(first_shutdown.await.unwrap(), Ok(()));
assert_eq!(second_shutdown.await.unwrap(), Ok(()));
assert_eq!(runtime.shutdown().await, Ok(()));
}
#[tokio::test]
async fn shutdown_rejects_cell_admission_queued_before_the_registry_lock() {
let (delegate, _) = BlockingCloseDelegate::new(FirstCloseOutcome::Succeeds);
let runtime = Arc::new(SessionRuntime::new(delegate));
let runtime = Arc::new(SessionRuntime::new(Arc::new(RecordingDelegate)));
let cells = runtime.inner.cells.lock().await;
let execution = runtime.execute(
@@ -374,8 +156,7 @@ async fn shutdown_rejects_cell_admission_queued_before_the_registry_lock() {
#[tokio::test]
async fn shutdown_rejects_cell_admission_queued_after_the_registry_lock() {
let (delegate, _) = BlockingCloseDelegate::new(FirstCloseOutcome::Succeeds);
let runtime = Arc::new(SessionRuntime::new(delegate));
let runtime = Arc::new(SessionRuntime::new(Arc::new(RecordingDelegate)));
let cells = runtime.inner.cells.lock().await;
let shutdown = runtime.shutdown();
@@ -409,8 +190,7 @@ async fn shutdown_rejects_cell_admission_queued_after_the_registry_lock() {
#[tokio::test]
async fn cancelling_shutdown_while_waiting_for_the_registry_lock_keeps_the_session_closed() {
let (delegate, _) = BlockingCloseDelegate::new(FirstCloseOutcome::Succeeds);
let runtime = SessionRuntime::new(delegate);
let runtime = SessionRuntime::new(Arc::new(RecordingDelegate));
let cells = runtime.inner.cells.lock().await;
let mut shutdown = Box::pin(runtime.shutdown());
@@ -439,18 +219,15 @@ async fn cancelling_shutdown_while_waiting_for_the_registry_lock_keeps_the_sessi
#[tokio::test]
async fn drop_terminates_cells_when_the_registry_is_locked() {
let (delegate, mut close_started_rx) = BlockingCloseDelegate::new(FirstCloseOutcome::Succeeds);
let runtime = SessionRuntime::new(Arc::clone(&delegate));
let started = runtime
.execute(
execute_request("while (true) {}"),
ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
)
.await
.unwrap();
let cell_id = started.cell_id.clone();
let runtime = SessionRuntime::new(Arc::new(RecordingDelegate));
assert_eq!(
started.initial_event().await.unwrap(),
runtime
.execute(
execute_request("while (true) {}"),
ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
)
.await
.unwrap(),
CellEvent::Yielded {
content_items: Vec::new(),
}
@@ -462,14 +239,6 @@ async fn drop_terminates_cells_when_the_registry_is_locked() {
drop(runtime);
drop(cells);
assert_eq!(
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), close_started_rx.recv())
.await
.unwrap(),
Some(cell_id)
);
delegate.close_release.add_permits(/*n*/ 1);
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), cell_count.changed())
.await
.unwrap()
@@ -481,20 +250,24 @@ async fn drop_terminates_cells_when_the_registry_is_locked() {
async fn drop_cancels_notifications_during_natural_completion_when_registry_is_locked() {
let (delegate, mut events_rx) = BlockingNotificationDelegate::new();
let runtime = SessionRuntime::new(Arc::clone(&delegate));
let started = runtime
.execute(
execute_request(r#"notify("pending");"#),
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
)
.await
.unwrap();
let cell_id = started.cell_id;
assert_eq!(
runtime
.execute(
execute_request(r#"yield_control(); notify("pending");"#),
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
)
.await
.unwrap(),
CellEvent::Yielded {
content_items: Vec::new(),
}
);
assert_eq!(
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), events_rx.recv())
.await
.unwrap(),
Some(NotificationEvent::Started)
Some(())
);
let inner = Arc::clone(&runtime.inner);
@@ -503,15 +276,43 @@ async fn drop_cancels_notifications_during_natural_completion_when_registry_is_l
drop(runtime);
drop(cells);
assert_eq!(
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), events_rx.recv())
.await
.unwrap(),
Some(NotificationEvent::Closed(cell_id))
);
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), cell_count.changed())
.await
.unwrap()
.unwrap();
assert_eq!(*cell_count.borrow_and_update(), 0);
}
#[tokio::test]
async fn client_owned_cell_id_is_preserved_and_active_duplicates_are_rejected() {
let runtime = SessionRuntime::new(Arc::new(RecordingDelegate));
let cell_id = CellId::new("client-generated");
assert_eq!(
runtime
.execute(
execute_request_with_id(cell_id.as_str(), "await new Promise(() => {});"),
ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
)
.await
.unwrap(),
CellEvent::Yielded {
content_items: Vec::new(),
}
);
assert_eq!(
runtime
.execute(
execute_request_with_id(cell_id.as_str(), "text('duplicate');"),
ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
)
.await,
Err(Error::DuplicateCell(cell_id.clone()))
);
assert_eq!(
runtime.terminate(&cell_id).await,
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
}
codex-rs/code-mode/src/session_runtime/types.rs
+1 -4
View File
@@ -74,6 +74,7 @@ pub enum ImageDetail {
/// Transport-neutral input for starting a cell.
pub struct ExecuteRequest {
pub cell_id: CellId,
pub tool_call_id: String,
pub enabled_tools: Vec<ToolDefinition>,
pub source: String,
@@ -114,8 +115,6 @@ pub struct NestedToolCall {
///
/// Implementations should forward callback cancellation tokens to downstream
/// work. The runtime stops awaiting callbacks once cancellation begins.
/// Implementations must not return from `cell_closed` until downstream routing
/// can no longer target the cell.
pub trait SessionRuntimeDelegate: Send + Sync + 'static {
fn invoke_tool(
&self,
@@ -130,8 +129,6 @@ pub trait SessionRuntimeDelegate: Send + Sync + 'static {
text: String,
cancellation_token: CancellationToken,
) -> impl Future<Output = Result<(), String>> + Send;
fn cell_closed(&self, cell_id: &CellId) -> impl Future<Output = Result<(), String>> + Send;
}
/// A failure reported by a session runtime operation.
codex-rs/core/src/tools/code_mode/delegate.rs
+10 -89
View File
@@ -1,6 +1,4 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use codex_code_mode::CellId;
use codex_code_mode::CodeModeNestedToolCall;
@@ -11,7 +9,6 @@ use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseItem;
use serde_json::Value as JsonValue;
use tokio::sync::oneshot;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use super::ExecContext;
@@ -24,7 +21,6 @@ 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>>>>,
}
impl CodeModeDispatchBroker {
@@ -33,18 +29,9 @@ impl CodeModeDispatchBroker {
Self {
dispatch_tx,
dispatch_rx,
dispatch_gates: Arc::new(Mutex::new(HashMap::new())),
}
}
pub(super) fn mark_cell_ready_for_dispatch(&self, cell_id: &CellId) {
dispatch_gate(&self.dispatch_gates, cell_id).send_replace(true);
}
pub(super) fn close_cell(&self, cell_id: &CellId) {
remove_dispatch_gate(&self.dispatch_gates, cell_id);
}
pub(super) fn start_turn_worker(
&self,
exec: ExecContext,
@@ -59,7 +46,6 @@ impl CodeModeDispatchBroker {
);
let host = Arc::new(CoreTurnHost { exec, tool_runtime });
let dispatch_rx = self.dispatch_rx.clone();
let dispatch_gates = Arc::clone(&self.dispatch_gates);
let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
tokio::spawn(async move {
loop {
@@ -78,17 +64,12 @@ impl CodeModeDispatchBroker {
cancellation_token,
response_tx,
} => {
let response = if wait_until_cell_ready_for_dispatch(
&dispatch_gates,
&cell_id,
&cancellation_token,
)
.await
{
host.notify(call_id, cell_id, text).await
} else {
remove_dispatch_gate(&dispatch_gates, &cell_id);
Err("code mode notification cancelled".to_string())
let response = tokio::select! {
biased;
_ = cancellation_token.cancelled() => {
Err("code mode notification cancelled".to_string())
}
response = host.notify(call_id, cell_id, text) => response,
};
let _ = response_tx.send(response);
}
@@ -97,17 +78,6 @@ impl CodeModeDispatchBroker {
cancellation_token,
response_tx,
} => {
let cell_id = invocation.cell_id.clone();
if !wait_until_cell_ready_for_dispatch(
&dispatch_gates,
&cell_id,
&cancellation_token,
)
.await
{
remove_dispatch_gate(&dispatch_gates, &cell_id);
continue;
}
let host = Arc::clone(&host);
tokio::spawn(async move {
let response = tokio::select! {
@@ -129,55 +99,6 @@ impl CodeModeDispatchBroker {
}
}
fn dispatch_gate(
dispatch_gates: &Mutex<HashMap<CellId, watch::Sender<bool>>>,
cell_id: &CellId,
) -> 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()
}
fn remove_dispatch_gate(
dispatch_gates: &Mutex<HashMap<CellId, watch::Sender<bool>>>,
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);
}
async fn wait_until_cell_ready_for_dispatch(
dispatch_gates: &Mutex<HashMap<CellId, watch::Sender<bool>>>,
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();
loop {
if *ready_rx.borrow_and_update() {
return true;
}
tokio::select! {
changed = ready_rx.changed() => {
if changed.is_err() {
return false;
}
}
_ = cancellation_token.cancelled() => return false,
}
}
}
impl CodeModeSessionDelegate for CodeModeDispatchBroker {
fn invoke_tool<'a>(
&'a self,
@@ -238,10 +159,6 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker {
}
})
}
fn cell_closed(&self, cell_id: &CellId) {
self.close_cell(cell_id);
}
}
enum DispatchMessage {
@@ -310,3 +227,7 @@ impl CoreTurnHost {
})
}
}
#[cfg(test)]
#[path = "delegate_tests.rs"]
mod tests;
codex-rs/core/src/tools/code_mode/delegate_tests.rs
+77
View File
@@ -0,0 +1,77 @@
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode::CellId;
use codex_code_mode::CodeModeSessionDelegate;
use pretty_assertions::assert_eq;
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use super::super::ExecContext;
use super::CodeModeDispatchBroker;
use crate::session::tests::make_session_and_context;
use crate::state::ActiveTurn;
use crate::tools::registry::ToolRegistry;
use crate::tools::router::ToolRouter;
use crate::turn_diff_tracker::TurnDiffTracker;
#[tokio::test]
async fn cancelled_notification_does_not_keep_the_turn_worker_alive() {
let (session, turn) = make_session_and_context().await;
let session = Arc::new(session);
let turn = Arc::new(turn);
let session_weak = Arc::downgrade(&session);
let active_turn = ActiveTurn::default();
let turn_state = Arc::clone(&active_turn.turn_state);
*session.active_turn.lock().await = Some(active_turn);
let _turn_state_guard = turn_state.lock().await;
let broker = CodeModeDispatchBroker::new();
let cell_id = CellId::new("cell-1".to_string());
let worker = broker.start_turn_worker(
ExecContext {
session: Arc::clone(&session),
turn: Arc::clone(&turn),
},
Arc::new(ToolRouter::from_parts(
ToolRegistry::empty_for_test(),
Vec::new(),
)),
Arc::new(Mutex::new(TurnDiffTracker::new())),
);
let cancellation_token = CancellationToken::new();
let notification = broker.notify(
"call-1".to_string(),
cell_id,
"pending".to_string(),
cancellation_token.clone(),
);
tokio::pin!(notification);
std::future::poll_fn(|context| match notification.as_mut().poll(context) {
std::task::Poll::Pending => std::task::Poll::Ready(()),
std::task::Poll::Ready(result) => {
panic!("notification returned before cancellation: {result:?}")
}
})
.await;
for _ in 0..8 {
tokio::task::yield_now().await;
}
cancellation_token.cancel();
assert_eq!(
notification.await,
Err("code mode notification cancelled".to_string())
);
drop(worker);
drop(turn);
drop(session);
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), async {
while session_weak.upgrade().is_some() {
tokio::task::yield_now().await;
}
})
.await
.expect("cancelled notification kept the turn worker alive");
}
codex-rs/core/src/tools/code_mode/execute_handler.rs
+24 -23
View File
@@ -39,20 +39,7 @@ impl CodeModeExecuteHandler {
let enabled_tools =
codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs);
let started_at = std::time::Instant::now();
let started_cell = exec
.session
.services
.code_mode_service
.execute(codex_code_mode::ExecuteRequest {
tool_call_id: call_id.clone(),
enabled_tools,
source: args.code.clone(),
yield_time_ms: args.yield_time_ms,
max_output_tokens: args.max_output_tokens,
})
.await
.map_err(FunctionCallError::RespondToModel)?;
let cell_id = started_cell.cell_id.clone();
let cell_id = codex_code_mode::CellId::new(uuid::Uuid::new_v4().to_string());
let runtime_cell_id = cell_id.to_string();
let code_cell_trace = exec
.session
@@ -64,14 +51,32 @@ impl CodeModeExecuteHandler {
call_id.as_str(),
args.code.as_str(),
);
exec.session
let response = match exec
.session
.services
.code_mode_service
.mark_cell_ready_for_dispatch(&cell_id);
let response = started_cell
.initial_response()
.execute(codex_code_mode::ExecuteRequest {
cell_id: cell_id.clone(),
tool_call_id: call_id.clone(),
enabled_tools,
source: args.code.clone(),
yield_time_ms: args.yield_time_ms,
max_output_tokens: args.max_output_tokens,
})
.await
.map_err(FunctionCallError::RespondToModel)?;
{
Ok(response) => response,
Err(error) => {
let response = codex_code_mode::RuntimeResponse::Result {
cell_id,
content_items: Vec::new(),
error_text: Some(error.clone()),
};
code_cell_trace.record_initial_response(&response);
code_cell_trace.record_ended(&response);
return Err(FunctionCallError::RespondToModel(error));
}
};
// Record the raw runtime boundary. The model-visible custom-tool output
// is produced by `handle_runtime_response` and later linked through
// `CodeCell.output_item_ids` in the reduced trace.
@@ -80,10 +85,6 @@ impl CodeModeExecuteHandler {
// here when the first response also ended the runtime.
if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) {
code_cell_trace.record_ended(&response);
exec.session
.services
.code_mode_service
.finish_cell_dispatch(&cell_id);
}
handle_runtime_response(&exec, response, args.max_output_tokens, started_at)
.await
codex-rs/core/src/tools/code_mode/mod.rs
+1 -9
View File
@@ -77,7 +77,7 @@ impl CodeModeService {
pub(crate) async fn execute(
&self,
request: codex_code_mode::ExecuteRequest,
) -> Result<codex_code_mode::StartedCell, String> {
) -> Result<codex_code_mode::RuntimeResponse, String> {
self.session()?.execute(request).await
}
@@ -102,14 +102,6 @@ impl CodeModeService {
}
}
pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode::CellId) {
self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id);
}
pub(crate) fn finish_cell_dispatch(&self, cell_id: &CellId) {
self.dispatch_broker.close_cell(cell_id);
}
pub(crate) fn start_turn_worker(
&self,
session: &Arc<Session>,
codex-rs/core/src/tools/code_mode/wait_handler.rs
-4
View File
@@ -115,10 +115,6 @@ impl CodeModeWaitHandler {
runtime_cell_id.as_str(),
)
.record_ended(response);
exec.session
.services
.code_mode_service
.finish_cell_dispatch(runtime_cell_id);
}
handle_runtime_response(&exec, wait_response.into(), args.max_tokens, started_at)
.await