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

code-mode: use client cell ids and linear observations

Browse Source
This commit is contained in:
Channing Conger
2026-06-22 08:49:51 +00:00
parent 28536e0bb5
commit e7f4fc2005
18 changed files with 349 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/code_mode/delegate_tests.rs
+4 -2
View File
@@ -5,6 +5,7 @@ use codex_code_mode::CellId;
use codex_code_mode::CodeModeService;
use codex_code_mode::CodeModeSessionDelegate;
use codex_code_mode::CreateCellRequest;
use codex_code_mode::ObservationGeneration;
use codex_code_mode::ObserveOutcome;
use codex_code_mode::ObserveRequest;
@@ -42,9 +43,10 @@ fn terminal_notification_before_readiness_does_not_reopen_the_dispatch_gate() {
async fn background_completion_removes_the_dispatch_gate_without_another_observation() {
let broker = Arc::new(CodeModeDispatchBroker::new());
let service = CodeModeService::with_delegate(broker.clone());
let requested_cell_id = CellId::new("cell-a7".to_string());
let cell_id = service
.create_cell(CreateCellRequest {
idempotency_key: "delegate-cleanup-cell".to_string(),
cell_id: requested_cell_id,
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: concat!(
@@ -60,8 +62,8 @@ async fn background_completion_removes_the_dispatch_gate_without_another_observa
assert_eq!(
service
.observe(ObserveRequest {
idempotency_key: "delegate-cleanup-observation".to_string(),
cell_id: cell_id.clone(),
generation: ObservationGeneration::INITIAL,
yield_time_ms: 1,
})
.await
codex-rs/core/src/tools/code_mode/execute_handler.rs
+20 -19
View File
@@ -40,19 +40,7 @@ impl CodeModeExecuteHandler {
let enabled_tools =
codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs);
let started_at = std::time::Instant::now();
let idempotency_key = format!("{}:{call_id}", exec.session.thread_id());
let cell_id = exec
.session
.services
.code_mode_service
.create_cell(codex_code_mode::CreateCellRequest {
idempotency_key: idempotency_key.clone(),
tool_call_id: call_id.clone(),
enabled_tools,
source: args.code.clone(),
})
.await
.map_err(FunctionCallError::RespondToModel)?;
let cell_id = super::cell_id_for_tool_call(&exec.session.thread_id().to_string(), &call_id);
let runtime_cell_id = cell_id.to_string();
let code_cell_trace = exec
.session
@@ -70,6 +58,21 @@ impl CodeModeExecuteHandler {
code_cell_trace.clone(),
)
.map_err(FunctionCallError::RespondToModel)?;
let cell_id = exec
.session
.services
.code_mode_service
.create_cell(codex_code_mode::CreateCellRequest {
cell_id,
tool_call_id: call_id.clone(),
enabled_tools,
source: args.code.clone(),
})
.await
.map_err(|error| {
initial_observation_guard.finish_with_error(&error);
FunctionCallError::RespondToModel(error)
})?;
exec.session
.services
.code_mode_service
@@ -78,13 +81,11 @@ impl CodeModeExecuteHandler {
.session
.services
.code_mode_service
.observe(codex_code_mode::ObserveRequest {
idempotency_key,
cell_id: cell_id.clone(),
yield_time_ms: args
.yield_time_ms
.observe(
cell_id.clone(),
args.yield_time_ms
.unwrap_or(codex_code_mode::DEFAULT_EXEC_YIELD_TIME_MS),
})
)
.await;
let response: codex_code_mode::RuntimeResponse = match observed {
Ok(outcome) => outcome.into(),
codex-rs/core/src/tools/code_mode/mod.rs
+97 -5
View File
@@ -5,13 +5,16 @@ mod response_adapter;
mod wait_handler;
pub(crate) mod wait_spec;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
use codex_code_mode::CellId;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeSession;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::ObservationGeneration;
use codex_code_mode::RuntimeResponse;
use codex_protocol::models::FunctionCallOutputContentItem;
use serde_json::Value as JsonValue;
@@ -61,6 +64,7 @@ pub(crate) struct ExecContext {
pub(crate) struct CodeModeService {
session: Option<Arc<dyn CodeModeSession>>,
dispatch_broker: Arc<CodeModeDispatchBroker>,
observation_generations: Mutex<HashMap<CellId, ObservationGeneration>>,
}
struct InitialObservationGuard {
@@ -128,6 +132,7 @@ impl CodeModeService {
dispatch_broker.clone(),
))),
dispatch_broker,
observation_generations: Mutex::new(HashMap::new()),
}
}
@@ -135,14 +140,50 @@ impl CodeModeService {
&self,
request: codex_code_mode::CreateCellRequest,
) -> Result<CellId, String> {
self.session()?.create_cell(request).await
let requested_cell_id = request.cell_id.clone();
let admitted_cell_id = self.session()?.create_cell(request).await?;
if admitted_cell_id != requested_cell_id {
return Err(format!(
"code-mode host admitted cell {admitted_cell_id} for requested cell {requested_cell_id}"
));
}
self.observation_generations
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.entry(admitted_cell_id.clone())
.or_insert(ObservationGeneration::INITIAL);
Ok(admitted_cell_id)
}
pub(crate) async fn observe(
&self,
request: codex_code_mode::ObserveRequest,
cell_id: CellId,
yield_time_ms: u64,
) -> Result<codex_code_mode::ObserveOutcome, String> {
self.session()?.observe(request).await
let generation = self
.observation_generations
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(&cell_id)
.copied()
.unwrap_or(ObservationGeneration::INITIAL);
let outcome = self
.session()?
.observe(codex_code_mode::ObserveRequest {
cell_id: cell_id.clone(),
generation,
yield_time_ms,
})
.await?;
let next_generation = generation.next().ok_or_else(|| {
format!("observation generation exhausted for code-mode cell {cell_id}")
})?;
let mut generations = self
.observation_generations
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
generations.insert(cell_id, next_generation);
Ok(outcome)
}
pub(crate) async fn terminate(
@@ -153,10 +194,15 @@ impl CodeModeService {
}
pub(crate) async fn shutdown(&self) -> Result<(), String> {
match &self.session {
let result = match &self.session {
Some(session) => session.shutdown().await,
None => Ok(()),
}
};
self.observation_generations
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clear();
result
}
pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode::CellId) {
@@ -198,6 +244,30 @@ impl CodeModeService {
}
}
fn cell_id_for_tool_call(thread_id: &str, call_id: &str) -> CellId {
use sha1::Digest;
const CELL_ID_ALPHABET: &[u8; 32] = b"0123456789abcdefghjkmnpqrstvwxyz";
const CELL_ID_LENGTH: usize = 16;
let digest = sha1::Sha1::new()
.chain_update(thread_id.as_bytes())
.chain_update([0])
.chain_update(call_id.as_bytes())
.finalize();
let mut encoded = String::with_capacity(CELL_ID_LENGTH);
for chunk in digest[..10].chunks_exact(5) {
let chunk_bits = chunk
.iter()
.fold(0_u64, |value, byte| (value << 8) | u64::from(*byte));
for shift in (0..40).step_by(5).rev() {
let index = ((chunk_bits >> shift) & 0x1f) as usize;
encoded.push(char::from(CELL_ID_ALPHABET[index]));
}
}
CellId::new(encoded)
}
pub(super) async fn handle_runtime_response(
exec: &ExecContext,
response: RuntimeResponse,
@@ -382,6 +452,7 @@ mod tests {
use super::CodeModeDispatchBroker;
use super::InitialObservationGuard;
use super::build_nested_tool_payload;
use super::cell_id_for_tool_call;
use super::truncate_code_mode_result;
use crate::tools::context::ToolPayload;
use codex_code_mode::CellId;
@@ -494,6 +565,27 @@ mod tests {
);
}
#[test]
fn cell_id_is_compact_and_stable_for_a_resumed_thread() {
let thread_id = "019f2f5a-7b6c-7000-8000-000000000001";
let first = cell_id_for_tool_call(thread_id, "call-42");
let resumed = cell_id_for_tool_call(thread_id, "call-42");
assert_eq!(resumed, first);
assert_eq!(first.as_str().len(), 16);
assert!(
first
.as_str()
.bytes()
.all(|byte| b"0123456789abcdefghjkmnpqrstvwxyz".contains(&byte))
);
assert_ne!(cell_id_for_tool_call(thread_id, "call-43"), first);
assert_ne!(
cell_id_for_tool_call("019f2f5a-7b6c-7000-8000-000000000002", "call-42"),
first
);
}
#[tokio::test]
async fn dropping_initial_observation_guard_terminates_the_cell() {
let session = Arc::new(RecordingCodeModeSession {
codex-rs/core/src/tools/code_mode/wait_handler.rs
+1 -7
View File
@@ -66,7 +66,6 @@ impl CodeModeWaitHandler {
let ToolInvocation {
session,
turn,
call_id,
tool_name,
payload,
..
@@ -77,7 +76,6 @@ impl CodeModeWaitHandler {
if tool_name.namespace.is_none() && tool_name.name.as_str() == WAIT_TOOL_NAME =>
{
let args: ExecWaitArgs = parse_arguments(&arguments)?;
let idempotency_key = format!("{}:{call_id}", session.thread_id());
let exec = ExecContext { session, turn };
let started_at = std::time::Instant::now();
let cell_id = codex_code_mode::CellId::new(args.cell_id);
@@ -100,11 +98,7 @@ impl CodeModeWaitHandler {
.session
.services
.code_mode_service
.observe(codex_code_mode::ObserveRequest {
idempotency_key,
cell_id,
yield_time_ms: args.yield_time_ms,
})
.observe(cell_id, args.yield_time_ms)
.await
.map_err(FunctionCallError::RespondToModel)?;
let cell_was_present =