openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-08-23 13:09:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Forward gRPC code-mode callbacks to session delegates (#38072)

Browse Source 
## What changed

- Subscribe each gRPC code-mode session to nested tool calls and forward tool and notification callbacks to its delegate.
- Complete tool calls through the host while bounding oversized results and errors.
- Track callback ownership and cancellation so completed cells drain notifications, terminated cells cancel them, and shutdown revokes outstanding work.
- Validate callback identifiers, cell ownership, enabled tools, and pending callback limits without serializing independent callbacks or sessions.

## Testing

- Add integration and state tests for callback forwarding, completion ordering, cancellation, malformed callbacks, delegate panics, oversized results, and concurrent work.

GitOrigin-RevId: 005afbb90eea0eb77d746b930a1a96ca6dfcd4e7
This commit is contained in:
Channing Conger
2026-08-11 20:35:11 +00:00
committed by copyberry
parent b43de77679
commit ba2fb48319
16 changed files with 1710 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
codex-rs/Cargo.lock generated
View File

@@ -2564,6 +2564,7 @@ dependencies = [
"codex-websocket-client",
"futures",
"pretty_assertions",
"prost",
"serde_json",
"tokio",
"tokio-tungstenite",

2
codex-rs/code-mode-host/src/grpc/validation.rs
View File

@@ -3,10 +3,10 @@ use tonic::Status;
use uuid::Uuid;
pub(super) use codex_code_mode_protocol::grpc::MAX_IDENTIFIER_BYTES;
pub(super) use codex_code_mode_protocol::grpc::MAX_TOOL_ERROR_BYTES;
pub(super) const MAX_TOOL_FILTERS: usize = 64;
pub(super) const MAX_TOOL_DEFINITIONS: usize = 1_024;
pub(super) const MAX_TOOL_DESCRIPTION_BYTES: usize = 16 * 1_024;
pub(super) const MAX_TOOL_ERROR_BYTES: usize = 64 * 1_024;
pub(super) fn identifier(value: &str, field: &str) -> Result<(), Status> {
if value.is_empty() {

262
codex-rs/code-mode-host/tests/grpc.rs
View File

@@ -4,34 +4,97 @@ use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
use codex_code_mode::CellId;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeSession;
use codex_code_mode::CodeModeSessionCellExecutionLimits;
use codex_code_mode::CodeModeSessionDelegate;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::ExecuteRequest;
use codex_code_mode::FunctionCallOutputContentItem;
use codex_code_mode::GrpcCodeModeSessionProvider;
use codex_code_mode::NoopCodeModeSessionDelegate;
use codex_code_mode::NotificationFuture;
use codex_code_mode::RuntimeResponse;
use codex_code_mode::ToolDefinition;
use codex_code_mode::ToolInvocationFuture;
use codex_code_mode::WaitOutcome;
use codex_code_mode::WaitRequest;
use codex_code_mode_protocol::grpc;
use codex_code_mode_protocol::grpc::code_mode_host_client::CodeModeHostClient;
use codex_protocol::ToolName;
use futures::FutureExt;
use pretty_assertions::assert_eq;
use serde_json::json;
use tokio::sync::Semaphore;
use tokio::time::timeout;
use tokio_util::sync::CancellationToken;
use tonic::Code;
#[path = "support/host.rs"]
mod host;
#[path = "support/large_tool_delegate.rs"]
mod large_tool_delegate;
#[path = "support/recording_delegate.rs"]
mod recording_delegate;
use host::HostHarness;
use large_tool_delegate::LargeToolResultDelegate;
use recording_delegate::RecordingDelegate;
use recording_delegate::cell_id;
const TEST_TIMEOUT: Duration = Duration::from_secs(20);
struct PanickingDelegate;
struct SelfCancellingToolDelegate;
impl CodeModeSessionDelegate for PanickingDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
_cancellation: CancellationToken,
) -> ToolInvocationFuture<'a> {
panic!("synchronous tool delegate panic")
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation: CancellationToken,
) -> NotificationFuture<'a> {
panic!("synchronous notification delegate panic")
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
impl CodeModeSessionDelegate for SelfCancellingToolDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
cancellation: CancellationToken,
) -> ToolInvocationFuture<'a> {
cancellation.cancel();
Box::pin(async { Err("tool delegate cancelled itself".to_string()) })
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
fn request(source: &str) -> ExecuteRequest {
ExecuteRequest {
tool_call_id: "call-1".to_string(),
@@ -42,6 +105,17 @@ fn request(source: &str) -> ExecuteRequest {
}
}
fn tool(name: &str) -> ToolDefinition {
ToolDefinition {
name: name.to_string(),
tool_name: ToolName::plain(name),
description: String::new(),
kind: CodeModeToolKind::Function,
input_schema: None,
output_schema: None,
}
}
fn text_response(cell: &str, value: &str) -> RuntimeResponse {
RuntimeResponse::Result {
cell_id: cell_id(cell),
@@ -92,7 +166,7 @@ async fn start_active_wait(
}
#[tokio::test]
async fn tcp_session_persists_values_and_reports_cell_closure() -> Result<()> {
async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closure() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
assert!(host.endpoint.starts_with("http://127.0.0.1:"));
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
@@ -111,9 +185,37 @@ async fn tcp_session_persists_values_and_reports_cell_closure() -> Result<()> {
}
);
let mut callback = request(
r#"const result = await tools.echo({ value: String(load("key")) }); notify("notice"); text(result.value);"#,
);
callback.tool_call_id = "call-2".to_string();
callback.enabled_tools = vec![tool("echo")];
assert_eq!(
execute(&session, request(r#"text(String(load("key")));"#)).await?,
text_response("2", "persisted")
execute(&session, callback).await?,
text_response("2", "output")
);
timeout(TEST_TIMEOUT, delegate.notification_delivered.notified())
.await
.context("notification was not delivered")?;
assert_eq!(
*delegate
.invocations
.lock()
.unwrap_or_else(PoisonError::into_inner),
vec![CodeModeNestedToolCall {
cell_id: cell_id("2"),
runtime_tool_call_id: "tool-1".to_string(),
tool_name: ToolName::plain("echo"),
tool_kind: CodeModeToolKind::Function,
input: Some(json!({ "value": "persisted" })),
}]
);
assert_eq!(
*delegate
.notifications
.lock()
.unwrap_or_else(PoisonError::into_inner),
vec![("call-2".to_string(), cell_id("2"), "notice".to_string())]
);
session.shutdown().await.map_err(anyhow::Error::msg)?;
@@ -344,6 +446,66 @@ async fn dropping_an_initial_response_terminates_its_pending_remote_execution()
Ok(())
}
#[tokio::test]
async fn synchronous_delegate_panics_do_not_orphan_callbacks_or_close_the_session() -> 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(PanickingDelegate))
.await
.map_err(anyhow::Error::msg)?;
let mut tool_panic = request(
r#"try { await tools.echo({}); text("unexpected"); } catch (_) { text("tool recovered"); }"#,
);
tool_panic.enabled_tools = vec![tool("echo")];
assert_eq!(
execute(&session, tool_panic).await?,
text_response("1", "tool recovered")
);
assert_eq!(
execute(
&session,
request(r#"notify("panic"); text("notification recovered");"#),
)
.await?,
text_response("2", "notification recovered")
);
assert_eq!(
execute(&session, request(r#"text("still alive");"#)).await?,
text_response("3", "still alive")
);
session.shutdown().await.map_err(anyhow::Error::msg)?;
Ok(())
}
#[tokio::test]
async fn tool_delegate_self_cancellation_returns_an_error_without_hanging() -> 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(SelfCancellingToolDelegate))
.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")];
assert_eq!(
execute(&session, callback).await?,
text_response("1", "tool recovered")
);
assert_eq!(
execute(&session, request(r#"text("still alive");"#)).await?,
text_response("2", "still alive")
);
session.shutdown().await.map_err(anyhow::Error::msg)?;
Ok(())
}
#[tokio::test]
async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
@@ -521,6 +683,100 @@ async fn dropping_a_session_off_runtime_retires_its_active_cells() -> Result<()>
Ok(())
}
#[tokio::test]
async fn large_unary_tool_completion_does_not_block_an_independent_session() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(LargeToolResultDelegate {
started: Semaphore::new(/*permits*/ 0),
release: Semaphore::new(/*permits*/ 0),
});
let slow_session = provider
.create_session(delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
let fast_session = provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.await
.map_err(anyhow::Error::msg)?;
let mut slow_request = request(
r#"const result = await tools.large({ value: "request" }); text(String(result.value.length));"#,
);
slow_request.enabled_tools = vec![tool("large")];
slow_request.yield_time_ms = Some(/*value*/ 20_000);
let slow_cell = slow_session
.execute(slow_request)
.await
.map_err(anyhow::Error::msg)?;
timeout(TEST_TIMEOUT, delegate.started.acquire())
.await
.context("large tool callback did not start")??
.forget();
assert_eq!(
execute(&fast_session, request(r#"text("fast-before");"#)).await?,
text_response("1", "fast-before")
);
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 (slow_response, fast_response) = tokio::join!(slow_response, fast_response);
assert_eq!(fast_response?, text_response("2", "fast-during"));
assert_eq!(
slow_response
.context("large unary tool response did not complete")?
.map_err(anyhow::Error::msg)?,
text_response("1", "8388608")
);
slow_session.shutdown().await.map_err(anyhow::Error::msg)?;
fast_session.shutdown().await.map_err(anyhow::Error::msg)?;
Ok(())
}
#[tokio::test]
async fn single_subscription_processes_slow_and_fast_tools_concurrently() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(LargeToolResultDelegate {
started: Semaphore::new(/*permits*/ 0),
release: Semaphore::new(/*permits*/ 0),
});
let session = provider
.create_session(delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
let mut slow =
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)?;
timeout(TEST_TIMEOUT, delegate.started.acquire())
.await
.context("slow tool did not start")??
.forget();
let mut fast = request(r#"const result = await tools.fast({}); text(result.value);"#);
fast.enabled_tools = vec![tool("fast")];
assert_eq!(
execute(&session, fast).await?,
text_response("2", "isolated")
);
delegate.release.add_permits(/*n*/ 1);
assert_eq!(
timeout(TEST_TIMEOUT, slow_cell.initial_response())
.await
.context("slow tool did not finish")?
.map_err(anyhow::Error::msg)?,
text_response("1", "8388608")
);
session.shutdown().await.map_err(anyhow::Error::msg)?;
Ok(())
}
#[tokio::test]
async fn sessions_enforce_independent_yield_limits() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;

313
codex-rs/code-mode-host/tests/grpc_notifications.rs Normal file
View File

@@ -0,0 +1,313 @@
use std::sync::Arc;
use std::sync::PoisonError;
use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
use codex_code_mode::CellId;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeSession;
use codex_code_mode::CodeModeSessionDelegate;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::ExecuteRequest;
use codex_code_mode::FunctionCallOutputContentItem;
use codex_code_mode::GrpcCodeModeSessionProvider;
use codex_code_mode::NotificationFuture;
use codex_code_mode::RuntimeResponse;
use codex_code_mode::ToolInvocationFuture;
use codex_code_mode::WaitOutcome;
use codex_code_mode::WaitRequest;
use pretty_assertions::assert_eq;
use tokio::sync::Semaphore;
use tokio::time::timeout;
use tokio_util::sync::CancellationToken;
#[path = "support/host.rs"]
mod host;
#[path = "support/recording_delegate.rs"]
mod recording_delegate;
use host::HostHarness;
use recording_delegate::RecordingDelegate;
use recording_delegate::cell_id;
const TEST_TIMEOUT: Duration = Duration::from_secs(20);
struct BlockingNotificationDelegate {
started: Semaphore,
release: Semaphore,
delivered: Semaphore,
cancelled: Semaphore,
closed: Semaphore,
}
impl BlockingNotificationDelegate {
fn new() -> Self {
Self {
started: Semaphore::new(/*permits*/ 0),
release: Semaphore::new(/*permits*/ 0),
delivered: Semaphore::new(/*permits*/ 0),
cancelled: Semaphore::new(/*permits*/ 0),
closed: Semaphore::new(/*permits*/ 0),
}
}
}
impl CodeModeSessionDelegate for BlockingNotificationDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
_cancellation: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async { Err("unexpected tool invocation".to_string()) })
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
cancellation: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async move {
self.started.add_permits(/*n*/ 1);
tokio::select! {
_ = cancellation.cancelled() => {
self.cancelled.add_permits(/*n*/ 1);
Err("notification cancelled".to_string())
}
permit = self.release.acquire() => {
permit
.map_err(|_| "notification release closed".to_string())?
.forget();
self.delivered.add_permits(/*n*/ 1);
Ok(())
}
}
})
}
fn cell_closed(&self, _cell_id: &CellId) {
self.closed.add_permits(/*n*/ 1);
}
}
fn request(source: &str) -> ExecuteRequest {
ExecuteRequest {
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
yield_time_ms: Some(/*value*/ 5_000),
max_output_tokens: Some(/*value*/ 1_000),
}
}
fn text_response(cell: &str, value: &str) -> RuntimeResponse {
RuntimeResponse::Result {
cell_id: cell_id(cell),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: value.to_string(),
}],
error_text: None,
}
}
async fn execute(
session: &Arc<dyn CodeModeSession>,
request: ExecuteRequest,
) -> Result<RuntimeResponse> {
timeout(TEST_TIMEOUT, async {
session
.execute(request)
.await
.map_err(anyhow::Error::msg)?
.initial_response()
.await
.map_err(anyhow::Error::msg)
})
.await
.context("timed out executing gRPC code-mode cell")?
}
#[tokio::test]
async fn completed_cells_drain_pending_notifications_before_completion() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(BlockingNotificationDelegate::new());
let session = provider
.create_session(delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
let executing = Arc::clone(&session);
let completion = tokio::spawn(async move {
execute(&executing, request(r#"notify("notice"); text("done");"#)).await
});
timeout(TEST_TIMEOUT, delegate.started.acquire())
.await
.context("notification did not start")??
.forget();
assert!(!completion.is_finished());
delegate.release.add_permits(/*n*/ 1);
assert_eq!(
timeout(TEST_TIMEOUT, completion)
.await
.context("completed cell did not finish after notification delivery")???,
text_response("1", "done")
);
timeout(TEST_TIMEOUT, delegate.delivered.acquire())
.await
.context("completed cell did not deliver its pending notification")??
.forget();
assert!(delegate.cancelled.try_acquire().is_err());
timeout(TEST_TIMEOUT, delegate.closed.acquire())
.await
.context("completed cell was not retired")??
.forget();
assert_eq!(
execute(&session, request(r#"text("still alive");"#)).await?,
text_response("2", "still alive")
);
session.shutdown().await.map_err(anyhow::Error::msg)?;
Ok(())
}
#[tokio::test]
async fn completed_waits_drain_pending_notifications_before_returning() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(BlockingNotificationDelegate::new());
let session = provider
.create_session(delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
let mut pending = request(
r#"await new Promise(resolve => setTimeout(resolve, 25)); notify("notice"); text("done");"#,
);
pending.yield_time_ms = Some(/*value*/ 1);
let cell = session.execute(pending).await.map_err(anyhow::Error::msg)?;
assert_eq!(
cell.initial_response().await.map_err(anyhow::Error::msg)?,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
let waiting = Arc::clone(&session);
let completion = tokio::spawn(async move {
waiting
.wait(WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 5_000,
})
.await
.map_err(anyhow::Error::msg)
});
timeout(TEST_TIMEOUT, delegate.started.acquire())
.await
.context("wait notification did not start")??
.forget();
assert!(!completion.is_finished());
delegate.release.add_permits(/*n*/ 1);
assert_eq!(
timeout(TEST_TIMEOUT, completion)
.await
.context("wait did not finish after notification delivery")???,
WaitOutcome::LiveCell(text_response("1", "done"))
);
timeout(TEST_TIMEOUT, delegate.delivered.acquire())
.await
.context("wait did not deliver its pending notification")??
.forget();
assert!(delegate.cancelled.try_acquire().is_err());
session.shutdown().await.map_err(anyhow::Error::msg)?;
Ok(())
}
#[tokio::test]
async fn termination_cancels_pending_notifications() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(BlockingNotificationDelegate::new());
let session = provider
.create_session(delegate.clone())
.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)?;
timeout(TEST_TIMEOUT, delegate.started.acquire())
.await
.context("notification did not start")??
.forget();
assert_eq!(
cell.initial_response().await.map_err(anyhow::Error::msg)?,
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
assert_eq!(
session
.terminate(cell_id("1"))
.await
.map_err(anyhow::Error::msg)?,
WaitOutcome::LiveCell(RuntimeResponse::Terminated {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
timeout(TEST_TIMEOUT, delegate.cancelled.acquire())
.await
.context("termination did not cancel notification delivery")??
.forget();
timeout(TEST_TIMEOUT, delegate.closed.acquire())
.await
.context("terminated cell was not retired")??
.forget();
session.shutdown().await.map_err(anyhow::Error::msg)?;
Ok(())
}
#[tokio::test]
async fn oversized_notification_text_is_truncated_at_a_utf8_boundary() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
let delegate = Arc::new(RecordingDelegate::default());
let session = provider
.create_session(delegate.clone())
.await
.map_err(anyhow::Error::msg)?;
assert_eq!(
execute(
&session,
request(r#"notify("🦀".repeat(512)); text("done");"#),
)
.await?,
text_response("1", "done")
);
timeout(TEST_TIMEOUT, delegate.notification_delivered.notified())
.await
.context("truncated notification was not delivered")?;
assert_eq!(
*delegate
.notifications
.lock()
.unwrap_or_else(PoisonError::into_inner),
vec![(
"call-1".to_string(),
cell_id("1"),
format!("{}... [truncated]", "🦀".repeat(252)),
)]
);
session.shutdown().await.map_err(anyhow::Error::msg)?;
Ok(())
}

3
codex-rs/code-mode-host/tests/support/recording_delegate.rs
View File

@@ -7,12 +7,14 @@ use codex_code_mode::CodeModeSessionDelegate;
use codex_code_mode::NotificationFuture;
use codex_code_mode::ToolInvocationFuture;
use serde_json::json;
use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;
#[derive(Default)]
pub(crate) struct RecordingDelegate {
pub(crate) invocations: Mutex<Vec<CodeModeNestedToolCall>>,
pub(crate) notifications: Mutex<Vec<(String, CellId, String)>>,
pub(crate) notification_delivered: Notify,
pub(crate) closed_cells: Mutex<Vec<CellId>>,
}
@@ -40,6 +42,7 @@ impl CodeModeSessionDelegate for RecordingDelegate {
.lock()
.unwrap_or_else(PoisonError::into_inner)
.push((call_id, cell_id, text));
self.notification_delivered.notify_one();
Box::pin(async { Ok(()) })
}

1
codex-rs/code-mode-protocol/src/grpc/mod.rs
View File

@@ -5,3 +5,4 @@ pub use code_mode_proto::codex::code_mode::v1::*;
tonic::include_proto!("codex.code_mode.v1");
pub const MAX_IDENTIFIER_BYTES: usize = 256;
pub const MAX_TOOL_ERROR_BYTES: usize = 64 * 1_024;

1
codex-rs/code-mode/Cargo.toml
View File

@@ -19,6 +19,7 @@ codex-install-context = { workspace = true }
codex-protocol = { workspace = true }
codex-websocket-client = { workspace = true }
futures = { workspace = true }
prost = "0.14.3"
serde_json = { workspace = true }
tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt", "sync", "time"] }
tokio-tungstenite = { workspace = true }

216
codex-rs/code-mode/src/grpc_session/callbacks.rs
View File

@@ -1,42 +1,69 @@
use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::sync::PoisonError;
use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::grpc;
use futures::FutureExt;
use tokio_util::sync::CancellationToken;
use tracing::warn;
use super::SessionInner;
use super::completion;
use super::conversion;
use super::deadline;
use super::state::CallbackAdmission;
const MAX_NOTIFICATION_BYTES: usize = 1_024;
const TRUNCATED_NOTIFICATION_SUFFIX: &str = "... [truncated]";
impl SessionInner {
pub(super) fn spawn_session_events(
self: &Arc<Self>,
mut events: tonic::Streaming<grpc::SessionEvent>,
events: tonic::Streaming<grpc::SessionEvent>,
) {
self.spawn_stream(events, "session lease", Self::handle_session_event);
}
pub(super) fn spawn_tool_subscription(
self: &Arc<Self>,
calls: tonic::Streaming<grpc::ToolCall>,
) {
self.spawn_stream(calls, "tool subscription", Self::handle_tool_call);
}
fn spawn_stream<T: Send + 'static>(
self: &Arc<Self>,
mut stream: tonic::Streaming<T>,
stream_name: &'static str,
handle: fn(&Arc<Self>, T) -> Result<(), String>,
) {
let inner = Arc::clone(self);
self.stream_tasks.spawn(async move {
loop {
let event = tokio::select! {
let message = tokio::select! {
biased;
_ = inner.stopped.cancelled() => return,
event = events.message() => event,
message = stream.message() => message,
};
match event {
Ok(Some(event)) => {
if let Err(error) = inner.handle_session_event(event) {
match message {
Ok(Some(message)) => {
if let Err(error) = handle(&inner, message) {
inner.fail(error);
return;
}
}
Ok(None) => {
if !inner.shutdown_requested.load(Ordering::Acquire) {
inner.fail(
"gRPC code-mode session lease closed unexpectedly".to_string(),
);
inner.fail(format!("gRPC code-mode {stream_name} closed unexpectedly"));
}
return;
}
Err(error) => {
if !inner.shutdown_requested.load(Ordering::Acquire) {
inner.fail(super::deadline::failure("session lease", error));
inner.fail(deadline::failure(stream_name, error));
}
return;
}
@@ -45,7 +72,7 @@ impl SessionInner {
});
}
fn handle_session_event(&self, event: grpc::SessionEvent) -> Result<(), String> {
fn handle_session_event(self: &Arc<Self>, event: grpc::SessionEvent) -> Result<(), String> {
match event
.event
.ok_or_else(|| "gRPC code-mode host sent an empty session event".to_string())?
@@ -53,9 +80,17 @@ impl SessionInner {
grpc::session_event::Event::Opened(_) => {
Err("gRPC code-mode host repeated the session opening event".to_string())
}
grpc::session_event::Event::ToolCallCancelled(_)
| grpc::session_event::Event::Notification(_)
| grpc::session_event::Event::NotificationCancelled(_) => Ok(()),
grpc::session_event::Event::ToolCallCancelled(cancelled) => {
self.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.cancel_invocation(&cancelled.invocation_id)?;
Ok(())
}
grpc::session_event::Event::Notification(notification) => {
self.handle_notification(notification)
}
grpc::session_event::Event::NotificationCancelled(_) => Ok(()),
grpc::session_event::Event::CellClosed(closed) => {
let cell = self
.state
@@ -67,4 +102,157 @@ impl SessionInner {
}
}
}
fn handle_tool_call(self: &Arc<Self>, call: grpc::ToolCall) -> Result<(), String> {
if call.session_id != self.id {
return Err(format!(
"gRPC code-mode tool invocation belongs to session {} instead of {}",
call.session_id, self.id
));
}
let admission = self
.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.admit_invocation(&call)?;
let invocation_id = call.invocation_id.clone();
let cancellation = match admission {
CallbackAdmission::Active(cancellation) => Ok(cancellation),
CallbackAdmission::Cancelled => return Ok(()),
CallbackAdmission::Closed => Err(format!("code-mode cell {} is closed", call.cell_id)),
CallbackAdmission::Rejected(error) => Err(error),
};
let cancellation = match cancellation {
Ok(cancellation) => cancellation,
Err(error) => {
let inner = Arc::clone(self);
tokio::spawn(async move {
inner
.complete_tool_call(invocation_id, CancellationToken::new(), Err(error))
.await;
});
return Ok(());
}
};
let invocation = conversion::tool_call(call);
let inner = Arc::clone(self);
tokio::spawn(async move {
let result = match invocation {
Ok(invocation) => {
let callback = AssertUnwindSafe(async {
inner
.delegate
.invoke_tool(invocation, cancellation.child_token())
.await
})
.catch_unwind();
tokio::select! {
biased;
_ = cancellation.cancelled() => return,
result = callback => match result {
Ok(result) => result,
Err(_) => Err("code-mode tool delegate panicked".to_string()),
},
}
}
Err(error) => Err(error),
};
inner
.complete_tool_call(invocation_id, cancellation, result)
.await;
});
Ok(())
}
async fn complete_tool_call(
&self,
invocation_id: String,
cancellation: CancellationToken,
result: Result<serde_json::Value, String>,
) {
let request = completion::request(&self.id, &invocation_id, result);
let mut client = self.client();
tokio::select! {
biased;
_ = cancellation.cancelled() => {}
result = deadline::request(
self,
"tool invocation completion",
Duration::ZERO,
client.complete_tool_call(request),
) => {
if let Err(error) = result
&& !cancellation.is_cancelled()
&& !self.stopped.is_cancelled()
{
self.fail(error);
}
}
}
self.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.finish_invocation(&invocation_id);
}
fn handle_notification(
self: &Arc<Self>,
mut notification: grpc::Notification,
) -> Result<(), String> {
if notification.text.len() > MAX_NOTIFICATION_BYTES {
let boundary = notification
.text
.floor_char_boundary(MAX_NOTIFICATION_BYTES - TRUNCATED_NOTIFICATION_SUFFIX.len());
notification.text.truncate(boundary);
notification.text.push_str(TRUNCATED_NOTIFICATION_SUFFIX);
}
let admission = self
.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.admit_notification(&notification)?;
let cancellation = match admission {
CallbackAdmission::Active(cancellation) => cancellation,
CallbackAdmission::Cancelled | CallbackAdmission::Closed => return Ok(()),
CallbackAdmission::Rejected(error) => {
warn!("code-mode notification was dropped: {error}");
return Ok(());
}
};
let execution_id = notification.execution_id;
let inner = Arc::clone(self);
// Delegate callbacks stay outside the tracked session tasks so shutdown can cancel
// them without waiting for arbitrary delegate work to complete.
tokio::spawn(async move {
let callback = AssertUnwindSafe(async {
inner
.delegate
.notify(
notification.call_id,
CellId::new(notification.cell_id),
notification.text,
cancellation,
)
.await
})
.catch_unwind();
let result = tokio::select! {
biased;
_ = inner.stopped.cancelled() => return,
result = callback => result,
};
match result {
Ok(Ok(())) => {}
Ok(Err(error)) => warn!("code-mode notification delegate failed: {error}"),
Err(_) => warn!("code-mode notification delegate panicked"),
}
let cell = inner
.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.finish_notification(&execution_id);
inner.report_closed_cell(cell);
});
Ok(())
}
}

70
codex-rs/code-mode/src/grpc_session/completion.rs Normal file
View File

@@ -0,0 +1,70 @@
use codex_code_mode_protocol::grpc;
use codex_code_mode_protocol::grpc::MAX_TOOL_ERROR_BYTES;
use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
use prost::Message;
const TRUNCATED_SUFFIX: &str = "... [truncated]";
pub(super) fn request(
session_id: &str,
invocation_id: &str,
result: Result<serde_json::Value, String>,
) -> grpc::CompleteToolCallRequest {
request_with_maximum(session_id, invocation_id, result, MAX_FRAME_BYTES)
}
fn request_with_maximum(
session_id: &str,
invocation_id: &str,
result: Result<serde_json::Value, String>,
maximum_message_bytes: usize,
) -> grpc::CompleteToolCallRequest {
let outcome = match result {
Ok(value) => match serde_json::to_vec(&value) {
Ok(output_json) => {
grpc::complete_tool_call_request::Outcome::Succeeded(grpc::ToolCallSucceeded {
output_json,
})
}
Err(error) => grpc::complete_tool_call_request::Outcome::Failed(grpc::ToolCallFailed {
message: bounded_error(format!("failed to encode code-mode tool result: {error}")),
}),
},
Err(message) => grpc::complete_tool_call_request::Outcome::Failed(grpc::ToolCallFailed {
message: bounded_error(message),
}),
};
let mut request = grpc::CompleteToolCallRequest {
session_id: session_id.to_string(),
invocation_id: invocation_id.to_string(),
outcome: Some(outcome),
};
let encoded_bytes = request.encoded_len();
if encoded_bytes > maximum_message_bytes {
request.outcome = Some(grpc::complete_tool_call_request::Outcome::Failed(
grpc::ToolCallFailed {
message: bounded_error(format!(
"code-mode tool result of {encoded_bytes} encoded bytes exceeds the gRPC message limit of {maximum_message_bytes} bytes"
)),
},
));
}
request
}
fn bounded_error(mut message: String) -> String {
if message.len() <= MAX_TOOL_ERROR_BYTES {
return message;
}
let mut boundary = MAX_TOOL_ERROR_BYTES - TRUNCATED_SUFFIX.len();
while !message.is_char_boundary(boundary) {
boundary -= 1;
}
message.truncate(boundary);
message.push_str(TRUNCATED_SUFFIX);
message
}
#[cfg(test)]
#[path = "completion_tests.rs"]
mod tests;

45
codex-rs/code-mode/src/grpc_session/completion_tests.rs Normal file
View File

@@ -0,0 +1,45 @@
use codex_code_mode_protocol::grpc;
use pretty_assertions::assert_eq;
use prost::Message;
use super::MAX_TOOL_ERROR_BYTES;
use super::TRUNCATED_SUFFIX;
use super::request;
use super::request_with_maximum;
#[test]
fn completion_size_includes_the_protobuf_envelope() {
let output = serde_json::Value::String("a".repeat(100));
let raw_json_bytes = serde_json::to_vec(&output).expect("valid JSON").len();
let completion = request_with_maximum("session", "invocation", Ok(output), raw_json_bytes);
assert!(matches!(
completion.outcome,
Some(grpc::complete_tool_call_request::Outcome::Failed(grpc::ToolCallFailed {
message,
})) if message.contains("encoded bytes exceeds the gRPC message limit")
));
}
#[test]
fn delegate_errors_are_truncated_at_a_utf8_boundary() {
let error = "🦀".repeat(MAX_TOOL_ERROR_BYTES);
let completion = request("session", "invocation", Err(error));
let Some(grpc::complete_tool_call_request::Outcome::Failed(failure)) = completion.outcome
else {
panic!("expected a failed tool completion");
};
assert!(failure.message.len() <= MAX_TOOL_ERROR_BYTES);
assert!(failure.message.ends_with(TRUNCATED_SUFFIX));
assert!(failure.message.starts_with('🦀'));
}
#[test]
fn completion_at_exact_message_limit_is_accepted() {
let value = serde_json::json!({ "ok": true });
let expected = request("session", "invocation", Ok(value.clone()));
let actual = request_with_maximum("session", "invocation", Ok(value), expected.encoded_len());
assert_eq!(actual, expected);
}

31
codex-rs/code-mode/src/grpc_session/conversion.rs
View File

@@ -1,4 +1,5 @@
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeNestedToolCall;
use codex_code_mode_protocol::CodeModeToolKind;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
@@ -7,6 +8,7 @@ use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::ToolDefinition;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::grpc;
use codex_protocol::ToolName;
pub(super) fn execute_request(
session_id: &str,
@@ -57,6 +59,35 @@ fn tool_definition(definition: ToolDefinition) -> Result<grpc::ToolDefinition, S
})
}
pub(super) fn tool_call(call: grpc::ToolCall) -> Result<CodeModeNestedToolCall, String> {
let name = call
.tool_name
.ok_or_else(|| "code-mode tool invocation omitted its tool name".to_string())?;
let kind = match grpc::ToolKind::try_from(call.tool_kind) {
Ok(grpc::ToolKind::Function) => CodeModeToolKind::Function,
Ok(grpc::ToolKind::Freeform) => CodeModeToolKind::Freeform,
Ok(grpc::ToolKind::Unspecified) | Err(_) => {
return Err(format!(
"code-mode tool invocation has invalid kind {}",
call.tool_kind
));
}
};
let input = call
.input_json
.map(|input| serde_json::from_slice(&input))
.transpose()
.map_err(|error| format!("code-mode tool invocation contains invalid JSON: {error}"))?;
Ok(CodeModeNestedToolCall {
cell_id: CellId::new(call.cell_id),
runtime_tool_call_id: call.runtime_tool_call_id,
tool_name: ToolName::new(name.namespace, name.name),
tool_kind: kind,
input,
})
}
pub(super) fn runtime_response(outcome: grpc::ExecutionOutcome) -> Result<RuntimeResponse, String> {
super::validate_identifier(&outcome.cell_id, "cell ID")?;
let cell_id = CellId::new(outcome.cell_id);

31
codex-rs/code-mode/src/grpc_session/conversion_tests.rs
View File

@@ -1,4 +1,5 @@
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeNestedToolCall;
use codex_code_mode_protocol::CodeModeToolKind;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
@@ -13,6 +14,7 @@ use serde_json::json;
use super::execute_request;
use super::runtime_response;
use super::tool_call;
use super::wait_outcome;
#[test]
@@ -56,6 +58,35 @@ fn execute_request_preserves_tool_schemas_namespaces_and_limits() {
);
}
#[test]
fn tool_call_decodes_structured_input_and_namespace() {
let call = grpc::ToolCall {
session_id: "session".to_string(),
execution_id: "execution".to_string(),
cell_id: "cell".to_string(),
invocation_id: "invocation".to_string(),
runtime_tool_call_id: "runtime-call".to_string(),
tool_name: Some(grpc::ToolName {
name: "search".to_string(),
namespace: Some("work".to_string()),
}),
tool_kind: grpc::ToolKind::Function as i32,
input_json: Some(br#"{"query":"hello"}"#.to_vec()),
sequence: 1,
};
assert_eq!(
tool_call(call),
Ok(CodeModeNestedToolCall {
cell_id: CellId::new("cell".to_string()),
runtime_tool_call_id: "runtime-call".to_string(),
tool_name: ToolName::namespaced("work", "search"),
tool_kind: CodeModeToolKind::Function,
input: Some(json!({"query": "hello"})),
})
);
}
#[test]
fn runtime_response_decodes_mixed_content_items() {
let outcome = grpc::ExecutionOutcome {

18
codex-rs/code-mode/src/grpc_session/mod.rs
View File

@@ -34,6 +34,7 @@ use crate::remote_session::ShutdownResultReceiver;
use crate::remote_session::wait_for_watch;
mod callbacks;
mod completion;
mod conversion;
mod deadline;
mod operations;
@@ -116,6 +117,23 @@ impl GrpcCodeModeSessionProvider {
inner: Some(Arc::clone(&inner)),
};
inner.spawn_session_events(lease);
let request = grpc::SubscribeToToolCallsRequest {
session_id: inner.id.clone(),
tool_names: Vec::new(),
};
let mut client = inner.client();
let response =
match deadline::startup("tool subscription", client.subscribe_to_tool_calls(request))
.await
{
Ok(response) => response,
Err(error) => {
let _ = wait_for_watch(inner.request_shutdown()).await;
return Err(error);
}
};
inner.spawn_tool_subscription(response.into_inner());
inner.require_open()?;
opening.inner = None;
Ok(Arc::new(GrpcCodeModeSession { inner }))

54
codex-rs/code-mode/src/grpc_session/operations.rs
View File

@@ -72,7 +72,7 @@ impl SessionInner {
self.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.begin_execution(execution_id.clone())?;
.begin_execution(&request)?;
let ownership = ExecutionOwnership {
session: Arc::clone(self),
execution_id,
@@ -202,18 +202,25 @@ impl SessionInner {
Err(error) => Err(error),
},
};
let outcome = outcome.and_then(|response| {
if runtime_response_cell_id(&response) != &cell_id {
let outcome = match outcome {
Ok(response) if runtime_response_cell_id(&response) != &cell_id => {
let error = format!(
"gRPC code-mode execution returned cell {} instead of {cell_id}",
runtime_response_cell_id(&response)
);
self.fail(error.clone());
Err(error)
} else {
}
Ok(response) => {
tokio::select! {
biased;
_ = response_tx.closed() => return,
_ = self.settle_notifications(&response) => {}
}
Ok(response)
}
});
Err(error) => Err(error),
};
let _ = response_tx.send(outcome);
}
@@ -276,7 +283,7 @@ impl SessionInner {
cancellation.disarm();
self.prune_wait_slots();
let outcome = conversion::wait_outcome(response?.into_inner())?;
self.validate_wait_cell(&expected_cell_id, outcome)
self.validate_wait_cell(&expected_cell_id, outcome).await
}
pub(super) async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, String> {
@@ -294,19 +301,18 @@ impl SessionInner {
.await?
.into_inner();
let outcome = conversion::wait_outcome(response)?;
self.validate_wait_cell(&cell_id, outcome)
self.validate_wait_cell(&cell_id, outcome).await
}
fn validate_wait_cell(
async fn validate_wait_cell(
&self,
expected_cell_id: &CellId,
outcome: WaitOutcome,
) -> Result<WaitOutcome, String> {
let actual_cell_id = match &outcome {
WaitOutcome::LiveCell(response) | WaitOutcome::MissingCell(response) => {
runtime_response_cell_id(response)
}
let response = match &outcome {
WaitOutcome::LiveCell(response) | WaitOutcome::MissingCell(response) => response,
};
let actual_cell_id = runtime_response_cell_id(response);
if actual_cell_id != expected_cell_id {
let error = format!(
"gRPC code-mode host returned cell {actual_cell_id} instead of {expected_cell_id}"
@@ -314,9 +320,33 @@ impl SessionInner {
self.fail(error.clone());
return Err(error);
}
self.settle_notifications(response).await;
Ok(outcome)
}
async fn settle_notifications(&self, response: &RuntimeResponse) {
match response {
RuntimeResponse::Yielded { .. } => {}
RuntimeResponse::Terminated { cell_id, .. } => self
.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.cancel_notifications(cell_id),
RuntimeResponse::Result { cell_id, .. } => {
let cancellation = self
.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.notification_cancellation(cell_id);
if let Some(cancellation) = cancellation {
// CellClosed follows notifications on the lease stream and cancels this
// token only after every admitted notification has finished.
cancellation.cancelled().await;
}
}
}
}
fn prune_wait_slots(&self) {
self.wait_slots
.lock()

271
codex-rs/code-mode/src/grpc_session/state.rs
View File

@@ -1,14 +1,38 @@
use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::grpc;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
use codex_protocol::ToolName;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
const MAX_RECENT_CALLBACK_IDS: usize = 4_096;
struct ActiveCallback {
execution_id: String,
cancellation: CancellationToken,
}
pub(super) enum CallbackAdmission {
Active(CancellationToken),
Cancelled,
Closed,
Rejected(String),
}
#[derive(Default)]
struct ExecutionRecord {
cell_id: Option<CellId>,
tool_call_id: String,
enabled_tools: HashMap<ToolName, i32>,
started: bool,
ready: bool,
closed: bool,
notifications: usize,
cancellation: CancellationToken,
}
impl ExecutionRecord {
@@ -27,9 +51,41 @@ impl ExecutionRecord {
}
}
#[derive(Default)]
struct RecentIds {
values: HashSet<Uuid>,
order: VecDeque<Uuid>,
}
impl RecentIds {
fn remember(&mut self, value: Uuid) {
if !self.values.insert(value) {
return;
}
self.order.push_back(value);
while self.order.len() > MAX_RECENT_CALLBACK_IDS {
if let Some(expired) = self.order.pop_front() {
self.values.remove(&expired);
}
}
}
fn remove(&mut self, value: &Uuid) -> bool {
self.values.remove(value)
}
fn contains(&self, value: &Uuid) -> bool {
self.values.contains(value)
}
}
#[derive(Default)]
pub(super) struct SessionState {
executions: HashMap<String, ExecutionRecord>,
invocations: HashMap<String, ActiveCallback>,
notifications: usize,
seen_invocations: RecentIds,
cancelled_invocations: RecentIds,
failure: Option<String>,
closed: bool,
}
@@ -45,13 +101,35 @@ impl SessionState {
Ok(())
}
pub(super) fn begin_execution(&mut self, execution_id: String) -> Result<(), String> {
pub(super) fn begin_execution(&mut self, request: &grpc::ExecuteRequest) -> Result<(), String> {
self.require_open()?;
if execution_id.is_empty() || self.executions.contains_key(&execution_id) {
if request.execution_id.is_empty() || self.executions.contains_key(&request.execution_id) {
return Err("code-mode execution ID was empty or reused".to_string());
}
self.executions
.insert(execution_id, ExecutionRecord::default());
super::validate_identifier(&request.tool_call_id, "tool call ID")?;
let enabled_tools = request
.enabled_tools
.iter()
.map(|definition| {
let name = definition
.tool_name
.as_ref()
.ok_or_else(|| "code-mode enabled tool omitted its tool name".to_string())?;
Ok((
ToolName::new(name.namespace.clone(), name.name.clone())
.with_default_namespace(),
definition.kind,
))
})
.collect::<Result<HashMap<_, _>, String>>()?;
self.executions.insert(
request.execution_id.clone(),
ExecutionRecord {
tool_call_id: request.tool_call_id.clone(),
enabled_tools,
..ExecutionRecord::default()
},
);
Ok(())
}
@@ -61,15 +139,7 @@ impl SessionState {
cell_id: &str,
) -> Result<(), String> {
self.require_open()?;
if self.executions.iter().any(|(id, execution)| {
id != execution_id
&& execution
.cell_id
.as_ref()
.is_some_and(|current| current.as_str() == cell_id)
}) {
return Err(format!("code-mode host reused active cell ID {cell_id}"));
}
self.check_cell_ownership(execution_id, cell_id)?;
let execution = self
.executions
.get_mut(execution_id)
@@ -100,11 +170,135 @@ impl SessionState {
Ok(self.close_execution_if_ready(execution_id))
}
pub(super) fn admit_invocation(
&mut self,
call: &grpc::ToolCall,
) -> Result<CallbackAdmission, String> {
self.require_open()?;
let invocation_id = Uuid::parse_str(&call.invocation_id)
.map_err(|_| "code-mode tool invocation ID must be a UUID".to_string())?;
if self.invocations.contains_key(&call.invocation_id)
|| self.seen_invocations.contains(&invocation_id)
{
return Err("code-mode tool invocation ID was reused".to_string());
}
self.check_cell_ownership(&call.execution_id, &call.cell_id)?;
let Some(execution) = self.executions.get_mut(&call.execution_id) else {
self.seen_invocations.remember(invocation_id);
self.cancelled_invocations.remove(&invocation_id);
return Ok(CallbackAdmission::Closed);
};
execution.accept_cell(&call.cell_id)?;
let execution_closed = execution.closed;
self.seen_invocations.remember(invocation_id);
let invocation_cancelled = self.cancelled_invocations.remove(&invocation_id);
if execution_closed {
return Ok(CallbackAdmission::Closed);
}
if invocation_cancelled {
return Ok(CallbackAdmission::Cancelled);
}
let Some(name) = call.tool_name.as_ref() else {
return Ok(CallbackAdmission::Rejected(
"code-mode tool invocation omitted its tool name".to_string(),
));
};
let tool_name =
ToolName::new(name.namespace.clone(), name.name.clone()).with_default_namespace();
if execution.enabled_tools.get(&tool_name) != Some(&call.tool_kind) {
return Ok(CallbackAdmission::Rejected(format!(
"code-mode tool {tool_name} is not enabled for this execution"
)));
}
if self.invocations.len() + self.notifications >= MAX_PENDING_DELEGATE_CALLS {
return Ok(CallbackAdmission::Rejected(
"code-mode host exceeded its pending delegate callback limit".to_string(),
));
}
let cancellation = CancellationToken::new();
self.invocations.insert(
call.invocation_id.clone(),
ActiveCallback {
execution_id: call.execution_id.clone(),
cancellation: cancellation.clone(),
},
);
Ok(CallbackAdmission::Active(cancellation))
}
pub(super) fn admit_notification(
&mut self,
notification: &grpc::Notification,
) -> Result<CallbackAdmission, String> {
self.require_open()?;
Uuid::parse_str(&notification.notification_id)
.map_err(|_| "code-mode notification ID must be a UUID".to_string())?;
super::validate_identifier(&notification.call_id, "notification call ID")?;
self.check_cell_ownership(&notification.execution_id, &notification.cell_id)?;
let Some(execution) = self.executions.get_mut(&notification.execution_id) else {
return Ok(CallbackAdmission::Closed);
};
execution.accept_cell(&notification.cell_id)?;
if notification.call_id != execution.tool_call_id {
return Err("code-mode notification call ID does not match its execution".to_string());
}
if execution.closed {
return Ok(CallbackAdmission::Closed);
}
if self.invocations.len() + self.notifications >= MAX_PENDING_DELEGATE_CALLS {
return Ok(CallbackAdmission::Rejected(
"code-mode host exceeded its pending delegate callback limit".to_string(),
));
}
execution.notifications += 1;
self.notifications += 1;
Ok(CallbackAdmission::Active(
execution.cancellation.child_token(),
))
}
pub(super) fn finish_notification(&mut self, execution_id: &str) -> Option<CellId> {
let execution = self.executions.get_mut(execution_id)?;
execution.notifications = execution.notifications.checked_sub(1)?;
self.notifications -= 1;
self.close_execution_if_ready(execution_id)
}
pub(super) fn cancel_notifications(&self, cell_id: &CellId) {
if let Some(cancellation) = self.notification_cancellation(cell_id) {
cancellation.cancel();
}
}
pub(super) fn notification_cancellation(&self, cell_id: &CellId) -> Option<CancellationToken> {
self.executions
.values()
.find(|execution| execution.cell_id.as_ref() == Some(cell_id))
.map(|execution| execution.cancellation.clone())
}
pub(super) fn cancel_invocation(&mut self, invocation_id: &str) -> Result<(), String> {
let parsed = Uuid::parse_str(invocation_id)
.map_err(|_| "code-mode tool invocation ID must be a UUID".to_string())?;
if let Some(callback) = self.invocations.remove(invocation_id) {
callback.cancellation.cancel();
} else if !self.seen_invocations.contains(&parsed) {
self.cancelled_invocations.remember(parsed);
}
Ok(())
}
pub(super) fn finish_invocation(&mut self, invocation_id: &str) {
self.invocations.remove(invocation_id);
}
pub(super) fn close_cell(
&mut self,
closed: grpc::CellClosed,
) -> Result<Option<CellId>, String> {
self.require_open()?;
self.check_cell_ownership(&closed.execution_id, &closed.cell_id)?;
let Some(execution) = self.executions.get_mut(&closed.execution_id) else {
return Ok(None);
};
@@ -116,6 +310,10 @@ impl SessionState {
));
}
execution.closed = true;
if execution.notifications == 0 {
execution.cancellation.cancel();
}
self.revoke_execution_callbacks(&closed.execution_id);
Ok(self.close_execution_if_ready(&closed.execution_id))
}
@@ -125,22 +323,63 @@ impl SessionState {
}
self.closed = true;
self.failure = failure;
self.notifications = 0;
for (_, callback) in self.invocations.drain() {
callback.cancellation.cancel();
}
self.executions
.drain()
.filter_map(|(_, execution)| execution.cell_id)
.filter_map(|(_, execution)| {
execution.cancellation.cancel();
execution.cell_id
})
.collect()
}
fn close_execution_if_ready(&mut self, execution_id: &str) -> Option<CellId> {
self.executions
.get(execution_id)
.is_some_and(|execution| execution.started && execution.ready && execution.closed)
.is_some_and(|execution| {
execution.started
&& execution.ready
&& execution.closed
&& execution.notifications == 0
})
.then(|| self.remove_execution(execution_id))
.flatten()
}
pub(super) fn remove_execution(&mut self, execution_id: &str) -> Option<CellId> {
self.executions.remove(execution_id)?.cell_id
let execution = self.executions.remove(execution_id)?;
self.notifications -= execution.notifications;
execution.cancellation.cancel();
self.revoke_execution_callbacks(execution_id);
execution.cell_id
}
fn check_cell_ownership(&self, execution_id: &str, cell_id: &str) -> Result<(), String> {
if self.executions.contains_key(execution_id)
&& self.executions.iter().any(|(id, execution)| {
id != execution_id
&& execution
.cell_id
.as_ref()
.is_some_and(|current| current.as_str() == cell_id)
})
{
return Err(format!("code-mode host reused active cell ID {cell_id}"));
}
Ok(())
}
fn revoke_execution_callbacks(&mut self, execution_id: &str) {
self.invocations.retain(|_, callback| {
if callback.execution_id != execution_id {
return true;
}
callback.cancellation.cancel();
false
});
}
}

442
codex-rs/code-mode/src/grpc_session/state_tests.rs
View File

@@ -1,14 +1,119 @@
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::grpc;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
use pretty_assertions::assert_eq;
use uuid::Uuid;
use super::CallbackAdmission;
use super::SessionState;
fn request(execution_id: &str) -> grpc::ExecuteRequest {
grpc::ExecuteRequest {
session_id: "session".to_string(),
execution_id: execution_id.to_string(),
tool_call_id: "call".to_string(),
source: String::new(),
enabled_tools: vec![grpc::ToolDefinition {
name: "tool".to_string(),
tool_name: Some(grpc::ToolName {
name: "tool".to_string(),
namespace: None,
}),
description: String::new(),
kind: grpc::ToolKind::Function as i32,
input_schema_json: None,
output_schema_json: None,
}],
yield_time_ms: None,
max_output_tokens: None,
}
}
fn tool_call(execution_id: &str, invocation_id: u128) -> grpc::ToolCall {
let invocation_id = Uuid::from_u128(invocation_id).to_string();
grpc::ToolCall {
session_id: "session".to_string(),
execution_id: execution_id.to_string(),
cell_id: "cell".to_string(),
invocation_id: invocation_id.clone(),
runtime_tool_call_id: format!("runtime-{invocation_id}"),
tool_name: Some(grpc::ToolName {
name: "tool".to_string(),
namespace: None,
}),
tool_kind: grpc::ToolKind::Function as i32,
input_json: None,
sequence: 1,
}
}
fn notification(execution_id: &str, notification_id: u128) -> grpc::Notification {
grpc::Notification {
notification_id: Uuid::from_u128(notification_id).to_string(),
execution_id: execution_id.to_string(),
cell_id: "cell".to_string(),
call_id: "call".to_string(),
text: "hello".to_string(),
}
}
#[test]
fn cell_closure_drains_notifications_and_cancels_tool_callbacks() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.expect("register execution");
let first = state
.admit_invocation(&tool_call("execution", /*invocation_id*/ 1))
.expect("accept early invocation");
let CallbackAdmission::Active(first_cancellation) = first else {
panic!("first callback was not admitted");
};
let CallbackAdmission::Active(notification_cancellation) = state
.admit_notification(&notification("execution", /*notification_id*/ 1))
.expect("admit notification")
else {
panic!("notification was not admitted");
};
assert_eq!(
state
.close_cell(grpc::CellClosed {
execution_id: "execution".to_string(),
cell_id: "cell".to_string(),
final_tool_call_sequence: 3,
})
.expect("record cell closure"),
None
);
assert!(first_cancellation.is_cancelled());
assert!(!notification_cancellation.is_cancelled());
state
.admit_execution("execution", "cell")
.expect("admit started cell");
assert!(matches!(
state
.admit_invocation(&tool_call("execution", /*invocation_id*/ 2))
.expect("reject invocation for a closed cell"),
CallbackAdmission::Closed
));
assert_eq!(
state
.mark_execution_ready("execution")
.expect("claim started cell"),
None
);
assert_eq!(
state.finish_notification("execution"),
Some(CellId::new("cell".to_string()))
);
assert!(notification_cancellation.is_cancelled());
}
#[test]
fn cell_closure_waits_until_the_started_cell_is_claimed() {
let mut state = SessionState::default();
state
.begin_execution("execution".to_string())
.begin_execution(&request("execution"))
.expect("register execution");
assert_eq!(
@@ -36,7 +141,7 @@ fn cell_closure_waits_until_the_started_cell_is_claimed() {
fn oversized_cell_ids_are_rejected_before_admission() {
let mut state = SessionState::default();
state
.begin_execution("execution".to_string())
.begin_execution(&request("execution"))
.expect("register execution");
assert_eq!(
@@ -49,11 +154,50 @@ fn oversized_cell_ids_are_rejected_before_admission() {
assert_eq!(state.remove_execution("execution"), None);
}
#[test]
fn callbacks_cannot_claim_another_executions_cell() {
let mut state = SessionState::default();
state
.begin_execution(&request("first"))
.expect("register first execution");
state
.begin_execution(&request("second"))
.expect("register second execution");
state
.admit_invocation(&tool_call("first", /*invocation_id*/ 1))
.expect("allow the first execution to claim its cell");
let expected = "code-mode host reused active cell ID cell".to_string();
assert_eq!(
state
.admit_invocation(&tool_call("second", /*invocation_id*/ 2))
.err(),
Some(expected.clone())
);
assert_eq!(
state
.admit_notification(&notification("second", /*notification_id*/ 1))
.err(),
Some(expected.clone())
);
assert_eq!(
state
.close_cell(grpc::CellClosed {
execution_id: "second".to_string(),
cell_id: "cell".to_string(),
final_tool_call_sequence: 0,
})
.err(),
Some(expected.clone())
);
assert_eq!(state.admit_execution("second", "cell"), Err(expected));
}
#[test]
fn abandonment_before_start_ignores_later_cell_closure() {
let mut state = SessionState::default();
state
.begin_execution("execution".to_string())
.begin_execution(&request("execution"))
.expect("register execution");
assert_eq!(state.remove_execution("execution"), None);
@@ -71,19 +215,307 @@ fn abandonment_before_start_ignores_later_cell_closure() {
}
#[test]
fn disconnect_returns_each_live_cell_once() {
fn abandonment_revokes_callbacks_and_ignores_late_events() {
let mut state = SessionState::default();
state
.begin_execution("execution".to_string())
.begin_execution(&request("execution"))
.expect("register execution");
let CallbackAdmission::Active(invocation) = state
.admit_invocation(&tool_call("execution", /*invocation_id*/ 1))
.expect("admit callback before execution starts")
else {
panic!("invocation was not admitted");
};
assert_eq!(
state.remove_execution("execution"),
Some(CellId::new("cell".to_string()))
);
assert!(invocation.is_cancelled());
assert!(matches!(
state
.admit_invocation(&tool_call("execution", /*invocation_id*/ 2))
.expect("reject delayed tool invocation"),
CallbackAdmission::Closed
));
assert_eq!(
state
.close_cell(grpc::CellClosed {
execution_id: "execution".to_string(),
cell_id: "cell".to_string(),
final_tool_call_sequence: 1,
})
.expect("ignore delayed cell closure"),
None
);
assert!(state.close(/*failure*/ None).is_empty());
}
#[test]
fn invocation_cancellation_revokes_delegate_and_late_completion() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.expect("register execution");
state
.admit_execution("execution", "cell")
.expect("admit execution");
state
.mark_execution_ready("execution")
.expect("claim execution");
let invocation = tool_call("execution", /*invocation_id*/ 1);
let CallbackAdmission::Active(cancellation) = state
.admit_invocation(&invocation)
.expect("accept invocation")
else {
panic!("invocation was not admitted");
};
state
.cancel_invocation(&invocation.invocation_id)
.expect("cancel invocation");
assert!(cancellation.is_cancelled());
state.finish_invocation(&invocation.invocation_id);
assert_eq!(
state
.close_cell(grpc::CellClosed {
execution_id: "execution".to_string(),
cell_id: "cell".to_string(),
final_tool_call_sequence: 1,
})
.expect("close cell"),
Some(CellId::new("cell".to_string()))
);
}
#[test]
fn duplicate_invocation_ids_are_rejected() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.expect("register execution");
state
.admit_execution("execution", "cell")
.expect("admit execution");
let invocation = tool_call("execution", /*invocation_id*/ 1);
state
.admit_invocation(&invocation)
.expect("accept invocation");
state.finish_invocation(&invocation.invocation_id);
assert!(state.admit_invocation(&invocation).is_err());
}
#[test]
fn tool_callbacks_must_match_the_executions_enabled_tools() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.expect("register execution");
let mut disabled = tool_call("execution", /*invocation_id*/ 1);
disabled.tool_name = Some(grpc::ToolName {
name: "hidden".to_string(),
namespace: None,
});
assert!(matches!(
state.admit_invocation(&disabled),
Ok(CallbackAdmission::Rejected(error))
if error == "code-mode tool hidden is not enabled for this execution"
));
let mut wrong_namespace = tool_call("execution", /*invocation_id*/ 2);
wrong_namespace.tool_name = Some(grpc::ToolName {
name: "tool".to_string(),
namespace: Some("private".to_string()),
});
assert!(matches!(
state.admit_invocation(&wrong_namespace),
Ok(CallbackAdmission::Rejected(_))
));
let mut wrong_kind = tool_call("execution", /*invocation_id*/ 3);
wrong_kind.tool_kind = grpc::ToolKind::Freeform as i32;
assert!(matches!(
state.admit_invocation(&wrong_kind),
Ok(CallbackAdmission::Rejected(_))
));
let mut explicit_default_namespace = tool_call("execution", /*invocation_id*/ 4);
explicit_default_namespace.tool_name = Some(grpc::ToolName {
name: "tool".to_string(),
namespace: Some("functions".to_string()),
});
assert!(matches!(
state.admit_invocation(&explicit_default_namespace),
Ok(CallbackAdmission::Active(_))
));
assert_eq!(state.require_open(), Ok(()));
}
#[test]
fn execution_call_ids_must_be_bounded() {
let mut state = SessionState::default();
let mut oversized = request("execution");
oversized.tool_call_id = "x".repeat(grpc::MAX_IDENTIFIER_BYTES + 1);
assert_eq!(
state.begin_execution(&oversized),
Err(format!(
"gRPC code-mode host returned tool call ID exceeding {} bytes",
grpc::MAX_IDENTIFIER_BYTES
))
);
assert_eq!(state.remove_execution("execution"), None);
}
#[test]
fn notification_call_ids_must_match_their_execution() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.expect("register execution");
let mut oversized = notification("execution", /*notification_id*/ 1);
oversized.call_id = "x".repeat(grpc::MAX_IDENTIFIER_BYTES + 1);
assert_eq!(
state.admit_notification(&oversized).err(),
Some(format!(
"gRPC code-mode host returned notification call ID exceeding {} bytes",
grpc::MAX_IDENTIFIER_BYTES
))
);
let mut mismatched = notification("execution", /*notification_id*/ 2);
mismatched.call_id = "other-call".to_string();
assert_eq!(
state.admit_notification(&mismatched).err(),
Some("code-mode notification call ID does not match its execution".to_string())
);
assert!(matches!(
state.admit_notification(&notification("execution", /*notification_id*/ 3)),
Ok(CallbackAdmission::Active(_))
));
}
#[test]
fn malformed_callback_ids_are_rejected_before_retention() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.expect("register execution");
let mut invalid_invocation = tool_call("execution", /*invocation_id*/ 1);
invalid_invocation.invocation_id = "not-a-uuid".to_string();
assert_eq!(
state.admit_invocation(&invalid_invocation).err(),
Some("code-mode tool invocation ID must be a UUID".to_string())
);
assert_eq!(
state.cancel_invocation("not-a-uuid"),
Err("code-mode tool invocation ID must be a UUID".to_string())
);
assert_eq!(
state.cancel_invocation(&"x".repeat(grpc::MAX_IDENTIFIER_BYTES + 1)),
Err("code-mode tool invocation ID must be a UUID".to_string())
);
let mut invalid_notification = notification("execution", /*notification_id*/ 1);
invalid_notification.notification_id = "not-a-uuid".to_string();
assert_eq!(
state.admit_notification(&invalid_notification).err(),
Some("code-mode notification ID must be a UUID".to_string())
);
let invocation = tool_call("execution", /*invocation_id*/ 2);
state
.cancel_invocation(&invocation.invocation_id)
.expect("remember valid cancellation");
assert!(matches!(
state.admit_invocation(&invocation),
Ok(CallbackAdmission::Cancelled)
));
}
#[test]
fn notifications_and_tools_share_the_pending_delegate_limit() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.expect("register execution");
for index in 0..MAX_PENDING_DELEGATE_CALLS {
assert!(matches!(
state.admit_notification(&notification("execution", index as u128 + 1)),
Ok(CallbackAdmission::Active(_))
));
}
assert!(matches!(
state.admit_notification(&notification("execution", /*notification_id*/ 2_000)),
Ok(CallbackAdmission::Rejected(error))
if error == "code-mode host exceeded its pending delegate callback limit"
));
assert!(matches!(
state.admit_invocation(&tool_call("execution", /*invocation_id*/ 1)),
Ok(CallbackAdmission::Rejected(error))
if error == "code-mode host exceeded its pending delegate callback limit"
));
assert_eq!(state.require_open(), Ok(()));
assert_eq!(state.finish_notification("execution"), None);
assert!(matches!(
state.admit_invocation(&tool_call("execution", /*invocation_id*/ 2)),
Ok(CallbackAdmission::Active(_))
));
}
#[test]
fn terminated_cells_cancel_pending_notifications() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.expect("register execution");
let CallbackAdmission::Active(cancellation) = state
.admit_notification(&notification("execution", /*notification_id*/ 1))
.expect("admit notification")
else {
panic!("notification was not admitted");
};
state.cancel_notifications(&CellId::new("cell".to_string()));
assert!(cancellation.is_cancelled());
assert_eq!(state.finish_notification("execution"), None);
}
#[test]
fn disconnect_revokes_callbacks_and_returns_each_live_cell_once() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.expect("register execution");
state
.admit_execution("execution", "cell")
.expect("admit execution");
let CallbackAdmission::Active(cancellation) = state
.admit_invocation(&tool_call("execution", /*invocation_id*/ 1))
.expect("accept invocation")
else {
panic!("invocation was not admitted");
};
let CallbackAdmission::Active(notification_cancellation) = state
.admit_notification(&notification("execution", /*notification_id*/ 1))
.expect("admit notification")
else {
panic!("notification was not admitted");
};
assert_eq!(
state.close(Some("lease closed".to_string())),
vec![CellId::new("cell".to_string())]
);
assert!(cancellation.is_cancelled());
assert!(notification_cancellation.is_cancelled());
assert!(state.close(/*failure*/ None).is_empty());
assert_eq!(state.require_open(), Err("lease closed".to_string()));
}