codex/codex-rs/code-mode/src/service.rs

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::time::Duration;

use async_trait::async_trait;
use serde_json::Value as JsonValue;
use tokio::sync::Mutex;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use tracing::warn;

use crate::FunctionCallOutputContentItem;
use crate::runtime::CodeModeNestedToolCall;
use crate::runtime::DEFAULT_EXEC_YIELD_TIME_MS;
use crate::runtime::ExecuteRequest;
use crate::runtime::ExecuteToPendingOutcome;
use crate::runtime::PendingRuntimeMode;
use crate::runtime::RuntimeCommand;
use crate::runtime::RuntimeControlCommand;
use crate::runtime::RuntimeEvent;
use crate::runtime::RuntimeResponse;
use crate::runtime::TurnMessage;
use crate::runtime::WaitOutcome;
use crate::runtime::WaitRequest;
use crate::runtime::WaitToPendingOutcome;
use crate::runtime::WaitToPendingRequest;
use crate::runtime::spawn_runtime;

#[async_trait]
pub trait CodeModeTurnHost: Send + Sync {
    async fn invoke_tool(
        &self,
        invocation: CodeModeNestedToolCall,
        cancellation_token: CancellationToken,
    ) -> Result<JsonValue, String>;

    async fn notify(&self, call_id: String, cell_id: String, text: String) -> Result<(), String>;
}

#[derive(Clone)]
struct SessionHandle {
    control_tx: mpsc::UnboundedSender<SessionControlCommand>,
    runtime_tx: std::sync::mpsc::Sender<RuntimeCommand>,
}

struct Inner {
    stored_values: Mutex<HashMap<String, JsonValue>>,
    sessions: Mutex<HashMap<String, SessionHandle>>,
    turn_message_tx: async_channel::Sender<TurnMessage>,
    turn_message_rx: async_channel::Receiver<TurnMessage>,
    next_cell_id: AtomicU64,
}

pub struct CodeModeService {
    inner: Arc<Inner>,
}

impl CodeModeService {
    pub fn new() -> Self {
        let (turn_message_tx, turn_message_rx) = async_channel::unbounded();

        Self {
            inner: Arc::new(Inner {
                stored_values: Mutex::new(HashMap::new()),
                sessions: Mutex::new(HashMap::new()),
                turn_message_tx,
                turn_message_rx,
                next_cell_id: AtomicU64::new(1),
            }),
        }
    }

    pub async fn stored_values(&self) -> HashMap<String, JsonValue> {
        self.inner.stored_values.lock().await.clone()
    }

    pub async fn replace_stored_values(&self, values: HashMap<String, JsonValue>) {
        *self.inner.stored_values.lock().await = values;
    }

    /// Reserves the runtime cell id for a future `execute` request.
    ///
    /// The runtime can issue nested tool calls before the first `execute`
    /// response is returned. Hosts that need a parent trace object for those
    /// nested calls should allocate the cell id up front and pass it back on the
    /// `ExecuteRequest`.
    pub fn allocate_cell_id(&self) -> String {
        self.inner
            .next_cell_id
            .fetch_add(1, Ordering::Relaxed)
            .to_string()
    }

    pub async fn execute(&self, request: ExecuteRequest) -> Result<RuntimeResponse, String> {
        let initial_yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS);
        let (response_tx, response_rx) = oneshot::channel();
        self.start_session(
            request,
            SessionResponseSender::Runtime(response_tx),
            Some(initial_yield_time_ms),
            PendingRuntimeMode::Continue,
        )
        .await?;

