nanocodex-tools 0.2.0

use std::{
    path::PathBuf,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::Duration,
};

use eyre::{Result, eyre};
use nanocodex_core::{ResponseItem, ToolDefinition};
use serde_json::Value;
use tokio::sync::Semaphore;

use super::{
    CellUpdate, CodeModeExecution, CodeModeObserver, CodeModeUpdate, LiveCell, NestedToolCall,
    ObservationMode, nested_tool_yield_after, observe_cell, observer_yield_timeout,
    parse_exec_source,
};
use crate::{
    Tool, ToolContext, ToolExecution, ToolInput, ToolOutputBody, ToolOutputContent, ToolResult,
    ToolRuntime, Tools, WebSearchConfig,
};

struct ConcurrencyProbe {
    state: Arc<ConcurrencyProbeState>,
}

struct ConcurrencyProbeState {
    active: AtomicUsize,
    maximum: AtomicUsize,
    release: Semaphore,
}

#[async_trait::async_trait]
impl Tool for ConcurrencyProbe {
    fn name(&self) -> &'static str {
        "concurrency_probe"
    }

    fn definition(&self) -> ToolDefinition {
        ToolDefinition::function(
            self.name(),
            "Waits until released by the test.",
            serde_json::json!({
                "type": "object",
                "properties": {},
                "additionalProperties": false
            }),
        )
    }

    async fn execute(&self, _input: ToolInput, _context: ToolContext<'_>) -> ToolResult {
        let active = self.state.active.fetch_add(1, Ordering::SeqCst) + 1;
        self.state.maximum.fetch_max(active, Ordering::SeqCst);
        let permit = self
            .state
            .release
            .acquire()
            .await
            .map_err(|error| std::io::Error::other(error.to_string()))?;
        permit.forget();
        self.state.active.fetch_sub(1, Ordering::SeqCst);
        Ok(ToolExecution::text("released"))
    }
}

