use std::collections::HashMap;use std::sync::Arc;use std::sync::atomic::AtomicBool;use std::sync::atomic::AtomicU64;use std::sync::atomic::Ordering;use std::time::Duration;use codex_protocol::ToolName;use pretty_assertions::assert_eq;use tokio::sync::Notify;use tokio::sync::mpsc;use tokio::sync::oneshot;use tokio_util::sync::CancellationToken;use super::*;use crate::CodeModeToolKind;use crate::ToolDefinition;#[derive(Debug, PartialEq)]enum DelegateEvent {    NotificationStarted,    NotificationCancelled,    ToolStarted,    ToolCancelled,    CellClosed(CellId),}struct BlockingDelegate {    events_tx: mpsc::UnboundedSender<DelegateEvent>,    notification_finished: AtomicBool,    tool_finished: AtomicBool,}struct HeldNotificationDelegate {    events_tx: mpsc::UnboundedSender<DelegateEvent>,    notification_release: Notify,}impl HeldNotificationDelegate {    fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {        let (events_tx, events_rx) = mpsc::unbounded_channel();        (            Arc::new(Self {                events_tx,                notification_release: Notify::new(),            }),            events_rx,        )    }    fn release_notification(&self) {        self.notification_release.notify_one();    }}impl CodeModeSessionDelegate for HeldNotificationDelegate {    fn invoke_tool<'a>(        &'a self,        _invocation: CodeModeNestedToolCall,        cancellation_token: CancellationToken,    ) -> ToolInvocationFuture<'a> {        Box::pin(async move {            cancellation_token.cancelled().await;            Err("cancelled".to_string())        })    }    fn notify<'a>(        &'a self,        _call_id: String,        _cell_id: CellId,        _text: String,        cancellation_token: CancellationToken,    ) -> NotificationFuture<'a> {        Box::pin(async move {            let _ = self.events_tx.send(DelegateEvent::NotificationStarted);            cancellation_token.cancelled().await;            let _ = self.events_tx.send(DelegateEvent::NotificationCancelled);            self.notification_release.notified().await;            Ok(())        })    }    fn cell_closed(&self, cell_id: &CellId) {        let _ = self            .events_tx            .send(DelegateEvent::CellClosed(cell_id.clone()));    }}struct CellControlHarness {    event_tx: mpsc::UnboundedSender<RuntimeEvent>,    control_tx: mpsc::UnboundedSender<CellControlCommand>,    initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,    task: tokio::task::JoinHandle<()>,    _runtime_event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,}fn spawn_cell_control_harness(    initial_yield_time_ms: Option<u64>,    delegate: Arc<dyn CodeModeSessionDelegate>,) -> CellControlHarness {    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(        HashMap::new(),        execute_request("await new Promise(() => {});"),        runtime_event_tx,        PendingRuntimeMode::Continue,    )    .unwrap();    let inner = Arc::new(Inner {        stored_values: Mutex::new(HashMap::new()),        cells: Mutex::new(HashMap::new()),        delegate,        shutting_down: AtomicBool::new(false),        next_cell_id: AtomicU64::new(1),    });    let task = tokio::spawn(run_cell_control(        inner,        CellControlContext {            cell_id: cell_id("1"),            runtime_tx,            runtime_control_tx,            pending_mode: PendingRuntimeMode::Continue,            runtime_terminate_handle,            cancellation_token: CancellationToken::new(),        },        event_rx,        control_rx,        CellResponseSender::Runtime(initial_response_tx),        initial_yield_time_ms,    ));    CellControlHarness {        event_tx,        control_tx,        initial_response_rx,        task,        _runtime_event_rx: runtime_event_rx,    }}impl BlockingDelegate {    fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {        let (events_tx, events_rx) = mpsc::unbounded_channel();        (            Arc::new(Self {                events_tx,                notification_finished: AtomicBool::new(false),                tool_finished: AtomicBool::new(false),            }),            events_rx,        )    }}impl CodeModeSessionDelegate for BlockingDelegate {    fn invoke_tool<'a>(        &'a self,        _invocation: CodeModeNestedToolCall,        cancellation_token: CancellationToken,    ) -> ToolInvocationFuture<'a> {        Box::pin(async move {            let _ = self.events_tx.send(DelegateEvent::ToolStarted);            cancellation_token.cancelled().await;            self.tool_finished.store(true, Ordering::Release);            let _ = self.events_tx.send(DelegateEvent::ToolCancelled);            Err("cancelled".to_string())        })    }    fn notify<'a>(        &'a self,        _call_id: String,        _cell_id: CellId,        _text: String,        cancellation_token: CancellationToken,    ) -> NotificationFuture<'a> {        Box::pin(async move {            let _ = self.events_tx.send(DelegateEvent::NotificationStarted);            cancellation_token.cancelled().await;            self.notification_finished.store(true, Ordering::Release);            let _ = self.events_tx.send(DelegateEvent::NotificationCancelled);            Err("cancelled".to_string())        })    }    fn cell_closed(&self, cell_id: &CellId) {        let _ = self            .events_tx            .send(DelegateEvent::CellClosed(cell_id.clone()));    }}fn cell_id(value: &str) -> CellId {    CellId::new(value.to_string())}fn execute_request(source: &str) -> ExecuteRequest {    ExecuteRequest {        tool_call_id: "call-1".to_string(),        enabled_tools: Vec::new(),        source: source.to_string(),        yield_time_ms: Some(1),        max_output_tokens: None,    }}fn blocking_tool() -> ToolDefinition {    ToolDefinition {        name: "block".to_string(),        tool_name: ToolName::plain("block"),        description: String::new(),        kind: CodeModeToolKind::Function,        input_schema: None,        output_schema: None,    }}async fn next_event(events_rx: &mut mpsc::UnboundedReceiver<DelegateEvent>) -> DelegateEvent {    tokio::time::timeout(Duration::from_secs(2), events_rx.recv())        .await        .expect("delegate event timeout")        .expect("delegate event channel closed")}#[tokio::test]async fn yield_timer_preempts_buffered_runtime_output() {    let harness = spawn_cell_control_harness(        Some(/*initial_yield_time_ms*/ 0),        Arc::new(NoopCodeModeSessionDelegate),    );    harness.event_tx.send(RuntimeEvent::Started).unwrap();    harness        .event_tx        .send(RuntimeEvent::ContentItem(            FunctionCallOutputContentItem::InputText {                text: "queued output".to_string(),            },        ))        .unwrap();    assert_eq!(        harness.initial_response_rx.await.unwrap(),        Ok(RuntimeResponse::Yielded {            cell_id: cell_id("1"),            content_items: Vec::new(),        })    );    let (termination_tx, termination_rx) = oneshot::channel();    harness        .control_tx        .send(CellControlCommand::Terminate {            response_tx: termination_tx,        })        .unwrap();    drop(harness.event_tx);    assert_eq!(        termination_rx.await.unwrap(),        Ok(RuntimeResponse::Terminated {            cell_id: cell_id("1"),            content_items: vec![FunctionCallOutputContentItem::InputText {                text: "queued output".to_string(),            }],        })    );    harness.task.await.unwrap();}#[tokio::test]async fn queued_termination_preempts_unobserved_runtime_completion() {    let harness = spawn_cell_control_harness(        Some(/*initial_yield_time_ms*/ 60_000),        Arc::new(NoopCodeModeSessionDelegate),    );    harness        .event_tx        .send(RuntimeEvent::Result {            stored_value_writes: HashMap::new(),            error_text: None,        })        .unwrap();    let (termination_tx, termination_rx) = oneshot::channel();    harness        .control_tx        .send(CellControlCommand::Terminate {            response_tx: termination_tx,        })        .unwrap();    let terminated = Ok(RuntimeResponse::Terminated {        cell_id: cell_id("1"),        content_items: Vec::new(),    });    assert_eq!(termination_rx.await.unwrap(), terminated.clone());    assert_eq!(harness.initial_response_rx.await.unwrap(), terminated);    harness.task.await.unwrap();}#[tokio::test]async fn yields_and_resumes() {    let service = CodeModeService::new();    let cell = service        .execute(execute_request(            r#"text("before"); yield_control(); text("after");"#,        ))        .await        .unwrap();    assert_eq!(        cell.initial_response().await.unwrap(),        RuntimeResponse::Yielded {            cell_id: cell_id("1"),            content_items: vec![FunctionCallOutputContentItem::InputText {                text: "before".to_string(),            }],        }    );    assert_eq!(        service            .wait(WaitRequest {                cell_id: cell_id("1"),                yield_time_ms: 1,            })            .await            .unwrap(),        WaitOutcome::LiveCell(RuntimeResponse::Result {            cell_id: cell_id("1"),            content_items: vec![FunctionCallOutputContentItem::InputText {                text: "after".to_string(),            }],            error_text: None,        })    );}#[tokio::test]async fn returns_and_resumes_from_the_pending_frontier() {    let service = CodeModeService::new();    assert_eq!(        service            .execute_to_pending(execute_request(                r#"await new Promise((resolve) => setTimeout(resolve, 60_000));text("after");"#,            ))            .await            .unwrap(),        ExecuteToPendingOutcome::Pending {            cell_id: cell_id("1"),            content_items: Vec::new(),            pending_tool_call_ids: Vec::new(),        }    );    service        .inner        .cells        .lock()        .await        .get(&cell_id("1"))        .unwrap()        .runtime_tx        .send(RuntimeCommand::TimeoutFired { id: 1 })        .unwrap();    assert_eq!(        service            .wait_to_pending(WaitToPendingRequest {                cell_id: cell_id("1"),            })            .await            .unwrap(),        WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Completed(            RuntimeResponse::Result {                cell_id: cell_id("1"),                content_items: vec![FunctionCallOutputContentItem::InputText {                    text: "after".to_string(),                }],                error_text: None,            }        ))    );}#[tokio::test]async fn observed_natural_completion_wins_over_termination() {    let (delegate, mut events_rx) = BlockingDelegate::new();    let harness =        spawn_cell_control_harness(Some(/*initial_yield_time_ms*/ 60_000), delegate.clone());    harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();    assert_eq!(        harness.initial_response_rx.await.unwrap(),        Ok(RuntimeResponse::Yielded {            cell_id: cell_id("1"),            content_items: Vec::new(),        })    );    harness        .event_tx        .send(RuntimeEvent::ContentItem(            FunctionCallOutputContentItem::InputText {                text: "done".to_string(),            },        ))        .unwrap();    harness        .event_tx        .send(RuntimeEvent::Result {            stored_value_writes: HashMap::new(),            error_text: None,        })        .unwrap();    harness        .event_tx        .send(RuntimeEvent::Notify {            call_id: "notify-1".to_string(),            text: "completion observed".to_string(),        })        .unwrap();    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::NotificationStarted    );    let (termination_tx, termination_rx) = oneshot::channel();    harness        .control_tx        .send(CellControlCommand::Terminate {            response_tx: termination_tx,        })        .unwrap();    assert_eq!(        termination_rx.await.unwrap(),        Ok(RuntimeResponse::Result {            cell_id: cell_id("1"),            content_items: vec![FunctionCallOutputContentItem::InputText {                text: "done".to_string(),            }],            error_text: None,        })    );    harness.task.await.unwrap();    assert!(delegate.notification_finished.load(Ordering::Acquire));    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::NotificationCancelled    );    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::CellClosed(cell_id("1"))    );}#[tokio::test]async fn termination_cancels_pending_callbacks_before_responding() {    let (delegate, mut events_rx) = BlockingDelegate::new();    let service = CodeModeService::with_delegate(delegate.clone());    let cell = service        .execute(execute_request(            r#"notify("pending"); await new Promise(() => {});"#,        ))        .await        .unwrap();    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::NotificationStarted    );    assert_eq!(        cell.initial_response().await.unwrap(),        RuntimeResponse::Yielded {            cell_id: cell_id("1"),            content_items: Vec::new(),        }    );    assert_eq!(        service.terminate(cell_id("1")).await.unwrap(),        WaitOutcome::LiveCell(RuntimeResponse::Terminated {            cell_id: cell_id("1"),            content_items: Vec::new(),        })    );    assert!(delegate.notification_finished.load(Ordering::Acquire));    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::NotificationCancelled    );    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::CellClosed(cell_id("1"))    );}#[tokio::test]async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {    let (delegate, mut events_rx) = HeldNotificationDelegate::new();    let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));    let cell = service        .execute(execute_request(            r#"notify("pending"); await new Promise(() => {});"#,        ))        .await        .unwrap();    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::NotificationStarted    );    assert_eq!(        cell.initial_response().await.unwrap(),        RuntimeResponse::Yielded {            cell_id: cell_id("1"),            content_items: Vec::new(),        }    );    let terminating_service = Arc::clone(&service);    let first_termination =        tokio::spawn(async move { terminating_service.terminate(cell_id("1")).await });    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::NotificationCancelled    );    let repeated_termination = service.terminate(cell_id("1")).await;    delegate.release_notification();    assert_eq!(        repeated_termination.unwrap_err(),        "exec cell 1 is already terminating"    );    assert_eq!(        first_termination.await.unwrap().unwrap(),        WaitOutcome::LiveCell(RuntimeResponse::Terminated {            cell_id: cell_id("1"),            content_items: Vec::new(),        })    );    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::CellClosed(cell_id("1"))    );}#[tokio::test]async fn second_observer_is_rejected_without_displacing_the_first() {    let service = CodeModeService::new();    let cell = service        .execute(execute_request("await new Promise(() => {});"))        .await        .unwrap();    assert_eq!(        cell.initial_response().await.unwrap(),        RuntimeResponse::Yielded {            cell_id: cell_id("1"),            content_items: Vec::new(),        }    );    let first_observer = service        .begin_wait(WaitRequest {            cell_id: cell_id("1"),            yield_time_ms: 60_000,        })        .await;    assert_eq!(        service            .wait(WaitRequest {                cell_id: cell_id("1"),                yield_time_ms: 60_000,            })            .await            .unwrap_err(),        "exec cell 1 already has an active observer"    );    let terminated = RuntimeResponse::Terminated {        cell_id: cell_id("1"),        content_items: Vec::new(),    };    assert_eq!(        service.terminate(cell_id("1")).await.unwrap(),        WaitOutcome::LiveCell(terminated.clone())    );    assert_eq!(        first_observer.await.unwrap(),        WaitOutcome::LiveCell(terminated)    );}#[tokio::test]async fn natural_completion_cleans_up_callbacks_before_responding() {    let (delegate, mut events_rx) = BlockingDelegate::new();    let service = CodeModeService::with_delegate(delegate.clone());    let cell = service        .execute(ExecuteRequest {            enabled_tools: vec![blocking_tool()],            source: r#"tools.block({}); text("done");"#.to_string(),            yield_time_ms: Some(60_000),            ..execute_request("")        })        .await        .unwrap();    assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);    assert_eq!(        cell.initial_response().await.unwrap(),        RuntimeResponse::Result {            cell_id: cell_id("1"),            content_items: vec![FunctionCallOutputContentItem::InputText {                text: "done".to_string(),            }],            error_text: None,        }    );    assert!(delegate.tool_finished.load(Ordering::Acquire));    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::ToolCancelled    );    assert_eq!(        next_event(&mut events_rx).await,        DelegateEvent::CellClosed(cell_id("1"))    );}