        response_rx
            .await
            .map_err(|_| "exec runtime ended unexpectedly".to_string())
    }

    pub async fn execute_to_pending(
        &self,
        request: ExecuteRequest,
    ) -> Result<ExecuteToPendingOutcome, String> {
        let (response_tx, response_rx) = oneshot::channel();
        self.start_session(
            request,
            SessionResponseSender::ExecuteToPending(response_tx),
            /*initial_yield_time_ms*/ None,
            PendingRuntimeMode::PauseUntilResumed,
        )
        .await?;

        response_rx
            .await
            .map_err(|_| "exec runtime ended unexpectedly".to_string())
    }

    async fn start_session(
        &self,
        request: ExecuteRequest,
        initial_response_tx: SessionResponseSender,
        initial_yield_time_ms: Option<u64>,
        pending_mode: PendingRuntimeMode,
    ) -> Result<(), String> {
        let cell_id = request.cell_id.clone();
        let (event_tx, event_rx) = mpsc::unbounded_channel();
        let (control_tx, control_rx) = mpsc::unbounded_channel();
        let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = {
            let mut sessions = self.inner.sessions.lock().await;
            if sessions.contains_key(&cell_id) {
                return Err(format!("exec cell {cell_id} already exists"));
            }

            let (runtime_tx, runtime_control_tx, runtime_terminate_handle) =
                spawn_runtime(request, event_tx, pending_mode)?;

            // Keep the session registry locked through insertion so a
            // caller-owned cell id cannot race with another execute and replace
            // a live runtime.
            sessions.insert(
                cell_id.clone(),
                SessionHandle {
                    control_tx,
                    runtime_tx: runtime_tx.clone(),
                },
            );
            (runtime_tx, runtime_control_tx, runtime_terminate_handle)
        };

        tokio::spawn(run_session_control(
            Arc::clone(&self.inner),
            SessionControlContext {
                cell_id: cell_id.clone(),
                runtime_tx,
                runtime_control_tx,
                pending_mode,
                runtime_terminate_handle,
            },
            event_rx,
            control_rx,
            initial_response_tx,
            initial_yield_time_ms,
        ));

        Ok(())
    }

    pub async fn wait(&self, request: WaitRequest) -> Result<WaitOutcome, String> {
        let cell_id = request.cell_id.clone();
        let handle = self
            .inner
            .sessions
            .lock()
            .await
            .get(&request.cell_id)
            .cloned();
        let Some(handle) = handle else {
            return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id)));
        };
        let (response_tx, response_rx) = oneshot::channel();
        let control_message = if request.terminate {
            SessionControlCommand::Terminate { response_tx }
        } else {
            SessionControlCommand::Poll {
                yield_time_ms: request.yield_time_ms,
                response_tx,
            }
        };
        if handle.control_tx.send(control_message).is_err() {
            return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id)));
        }
        match response_rx.await {
            Ok(response) => Ok(WaitOutcome::LiveCell(response)),
            Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(
                request.cell_id,
            ))),
        }
    }

    pub async fn wait_to_pending(
        &self,
        request: WaitToPendingRequest,
    ) -> Result<WaitToPendingOutcome, String> {
        let cell_id = request.cell_id.clone();
        let handle = self
            .inner
            .sessions
            .lock()
            .await
            .get(&request.cell_id)
            .cloned();
        let Some(handle) = handle else {
            return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response(
                cell_id,
            )));
        };
        let (response_tx, response_rx) = oneshot::channel();
        if handle
            .control_tx
            .send(SessionControlCommand::PollToPending { response_tx })
            .is_err()
        {
            return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response(
                cell_id,
            )));
        }
        match response_rx.await {
            Ok(response) => Ok(WaitToPendingOutcome::LiveCell(response)),
            Err(_) => Ok(WaitToPendingOutcome::MissingCell(missing_cell_response(
                request.cell_id,
            ))),
        }
    }

    pub fn start_turn_worker(&self, host: Arc<dyn CodeModeTurnHost>) -> CodeModeTurnWorker {
        let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
        let inner = Arc::clone(&self.inner);
        let turn_message_rx = self.inner.turn_message_rx.clone();

        tokio::spawn(async move {
            loop {
                let next_message = tokio::select! {
                    _ = &mut shutdown_rx => break,
                    message = turn_message_rx.recv() => message.ok(),
                };
                let Some(next_message) = next_message else {
                    break;
                };
                match next_message {
                    TurnMessage::Notify {
                        cell_id,
                        call_id,
                        text,
                    } => {
                        if let Err(err) = host.notify(call_id, cell_id.clone(), text).await {
                            warn!(
                                "failed to deliver code mode notification for cell {cell_id}: {err}"
                            );
                        }
                    }
                    TurnMessage::ToolCall(invocation) => {
                        let host = Arc::clone(&host);
                        let inner = Arc::clone(&inner);
                        tokio::spawn(async move {
                            let cell_id = invocation.cell_id.clone();
                            let runtime_tool_call_id = invocation.runtime_tool_call_id.clone();
                            let response =
                                host.invoke_tool(invocation, CancellationToken::new()).await;
                            let runtime_tx = inner
                                .sessions
                                .lock()
                                .await
                                .get(&cell_id)
                                .map(|handle| handle.runtime_tx.clone());
                            let Some(runtime_tx) = runtime_tx else {
                                return;
                            };
                            let command = match response {
                                Ok(result) => RuntimeCommand::ToolResponse {
                                    id: runtime_tool_call_id,
                                    result,
                                },
                                Err(error_text) => RuntimeCommand::ToolError {
                                    id: runtime_tool_call_id,
                                    error_text,
                                },
                            };
                            let _ = runtime_tx.send(command);
                        });
                    }
                }
            }
        });

        CodeModeTurnWorker {
            shutdown_tx: Some(shutdown_tx),
        }
    }
}

impl Default for CodeModeService {
    fn default() -> Self {
        Self::new()
    }
}

pub struct CodeModeTurnWorker {
    shutdown_tx: Option<oneshot::Sender<()>>,
}

impl Drop for CodeModeTurnWorker {
    fn drop(&mut self) {
        if let Some(shutdown_tx) = self.shutdown_tx.take() {
            let _ = shutdown_tx.send(());
        }
    }
}

enum SessionControlCommand {
    Poll {
        yield_time_ms: u64,
        response_tx: oneshot::Sender<RuntimeResponse>,
    },
    PollToPending {
        response_tx: oneshot::Sender<ExecuteToPendingOutcome>,
    },
    Terminate {
        response_tx: oneshot::Sender<RuntimeResponse>,
    },
}

