| use std::sync::Arc; |
|
|
| use codex_code_mode_protocol::CellId; |
| use codex_code_mode_protocol::CodeModeNestedToolCall; |
| use codex_code_mode_protocol::CodeModeSessionDelegate; |
| use codex_code_mode_protocol::NotificationFuture; |
| use codex_code_mode_protocol::RuntimeResponse; |
| use codex_code_mode_protocol::StartedCell; |
| use codex_code_mode_protocol::ToolInvocationFuture; |
| use codex_code_mode_protocol::WaitOutcome; |
| use tokio_util::sync::CancellationToken; |
|
|
| pub(super) struct GenerationDelegate { |
| pub(super) delegate: Arc<dyn CodeModeSessionDelegate>, |
| pub(super) generation: u64, |
| } |
|
|
| impl CodeModeSessionDelegate for GenerationDelegate { |
| fn invoke_tool<'a>( |
| &'a self, |
| mut invocation: CodeModeNestedToolCall, |
| cancellation_token: CancellationToken, |
| ) -> ToolInvocationFuture<'a> { |
| invocation.cell_id = public_cell_id(self.generation, &invocation.cell_id); |
| self.delegate.invoke_tool(invocation, cancellation_token) |
| } |
|
|
| fn notify<'a>( |
| &'a self, |
| call_id: String, |
| cell_id: CellId, |
| text: String, |
| cancellation_token: CancellationToken, |
| ) -> NotificationFuture<'a> { |
| self.delegate.notify( |
| call_id, |
| public_cell_id(self.generation, &cell_id), |
| text, |
| cancellation_token, |
| ) |
| } |
|
|
| fn cell_closed(&self, cell_id: &CellId) { |
| self.delegate |
| .cell_closed(&public_cell_id(self.generation, cell_id)); |
| } |
| } |
|
|
| fn public_cell_id(generation: u64, cell_id: &CellId) -> CellId { |
| if generation == 1 { |
| cell_id.clone() |
| } else { |
| CellId::new(format!("g{generation}:{cell_id}")) |
| } |
| } |
|
|
| pub(super) fn remote_cell_id(generation: u64, cell_id: &CellId) -> Result<CellId, String> { |
| if generation == 1 { |
| return Ok(cell_id.clone()); |
| } |
|
|
| let prefix = format!("g{generation}:"); |
| cell_id |
| .as_str() |
| .strip_prefix(&prefix) |
| .map(|cell_id| CellId::new(cell_id.to_string())) |
| .ok_or_else(|| "cell belongs to a stale code-mode host generation".to_string()) |
| } |
|
|
| pub(super) fn public_started_cell(generation: u64, started: StartedCell) -> StartedCell { |
| if generation == 1 { |
| return started; |
| } |
| let cell_id = public_cell_id(generation, &started.cell_id); |
| StartedCell::from_future(cell_id, async move { |
| started |
| .initial_response() |
| .await |
| .map(|response| public_runtime_response(generation, response)) |
| }) |
| } |
|
|
| fn public_runtime_response(generation: u64, response: RuntimeResponse) -> RuntimeResponse { |
| match response { |
| RuntimeResponse::Yielded { |
| cell_id, |
| content_items, |
| code_mode_host_duration, |
| } => RuntimeResponse::Yielded { |
| cell_id: public_cell_id(generation, &cell_id), |
| content_items, |
| code_mode_host_duration, |
| }, |
| RuntimeResponse::Terminated { |
| cell_id, |
| content_items, |
| code_mode_host_duration, |
| } => RuntimeResponse::Terminated { |
| cell_id: public_cell_id(generation, &cell_id), |
| content_items, |
| code_mode_host_duration, |
| }, |
| RuntimeResponse::Result { |
| cell_id, |
| content_items, |
| error_text, |
| code_mode_host_duration, |
| } => RuntimeResponse::Result { |
| cell_id: public_cell_id(generation, &cell_id), |
| content_items, |
| error_text, |
| code_mode_host_duration, |
| }, |
| } |
| } |
|
|
| pub(super) fn public_wait_outcome(generation: u64, outcome: WaitOutcome) -> WaitOutcome { |
| match outcome { |
| WaitOutcome::LiveCell(response) => { |
| WaitOutcome::LiveCell(public_runtime_response(generation, response)) |
| } |
| WaitOutcome::MissingCell(response) => { |
| WaitOutcome::MissingCell(public_runtime_response(generation, response)) |
| } |
| } |
| } |
|
|
| #[cfg(test)] |
| #[path = "generation_tests.rs"] |
| mod tests; |
|
|
