code-mode: expose create and observe operations

This commit is contained in:
Channing Conger
2026-06-21 07:20:01 +00:00
parent d969450436
commit e81e125cb1
17 changed files with 964 additions and 436 deletions
+6 -7
View File
@@ -20,17 +20,17 @@ pub use description::render_json_schema_to_typescript;
pub use response::DEFAULT_IMAGE_DETAIL;
pub use response::FunctionCallOutputContentItem;
pub use response::ImageDetail;
pub use runtime::CellOutcome;
pub use runtime::CodeModeNestedToolCall;
pub use runtime::CreateCellRequest;
pub use runtime::DEFAULT_EXEC_YIELD_TIME_MS;
pub use runtime::DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL;
pub use runtime::DEFAULT_WAIT_YIELD_TIME_MS;
pub use runtime::ExecuteRequest;
pub use runtime::ExecuteToPendingOutcome;
pub use runtime::ObserveRequest;
pub use runtime::ObserveToPendingOutcome;
pub use runtime::ObserveToPendingRequest;
pub use runtime::PendingOutcome;
pub use runtime::RuntimeResponse;
pub use runtime::WaitOutcome;
pub use runtime::WaitRequest;
pub use runtime::WaitToPendingOutcome;
pub use runtime::WaitToPendingRequest;
pub use session::CellId;
pub use session::CodeModeSession;
pub use session::CodeModeSessionDelegate;
@@ -38,7 +38,6 @@ pub use session::CodeModeSessionProvider;
pub use session::CodeModeSessionProviderFuture;
pub use session::CodeModeSessionResultFuture;
pub use session::NotificationFuture;
pub use session::StartedCell;
pub use session::ToolInvocationFuture;
pub const PUBLIC_TOOL_NAME: &str = "exec";
+14 -12
View File
@@ -13,33 +13,31 @@ pub const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000;
pub const DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL: usize = 10_000;
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct ExecuteRequest {
pub struct CreateCellRequest {
pub tool_call_id: String,
pub enabled_tools: Vec<ToolDefinition>,
pub source: String,
pub yield_time_ms: Option<u64>,
pub max_output_tokens: Option<usize>,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct WaitRequest {
pub struct ObserveRequest {
pub cell_id: CellId,
pub yield_time_ms: u64,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct WaitToPendingRequest {
pub struct ObserveToPendingRequest {
pub cell_id: CellId,
}
#[derive(Debug, Deserialize, PartialEq, Serialize)]
pub enum WaitOutcome {
pub enum CellOutcome {
LiveCell(RuntimeResponse),
MissingCell(RuntimeResponse),
}
#[derive(Debug, Deserialize, PartialEq, Serialize)]
pub enum ExecuteToPendingOutcome {
pub enum PendingOutcome {
Pending {
cell_id: CellId,
content_items: Vec<FunctionCallOutputContentItem>,
@@ -49,15 +47,15 @@ pub enum ExecuteToPendingOutcome {
}
#[derive(Debug, Deserialize, PartialEq, Serialize)]
pub enum WaitToPendingOutcome {
LiveCell(ExecuteToPendingOutcome),
pub enum ObserveToPendingOutcome {
LiveCell(PendingOutcome),
MissingCell(RuntimeResponse),
}
impl From<WaitOutcome> for RuntimeResponse {
fn from(outcome: WaitOutcome) -> Self {
impl From<CellOutcome> for RuntimeResponse {
fn from(outcome: CellOutcome) -> Self {
match outcome {
WaitOutcome::LiveCell(response) | WaitOutcome::MissingCell(response) => response,
CellOutcome::LiveCell(response) | CellOutcome::MissingCell(response) => response,
}
}
}
@@ -87,3 +85,7 @@ pub struct CodeModeNestedToolCall {
pub tool_kind: CodeModeToolKind,
pub input: Option<JsonValue>,
}
#[cfg(test)]
#[path = "runtime_tests.rs"]
mod tests;
@@ -0,0 +1,301 @@
use std::fmt::Debug;
use codex_protocol::ToolName;
use pretty_assertions::assert_eq;
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::Value as JsonValue;
use serde_json::json;
use super::*;
use crate::ImageDetail;
fn assert_json_round_trip<T>(value: &T, expected: JsonValue)
where
T: Debug + DeserializeOwned + PartialEq + Serialize,
{
assert_eq!(serde_json::to_value(value).unwrap(), expected);
assert_eq!(&serde_json::from_value::<T>(expected).unwrap(), value);
}
fn text(value: &str) -> FunctionCallOutputContentItem {
FunctionCallOutputContentItem::InputText {
text: value.to_string(),
}
}
fn image(image_url: &str, detail: Option<ImageDetail>) -> FunctionCallOutputContentItem {
FunctionCallOutputContentItem::InputImage {
image_url: image_url.to_string(),
detail,
}
}
#[test]
fn requests_round_trip_with_exact_tool_and_namespace_fields() {
let create = CreateCellRequest {
tool_call_id: "response-call-7".to_string(),
enabled_tools: vec![
ToolDefinition {
name: "mcp__search__query".to_string(),
tool_name: ToolName::namespaced("search", "query"),
description: "Search indexed documents".to_string(),
kind: CodeModeToolKind::Function,
input_schema: Some(json!({
"type": "object",
"properties": {"q": {"type": "string"}},
})),
output_schema: None,
},
ToolDefinition {
name: "apply_patch".to_string(),
tool_name: ToolName::plain("apply_patch"),
description: "Apply a patch".to_string(),
kind: CodeModeToolKind::Freeform,
input_schema: None,
output_schema: Some(json!({"type": "string"})),
},
],
source: "await tools.mcp__search__query({ q: 'rust' });".to_string(),
};
assert_json_round_trip(
&create,
json!({
"tool_call_id": "response-call-7",
"enabled_tools": [
{
"name": "mcp__search__query",
"tool_name": {"name": "query", "namespace": "search"},
"description": "Search indexed documents",
"kind": "function",
"input_schema": {
"type": "object",
"properties": {"q": {"type": "string"}},
},
"output_schema": null,
},
{
"name": "apply_patch",
"tool_name": {"name": "apply_patch", "namespace": null},
"description": "Apply a patch",
"kind": "freeform",
"input_schema": null,
"output_schema": {"type": "string"},
},
],
"source": "await tools.mcp__search__query({ q: 'rust' });",
}),
);
assert_json_round_trip(
&ObserveRequest {
cell_id: CellId::new("cell-a7".to_string()),
yield_time_ms: 250,
},
json!({"cell_id": "cell-a7", "yield_time_ms": 250}),
);
assert_json_round_trip(
&ObserveToPendingRequest {
cell_id: CellId::new("cell-b9".to_string()),
},
json!({"cell_id": "cell-b9"}),
);
}
#[test]
fn nested_tool_calls_round_trip_with_optional_input_and_tool_kind() {
assert_json_round_trip(
&CodeModeNestedToolCall {
cell_id: CellId::new("cell-a7".to_string()),
runtime_tool_call_id: "runtime-call-1".to_string(),
tool_name: ToolName::namespaced("search", "query"),
tool_kind: CodeModeToolKind::Function,
input: Some(json!({"q": "rust"})),
},
json!({
"cell_id": "cell-a7",
"runtime_tool_call_id": "runtime-call-1",
"tool_name": {"name": "query", "namespace": "search"},
"tool_kind": "function",
"input": {"q": "rust"},
}),
);
assert_json_round_trip(
&CodeModeNestedToolCall {
cell_id: CellId::new("cell-b9".to_string()),
runtime_tool_call_id: "runtime-call-2".to_string(),
tool_name: ToolName::plain("apply_patch"),
tool_kind: CodeModeToolKind::Freeform,
input: None,
},
json!({
"cell_id": "cell-b9",
"runtime_tool_call_id": "runtime-call-2",
"tool_name": {"name": "apply_patch", "namespace": null},
"tool_kind": "freeform",
"input": null,
}),
);
}
#[test]
fn runtime_responses_round_trip_with_every_output_and_terminal_variant() {
let yielded = RuntimeResponse::Yielded {
cell_id: CellId::new("cell-a7".to_string()),
content_items: vec![
text("before"),
image("data:image/png;base64,auto", Some(ImageDetail::Auto)),
image("data:image/png;base64,low", Some(ImageDetail::Low)),
image("data:image/png;base64,high", Some(ImageDetail::High)),
image(
"data:image/png;base64,original",
Some(ImageDetail::Original),
),
image("data:image/png;base64,default", None),
],
};
assert_json_round_trip(
&yielded,
json!({
"Yielded": {
"cell_id": "cell-a7",
"content_items": [
{"type": "input_text", "text": "before"},
{"type": "input_image", "image_url": "data:image/png;base64,auto", "detail": "auto"},
{"type": "input_image", "image_url": "data:image/png;base64,low", "detail": "low"},
{"type": "input_image", "image_url": "data:image/png;base64,high", "detail": "high"},
{"type": "input_image", "image_url": "data:image/png;base64,original", "detail": "original"},
{"type": "input_image", "image_url": "data:image/png;base64,default"},
],
},
}),
);
assert_json_round_trip(
&RuntimeResponse::Terminated {
cell_id: CellId::new("cell-b9".to_string()),
content_items: vec![text("partial")],
},
json!({
"Terminated": {
"cell_id": "cell-b9",
"content_items": [{"type": "input_text", "text": "partial"}],
},
}),
);
assert_json_round_trip(
&RuntimeResponse::Result {
cell_id: CellId::new("cell-c3".to_string()),
content_items: vec![text("failed")],
error_text: Some("tool failed".to_string()),
},
json!({
"Result": {
"cell_id": "cell-c3",
"content_items": [{"type": "input_text", "text": "failed"}],
"error_text": "tool failed",
},
}),
);
assert_json_round_trip(
&RuntimeResponse::Result {
cell_id: CellId::new("cell-d4".to_string()),
content_items: Vec::new(),
error_text: None,
},
json!({
"Result": {
"cell_id": "cell-d4",
"content_items": [],
"error_text": null,
},
}),
);
}
#[test]
fn observation_outcomes_round_trip_with_live_missing_pending_and_completed_cells() {
let result = RuntimeResponse::Result {
cell_id: CellId::new("cell-a7".to_string()),
content_items: vec![text("done")],
error_text: None,
};
assert_json_round_trip(
&CellOutcome::LiveCell(result.clone()),
json!({
"LiveCell": {
"Result": {
"cell_id": "cell-a7",
"content_items": [{"type": "input_text", "text": "done"}],
"error_text": null,
},
},
}),
);
assert_json_round_trip(
&CellOutcome::MissingCell(RuntimeResponse::Terminated {
cell_id: CellId::new("missing".to_string()),
content_items: Vec::new(),
}),
json!({
"MissingCell": {
"Terminated": {"cell_id": "missing", "content_items": []},
},
}),
);
let pending = PendingOutcome::Pending {
cell_id: CellId::new("cell-b9".to_string()),
content_items: vec![text("waiting")],
pending_tool_call_ids: vec!["runtime-call-1".to_string()],
};
assert_json_round_trip(
&pending,
json!({
"Pending": {
"cell_id": "cell-b9",
"content_items": [{"type": "input_text", "text": "waiting"}],
"pending_tool_call_ids": ["runtime-call-1"],
},
}),
);
assert_json_round_trip(
&PendingOutcome::Completed(result),
json!({
"Completed": {
"Result": {
"cell_id": "cell-a7",
"content_items": [{"type": "input_text", "text": "done"}],
"error_text": null,
},
},
}),
);
assert_json_round_trip(
&ObserveToPendingOutcome::LiveCell(pending),
json!({
"LiveCell": {
"Pending": {
"cell_id": "cell-b9",
"content_items": [{"type": "input_text", "text": "waiting"}],
"pending_tool_call_ids": ["runtime-call-1"],
},
},
}),
);
assert_json_round_trip(
&ObserveToPendingOutcome::MissingCell(RuntimeResponse::Result {
cell_id: CellId::new("missing".to_string()),
content_items: Vec::new(),
error_text: Some("cell not found".to_string()),
}),
json!({
"MissingCell": {
"Result": {
"cell_id": "missing",
"content_items": [],
"error_text": "cell not found",
},
},
}),
);
}
+15 -52
View File
@@ -6,14 +6,12 @@ use std::sync::Arc;
use serde::Deserialize;
use serde::Serialize;
use serde_json::Value as JsonValue;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use crate::CellOutcome;
use crate::CodeModeNestedToolCall;
use crate::ExecuteRequest;
use crate::RuntimeResponse;
use crate::WaitOutcome;
use crate::WaitRequest;
use crate::CreateCellRequest;
use crate::ObserveRequest;
pub type CodeModeSessionResultFuture<'a, T> =
Pin<Box<dyn Future<Output = Result<T, String>> + Send + 'a>>;
@@ -48,42 +46,6 @@ impl fmt::Display for CellId {
}
}
pub struct StartedCell {
pub cell_id: CellId,
initial_response: CodeModeSessionResultFuture<'static, RuntimeResponse>,
}
impl StartedCell {
pub fn new(cell_id: CellId, initial_response_rx: oneshot::Receiver<RuntimeResponse>) -> Self {
Self {
cell_id,
initial_response: Box::pin(async move {
initial_response_rx
.await
.map_err(|_| "exec runtime ended unexpectedly".to_string())
}),
}
}
pub fn from_result_receiver(
cell_id: CellId,
initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,
) -> Self {
Self {
cell_id,
initial_response: Box::pin(async move {
initial_response_rx
.await
.map_err(|_| "exec runtime ended unexpectedly".to_string())?
}),
}
}
pub async fn initial_response(self) -> Result<RuntimeResponse, String> {
self.initial_response.await
}
}
/// Host callbacks used by a code-mode session while cells are executing.
pub trait CodeModeSessionDelegate: Send + Sync {
fn invoke_tool<'a>(
@@ -104,11 +66,13 @@ pub trait CodeModeSessionDelegate: Send + Sync {
fn cell_closed(&self, cell_id: &CellId);
}
/// A durable code-mode session owned by one Codex thread.
/// A stateful code-mode session owned by one Codex thread.
///
/// Cells executed in the same session share stored values. Separate sessions
/// must keep those values isolated. Implementations may execute cells
/// in-process or remotely.
/// in-process or remotely. Implementations should surface lost connections or
/// protocol desynchronization, but do not need to preserve cells or stored
/// values across process failure or restart.
pub trait CodeModeSession: Send + Sync {
/// Returns whether the session can still accept requests.
///
@@ -116,14 +80,17 @@ pub trait CodeModeSession: Send + Sync {
/// connection fails so callers can create a fresh session for later work.
fn is_alive(&self) -> bool;
fn execute<'a>(
fn create_cell<'a>(
&'a self,
request: ExecuteRequest,
) -> CodeModeSessionResultFuture<'a, StartedCell>;
request: CreateCellRequest,
) -> CodeModeSessionResultFuture<'a, CellId>;
fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
fn observe<'a>(
&'a self,
request: ObserveRequest,
) -> CodeModeSessionResultFuture<'a, CellOutcome>;
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, CellOutcome>;
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>;
}
@@ -138,7 +105,3 @@ pub trait CodeModeSessionProvider: Send + Sync {
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionProviderFuture<'a>;
}
#[cfg(test)]
#[path = "session_tests.rs"]
mod tests;
@@ -1,19 +0,0 @@
use pretty_assertions::assert_eq;
use tokio::sync::oneshot;
use super::CellId;
use super::StartedCell;
#[tokio::test]
async fn started_cell_preserves_remote_initial_response_errors() {
let (response_tx, response_rx) = oneshot::channel();
response_tx
.send(Err("remote runtime failed".to_string()))
.expect("initial response receiver should be open");
let started = StartedCell::from_result_receiver(CellId::new("1".to_string()), response_rx);
assert_eq!(
started.initial_response().await,
Err("remote runtime failed".to_string())
);
}