enum SessionResponseSender {
    Runtime(oneshot::Sender<RuntimeResponse>),
    ExecuteToPending(oneshot::Sender<ExecuteToPendingOutcome>),
}

struct PendingResult {
    content_items: Vec<FunctionCallOutputContentItem>,
    stored_values: HashMap<String, JsonValue>,
    error_text: Option<String>,
}

struct SessionControlContext {
    cell_id: String,
    runtime_tx: std::sync::mpsc::Sender<RuntimeCommand>,
    runtime_control_tx: std::sync::mpsc::Sender<RuntimeControlCommand>,
    pending_mode: PendingRuntimeMode,
    runtime_terminate_handle: v8::IsolateHandle,
}

fn missing_cell_response(cell_id: String) -> RuntimeResponse {
    RuntimeResponse::Result {
        error_text: Some(format!("exec cell {cell_id} not found")),
        cell_id,
        content_items: Vec::new(),
        stored_values: HashMap::new(),
    }
}

fn pending_result_response(cell_id: &str, result: PendingResult) -> RuntimeResponse {
    RuntimeResponse::Result {
        cell_id: cell_id.to_string(),
        content_items: result.content_items,
        stored_values: result.stored_values,
        error_text: result.error_text,
    }
}

fn send_terminal_response(response_tx: SessionResponseSender, response: RuntimeResponse) {
    match response_tx {
        SessionResponseSender::Runtime(response_tx) => {
            let _ = response_tx.send(response);
        }
        SessionResponseSender::ExecuteToPending(response_tx) => {
            let _ = response_tx.send(ExecuteToPendingOutcome::Completed(response));
        }
    }
}

fn send_or_buffer_result(
    cell_id: &str,
    result: PendingResult,
    response_tx: &mut Option<SessionResponseSender>,
    pending_result: &mut Option<PendingResult>,
) -> bool {
    if let Some(response_tx) = response_tx.take() {
        let response = pending_result_response(cell_id, result);
        send_terminal_response(response_tx, response);
        return true;
    }

    *pending_result = Some(result);
    false
}

fn send_yield_response(
    cell_id: &str,
    content_items: &mut Vec<FunctionCallOutputContentItem>,
    response_tx: &mut Option<SessionResponseSender>,
) {
    let Some(current_response_tx) = response_tx.take() else {
        return;
    };
    match current_response_tx {
        SessionResponseSender::Runtime(response_tx) => {
            let _ = response_tx.send(RuntimeResponse::Yielded {
                cell_id: cell_id.to_string(),
                content_items: std::mem::take(content_items),
            });
        }
        SessionResponseSender::ExecuteToPending(execute_to_pending_tx) => {
            *response_tx = Some(SessionResponseSender::ExecuteToPending(
                execute_to_pending_tx,
            ));
        }
    }
}

