openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Scope code mode callback delegates to individual executions (#44865)

Browse Source 
## Why

A delegate bound at session creation cannot provide different callbacks for cells sharing that session. Each execution needs to retain its own delegate across yields and release it when the cell is cleaned up.

## What changed

- Pass `CodeModeSessionDelegate` to `execute` instead of session creation.
- Route tool calls, notifications, and cell closure callbacks through the execution's delegate in the in-process runtime and the gRPC and stdio transports.
- Retain delegates with pending executions and live cells, releasing them through closure and cancellation cleanup.

## Testing

Add coverage for distinct delegates across yielded cells, gRPC callbacks before cell admission, and delegate release after completion or abandoned execution cleanup. Update transport tests to verify callbacks reach the owning cell's delegate.

GitOrigin-RevId: 7469f52104b993e790a40c65fb800ad02b7fd606
This commit is contained in:
Abhinav
2026-09-11 17:00:07 +00:00
committed by copyberry
parent bc5957eac9
commit 3305c4f31d
30 changed files with 1027 additions and 637 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-host/src/grpc/mod.rs
+1 -1
View File
@@ -148,7 +148,7 @@ impl GrpcCodeModeHost {
_ = session.closed.cancelled() => {
return Err(Status::cancelled("code-mode session is closed"));
}
result = session.runtime.execute(request) => {
result = session.runtime.execute(request, Arc::new(delegate::GrpcDelegate::new(Arc::downgrade(&session)))) => {
result.map_err(Status::failed_precondition)?
}
};
codex-rs/code-mode-host/src/grpc/session.rs
+4 -9
View File
@@ -27,7 +27,6 @@ use tonic::Status;
use uuid::Uuid;
use super::GrpcStream;
use super::delegate::GrpcDelegate;
use super::events::EventSender;
use super::validation;
use super::waits::ActiveWait;
@@ -213,7 +212,6 @@ impl GrpcSession {
limits: CodeModeSessionCellExecutionLimits,
) -> Arc<Self> {
Arc::new_cyclic(|weak: &Weak<Self>| {
let delegate = Arc::new(GrpcDelegate::new(weak.clone()));
let failure_session = weak.clone();
let failure_handler = Arc::new(move |reason: String| {
if let Some(session) = failure_session.upgrade() {
@@ -223,13 +221,10 @@ impl GrpcSession {
});
Self {
id,
runtime: Arc::new(
InProcessCodeModeSession::with_delegate_and_task_failure_handler(
delegate,
failure_handler,
limits,
),
),
runtime: Arc::new(InProcessCodeModeSession::with_task_failure_handler(
failure_handler,
limits,
)),
closed,
state: Mutex::new(SessionState::default()),
events,
codex-rs/code-mode-host/src/lib.rs
+9 -12
View File
@@ -412,7 +412,11 @@ impl HostState {
);
return;
};
let result = session.execute(request).await;
let delegate = Arc::new(RemoteDelegate::new(
session_id.clone(),
Arc::clone(&self.peer),
));
let result = session.execute(request, delegate).await;
match result {
Ok(started) => {
let cell_id = started.cell_id.clone();
@@ -517,10 +521,6 @@ impl HostState {
{
return Err(format!("code-mode session ID `{session_id}` was reused"));
}
let delegate = Arc::new(RemoteDelegate::new(
session_id.clone(),
Arc::clone(&self.peer),
));
let peer = Arc::downgrade(&self.peer);
let task_failure_handler = Arc::new(move |reason| {
if let Some(peer) = peer.upgrade() {
@@ -529,13 +529,10 @@ impl HostState {
});
sessions.insert(
session_id,
Arc::new(
InProcessCodeModeSession::with_delegate_and_task_failure_handler(
delegate,
task_failure_handler,
cell_execution_limits,
),
),
Arc::new(InProcessCodeModeSession::with_task_failure_handler(
task_failure_handler,
cell_execution_limits,
)),
);
Ok(())
}
codex-rs/code-mode-host/tests/grpc.rs
+146 -60
View File
@@ -144,10 +144,11 @@ fn text_response(
async fn execute(
session: &Arc<dyn CodeModeSession>,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> Result<RuntimeResponse> {
timeout(TEST_TIMEOUT, async {
session
.execute(request)
.execute(request, delegate.clone())
.await
.map_err(anyhow::Error::msg)?
.initial_response()
@@ -190,7 +191,7 @@ async fn grpc_endpoints_reject_credentials_without_disclosing_them() {
] {
let provider = GrpcCodeModeSessionProvider::new(endpoint);
let error = provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.create_session()
.await
.err()
.expect("gRPC credentials should be rejected");
@@ -207,12 +208,18 @@ async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closur
assert!(host.endpoint.starts_with("http://127.0.0.1:"));
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(RecordingDelegate::default());
let first_delegate = Arc::new(RecordingDelegate::default());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let actual = execute(&session, request(r#"store("key", "persisted");"#)).await?;
let actual = execute(
&session,
request(r#"store("key", "persisted");"#),
first_delegate.clone(),
)
.await?;
assert_eq!(
actual,
RuntimeResponse::Result {
@@ -228,7 +235,7 @@ async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closur
);
callback.tool_call_id = "call-2".to_string();
callback.enabled_tools = vec![tool("echo")];
let actual = execute(&session, callback).await?;
let actual = execute(&session, callback, delegate.clone()).await?;
assert_eq!(
actual,
text_response("2", "output", actual.code_mode_host_duration())
@@ -263,7 +270,14 @@ async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closur
.closed_cells
.lock()
.unwrap_or_else(PoisonError::into_inner),
vec![cell_id("1"), cell_id("2")]
vec![cell_id("2")]
);
assert_eq!(
*first_delegate
.closed_cells
.lock()
.unwrap_or_else(PoisonError::into_inner),
vec![cell_id("1")]
);
Ok(())
}
@@ -273,14 +287,20 @@ async fn shutdown_immediately_rejects_new_operations() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let session = provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let shutdown = session.shutdown();
let expected = "code mode session is shutting down".to_string();
assert_eq!(
session.execute(request("text('too late');")).await.err(),
session
.execute(
request("text('too late');"),
Arc::new(NoopCodeModeSessionDelegate)
)
.await
.err(),
Some(expected.clone())
);
assert_eq!(
@@ -304,13 +324,18 @@ async fn cancelling_execution_before_admission_keeps_the_session_usable() -> Res
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(RecordingDelegate::default());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let mut pending = request("await new Promise(() => {});");
pending.yield_time_ms = Some(/*value*/ 1);
assert!(session.execute(pending).now_or_never().is_none());
assert!(
session
.execute(pending, delegate.clone())
.now_or_never()
.is_none()
);
let abandoned_cell = cell_id("1");
timeout(TEST_TIMEOUT, async {
@@ -354,7 +379,12 @@ async fn cancelling_execution_before_admission_keeps_the_session_usable() -> Res
.await
.context("cancelled execution leaked its remote cell")??;
let actual = execute(&session, request(r#"text("still alive");"#)).await?;
let actual = execute(
&session,
request(r#"text("still alive");"#),
delegate.clone(),
)
.await?;
assert_eq!(
actual,
text_response("2", "still alive", actual.code_mode_host_duration())
@@ -370,12 +400,15 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(RecordingDelegate::default());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let mut pending = request("await new Promise(() => {});");
pending.yield_time_ms = Some(/*value*/ 1);
let started = session.execute(pending).await.map_err(anyhow::Error::msg)?;
let started = session
.execute(pending, delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
let abandoned_cell = started.cell_id.clone();
timeout(TEST_TIMEOUT, async {
@@ -436,7 +469,12 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec
.contains(&abandoned_cell)
);
let actual = execute(&session, request(r#"text("still alive");"#)).await?;
let actual = execute(
&session,
request(r#"text("still alive");"#),
delegate.clone(),
)
.await?;
assert_eq!(
actual,
text_response("2", "still alive", actual.code_mode_host_duration())
@@ -451,12 +489,15 @@ async fn dropping_an_initial_response_terminates_its_pending_remote_execution()
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(RecordingDelegate::default());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let mut pending = request("await new Promise(() => {});");
pending.yield_time_ms = Some(/*value*/ 60_000);
let started = session.execute(pending).await.map_err(anyhow::Error::msg)?;
let started = session
.execute(pending, delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
let abandoned_cell = started.cell_id.clone();
let initial_response = tokio::spawn(started.initial_response());
tokio::task::yield_now().await;
@@ -479,7 +520,12 @@ async fn dropping_an_initial_response_terminates_its_pending_remote_execution()
.await
.context("dropping an initial response did not terminate its pending remote execution")?;
let actual = execute(&session, request(r#"text("still alive");"#)).await?;
let actual = execute(
&session,
request(r#"text("still alive");"#),
delegate.clone(),
)
.await?;
assert_eq!(
actual,
text_response("2", "still alive", actual.code_mode_host_duration())
@@ -493,7 +539,7 @@ async fn synchronous_delegate_panics_do_not_orphan_callbacks_or_close_the_sessio
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let session = provider
.create_session(Arc::new(PanickingDelegate))
.create_session()
.await
.map_err(anyhow::Error::msg)?;
@@ -501,7 +547,7 @@ async fn synchronous_delegate_panics_do_not_orphan_callbacks_or_close_the_sessio
r#"try { await tools.echo({}); text("unexpected"); } catch (_) { text("tool recovered"); }"#,
);
tool_panic.enabled_tools = vec![tool("echo")];
let actual = execute(&session, tool_panic).await?;
let actual = execute(&session, tool_panic, Arc::new(PanickingDelegate)).await?;
assert_eq!(
actual,
text_response("1", "tool recovered", actual.code_mode_host_duration())
@@ -510,6 +556,7 @@ async fn synchronous_delegate_panics_do_not_orphan_callbacks_or_close_the_sessio
let actual = execute(
&session,
request(r#"notify("panic"); text("notification recovered");"#),
Arc::new(PanickingDelegate),
)
.await?;
assert_eq!(
@@ -520,7 +567,12 @@ async fn synchronous_delegate_panics_do_not_orphan_callbacks_or_close_the_sessio
actual.code_mode_host_duration()
)
);
let actual = execute(&session, request(r#"text("still alive");"#)).await?;
let actual = execute(
&session,
request(r#"text("still alive");"#),
Arc::new(PanickingDelegate),
)
.await?;
assert_eq!(
actual,
text_response("3", "still alive", actual.code_mode_host_duration())
@@ -535,19 +587,24 @@ async fn tool_delegate_self_cancellation_returns_an_error_without_hanging() -> R
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let session = provider
.create_session(Arc::new(SelfCancellingToolDelegate))
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let mut callback =
request(r#"try { await tools.echo({}); } catch (_) { text("tool recovered"); }"#);
callback.enabled_tools = vec![tool("echo")];
let actual = execute(&session, callback).await?;
let actual = execute(&session, callback, Arc::new(SelfCancellingToolDelegate)).await?;
assert_eq!(
actual,
text_response("1", "tool recovered", actual.code_mode_host_duration())
);
let actual = execute(&session, request(r#"text("still alive");"#)).await?;
let actual = execute(
&session,
request(r#"text("still alive");"#),
Arc::new(SelfCancellingToolDelegate),
)
.await?;
assert_eq!(
actual,
text_response("2", "still alive", actual.code_mode_host_duration())
@@ -562,12 +619,15 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let session = provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let mut pending = request("await new Promise(() => {});");
pending.yield_time_ms = Some(/*value*/ 1);
let started = session.execute(pending).await.map_err(anyhow::Error::msg)?;
let started = session
.execute(pending, Arc::new(NoopCodeModeSessionDelegate))
.await
.map_err(anyhow::Error::msg)?;
let running_cell = started.cell_id.clone();
let actual = started
.initial_response()
@@ -639,12 +699,15 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<(
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let session = provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let mut pending = request("await new Promise(() => {});");
pending.yield_time_ms = Some(/*value*/ 1);
let started = session.execute(pending).await.map_err(anyhow::Error::msg)?;
let started = session
.execute(pending, Arc::new(NoopCodeModeSessionDelegate))
.await
.map_err(anyhow::Error::msg)?;
let running_cell = started.cell_id.clone();
let actual = started
.initial_response()
@@ -710,12 +773,12 @@ async fn dropping_a_session_off_runtime_retires_its_active_cells() -> Result<()>
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(RecordingDelegate::default());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let mut pending = request("await new Promise(() => {});");
pending.yield_time_ms = Some(/*value*/ 1);
let actual = execute(&session, pending).await?;
let actual = execute(&session, pending, delegate.clone()).await?;
assert_eq!(
actual,
RuntimeResponse::Yielded {
@@ -757,11 +820,11 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() ->
release: Semaphore::new(/*permits*/ 0),
});
let slow_session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let fast_session = provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let mut slow_request = request(
@@ -770,7 +833,7 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() ->
slow_request.enabled_tools = vec![tool("large")];
slow_request.yield_time_ms = Some(/*value*/ 20_000);
let slow_cell = slow_session
.execute(slow_request)
.execute(slow_request, delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
timeout(TEST_TIMEOUT, delegate.started.acquire())
@@ -778,7 +841,12 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() ->
.context("large tool callback did not start")??
.forget();
let actual = execute(&fast_session, request(r#"text("fast-before");"#)).await?;
let actual = execute(
&fast_session,
request(r#"text("fast-before");"#),
Arc::new(NoopCodeModeSessionDelegate),
)
.await?;
assert_eq!(
actual,
text_response("1", "fast-before", actual.code_mode_host_duration())
@@ -786,7 +854,11 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() ->
delegate.release.add_permits(/*n*/ 1);
let slow_response = timeout(TEST_TIMEOUT, slow_cell.initial_response());
let fast_response = execute(&fast_session, request(r#"text("fast-during");"#));
let fast_response = execute(
&fast_session,
request(r#"text("fast-during");"#),
Arc::new(NoopCodeModeSessionDelegate),
);
let (slow_response, fast_response) = tokio::join!(slow_response, fast_response);
let actual = fast_response?;
@@ -822,7 +894,7 @@ async fn single_subscription_processes_slow_and_fast_tools_concurrently() -> Res
release: Semaphore::new(/*permits*/ 0),
});
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
@@ -830,7 +902,10 @@ async fn single_subscription_processes_slow_and_fast_tools_concurrently() -> Res
request(r#"const result = await tools.large({}); text(String(result.value.length));"#);
slow.enabled_tools = vec![tool("large")];
slow.yield_time_ms = Some(/*value*/ 20_000);
let slow_cell = session.execute(slow).await.map_err(anyhow::Error::msg)?;
let slow_cell = session
.execute(slow, delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
timeout(TEST_TIMEOUT, delegate.started.acquire())
.await
.context("slow tool did not start")??
@@ -838,7 +913,7 @@ async fn single_subscription_processes_slow_and_fast_tools_concurrently() -> Res
let mut fast = request(r#"const result = await tools.fast({}); text(result.value);"#);
fast.enabled_tools = vec![tool("fast")];
let actual = execute(&session, fast).await?;
let actual = execute(&session, fast, delegate.clone()).await?;
assert_eq!(
actual,
text_response("2", "isolated", actual.code_mode_host_duration())
@@ -869,29 +944,23 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let limited = provider
.create_session_with_limits(
Arc::new(NoopCodeModeSessionDelegate),
CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(/*value*/ 1),
max_heap_size_bytes: Some(/*value*/ 16 * 1024 * 1024),
},
)
.create_session_with_limits(CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(/*value*/ 1),
max_heap_size_bytes: Some(/*value*/ 16 * 1024 * 1024),
})
.await
.map_err(anyhow::Error::msg)?;
let other = provider
.create_session_with_limits(
Arc::new(NoopCodeModeSessionDelegate),
CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(/*value*/ 1_000),
max_heap_size_bytes: None,
},
)
.create_session_with_limits(CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(/*value*/ 1_000),
max_heap_size_bytes: None,
})
.await
.map_err(anyhow::Error::msg)?;
let mut pending = request("await new Promise(() => {});");
pending.yield_time_ms = Some(/*value*/ 60_000);
let actual = execute(&limited, pending).await?;
let actual = execute(&limited, pending, Arc::new(NoopCodeModeSessionDelegate)).await?;
assert_eq!(
actual,
RuntimeResponse::Yielded {
@@ -921,6 +990,7 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> {
let actual = execute(
&other,
request(r#"await new Promise(resolve => setTimeout(resolve, 25)); text("isolated");"#),
Arc::new(NoopCodeModeSessionDelegate),
)
.await?;
assert_eq!(
@@ -939,7 +1009,12 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> {
content_items: Vec::new(),
})
);
let actual = execute(&limited, request("await new Promise(() => {});")).await?;
let actual = execute(
&limited,
request("await new Promise(() => {});"),
Arc::new(NoopCodeModeSessionDelegate),
)
.await?;
assert_eq!(
actual,
RuntimeResponse::Yielded {
@@ -1008,14 +1083,17 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
let provider = GrpcCodeModeSessionProvider::new(original.endpoint.clone());
let delegate = Arc::new(RecordingDelegate::default());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let mut pending = request("await tools.echo({generation: 1}); await new Promise(() => {});");
pending.enabled_tools = vec![tool("echo")];
pending.yield_time_ms = Some(/*value*/ 1);
let started = session.execute(pending).await.map_err(anyhow::Error::msg)?;
let started = session
.execute(pending, delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
let old_cell_id = started.cell_id.clone();
assert_eq!(old_cell_id, cell_id("1"));
assert!(matches!(
@@ -1078,8 +1156,8 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
callback.tool_call_id = "reconnected-call".to_string();
callback.enabled_tools = vec![tool("echo")];
let (callback_response, concurrent_response) = tokio::join!(
execute(&session, callback),
execute(&session, request(r#"text("concurrent")"#)),
execute(&session, callback, delegate.clone()),
execute(&session, request(r#"text("concurrent")"#), delegate.clone()),
);
let callback_response = callback_response?;
let RuntimeResponse::Result {
@@ -1145,7 +1223,10 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
let mut pending = request("await new Promise(() => {});");
pending.yield_time_ms = Some(/*value*/ 1);
let started = session.execute(pending).await.map_err(anyhow::Error::msg)?;
let started = session
.execute(pending, delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
let replacement_cell_id = started.cell_id.clone();
assert_eq!(replacement_cell_id, cell_id("g2:3"));
let actual = started
@@ -1219,10 +1300,15 @@ async fn unix_socket_endpoints_execute_code_mode_cells() -> Result<()> {
format!("unix:{}", socket_path.display()),
] {
let session = GrpcCodeModeSessionProvider::new(endpoint)
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let actual = execute(&session, request(r#"text("unix socket")"#)).await?;
let actual = execute(
&session,
request(r#"text("unix socket")"#),
Arc::new(NoopCodeModeSessionDelegate),
)
.await?;
assert_eq!(
actual,
text_response("1", "unix socket", actual.code_mode_host_duration())
codex-rs/code-mode-host/tests/grpc_notifications.rs
+28 -9
View File
@@ -120,10 +120,11 @@ fn text_response(
async fn execute(
session: &Arc<dyn CodeModeSession>,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> Result<RuntimeResponse> {
timeout(TEST_TIMEOUT, async {
session
.execute(request)
.execute(request, delegate.clone())
.await
.map_err(anyhow::Error::msg)?
.initial_response()
@@ -140,13 +141,19 @@ async fn completed_cells_drain_pending_notifications_before_completion() -> Resu
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(BlockingNotificationDelegate::new());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let executing = Arc::clone(&session);
let execution_delegate = delegate.clone();
let completion = tokio::spawn(async move {
execute(&executing, request(r#"notify("notice"); text("done");"#)).await
execute(
&executing,
request(r#"notify("notice"); text("done");"#),
execution_delegate,
)
.await
});
timeout(TEST_TIMEOUT, delegate.started.acquire())
.await
@@ -171,7 +178,12 @@ async fn completed_cells_drain_pending_notifications_before_completion() -> Resu
.context("completed cell was not retired")??
.forget();
let actual = execute(&session, request(r#"text("still alive");"#)).await?;
let actual = execute(
&session,
request(r#"text("still alive");"#),
delegate.clone(),
)
.await?;
assert_eq!(
actual,
text_response("2", "still alive", actual.code_mode_host_duration())
@@ -187,11 +199,14 @@ async fn completed_waits_drain_pending_notifications_before_returning() -> Resul
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(BlockingNotificationDelegate::new());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let pending = request(r#"yield_control(); notify("notice"); text("done");"#);
let cell = session.execute(pending).await.map_err(anyhow::Error::msg)?;
let cell = session
.execute(pending, delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
let actual = cell.initial_response().await.map_err(anyhow::Error::msg)?;
assert_eq!(
actual,
@@ -248,12 +263,15 @@ async fn termination_cancels_pending_notifications() -> Result<()> {
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(BlockingNotificationDelegate::new());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let mut pending = request(r#"notify("notice"); await new Promise(() => {});"#);
pending.yield_time_ms = Some(/*value*/ 1);
let cell = session.execute(pending).await.map_err(anyhow::Error::msg)?;
let cell = session
.execute(pending, delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
timeout(TEST_TIMEOUT, delegate.started.acquire())
.await
@@ -299,13 +317,14 @@ async fn oversized_notification_text_is_delivered_unchanged() -> Result<()> {
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(RecordingDelegate::default());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let actual = execute(
&session,
request(r#"notify("🦀".repeat(512)); text("done");"#),
delegate.clone(),
)
.await?;
assert_eq!(
codex-rs/code-mode-host/tests/grpc_timing.rs
+24 -18
View File
@@ -119,7 +119,7 @@ async fn stdio_execution_timing_includes_javascript_but_excludes_delayed_reads()
);
let delegate = Arc::new(RecordingDelegate::default());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.map_err(anyhow::Error::msg)?;
@@ -127,17 +127,20 @@ async fn stdio_execution_timing_includes_javascript_but_excludes_delayed_reads()
for (index, suffix) in endings.into_iter().enumerate() {
let observed = Instant::now();
let started = session
.execute(ExecuteRequest {
tool_call_id: format!("call-{index}"),
source: format!(
"const until = Date.now() + 150; while (Date.now() < until) {{}} \
.execute(
ExecuteRequest {
tool_call_id: format!("call-{index}"),
source: format!(
"const until = Date.now() + 150; while (Date.now() < until) {{}} \
await new Promise(resolve => setTimeout(resolve, 150)); \
notify('finished'); {suffix}"
),
enabled_tools: Vec::new(),
yield_time_ms: Some(/*value*/ 5_000),
max_output_tokens: Some(/*value*/ 1_000),
})
),
enabled_tools: Vec::new(),
yield_time_ms: Some(/*value*/ 5_000),
max_output_tokens: Some(/*value*/ 1_000),
},
delegate.clone(),
)
.await
.map_err(anyhow::Error::msg)?;
timeout(TEST_TIMEOUT, delegate.notification_delivered.notified())
@@ -184,17 +187,20 @@ async fn observation_timing_excludes_previous_requests_and_background_time() ->
];
for provider in providers {
let session = provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.create_session()
.await
.map_err(anyhow::Error::msg)?;
let started = session
.execute(ExecuteRequest {
tool_call_id: "call-1".to_string(),
source: "await new Promise(() => {});".to_string(),
enabled_tools: Vec::new(),
yield_time_ms: Some(/*value*/ 200),
max_output_tokens: Some(/*value*/ 1_000),
})
.execute(
ExecuteRequest {
tool_call_id: "call-1".to_string(),
source: "await new Promise(() => {});".to_string(),
enabled_tools: Vec::new(),
yield_time_ms: Some(/*value*/ 200),
max_output_tokens: Some(/*value*/ 1_000),
},
Arc::new(NoopCodeModeSessionDelegate),
)
.await
.map_err(anyhow::Error::msg)?;
let cell_id = started.cell_id.clone();
codex-rs/code-mode-host/tests/stdio.rs
+103 -44
View File
@@ -174,9 +174,13 @@ fn execute_request(source: &str) -> ExecuteRequest {
}
}
async fn execute(session: &Arc<dyn CodeModeSession>, request: ExecuteRequest) -> RuntimeResponse {
async fn execute(
session: &Arc<dyn CodeModeSession>,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> RuntimeResponse {
session
.execute(request)
.execute(request, delegate.clone())
.await
.expect("start execution")
.initial_response()
@@ -187,8 +191,12 @@ async fn execute(session: &Arc<dyn CodeModeSession>, request: ExecuteRequest) ->
async fn execute_to_terminal(
session: &Arc<dyn CodeModeSession>,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> RuntimeResponse {
let started = session.execute(request).await.expect("start execution");
let started = session
.execute(request, delegate.clone())
.await
.expect("start execution");
let mut response = started.initial_response().await.expect("initial response");
loop {
match response {
@@ -227,29 +235,27 @@ async fn session_execution_limits_are_isolated_on_a_shared_process_host() {
codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"),
));
let limited = provider
.create_session_with_limits(
Arc::new(RecordingDelegate::default()),
CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(1),
max_heap_size_bytes: None,
},
)
.create_session_with_limits(CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(1),
max_heap_size_bytes: None,
})
.await
.expect("create limited session");
let other = provider
.create_session_with_limits(
Arc::new(RecordingDelegate::default()),
CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(1_000),
max_heap_size_bytes: None,
},
)
.create_session_with_limits(CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(1_000),
max_heap_size_bytes: None,
})
.await
.expect("create independently limited session");
let response = tokio::time::timeout(
Duration::from_secs(5),
execute(&limited, execute_request("await new Promise(() => {});")),
execute(
&limited,
execute_request("await new Promise(() => {});"),
Arc::new(RecordingDelegate::default()),
),
)
.await
.expect("session limit should bound the default execution wait");
@@ -297,12 +303,18 @@ async fn remote_session_persists_values_forwards_delegates_and_controls_cells()
codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"),
);
let delegate = Arc::new(RecordingDelegate::default());
let first_delegate = Arc::new(RecordingDelegate::default());
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.expect("create remote session");
let actual = execute(&session, execute_request(r#"store("key", "persisted");"#)).await;
let actual = execute(
&session,
execute_request(r#"store("key", "persisted");"#),
first_delegate.clone(),
)
.await;
assert_eq!(
actual,
RuntimeResponse::Result {
@@ -329,7 +341,7 @@ text(result.value);
input_schema: None,
output_schema: None,
}];
let actual = execute(&session, callback_request).await;
let actual = execute(&session, callback_request, delegate.clone()).await;
assert_eq!(
actual,
RuntimeResponse::Result {
@@ -359,7 +371,7 @@ text(result.value);
let mut pending_request = execute_request("await new Promise(() => {});");
pending_request.tool_call_id = "call-3".to_string();
pending_request.yield_time_ms = Some(1);
let actual = execute(&session, pending_request).await;
let actual = execute(&session, pending_request, delegate.clone()).await;
assert_eq!(
actual,
RuntimeResponse::Yielded {
@@ -399,7 +411,14 @@ text(result.value);
session.shutdown().await.expect("shutdown remote session");
assert_eq!(
*delegate.closed_cells.lock().expect("closed cells lock"),
vec![cell_id("1"), cell_id("2"), cell_id("3")]
vec![cell_id("2"), cell_id("3")]
);
assert_eq!(
*first_delegate
.closed_cells
.lock()
.expect("closed cells lock"),
vec![cell_id("1")]
);
}
@@ -409,12 +428,15 @@ async fn dropping_long_wait_releases_observer_before_next_wait() {
codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"),
);
let session = provider
.create_session(Arc::new(RecordingDelegate::default()))
.create_session()
.await
.expect("create remote session");
let mut request = execute_request("await new Promise(() => {});");
request.yield_time_ms = Some(1);
let started = session.execute(request).await.expect("start execution");
let started = session
.execute(request, Arc::new(RecordingDelegate::default()))
.await
.expect("start execution");
let running_cell_id = started.cell_id.clone();
let actual = started.initial_response().await.expect("initial response");
assert_eq!(
@@ -472,7 +494,7 @@ async fn unawaited_slow_tool_is_cancelled_after_parallel_tools_complete() {
);
let (delegate, mut events_rx) = CancellationDelegate::new();
let session = provider
.create_session(delegate.clone())
.create_session()
.await
.expect("create remote session");
let mut request = execute_request(
@@ -508,7 +530,10 @@ return;
})
.collect();
let started = session.execute(request).await.expect("start execution");
let started = session
.execute(request, delegate.clone())
.await
.expect("start execution");
let running_cell_id = started.cell_id.clone();
let actual = started.initial_response().await.expect("initial response");
assert_eq!(
@@ -587,11 +612,14 @@ async fn oversized_execute_request_does_not_close_the_shared_host() {
codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"),
);
let session = provider
.create_session(Arc::new(RecordingDelegate::default()))
.create_session()
.await
.expect("create remote session");
let error = session
.execute(execute_request(&"x".repeat(MAX_FRAME_BYTES)))
.execute(
execute_request(&"x".repeat(MAX_FRAME_BYTES)),
Arc::new(RecordingDelegate::default()),
)
.await
.err()
.expect("oversized execute should fail");
@@ -600,7 +628,12 @@ async fn oversized_execute_request_does_not_close_the_shared_host() {
"unexpected error: {error}"
);
let actual = execute(&session, execute_request(r#"text("still alive");"#)).await;
let actual = execute(
&session,
execute_request(r#"text("still alive");"#),
Arc::new(RecordingDelegate::default()),
)
.await;
assert_eq!(
actual,
RuntimeResponse::Result {
@@ -621,7 +654,7 @@ async fn oversized_delegate_payloads_fail_only_the_tool_call() {
codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"),
);
let session = provider
.create_session(Arc::new(OversizedResultDelegate))
.create_session()
.await
.expect("create remote session");
let tool = |name: &str| ToolDefinition {
@@ -644,7 +677,12 @@ try {{
));
oversized_argument.enabled_tools = vec![tool("big_argument")];
oversized_argument.yield_time_ms = Some(60_000);
let actual = execute_to_terminal(&session, oversized_argument).await;
let actual = execute_to_terminal(
&session,
oversized_argument,
Arc::new(OversizedResultDelegate),
)
.await;
assert_eq!(
actual,
RuntimeResponse::Result {
@@ -668,7 +706,12 @@ try {
);
oversized_result.enabled_tools = vec![tool("big_result")];
oversized_result.yield_time_ms = Some(60_000);
let actual = execute_to_terminal(&session, oversized_result).await;
let actual = execute_to_terminal(
&session,
oversized_result,
Arc::new(OversizedResultDelegate),
)
.await;
assert_eq!(
actual,
RuntimeResponse::Result {
@@ -681,7 +724,12 @@ try {
}
);
let actual = execute(&session, execute_request(r#"text("still alive");"#)).await;
let actual = execute(
&session,
execute_request(r#"text("still alive");"#),
Arc::new(OversizedResultDelegate),
)
.await;
assert_eq!(
actual,
RuntimeResponse::Result {
@@ -702,13 +750,14 @@ async fn oversized_initial_response_does_not_close_the_shared_host() {
codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"),
);
let session = provider
.create_session(Arc::new(RecordingDelegate::default()))
.create_session()
.await
.expect("create remote session");
let started = session
.execute(execute_request(&format!(
r#"text("x".repeat({MAX_FRAME_BYTES}));"#
)))
.execute(
execute_request(&format!(r#"text("x".repeat({MAX_FRAME_BYTES}));"#)),
Arc::new(RecordingDelegate::default()),
)
.await
.expect("start oversized response");
let error = started
@@ -720,7 +769,12 @@ async fn oversized_initial_response_does_not_close_the_shared_host() {
"unexpected error: {error}"
);
let actual = execute(&session, execute_request(r#"text("still alive");"#)).await;
let actual = execute(
&session,
execute_request(r#"text("still alive");"#),
Arc::new(RecordingDelegate::default()),
)
.await;
assert_eq!(
actual,
RuntimeResponse::Result {
@@ -763,11 +817,11 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
delegate_a.hold_slow_cleanup();
let delegate_b = Arc::new(RecordingDelegate::default());
let session_a = provider
.create_session(delegate_a.clone())
.create_session()
.await
.expect("create first remote session");
let session_b = provider
.create_session(delegate_b.clone())
.create_session()
.await
.expect("create second remote session");
@@ -782,7 +836,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
output_schema: None,
}];
let started_a = session_a
.execute(request_a)
.execute(request_a, delegate_a.clone())
.await
.expect("start first cell");
let cell_a = started_a.cell_id.clone();
@@ -806,7 +860,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
let mut request_b = execute_request("await new Promise(() => {});");
request_b.yield_time_ms = Some(1);
let started_b = session_b
.execute(request_b)
.execute(request_b, delegate_b.clone())
.await
.expect("start second cell");
let cell_b = started_b.cell_id.clone();
@@ -899,7 +953,12 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
.await
.expect("unrelated session cleanup timeout");
let actual = execute(&session_b, execute_request(r#"text("replacement");"#)).await;
let actual = execute(
&session_b,
execute_request(r#"text("replacement");"#),
delegate_b.clone(),
)
.await;
assert_eq!(
actual,
RuntimeResponse::Result {