#[tokio::test]
async fn nested_tool_calls_are_bounded_at_128() -> Result<()> {
    const CALLS: usize = super::MAX_CONCURRENT_NESTED_CALLS + 1;

    let workspace = temporary_workspace("bounded-nested-tools")?;
    let state = Arc::new(ConcurrencyProbeState {
        active: AtomicUsize::new(0),
        maximum: AtomicUsize::new(0),
        release: Semaphore::new(0),
    });
    let tools = Tools::builder()
        .without_defaults()
        .tool(ConcurrencyProbe {
            state: Arc::clone(&state),
        })
        .build()?;
    let runtime = ToolRuntime::new_with_tools(&workspace, None, None, &tools);
    let execution = tokio::spawn(async move {
        let history = Vec::new();
        runtime
            .execute_code(
                &format!(
                    "await Promise.all(Array.from({{ length: {CALLS} }}, () => \
                     tools.concurrency_probe({{}})));"
                ),
                test_context(&history),
            )
            .await
    });

    tokio::time::timeout(Duration::from_secs(2), async {
        while state.active.load(Ordering::SeqCst) < super::MAX_CONCURRENT_NESTED_CALLS {
            tokio::task::yield_now().await;
        }
    })
    .await?;
    tokio::time::sleep(Duration::from_millis(25)).await;
    assert_eq!(
        state.active.load(Ordering::SeqCst),
        super::MAX_CONCURRENT_NESTED_CALLS
    );
    assert_eq!(
        state.maximum.load(Ordering::SeqCst),
        super::MAX_CONCURRENT_NESTED_CALLS
    );

    state.release.add_permits(CALLS);
    let execution = execution.await?;
    assert!(execution.success, "{}", execution_output(&execution));
    assert_eq!(execution.nested_calls.len(), CALLS);
    assert_eq!(
        state.maximum.load(Ordering::SeqCst),
        super::MAX_CONCURRENT_NESTED_CALLS
    );

    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[test]
fn long_observer_yields_include_completion_grace() {
    assert_eq!(
        observer_yield_timeout(Duration::from_secs(10)),
        Duration::from_secs(11)
    );
    assert_eq!(
        observer_yield_timeout(Duration::from_millis(9_999)),
        Duration::from_millis(9_999)
    );
}

#[tokio::test]
async fn prewarms_embedded_quickjs_host() -> Result<()> {
    let workspace = temporary_workspace("prewarmed-quickjs-host")?;
    let runtime = super::CodeModeRuntime::new(workspace.clone());

    assert!(runtime.host.lock().await.host.is_some());

    runtime.control().terminate_all().await;
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn execution_globals_do_not_leak_across_quickjs_contexts() -> Result<()> {
    let workspace = temporary_workspace("isolated-quickjs-contexts")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let context = test_context(&history);

    let source = r"
const previous = globalThis.__nanocodexContextGeneration;
globalThis.__nanocodexContextGeneration = (previous || 0) + 1;
text({ previous: previous ?? null, current: globalThis.__nanocodexContextGeneration });
";
    let first = tools.execute_code(source, context).await;
    let second = tools.execute_code(source, context).await;

    assert!(first.success);
    assert!(second.success);
    assert_eq!(emitted_text(&first)?, r#"{"previous":null,"current":1}"#);
    assert_eq!(emitted_text(&second)?, r#"{"previous":null,"current":1}"#);
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn execution_prototype_mutations_do_not_leak_across_quickjs_contexts() -> Result<()> {
    let workspace = temporary_workspace("isolated-quickjs-prototypes")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();

    let first = tools
        .execute_code(
            r#"
Object.prototype.__nanocodexPoisoned = "yes";
text(({}).__nanocodexPoisoned);
"#,
            test_context(&history),
        )
        .await;
    let second = tools
        .execute_code(
            r#"text(({}).__nanocodexPoisoned ?? "clean");"#,
            test_context(&history),
        )
        .await;

    assert!(first.success, "{}", execution_output(&first));
    assert!(second.success, "{}", execution_output(&second));
    assert_eq!(emitted_text(&first)?, "yes");
    assert_eq!(emitted_text(&second)?, "clean");
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn execution_local_bindings_do_not_leak_across_quickjs_calls() -> Result<()> {
    let workspace = temporary_workspace("scoped-quickjs-bindings")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let context = test_context(&history);

    let source = r"
const executionLocal = 1;
text(executionLocal);
";
    let first = tools.execute_code(source, context).await;
    let second = tools.execute_code(source, context).await;

    assert!(first.success, "{}", execution_output(&first));
    assert!(second.success, "{}", execution_output(&second));
    assert_eq!(emitted_text(&first)?, "1");
    assert_eq!(emitted_text(&second)?, "1");
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn embedded_quickjs_does_not_expose_node_or_host_callback_globals() -> Result<()> {
    let workspace = temporary_workspace("embedded-quickjs-globals")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
text({
  process: typeof process,
  require: typeof require,
  hostCallback: typeof __nanocodexTool,
  console: Object.hasOwn(globalThis, "console"),
  atomics: Object.hasOwn(globalThis, "Atomics"),
  sharedArrayBuffer: Object.hasOwn(globalThis, "SharedArrayBuffer"),
  webAssembly: Object.hasOwn(globalThis, "WebAssembly"),
});
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert_eq!(
        emitted_text(&execution)?,
        r#"{"process":"undefined","require":"undefined","hostCallback":"undefined","console":false,"atomics":false,"sharedArrayBuffer":false,"webAssembly":false}"#
    );
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn all_tools_exposes_only_codex_metadata_fields() -> Result<()> {
    let workspace = temporary_workspace("all-tools-metadata")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r"
text(ALL_TOOLS.map((tool) => ({
  keys: Object.keys(tool).sort(),
  frozen: Object.isFrozen(tool),
})));
",
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    let metadata = serde_json::from_str::<Value>(emitted_text(&execution)?)?;
    let tools = metadata
        .as_array()
        .ok_or_else(|| eyre!("ALL_TOOLS output was not an array"))?;
    assert!(!tools.is_empty());
    assert!(tools.iter().all(|tool| {
        tool["keys"] == serde_json::json!(["description", "name"]) && tool["frozen"] == true
    }));
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn function_tools_receive_an_empty_object_when_called_without_arguments() -> Result<()> {
    let workspace = temporary_workspace("empty-function-arguments")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r"
try {
  await tools.update_plan();
} catch (_) {}
",
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert_eq!(execution.nested_calls.len(), 1);
    assert_eq!(execution.nested_calls[0].input, serde_json::json!({}));
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn text_propagates_json_stringification_errors() -> Result<()> {
    let workspace = temporary_workspace("text-stringification-error")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r"
const value = {};
value.self = value;
text(value);
",
            test_context(&history),
        )
        .await;

    assert!(!execution.success);
    assert!(execution_output(&execution).contains("Script error:"));
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[cfg(unix)]
#[tokio::test]
async fn multiple_yielded_cells_continue_and_complete_independently() -> Result<()> {
    let workspace = temporary_workspace("multiple-live-cells")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let first = tools
        .execute_code(
            r#"
await yield_control();
const result = await tools.exec_command({ cmd: "sleep 0.04; printf 'first done'", login: false });
text(result.output);
"#,
            test_context_with_call(&history, "call-first"),
        )
        .await;
    let second = tools
        .execute_code(
            r#"
await yield_control();
const result = await tools.exec_command({ cmd: "sleep 0.01; printf 'second done'", login: false });
text(result.output);
"#,
            test_context_with_call(&history, "call-second"),
        )
        .await;

    assert!(execution_output(&first).contains("Script running with cell ID 1"));
    assert!(execution_output(&second).contains("Script running with cell ID 2"));

    let second = tools
        .wait_for_code(
            r#"{"cell_id":"2","yield_time_ms":5000}"#,
            test_context_with_call(&history, "call-wait-second"),
        )
        .await;
    let first = tools
        .wait_for_code(
            r#"{"cell_id":"1","yield_time_ms":5000}"#,
            test_context_with_call(&history, "call-wait-first"),
        )
        .await;

    assert!(second.success, "{}", execution_output(&second));
    assert!(
        execution_output(&second).contains("second done"),
        "{}",
        execution_output(&second)
    );
    assert_eq!(second.nested_calls.len(), 1);
    assert!(first.success, "{}", execution_output(&first));
    assert!(
        execution_output(&first).contains("first done"),
        "{}",
        execution_output(&first)
    );
    assert_eq!(first.nested_calls.len(), 1);
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[cfg(unix)]
#[tokio::test]
async fn promise_all_runs_nested_tools_concurrently() -> Result<()> {
    let workspace = temporary_workspace("parallel-nested-tools")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
const [first, second] = await Promise.all([
  tools.exec_command({
    cmd: "touch first.started; i=0; while [ \"$i\" -lt 100 ]; do [ -f second.started ] && exit 0; i=$((i + 1)); sleep 0.01; done; exit 91",
    login: false,
  }),
  tools.exec_command({
    cmd: "touch second.started; i=0; while [ \"$i\" -lt 100 ]; do [ -f first.started ] && exit 0; i=$((i + 1)); sleep 0.01; done; exit 92",
    login: false,
  }),
]);
text({ first: first.exit_code, second: second.exit_code });
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success);
    assert_eq!(
        call_ids(&execution.nested_calls),
        ["call-exec/code-1", "call-exec/code-2"]
    );
    let result = serde_json::from_str::<Value>(emitted_text(&execution)?)?;
    assert_eq!(result, serde_json::json!({ "first": 0, "second": 0 }));
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[cfg(unix)]
#[tokio::test]
async fn nested_call_updates_stream_in_start_and_resolution_order() -> Result<()> {
    #[derive(Default)]
    struct Timeline(Vec<String>);

    impl CodeModeObserver for Timeline {
        fn update(&mut self, update: CodeModeUpdate<'_>) {
            match update {
                CodeModeUpdate::NestedCallStarted { call_id, .. } => {
                    self.0.push(format!("start:{call_id}"));
                }
                CodeModeUpdate::NestedCallCompleted(call) => {
                    self.0.push(format!("done:{}", call.call_id));
                }
            }
        }
    }

    let workspace = temporary_workspace("streaming-nested-tools")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let mut timeline = Timeline::default();
    let execution = tools
        .execute_code_with_updates(
            r#"
await Promise.all([
  tools.exec_command({
    cmd: "i=0; while [ \"$i\" -lt 500 ]; do [ -f second.done ] && exit 0; i=$((i + 1)); sleep 0.01; done; exit 91",
    login: false,
  }),
  tools.exec_command({ cmd: "touch second.done", login: false }),
]);
"#,
            test_context(&history),
            &mut timeline,
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert_eq!(
        timeline.0,
        [
            "start:call-exec/code-1",
            "start:call-exec/code-2",
            "done:call-exec/code-2",
            "done:call-exec/code-1",
        ]
    );
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn failed_nested_tool_rejects_its_javascript_promise() -> Result<()> {
    let workspace = temporary_workspace("nested-tool-rejection")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
try {
  await tools.view_image({ path: "missing.png" });
  text("unexpected success");
} catch (error) {
  text(error);
}
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success);
    assert!(emitted_text(&execution)?.contains("unable to locate image"));
    assert_eq!(execution.nested_calls.len(), 1);
    assert!(!execution.nested_calls[0].success);
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn image_helper_requires_data_urls() -> Result<()> {
    let workspace = temporary_workspace("code-mode-image-urls")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();

    let remote = tools
        .execute_code(
            r#"image("https://example.com/image.png");"#,
            test_context(&history),
        )
        .await;
    assert!(!remote.success);
    let remote_output = execution_output(&remote);
    assert!(remote_output.contains(
        "Script error:\nTool call failed: remote image URLs are not supported in tool outputs. Pass a base64 data URI instead"
    ));
    assert!(!remote_output.contains("at image"));

    let invalid = tools
        .execute_code(r#"image("not-an-image");"#, test_context(&history))
        .await;
    assert!(!invalid.success);
    let invalid_output = execution_output(&invalid);
    assert!(invalid_output.contains(
        "Script error:\nTool call failed: invalid image output. Pass a base64 data URI instead"
    ));
    assert!(!invalid_output.contains("at image"));

    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn failed_cell_preserves_accumulated_output() -> Result<()> {
    let workspace = temporary_workspace("failed-cell-output")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
text("before crash");
image("data:image/png;base64,a", "original");
throw new Error("boom");
"#,
            test_context(&history),
        )
        .await;

    assert!(!execution.success);
    let ToolOutputBody::Content(content) = &execution.output else {
        return Err(eyre!("code-mode execution did not emit content"));
    };
    assert!(matches!(
        content.get(1),
        Some(ToolOutputContent::InputText { text }) if text == "before crash"
    ));
    assert!(matches!(
        content.get(2),
        Some(ToolOutputContent::InputImage {
            image_url,
            detail: crate::ImageDetail::Original,
        }) if image_url == "data:image/png;base64,a"
    ));
    assert!(matches!(
        content.get(3),
        Some(ToolOutputContent::InputText { text })
            if text.starts_with("Script error:\nError: boom\n")
    ));

    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn image_helper_normalizes_detail_and_honors_override() -> Result<()> {
    let workspace = temporary_workspace("code-mode-image-detail")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"image({ image_url: "data:image/png;base64,a", detail: "low" }, "ORIGINAL");"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    let ToolOutputBody::Content(content) = &execution.output else {
        return Err(eyre!("code-mode execution did not emit content"));
    };
    assert!(matches!(
        content.last(),
        Some(ToolOutputContent::InputImage {
            image_url,
            detail: crate::ImageDetail::Original,
        }) if image_url == "data:image/png;base64,a"
    ));

    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn view_image_rejects_unsupported_detail_with_codex_diagnostics() -> Result<()> {
    let workspace = temporary_workspace("view-image-detail-error")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
try {
  await tools.view_image({ path: "missing.png", detail: "low" });
} catch (error) {
  text(error);
}
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert_eq!(
        emitted_text(&execution)?,
        "view_image.detail only supports `high` or `original`; omit `detail` for default high resized behavior, got `low`"
    );
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn output_helpers_accept_raw_mcp_image_and_audio_blocks() -> Result<()> {
    let workspace = temporary_workspace("code-mode-mcp-media")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
const returnsUndefined = [
  image({
    type: "image",
    data: "a",
    mimeType: "image/png",
    _meta: { "codex/imageDetail": "original" },
  }),
  audio({
    type: "audio",
    data: "YXVkaW8=",
    mimeType: "audio/wav",
  }),
].map((value) => value === undefined);
text(returnsUndefined);
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    let ToolOutputBody::Content(content) = &execution.output else {
        return Err(eyre!("code-mode execution did not emit content"));
    };
    assert!(matches!(
        content.get(1),
        Some(ToolOutputContent::InputImage {
            image_url,
            detail: crate::ImageDetail::Original,
        }) if image_url == "data:image/png;base64,a"
    ));
    assert_eq!(emitted_text(&execution)?, "[true,true]");
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn generated_image_helper_appends_high_detail_image_and_hint() -> Result<()> {
    let workspace = temporary_workspace("code-mode-generated-image")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
generatedImage({
  image_url: "data:image/png;base64,a",
  output_hint: "generated image save hint",
});
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    let ToolOutputBody::Content(content) = &execution.output else {
        return Err(eyre!("code-mode execution did not emit content"));
    };
    assert!(matches!(
        content.get(1),
        Some(ToolOutputContent::InputImage {
            image_url,
            detail: crate::ImageDetail::High,
        }) if image_url == "data:image/png;base64,a"
    ));
    assert!(matches!(
        content.get(2),
        Some(ToolOutputContent::InputText { text }) if text == "generated image save hint"
    ));

    let invalid = tools
        .execute_code(
            r#"generatedImage({ image_url: "data:image/png;base64,a", output_hint: 1 });"#,
            test_context(&history),
        )
        .await;
    assert!(!invalid.success);
    assert!(
        execution_output(&invalid)
            .contains("generatedImage output_hint must be a string when provided")
    );

    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn notify_serializes_values_and_rejects_empty_text() -> Result<()> {
    let workspace = temporary_workspace("code-mode-notify")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"notify({ phase: "working" }); text("done");"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert_eq!(execution.notifications.len(), 1);
    assert_eq!(execution.notifications[0].call_id, "call-exec");
    assert_eq!(execution.notifications[0].text, r#"{"phase":"working"}"#);

    let empty = tools
        .execute_code(r#"notify("  ");"#, test_context(&history))
        .await;
    assert!(!empty.success);
    assert!(execution_output(&empty).contains("Script error:\nnotify expects non-empty text"));
    assert!(!execution_output(&empty).contains("at notify"));

    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn store_normalizes_json_values_and_coerces_keys() -> Result<()> {
    let workspace = temporary_workspace("code-mode-store-json")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let write = tools
        .execute_code(
            r"
const value = { kept: 1, dropped: undefined, array: [undefined, NaN] };
store(42, value);
value.kept = 99;
",
            test_context(&history),
        )
        .await;
    assert!(write.success, "{}", execution_output(&write));

    let read = tools
        .execute_code(
            r"text(load(42));",
            test_context_with_call(&history, "call-read"),
        )
        .await;
    assert!(read.success, "{}", execution_output(&read));
    assert_eq!(
        serde_json::from_str::<Value>(emitted_text(&read)?)?,
        serde_json::json!({ "kept": 1, "array": [null, null] })
    );

    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn store_rejects_non_serializable_values_at_the_call_boundary() -> Result<()> {
    let workspace = temporary_workspace("code-mode-store-errors")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(r#"store("candidate", undefined);"#, test_context(&history))
        .await;

    assert!(!execution.success);
    let output = execution_output(&execution);
    assert!(output.contains(
        "Script error:\nUnable to store \"candidate\". Only plain serializable objects can be stored."
    ));
    assert!(!output.contains("at store"));

    let read = tools
        .execute_code(
            r#"text(load("candidate"));"#,
            test_context_with_call(&history, "call-read"),
        )
        .await;
    assert!(read.success, "{}", execution_output(&read));
    assert_eq!(emitted_text(&read)?, "undefined");

    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn yielded_cell_completes_through_wait() -> Result<()> {
    let workspace = temporary_workspace("yielded-cell")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
text("before");
await yield_control();
await new Promise((resolve) => setTimeout(resolve, 10));
text("after");
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success);
    assert!(execution_output(&execution).contains("Script running with cell ID 1"));
    assert!(execution_output(&execution).contains("before"));

    let completed = tools
        .wait_for_code(
            r#"{"cell_id":"1","yield_time_ms":5000}"#,
            test_context(&history),
        )
        .await;
    assert!(completed.success);
    assert!(execution_output(&completed).contains("Script completed"));
    assert!(execution_output(&completed).contains("after"));
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[cfg(unix)]
#[tokio::test]
async fn yield_control_does_not_require_pending_tools_to_finish() -> Result<()> {
    let workspace = temporary_workspace("yield-with-pending-tool")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
const pending = tools.exec_command({
  cmd: "sleep 0.05; printf 'finished'",
  login: false,
});
text("before");
yield_control();
const result = await pending;
text(result.output);
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert!(execution_output(&execution).contains("Script running with cell ID 1"));
    assert!(execution_output(&execution).contains("before"));
    let completed = tools
        .wait_for_code(
            r#"{"cell_id":"1","yield_time_ms":5000}"#,
            test_context(&history),
        )
        .await;
    assert!(completed.success, "{}", execution_output(&completed));
    assert!(execution_output(&completed).contains("finished"));
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn pending_timeouts_do_not_keep_a_cell_alive() -> Result<()> {
    let workspace = temporary_workspace("unawaited-timeout")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let completed = tokio::time::timeout(
        Duration::from_secs(2),
        tools.execute_code(
            r#"
setTimeout(() => text("late"), 60_000);
text("done");
"#,
            test_context(&history),
        ),
    )
    .await?;

    assert!(completed.success, "{}", execution_output(&completed));
    assert_eq!(emitted_text(&completed)?, "done");
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn termination_returns_unobserved_output_since_the_last_yield() -> Result<()> {
    let workspace = temporary_workspace("terminated-cell-output")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let after_ready = workspace.join("after-ready");
    let yielded = tools
        .execute_code(
            r#"
text("before");
yield_control();
text("after");
await tools.exec_command({cmd: "touch after-ready", login: false});
await new Promise(() => {});
"#,
            test_context(&history),
        )
        .await;

    assert!(yielded.success, "{}", execution_output(&yielded));
    assert!(execution_output(&yielded).contains("before"));
    tokio::time::timeout(Duration::from_secs(5), async {
        while !after_ready.exists() {
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await?;
    let terminated = tools
        .wait_for_code(
            r#"{"cell_id":"1","terminate":true}"#,
            test_context(&history),
        )
        .await;
    assert!(terminated.success, "{}", execution_output(&terminated));
    assert!(execution_output(&terminated).contains("Script terminated"));
    assert!(execution_output(&terminated).contains("after"));
    assert!(!execution_output(&terminated).contains("was terminated"));
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn natural_completion_wins_over_a_late_termination_request() -> Result<()> {
    let workspace = temporary_workspace("completion-before-termination")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let yielded = tools
        .execute_code(
            r#"
yield_control();
text("done");
"#,
            test_context(&history),
        )
        .await;
    assert!(yielded.success, "{}", execution_output(&yielded));

    tokio::time::sleep(Duration::from_millis(20)).await;
    let completed = tools
        .wait_for_code(
            r#"{"cell_id":"1","terminate":true}"#,
            test_context(&history),
        )
        .await;
    assert!(completed.success, "{}", execution_output(&completed));
    assert!(execution_output(&completed).contains("Script completed"));
    assert!(execution_output(&completed).contains("done"));
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn running_shell_session_survives_output_only_javascript() -> Result<()> {
    let workspace = temporary_workspace("running-shell-session-output")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
const result = await tools.exec_command({ cmd: "sleep 5", yield_time_ms: 250 });
text(result.output);
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert!(
        execution_output(&execution)
            .contains("Nested shell process is still running with session ID 1")
    );
    tools.control().cancel().await;
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn running_shell_session_notice_is_not_duplicated_for_full_results() -> Result<()> {
    let workspace = temporary_workspace("visible-running-shell-session")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
const result = await tools.exec_command({ cmd: "sleep 5", yield_time_ms: 250 });
text(result);
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    let output = execution_output(&execution);
    assert!(output.contains(r#""session_id":1"#));
    assert!(!output.contains("Nested shell process is still running"));
    tools.control().cancel().await;
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn completed_shell_session_is_not_reported_as_running() -> Result<()> {
    let workspace = temporary_workspace("completed-shell-session")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
const command = await tools.exec_command({
  cmd: "read value; printf 'received:%s' \"$value\"",
  login: false,
  tty: true,
  yield_time_ms: 250,
});
const completed = await tools.write_stdin({
  session_id: command.session_id,
  chars: "parity\n",
  yield_time_ms: 5000,
});
text(completed.output);
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    let output = execution_output(&execution);
    assert!(output.contains("received:parity"));
    assert!(!output.contains("Nested shell process is still running"));
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn cancellation_terminates_yielded_code_cells() -> Result<()> {
    let workspace = temporary_workspace("cancelled-cell")?;
    let tools = test_tools(&workspace);
    let control = tools.control();
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r"
await yield_control();
await new Promise(() => {});
",
            test_context(&history),
        )
        .await;
    assert!(execution_output(&execution).contains("Script running with cell ID 1"));

    tokio::time::timeout(std::time::Duration::from_secs(2), control.cancel()).await?;
    let missing = tools
        .wait_for_code(r#"{"cell_id":"1"}"#, test_context(&history))
        .await;
    assert!(!missing.success);
    assert!(execution_output(&missing).contains("exec cell 1 not found"));

    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn cancellation_interrupts_busy_javascript_and_recreates_the_host() -> Result<()> {
    let workspace = temporary_workspace("cancelled-busy-cell")?;
    let tools = test_tools(&workspace);
    let control = tools.control();
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r"
await yield_control();
while (true) {}
",
            test_context(&history),
        )
        .await;
    assert!(execution_output(&execution).contains("Script running with cell ID 1"));

    tokio::time::timeout(std::time::Duration::from_secs(2), control.cancel()).await?;
    let recovered = tools
        .execute_code(r#"text("recovered")"#, test_context(&history))
        .await;

    assert!(recovered.success, "{}", execution_output(&recovered));
    assert_eq!(emitted_text(&recovered)?, "recovered");
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[cfg(unix)]
#[tokio::test]
async fn cancellation_drops_pending_tool_promises_before_recreating_the_host() -> Result<()> {
    let workspace = temporary_workspace("cancelled-pending-tool")?;
    let tools = test_tools(&workspace);
    let control = tools.control();
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"// @exec: {"yield_time_ms": 10}
await tools.exec_command({ cmd: "sleep 5", login: false });
"#,
            test_context(&history),
        )
        .await;
    assert!(execution_output(&execution).contains("Script running with cell ID 1"));

    tokio::time::timeout(std::time::Duration::from_secs(2), control.cancel()).await?;
    let recovered = tools
        .execute_code(r#"text("recovered")"#, test_context(&history))
        .await;

    assert!(recovered.success, "{}", execution_output(&recovered));
    assert_eq!(emitted_text(&recovered)?, "recovered");
    tools.control().cancel().await;
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn resumed_cell_notifications_keep_the_original_exec_call_id() -> Result<()> {
    let workspace = temporary_workspace("resumed-cell-notify")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
await yield_control();
notify("after yield");
text("done");
"#,
            test_context_with_call(&history, "call-original-exec"),
        )
        .await;

    assert!(execution.success);
    assert!(execution.notifications.is_empty());

    let completed = tools
        .wait_for_code(
            r#"{"cell_id":"1","yield_time_ms":1000}"#,
            test_context_with_call(&history, "call-wait"),
        )
        .await;
    assert!(completed.success, "{}", execution_output(&completed));
    assert_eq!(completed.notifications.len(), 1);
    assert_eq!(completed.notifications[0].call_id, "call-original-exec");
    assert_eq!(completed.notifications[0].text, "after yield");

    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn freeform_apply_patch_accepts_a_string() -> Result<()> {
    let workspace = temporary_workspace("freeform-apply-patch")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
const result = await tools.apply_patch("*** Begin Patch\n*** Add File: created.txt\n+created by patch\n*** End Patch");
text(result);
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert_eq!(emitted_text(&execution)?, "{}");
    assert_eq!(execution.nested_calls.len(), 1);
    assert_eq!(
        execution.nested_calls[0].input,
        Value::String(
            "*** Begin Patch\n*** Add File: created.txt\n+created by patch\n*** End Patch"
                .to_owned()
        )
    );
    assert_eq!(
        std::fs::read_to_string(workspace.join("created.txt"))?,
        "created by patch\n"
    );
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn freeform_apply_patch_rejects_with_codex_diagnostics() -> Result<()> {
    let workspace = temporary_workspace("freeform-apply-patch-error")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
try {
  await tools.apply_patch("*** Begin Patch\n*** Update File: missing.txt\n@@\n-before\n+after\n*** End Patch");
} catch (error) {
  text(error);
}
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert!(
        emitted_text(&execution)?
            .starts_with("apply_patch verification failed: Failed to read file to update "),
        "{}",
        execution_output(&execution)
    );
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn update_plan_matches_codex_handler_acceptance() -> Result<()> {
    let workspace = temporary_workspace("update-plan-acceptance")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
const result = await tools.update_plan({
  plan: [
    { step: "", status: "in_progress" },
    { step: "also active", status: "in_progress" },
  ],
});
text(result);
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert_eq!(emitted_text(&execution)?, "{}");
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn exec_pragma_and_wait_limit_direct_output() -> Result<()> {
    let workspace = temporary_workspace("code-output-limits")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            "// @exec: {\"max_output_tokens\": 2}\ntext(\"abcdefghijklmnop\")",
            test_context(&history),
        )
        .await;
    assert!(execution.success);
    assert!(execution_output(&execution).contains("Warning: truncated output"));

    let yielded = tools
        .execute_code(
            r#"
await yield_control();
text("abcdefghijklmnop");
"#,
            test_context(&history),
        )
        .await;
    assert!(yielded.success);
    let completed = tools
        .wait_for_code(
            r#"{"cell_id":"2","yield_time_ms":1000,"max_tokens":2}"#,
            test_context(&history),
        )
        .await;
    assert!(completed.success);
    assert!(execution_output(&completed).contains("Warning: truncated output"));
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[test]
fn exec_pragma_rejects_unknown_fields() {
    let error = parse_exec_source("// @exec: {\"unknown\": 1}\ntext('hi')")
        .err()
        .expect("unknown pragma fields should fail");
    assert!(error.contains("only supports"));
}

#[test]
fn nested_shell_yields_follow_the_handlers_bounds() {
    assert_eq!(
        nested_tool_yield_after("exec_command", &serde_json::json!({ "cmd": "sleep 1" })),
        Some(Duration::from_secs(10))
    );
    assert_eq!(
        nested_tool_yield_after("write_stdin", &serde_json::json!({ "session_id": 1 }),),
        Some(Duration::from_secs(5))
    );
    assert_eq!(
        nested_tool_yield_after(
            "write_stdin",
            &serde_json::json!({ "session_id": 1, "chars": "x" }),
        ),
        Some(Duration::from_millis(250))
    );
    assert_eq!(
        nested_tool_yield_after(
            "exec_command",
            &serde_json::json!({ "yield_time_ms": 45_000 }),
        ),
        Some(Duration::from_secs(30))
    );
    assert_eq!(
        nested_tool_yield_after(
            "write_stdin",
            &serde_json::json!({ "session_id": 1, "yield_time_ms": 120_000 }),
        ),
        Some(Duration::from_mins(2))
    );
    assert_eq!(
        nested_tool_yield_after(
            "write_stdin",
            &serde_json::json!({
                "session_id": 1,
                "chars": "x",
                "yield_time_ms": 120_000,
            }),
        ),
        Some(Duration::from_secs(30))
    );
    assert_eq!(
        nested_tool_yield_after(
            "apply_patch",
            &serde_json::json!({ "yield_time_ms": 30_000 })
        ),
        None
    );
}

#[tokio::test]
async fn negative_shell_limits_fail_during_codex_compatible_argument_parsing() -> Result<()> {
    let workspace = temporary_workspace("negative-shell-limits")?;
    let tools = test_tools(&workspace);
    let history = Vec::new();
    let execution = tools
        .execute_code(
            r#"
try {
  await tools.exec_command({ cmd: "true", yield_time_ms: -1 });
} catch (error) {
  text(error);
}
"#,
            test_context(&history),
        )
        .await;

    assert!(execution.success, "{}", execution_output(&execution));
    assert!(
        emitted_text(&execution)?
            .starts_with("failed to parse function arguments: invalid value: integer `-1`"),
        "{}",
        execution_output(&execution)
    );
    std::fs::remove_dir_all(workspace)?;
    Ok(())
}

#[tokio::test]
async fn default_cell_yield_extends_for_a_longer_nested_shell_wait() {
    let (updates_tx, updates) = tokio::sync::mpsc::unbounded_channel();
    let (terminate, _terminate_rx) = tokio::sync::oneshot::channel();
    let task = tokio::spawn(async move {
        updates_tx
            .send(CellUpdate::NestedCallStarted {
                call_id: "call/code-1".to_owned(),
                name: "write_stdin".to_owned(),
                input: serde_json::Value::Null,
                yield_after: Some(Duration::from_millis(40)),
            })
            .expect("observer should receive the nested call");
        tokio::time::sleep(Duration::from_millis(15)).await;
        updates_tx
            .send(CellUpdate::Completed)
            .expect("observer should receive cell completion");
    });
    let mut cell = LiveCell {
        id: 1,
        output_token_budget: crate::DEFAULT_TOOL_OUTPUT_TOKENS,
        updates,
        terminate: Some(terminate),
        task: Some(task),
    };

    let (execution, running) = observe_cell(
        &mut cell,
        std::time::Instant::now(),
        ObservationMode::YieldAfter(Duration::from_millis(5)),
        None,
        true,
        &mut super::IgnoreCodeModeUpdates,
    )
    .await;

    assert!(!running);
    assert!(execution.success);
    assert!(execution_output(&execution).contains("Script completed"));
    cell.join().await;
}

#[tokio::test]
async fn explicit_cell_yield_is_not_extended_by_a_nested_shell_wait() {
    let (updates_tx, updates) = tokio::sync::mpsc::unbounded_channel();
    let (terminate, _terminate_rx) = tokio::sync::oneshot::channel();
    let task = tokio::spawn(async move {
        updates_tx
            .send(CellUpdate::NestedCallStarted {
                call_id: "call/code-1".to_owned(),
                name: "write_stdin".to_owned(),
                input: serde_json::Value::Null,
                yield_after: Some(Duration::from_millis(40)),
            })
            .expect("observer should receive the nested call");
        tokio::time::sleep(Duration::from_millis(15)).await;
        let _ = updates_tx.send(CellUpdate::Completed);
    });
    let mut cell = LiveCell {
        id: 1,
        output_token_budget: crate::DEFAULT_TOOL_OUTPUT_TOKENS,
        updates,
        terminate: Some(terminate),
        task: Some(task),
    };

    let (execution, running) = observe_cell(
        &mut cell,
        std::time::Instant::now(),
        ObservationMode::YieldAfter(Duration::from_millis(5)),
        None,
        false,
        &mut super::IgnoreCodeModeUpdates,
    )
    .await;

    assert!(running);
    assert!(execution.success);
    assert!(execution_output(&execution).contains("Script running with cell ID 1"));
    cell.join().await;
}

#[test]
fn model_description_uses_codex_style_declarations() {
    let workspace = temporary_workspace("code-mode-description")
        .expect("temporary test workspace should be available");
    let tools = test_tools(&workspace);
    let specs = tools
        .model_specs()
        .into_iter()
        .map(|spec| serde_json::to_value(spec).unwrap())
        .collect::<Vec<_>>();
    let description = specs[0]["description"]
        .as_str()
        .expect("exec should have a description");
    assert!(description.contains("// @exec:"));
    assert!(description.contains("should be a base64-encoded `data:` URL"));
    assert!(description.contains("apply_patch(input: string): Promise<unknown>"));
    assert!(description.contains("exec_command(args: {"));
    assert!(!description.contains("Input schema:"));
    assert_eq!(
        specs[1]["parameters"]["properties"]["max_tokens"]["type"],
        "number"
    );
    std::fs::remove_dir_all(workspace).expect("temporary workspace should be removable");
}

fn emitted_text(execution: &CodeModeExecution) -> Result<&str> {
    let ToolOutputBody::Content(content) = &execution.output else {
        return Err(eyre!("code-mode execution did not emit content"));
    };
    content
        .iter()
        .rev()
        .find_map(|item| match item {
            ToolOutputContent::InputText { text } => Some(text.as_str()),
            ToolOutputContent::InputImage { .. } | ToolOutputContent::InputAudio { .. } => None,
        })
        .ok_or_else(|| eyre!("code-mode execution did not emit text"))
}

fn execution_output(execution: &CodeModeExecution) -> String {
    match &execution.output {
        ToolOutputBody::Text(text) => text.clone(),
        ToolOutputBody::Content(content) => content
            .iter()
            .filter_map(|item| match item {
                ToolOutputContent::InputText { text } => Some(text.as_str()),
                ToolOutputContent::InputImage { .. } | ToolOutputContent::InputAudio { .. } => None,
            })
            .collect::<Vec<_>>()
            .join("\n"),
    }
}

fn call_ids(calls: &[NestedToolCall]) -> Vec<&str> {
    calls.iter().map(|call| call.call_id.as_str()).collect()
}

fn test_tools(workspace: &std::path::Path) -> ToolRuntime {
    ToolRuntime::new(
        workspace,
        Some(WebSearchConfig {
            endpoint: "http://127.0.0.1:1/v1/alpha/search".to_owned(),
            auth: nanocodex_core::OpenAiAuth::api_key("test-key"),
        }),
        Some(super::super::ImageGenerationConfig {
            api_base_url: "http://127.0.0.1:1/v1".to_owned(),
            auth: nanocodex_core::OpenAiAuth::api_key("test-key"),
            save_root: workspace.to_path_buf(),
        }),
    )
}

fn test_context(history: &[ResponseItem]) -> ToolContext<'_> {
    test_context_with_call(history, "call-exec")
}

fn test_context_with_call<'a>(history: &'a [ResponseItem], call_id: &'a str) -> ToolContext<'a> {
    ToolContext {
        model: "test-model",
        session_id: "test-session",
        call_id,
        history,
        output_token_budget: crate::DEFAULT_TOOL_OUTPUT_TOKENS,
    }
}

fn temporary_workspace(label: &str) -> Result<PathBuf> {
    let path = std::env::temp_dir().join(format!(
        "nanocodex-{label}-{}-{}",
        std::process::id(),
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)?
            .as_nanos()
    ));
    std::fs::create_dir_all(&path)?;
    Ok(path)
}