async fn run_session_control(
    inner: Arc<Inner>,
    context: SessionControlContext,
    mut event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
    mut control_rx: mpsc::UnboundedReceiver<SessionControlCommand>,
    initial_response_tx: SessionResponseSender,
    initial_yield_time_ms: Option<u64>,
) {
    let SessionControlContext {
        cell_id,
        runtime_tx,
        runtime_control_tx,
        pending_mode,
        runtime_terminate_handle,
    } = context;
    let mut content_items = Vec::new();
    let mut pending_tool_call_ids = Vec::new();
    let mut pending_result: Option<PendingResult> = None;
    let mut response_tx = Some(initial_response_tx);
    let mut termination_requested = false;
    let mut runtime_closed = false;
    let mut yield_timer: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;

    loop {
        tokio::select! {
            maybe_event = async {
                if runtime_closed {
                    std::future::pending::<Option<RuntimeEvent>>().await
                } else {
                    event_rx.recv().await
                }
            } => {
                let Some(event) = maybe_event else {
                    runtime_closed = true;
                    if termination_requested {
                        if let Some(response_tx) = response_tx.take() {
                            let response = RuntimeResponse::Terminated {
                                cell_id: cell_id.clone(),
                                content_items: std::mem::take(&mut content_items),
                            };
                            send_terminal_response(response_tx, response);
                        }
                        break;
                    }
                    if pending_result.is_none() {
                        let result = PendingResult {
                            content_items: std::mem::take(&mut content_items),
                            stored_values: HashMap::new(),
                            error_text: Some("exec runtime ended unexpectedly".to_string()),
                        };
                        if send_or_buffer_result(
                            &cell_id,
                            result,
                            &mut response_tx,
                            &mut pending_result,
                        ) {
                            break;
                        }
                    }
                    continue;
                };
                match event {
                    RuntimeEvent::Started => {
                        yield_timer = initial_yield_time_ms.map(|initial_yield_time_ms| {
                            Box::pin(tokio::time::sleep(Duration::from_millis(initial_yield_time_ms)))
                        });
                    }
                    RuntimeEvent::Pending => {
                        if let Some(current_response_tx) = response_tx.take() {
                            match current_response_tx {
                                SessionResponseSender::Runtime(runtime_response_tx) => {
                                    response_tx =
                                        Some(SessionResponseSender::Runtime(runtime_response_tx));
                                }
                                SessionResponseSender::ExecuteToPending(response_tx) => {
                                    let _ = response_tx.send(ExecuteToPendingOutcome::Pending {
                                        cell_id: cell_id.clone(),
                                        content_items: std::mem::take(&mut content_items),
                                        pending_tool_call_ids: std::mem::take(
                                            &mut pending_tool_call_ids,
                                        ),
                                    });
                                }
                            }
                        }
                    }
                    RuntimeEvent::ContentItem(item) => {
                        content_items.push(item);
                    }
                    RuntimeEvent::YieldRequested => {
                        yield_timer = None;
                        send_yield_response(&cell_id, &mut content_items, &mut response_tx);
                    }
                    RuntimeEvent::Notify { call_id, text } => {
                        let _ = inner.turn_message_tx.send(TurnMessage::Notify {
                            cell_id: cell_id.clone(),
                            call_id,
                            text,
                        }).await;
                    }
                    RuntimeEvent::ToolCall {
                        id,
                        name,
                        kind,
                        input,
                    } => {
                        if pending_mode == PendingRuntimeMode::PauseUntilResumed {
                            pending_tool_call_ids.push(id.clone());
                        }
                        let tool_call = CodeModeNestedToolCall {
                            cell_id: cell_id.clone(),
                            runtime_tool_call_id: id,
                            tool_name: name,
                            tool_kind: kind,
                            input,
                        };
                        let _ = inner
                            .turn_message_tx
                            .send(TurnMessage::ToolCall(tool_call))
                            .await;
                    }
                    RuntimeEvent::Result {
                        stored_values,
                        error_text,
                    } => {
                        yield_timer = None;
                        if termination_requested {
                            if let Some(response_tx) = response_tx.take() {
                                let response = RuntimeResponse::Terminated {
                                    cell_id: cell_id.clone(),
                                    content_items: std::mem::take(&mut content_items),
                                };
                                send_terminal_response(response_tx, response);
                            }
                            break;
                        }
                        let result = PendingResult {
                            content_items: std::mem::take(&mut content_items),
                            stored_values,
                            error_text,
                        };
                        if send_or_buffer_result(
                            &cell_id,
                            result,
                            &mut response_tx,
                            &mut pending_result,
                        ) {
                            break;
                        }
                    }
                }
            }
            maybe_command = control_rx.recv() => {
                let Some(command) = maybe_command else {
                    break;
                };
                match command {
                    SessionControlCommand::Poll {
                        yield_time_ms,
                        response_tx: next_response_tx,
                    } => {
                        if let Some(result) = pending_result.take() {
                            let _ = next_response_tx.send(pending_result_response(&cell_id, result));
                            break;
                        }
                        response_tx = Some(SessionResponseSender::Runtime(next_response_tx));
                        yield_timer = Some(Box::pin(tokio::time::sleep(Duration::from_millis(yield_time_ms))));
                        resume_paused_runtime(&runtime_control_tx, pending_mode);
                    }
                    SessionControlCommand::PollToPending {
                        response_tx: next_response_tx,
                    } => {
                        if let Some(result) = pending_result.take() {
                            let response = pending_result_response(&cell_id, result);
                            let _ = next_response_tx
                                .send(ExecuteToPendingOutcome::Completed(response));
                            break;
                        }
                        response_tx =
                            Some(SessionResponseSender::ExecuteToPending(next_response_tx));
                        yield_timer = None;
                        resume_paused_runtime(&runtime_control_tx, pending_mode);
                    }
                    SessionControlCommand::Terminate { response_tx: next_response_tx } => {
                        if let Some(result) = pending_result.take() {
                            let _ = next_response_tx.send(pending_result_response(&cell_id, result));
                            break;
                        }

                        response_tx = Some(SessionResponseSender::Runtime(next_response_tx));
                        termination_requested = true;
                        yield_timer = None;
                        let _ = runtime_tx.send(RuntimeCommand::Terminate);
                        terminate_paused_runtime(&runtime_control_tx, pending_mode);
                        let _ = runtime_terminate_handle.terminate_execution();
                        if runtime_closed {
                            if let Some(response_tx) = response_tx.take() {
                                let response = RuntimeResponse::Terminated {
                                    cell_id: cell_id.clone(),
                                    content_items: std::mem::take(&mut content_items),
                                };
                                send_terminal_response(response_tx, response);
                            }
                            break;
                        } else {
                            continue;
                        }
                    }
                }
            }
            _ = async {
                if let Some(yield_timer) = yield_timer.as_mut() {
                    yield_timer.await;
                } else {
                    std::future::pending::<()>().await;
                }
            } => {
                yield_timer = None;
                send_yield_response(&cell_id, &mut content_items, &mut response_tx);
            }
        }
    }

    let _ = runtime_tx.send(RuntimeCommand::Terminate);
    terminate_paused_runtime(&runtime_control_tx, pending_mode);
    inner.sessions.lock().await.remove(&cell_id);
}

