diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index 6b35d92021..b1e8735155 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -49,6 +49,7 @@ async fn run() -> Result<(), String> { sessions: Mutex::new(HashMap::new()), next_session_id: AtomicU64::new(1), peer, + host_id: std::process::id().to_string(), }); let writer = tokio::spawn(async move { @@ -110,8 +111,12 @@ impl HostState { }); self.sessions.lock().await.insert( session_id, - Arc::new(CodeModeService::with_delegate(delegate)), + Arc::new(CodeModeService::with_delegate_and_cell_id_prefix( + delegate, + self.host_id.clone(), + )), ); + host_id: String, self.respond(request_id, Ok(HostResponse::SessionCreated { session_id })) .await; } diff --git a/codex-rs/code-mode-host/tests/host.rs b/codex-rs/code-mode-host/tests/host.rs index aa32957cd0..ce31eb8a96 100644 --- a/codex-rs/code-mode-host/tests/host.rs +++ b/codex-rs/code-mode-host/tests/host.rs @@ -25,6 +25,7 @@ async fn serves_code_mode_sessions_over_stdio() { .kill_on_drop(true) .spawn() .expect("spawn codex-code-mode-host"); + let host_id = child.id().expect("host process id"); let mut stdin = child.stdin.take().expect("host stdin"); let mut stdout = child.stdout.take().expect("host stdout"); @@ -70,6 +71,7 @@ async fn serves_code_mode_sessions_over_stdio() { }) => cell_id, message => panic!("unexpected execute response: {message:?}"), }; + assert_eq!(cell_id.as_str(), format!("{host_id}_1")); let response = match read_frame(&mut stdout) .await .expect("initial response frame") diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index ca53edbe54..d0856f484b 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -93,6 +93,7 @@ struct Inner { cells: Mutex>, delegate: Arc, shutting_down: AtomicBool, + cell_id_prefix: Option, next_cell_id: AtomicU64, } @@ -106,24 +107,38 @@ impl CodeModeService { } pub fn with_delegate(delegate: Arc) -> Self { + Self::with_inner(delegate, None) + } + + pub fn with_delegate_and_cell_id_prefix( + delegate: Arc, + cell_id_prefix: String, + ) -> Self { + Self::with_inner(delegate, Some(cell_id_prefix)) + } + + fn with_inner( + delegate: Arc, + cell_id_prefix: Option, + ) -> Self { Self { inner: Arc::new(Inner { stored_values: Mutex::new(HashMap::new()), cells: Mutex::new(HashMap::new()), delegate, shutting_down: AtomicBool::new(false), + cell_id_prefix, next_cell_id: AtomicU64::new(1), }), } } fn allocate_cell_id(&self) -> CellId { - CellId::new( - self.inner - .next_cell_id - .fetch_add(1, Ordering::Relaxed) - .to_string(), - ) + let cell_id = self.inner.next_cell_id.fetch_add(1, Ordering::Relaxed); + CellId::new(match &self.inner.cell_id_prefix { + Some(prefix) => format!("{prefix}_{cell_id}"), + None => cell_id.to_string(), + }) } pub async fn execute(&self, request: ExecuteRequest) -> Result { @@ -823,6 +838,7 @@ mod tests { }) } + cell_id_prefix: None, #[tokio::test] async fn synchronous_exit_returns_successfully() { let service = CodeModeService::new(); @@ -853,6 +869,27 @@ mod tests { async fn stored_values_are_shared_between_cells_but_not_sessions() { let first_session = CodeModeService::new(); let second_session = CodeModeService::new(); + #[tokio::test] + async fn cell_ids_include_the_configured_host_prefix() { + let service = CodeModeService::with_delegate_and_cell_id_prefix( + Arc::new(NoopCodeModeSessionDelegate), + "host7".to_string(), + ); + + let response = execute(&service, execute_request("text('done');")).await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("host7_1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "done".to_string(), + }], + error_text: None, + } + ); + } + let write_response = execute( &first_session,