fn resume_paused_runtime(
    runtime_control_tx: &std::sync::mpsc::Sender<RuntimeControlCommand>,
    pending_mode: PendingRuntimeMode,
) {
    if pending_mode == PendingRuntimeMode::PauseUntilResumed {
        let _ = runtime_control_tx.send(RuntimeControlCommand::Resume);
    }
}

fn terminate_paused_runtime(
    runtime_control_tx: &std::sync::mpsc::Sender<RuntimeControlCommand>,
    pending_mode: PendingRuntimeMode,
) {
    if pending_mode == PendingRuntimeMode::PauseUntilResumed {
        let _ = runtime_control_tx.send(RuntimeControlCommand::Terminate);
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::sync::Arc;
    use std::sync::atomic::AtomicU64;
    use std::time::Duration;

    use codex_protocol::ToolName;
    use pretty_assertions::assert_eq;
    use tokio::sync::Mutex;
    use tokio::sync::mpsc;
    use tokio::sync::oneshot;

    use super::CodeModeService;
    use super::Inner;
    use super::PendingRuntimeMode;
    use super::RuntimeCommand;
    use super::RuntimeResponse;
    use super::SessionControlCommand;
    use super::SessionControlContext;
    use super::SessionResponseSender;
    use super::WaitOutcome;
    use super::WaitRequest;
    use super::WaitToPendingOutcome;
    use super::WaitToPendingRequest;
    use super::run_session_control;
    use crate::CodeModeToolKind;
    use crate::FunctionCallOutputContentItem;
    use crate::ToolDefinition;
    use crate::runtime::ExecuteRequest;
    use crate::runtime::ExecuteToPendingOutcome;
    use crate::runtime::RuntimeEvent;
    use crate::runtime::spawn_runtime;

    fn execute_request(source: &str) -> ExecuteRequest {
        ExecuteRequest {
            cell_id: "1".to_string(),
            tool_call_id: "call_1".to_string(),
            enabled_tools: Vec::new(),
            source: source.to_string(),
            stored_values: HashMap::new(),
            yield_time_ms: Some(1),
            max_output_tokens: None,
        }
    }

    fn test_inner() -> Arc<Inner> {
        let (turn_message_tx, turn_message_rx) = async_channel::unbounded();
        Arc::new(Inner {
            stored_values: Mutex::new(HashMap::new()),
            sessions: Mutex::new(HashMap::new()),
            turn_message_tx,
            turn_message_rx,
            next_cell_id: AtomicU64::new(1),
        })
    }

    #[tokio::test]
    async fn synchronous_exit_returns_successfully() {
        let service = CodeModeService::new();

        let response = service
            .execute(ExecuteRequest {
                source: r#"text("before"); exit(); text("after");"#.to_string(),
                yield_time_ms: None,
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            RuntimeResponse::Result {
                cell_id: "1".to_string(),
                content_items: vec![FunctionCallOutputContentItem::InputText {
                    text: "before".to_string(),
                }],
                stored_values: HashMap::new(),
                error_text: None,
            }
        );
    }

    #[tokio::test]
    async fn execute_to_pending_returns_completed_for_synchronous_results() {
        let service = CodeModeService::new();

        let response = service
            .execute_to_pending(ExecuteRequest {
                source: r#"text("done");"#.to_string(),
                yield_time_ms: Some(60_000),
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            ExecuteToPendingOutcome::Completed(RuntimeResponse::Result {
                cell_id: "1".to_string(),
                content_items: vec![FunctionCallOutputContentItem::InputText {
                    text: "done".to_string(),
                }],
                stored_values: HashMap::new(),
                error_text: None,
            })
        );
    }

    #[tokio::test]
    async fn execute_to_pending_returns_once_the_runtime_is_quiescent() {
        let service = CodeModeService::new();

        let response = tokio::time::timeout(
            Duration::from_secs(1),
            service.execute_to_pending(ExecuteRequest {
                source: r#"text("before"); await new Promise(() => {});"#.to_string(),
                yield_time_ms: Some(60_000),
                ..execute_request("")
            }),
        )
        .await
        .unwrap()
        .unwrap();

        assert_eq!(
            response,
            ExecuteToPendingOutcome::Pending {
                cell_id: "1".to_string(),
                content_items: vec![FunctionCallOutputContentItem::InputText {
                    text: "before".to_string(),
                }],
                pending_tool_call_ids: Vec::new(),
            }
        );

        let termination = service
            .wait(WaitRequest {
                cell_id: "1".to_string(),
                yield_time_ms: 1,
                terminate: true,
            })
            .await
            .unwrap();

        assert_eq!(
            termination,
            WaitOutcome::LiveCell(RuntimeResponse::Terminated {
                cell_id: "1".to_string(),
                content_items: Vec::new(),
            })
        );
    }

    #[tokio::test]
    async fn execute_to_pending_identifies_tool_calls_in_paused_frontier() {
        let service = CodeModeService::new();

        let response = service
            .execute_to_pending(ExecuteRequest {
                enabled_tools: vec![ToolDefinition {
                    name: "echo".to_string(),
                    tool_name: ToolName::plain("echo"),
                    description: String::new(),
                    kind: CodeModeToolKind::Function,
                    input_schema: None,
                    output_schema: None,
                }],
                source: r#"
await Promise.all([
  tools.echo({ value: "first" }),
  tools.echo({ value: "second" }),
]);
"#
                .to_string(),
                yield_time_ms: Some(60_000),
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            ExecuteToPendingOutcome::Pending {
                cell_id: "1".to_string(),
                content_items: Vec::new(),
                pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()],
            }
        );

        let termination = service
            .wait(WaitRequest {
                cell_id: "1".to_string(),
                yield_time_ms: 1,
                terminate: true,
            })
            .await
            .unwrap();

        assert_eq!(
            termination,
            WaitOutcome::LiveCell(RuntimeResponse::Terminated {
                cell_id: "1".to_string(),
                content_items: Vec::new(),
            })
        );
    }

    #[tokio::test]
    async fn execute_to_pending_excludes_delayed_timeout_tool_calls_until_wait() {
        let service = CodeModeService::new();

        let initial_response = service
            .execute_to_pending(ExecuteRequest {
                enabled_tools: vec![ToolDefinition {
                    name: "echo".to_string(),
                    tool_name: ToolName::plain("echo"),
                    description: String::new(),
                    kind: CodeModeToolKind::Function,
                    input_schema: None,
                    output_schema: None,
                }],
                source: r#"
setTimeout(() => {
  tools.echo({ value: "delayed" });
}, 1000);
await Promise.all([
  tools.echo({ value: "second" }),
  tools.echo({ value: "third" }),
]);
"#
                .to_string(),
                yield_time_ms: Some(60_000),
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            initial_response,
            ExecuteToPendingOutcome::Pending {
                cell_id: "1".to_string(),
                content_items: Vec::new(),
                pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()],
            }
        );

        let runtime_tx = service
            .inner
            .sessions
            .lock()
            .await
            .get("1")
            .unwrap()
            .runtime_tx
            .clone();
        runtime_tx
            .send(RuntimeCommand::TimeoutFired { id: 1 })
            .unwrap();

        let resumed_response = tokio::time::timeout(
            Duration::from_secs(1),
            service.wait_to_pending(WaitToPendingRequest {
                cell_id: "1".to_string(),
            }),
        )
        .await
        .unwrap()
        .unwrap();

        assert_eq!(
            resumed_response,
            WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending {
                cell_id: "1".to_string(),
                content_items: Vec::new(),
                pending_tool_call_ids: vec!["tool-3".to_string()],
            })
        );

        let termination = service
            .wait(WaitRequest {
                cell_id: "1".to_string(),
                yield_time_ms: 1,
                terminate: true,
            })
            .await
            .unwrap();

        assert_eq!(
            termination,
            WaitOutcome::LiveCell(RuntimeResponse::Terminated {
                cell_id: "1".to_string(),
                content_items: Vec::new(),
            })
        );
    }

    #[tokio::test]
    async fn wait_to_pending_returns_after_resumed_runtime_becomes_quiescent_again() {
        let service = CodeModeService::new();

        let initial_response = service
            .execute_to_pending(ExecuteRequest {
                source: r#"
await new Promise((resolve) => setTimeout(resolve, 60_000));
text("after");
await new Promise(() => {});
"#
                .to_string(),
                yield_time_ms: Some(60_000),
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            initial_response,
            ExecuteToPendingOutcome::Pending {
                cell_id: "1".to_string(),
                content_items: Vec::new(),
                pending_tool_call_ids: Vec::new(),
            }
        );

        let runtime_tx = service
            .inner
            .sessions
            .lock()
            .await
            .get("1")
            .unwrap()
            .runtime_tx
            .clone();
        runtime_tx
            .send(RuntimeCommand::TimeoutFired { id: 1 })
            .unwrap();

        let resumed_response = tokio::time::timeout(
            Duration::from_secs(1),
            service.wait_to_pending(WaitToPendingRequest {
                cell_id: "1".to_string(),
            }),
        )
        .await
        .unwrap()
        .unwrap();

        assert_eq!(
            resumed_response,
            WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending {
                cell_id: "1".to_string(),
                content_items: vec![FunctionCallOutputContentItem::InputText {
                    text: "after".to_string(),
                }],
                pending_tool_call_ids: Vec::new(),
            })
        );

        let termination = service
            .wait(WaitRequest {
                cell_id: "1".to_string(),
                yield_time_ms: 1,
                terminate: true,
            })
            .await
            .unwrap();

        assert_eq!(
            termination,
            WaitOutcome::LiveCell(RuntimeResponse::Terminated {
                cell_id: "1".to_string(),
                content_items: Vec::new(),
            })
        );
    }

    #[tokio::test]
    async fn wait_to_pending_returns_completed_after_resumed_runtime_finishes() {
        let service = CodeModeService::new();

        let initial_response = service
            .execute_to_pending(ExecuteRequest {
                source: r#"
await new Promise((resolve) => setTimeout(resolve, 60_000));
text("done");
"#
                .to_string(),
                yield_time_ms: Some(60_000),
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            initial_response,
            ExecuteToPendingOutcome::Pending {
                cell_id: "1".to_string(),
                content_items: Vec::new(),
                pending_tool_call_ids: Vec::new(),
            }
        );

        let runtime_tx = service
            .inner
            .sessions
            .lock()
            .await
            .get("1")
            .unwrap()
            .runtime_tx
            .clone();
        runtime_tx
            .send(RuntimeCommand::TimeoutFired { id: 1 })
            .unwrap();

        let resumed_response = tokio::time::timeout(
            Duration::from_secs(1),
            service.wait_to_pending(WaitToPendingRequest {
                cell_id: "1".to_string(),
            }),
        )
        .await
        .unwrap()
        .unwrap();

        assert_eq!(
            resumed_response,
            WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Completed(
                RuntimeResponse::Result {
                    cell_id: "1".to_string(),
                    content_items: vec![FunctionCallOutputContentItem::InputText {
                        text: "done".to_string(),
                    }],
                    stored_values: HashMap::new(),
                    error_text: None,
                }
            ))
        );
    }

    #[tokio::test]
    async fn v8_console_is_not_exposed_on_global_this() {
        let service = CodeModeService::new();

        let response = service
            .execute(ExecuteRequest {
                source: r#"text(String(Object.hasOwn(globalThis, "console")));"#.to_string(),
                yield_time_ms: None,
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            RuntimeResponse::Result {
                cell_id: "1".to_string(),
                content_items: vec![FunctionCallOutputContentItem::InputText {
                    text: "false".to_string(),
                }],
                stored_values: HashMap::new(),
                error_text: None,
            }
        );
    }

    #[tokio::test]
    async fn date_locale_string_formats_with_icu_data() {
        let service = CodeModeService::new();

        let response = service
            .execute(ExecuteRequest {
                source: r#"
const value = new Date("2025-01-02T03:04:05Z")
  .toLocaleString("fr-FR", {
    weekday: "long",
    month: "long",
    day: "numeric",
    hour: "2-digit",
    minute: "2-digit",
    second: "2-digit",
    hour12: false,
    timeZone: "UTC",
  });
text(value);
"#
                .to_string(),
                yield_time_ms: None,
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            RuntimeResponse::Result {
                cell_id: "1".to_string(),
                content_items: vec![FunctionCallOutputContentItem::InputText {
                    text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(),
                }],
                stored_values: HashMap::new(),
                error_text: None,
            }
        );
    }

    #[tokio::test]
    async fn intl_date_time_format_formats_with_icu_data() {
        let service = CodeModeService::new();

        let response = service
            .execute(ExecuteRequest {
                source: r#"
const formatter = new Intl.DateTimeFormat("fr-FR", {
  weekday: "long",
  month: "long",
  day: "numeric",
  hour: "2-digit",
  minute: "2-digit",
  second: "2-digit",
  hour12: false,
  timeZone: "UTC",
});
text(formatter.format(new Date("2025-01-02T03:04:05Z")));
"#
                .to_string(),
                yield_time_ms: None,
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            RuntimeResponse::Result {
                cell_id: "1".to_string(),
                content_items: vec![FunctionCallOutputContentItem::InputText {
                    text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(),
                }],
                stored_values: HashMap::new(),
                error_text: None,
            }
        );
    }

    #[tokio::test]
    async fn output_helpers_return_undefined() {
        let service = CodeModeService::new();

        let response = service
            .execute(ExecuteRequest {
                source: r#"
const returnsUndefined = [
  text("first"),
  image("https://example.com/image.jpg"),
  notify("ping"),
].map((value) => value === undefined);
text(JSON.stringify(returnsUndefined));
"#
                .to_string(),
                yield_time_ms: None,
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            RuntimeResponse::Result {
                cell_id: "1".to_string(),
                content_items: vec![
                    FunctionCallOutputContentItem::InputText {
                        text: "first".to_string(),
                    },
                    FunctionCallOutputContentItem::InputImage {
                        image_url: "https://example.com/image.jpg".to_string(),
                        detail: Some(crate::DEFAULT_IMAGE_DETAIL),
                    },
                    FunctionCallOutputContentItem::InputText {
                        text: "[true,true,true]".to_string(),
                    },
                ],
                stored_values: HashMap::new(),
                error_text: None,
            }
        );
    }

    #[tokio::test]
    async fn image_helper_accepts_raw_mcp_image_block_with_original_detail() {
        let service = CodeModeService::new();

        let response = service
            .execute(ExecuteRequest {
                source: r#"
image({
  type: "image",
  data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==",
  mimeType: "image/png",
  _meta: { "codex/imageDetail": "original" },
});
"#
                .to_string(),
                yield_time_ms: None,
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            RuntimeResponse::Result {
                cell_id: "1".to_string(),
                content_items: vec![FunctionCallOutputContentItem::InputImage {
                    image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(),
                    detail: Some(crate::ImageDetail::Original),
                }],
                stored_values: HashMap::new(),
                error_text: None,
            }
        );
    }

    #[tokio::test]
    async fn image_helper_second_arg_overrides_explicit_object_detail() {
        let service = CodeModeService::new();

        let response = service
            .execute(ExecuteRequest {
                source: r#"
image(
  {
    image_url: "https://example.com/image.jpg",
    detail: "low",
  },
  "original",
);
"#
                .to_string(),
                yield_time_ms: None,
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            RuntimeResponse::Result {
                cell_id: "1".to_string(),
                content_items: vec![FunctionCallOutputContentItem::InputImage {
                    image_url: "https://example.com/image.jpg".to_string(),
                    detail: Some(crate::ImageDetail::Original),
                }],
                stored_values: HashMap::new(),
                error_text: None,
            }
        );
    }

    #[tokio::test]
    async fn image_helper_second_arg_overrides_raw_mcp_image_detail() {
        let service = CodeModeService::new();

        let response = service
            .execute(ExecuteRequest {
                source: r#"
image(
  {
    type: "image",
    data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==",
    mimeType: "image/png",
    _meta: { "codex/imageDetail": "original" },
  },
  "low",
);
"#
                .to_string(),
                yield_time_ms: None,
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            RuntimeResponse::Result {
                cell_id: "1".to_string(),
                content_items: vec![FunctionCallOutputContentItem::InputImage {
                    image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(),
                    detail: Some(crate::ImageDetail::Low),
                }],
                stored_values: HashMap::new(),
                error_text: None,
            }
        );
    }

    #[tokio::test]
    async fn image_helper_rejects_raw_mcp_result_container() {
        let service = CodeModeService::new();

        let response = service
            .execute(ExecuteRequest {
                source: r#"
image({
  content: [
    {
      type: "image",
      data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==",
      mimeType: "image/png",
      _meta: { "codex/imageDetail": "original" },
    },
  ],
  isError: false,
});
"#
                .to_string(),
                yield_time_ms: None,
                ..execute_request("")
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            RuntimeResponse::Result {
                cell_id: "1".to_string(),
                content_items: Vec::new(),
                stored_values: HashMap::new(),
                error_text: Some(
                    "image expects a non-empty image URL string, an object with image_url and optional detail, or a raw MCP image block".to_string(),
                ),
            }
        );
    }

    #[tokio::test]
    async fn wait_reports_missing_cell_separately_from_runtime_results() {
        let service = CodeModeService::new();

        let response = service
            .wait(WaitRequest {
                cell_id: "missing".to_string(),
                yield_time_ms: 1,
                terminate: false,
            })
            .await
            .unwrap();

        assert_eq!(
            response,
            WaitOutcome::MissingCell(RuntimeResponse::Result {
                cell_id: "missing".to_string(),
                content_items: Vec::new(),
                stored_values: HashMap::new(),
                error_text: Some("exec cell missing not found".to_string()),
            })
        );
    }

    #[tokio::test]
    async fn terminate_waits_for_runtime_shutdown_before_responding() {
        let inner = test_inner();
        let (event_tx, event_rx) = mpsc::unbounded_channel();
        let (control_tx, control_rx) = mpsc::unbounded_channel();
        let (initial_response_tx, initial_response_rx) = oneshot::channel();
        let (runtime_event_tx, _runtime_event_rx) = mpsc::unbounded_channel();
        let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = spawn_runtime(
            ExecuteRequest {
                source: "await new Promise(() => {})".to_string(),
                yield_time_ms: None,
                ..execute_request("")
            },
            runtime_event_tx,
            PendingRuntimeMode::Continue,
        )
        .unwrap();

        tokio::spawn(run_session_control(
            inner,
            SessionControlContext {
                cell_id: "cell-1".to_string(),
                runtime_tx: runtime_tx.clone(),
                runtime_control_tx,
                pending_mode: PendingRuntimeMode::Continue,
                runtime_terminate_handle,
            },
            event_rx,
            control_rx,
            SessionResponseSender::Runtime(initial_response_tx),
            Some(/*initial_yield_time_ms*/ 60_000),
        ));

        event_tx.send(RuntimeEvent::Started).unwrap();
        event_tx.send(RuntimeEvent::YieldRequested).unwrap();
        assert_eq!(
            initial_response_rx.await.unwrap(),
            RuntimeResponse::Yielded {
                cell_id: "cell-1".to_string(),
                content_items: Vec::new(),
            }
        );

        let (terminate_response_tx, terminate_response_rx) = oneshot::channel();
        control_tx
            .send(SessionControlCommand::Terminate {
                response_tx: terminate_response_tx,
            })
            .unwrap();
        let terminate_response = async { terminate_response_rx.await.unwrap() };
        tokio::pin!(terminate_response);
        assert!(
            tokio::time::timeout(Duration::from_millis(100), terminate_response.as_mut())
                .await
                .is_err()
        );

        drop(event_tx);

        assert_eq!(
            terminate_response.await,
            RuntimeResponse::Terminated {
                cell_id: "cell-1".to_string(),
                content_items: Vec::new(),
            }
        );

        let _ = runtime_tx.send(RuntimeCommand::Terminate);
    }
}