llm-tool 0.10.5

use super::*;

#[test]
fn tool_definition_serde_roundtrip() {
    let def = definition_of(&SampleTool).expect("schema");
    let json = serde_json::to_string(&def).expect("serialize");
    let parsed: ToolDefinition = serde_json::from_str(&json).expect("deserialize");
    assert_eq!(parsed, def);
}

struct EmptyParamTool;
impl RustTool for EmptyParamTool {
    type Params = EmptyParams;
    const NAME: &'static str = "empty";
    const DESCRIPTION: &'static str = "No params";
    async fn call(
        &self,
        _params: Self::Params,
        _ctx: &ToolContext,
    ) -> Result<ToolOutput, ToolError> {
        Ok("ok".into())
    }
}

#[test]
fn tool_definition_with_empty_schema() {
    let tool = definition_of(&EmptyParamTool).expect("schema");
    let json = serde_json::to_string(&tool).expect("serialize");
    let parsed: ToolDefinition = serde_json::from_str(&json).expect("deserialize");
    // Compare via JSON to handle serde normalization (None vs empty struct).
    let orig_json = serde_json::to_value(&tool.parameter_schema).unwrap();
    let parsed_json = serde_json::to_value(&parsed.parameter_schema).unwrap();
    assert_eq!(orig_json, parsed_json);
}

#[test]
fn tool_definition_with_complex_schema() {
    let tool = definition_of(&RunCommandTool).expect("schema");
    let schema_json = serde_json::to_value(&tool.parameter_schema).expect("schema to json");
    // The schema should have 'command' as a required field.
    let required = schema_json["required"]
        .as_array()
        .expect("required should be an array");
    assert!(
        required.iter().any(|v| v == "command"),
        "'command' should be required, got: {required:?}"
    );
}

// ── ToolRegistry tests ────────────────────────────────────────

#[tokio::test]
async fn registry_dispatch_valid_tool() {
    let mut d = ToolRegistry::new();
    d.register(SampleTool);
    let result = d
        .dispatch(
            "sample",
            serde_json::json!({"path": "/tmp/foo"}),
            &test_ctx(),
        )
        .await;
    assert_eq!(result.unwrap().content(), "/tmp/foo");
}

#[tokio::test]
async fn registry_dispatch_unknown_tool() {
    let d = ToolRegistry::new();
    // Dispatching an unknown tool yields a not_found error, not an execution error.
    let err = d
        .dispatch("nonexistent", serde_json::json!({}), &test_ctx())
        .await
        .expect_err("dispatching an unknown tool should return a not_found error");
    assert_eq!(err.metadata()["error_kind"], "not_registered");
    assert!(
        err.to_string().contains("nonexistent"),
        "error should name the missing tool, got: {err}"
    );
}

#[tokio::test]
async fn registry_dispatch_invalid_args() {
    let mut d = ToolRegistry::new();
    d.register(SampleTool);
    // SampleTool expects {"path": String}, not an integer.
    let result = d
        .dispatch("sample", serde_json::json!({"path": 42}), &test_ctx())
        .await;
    let err = result.unwrap_err();
    assert!(
        err.message.contains("deserialize"),
        "Error should mention deserialization, got: {err}"
    );
}

#[tokio::test]
async fn registry_dispatch_missing_required_field() {
    let mut d = ToolRegistry::new();
    d.register(SampleTool);
    // Missing the required "path" field entirely.
    let err = d
        .dispatch("sample", serde_json::json!({}), &test_ctx())
        .await
        .expect_err("Expected error for missing required field");
    assert!(
        err.message.contains("missing field"),
        "Error should mention missing field, got: {err}"
    );
}

#[test]
fn registry_definitions_returns_all() {
    let mut d = ToolRegistry::new();
    d.register(SampleTool);
    d.register(RunCommandTool);

    let defs = d.definitions();
    assert_eq!(defs.len(), 2);

    let mut names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect();
    names.sort_unstable();
    assert_eq!(names, vec!["run_command", "sample"]);
}

#[test]
fn registry_register_chaining() {
    let mut d = ToolRegistry::new();
    d.register(SampleTool).register(RunCommandTool);
    assert_eq!(d.len(), 2);
    assert!(!d.is_empty());
}

#[test]
fn registry_try_register_succeeds_and_chains() {
    let mut d = ToolRegistry::new();
    d.try_register(SampleTool)
        .expect("sample schema builds")
        .try_register(RunCommandTool)
        .expect("run_command schema builds");
    assert_eq!(d.len(), 2);
}

#[test]
fn registry_try_register_replaces_same_name() {
    let mut d = ToolRegistry::new();
    d.try_register(SampleTool).expect("first registration");
    d.try_register(SampleTool).expect("second registration");
    // Re-registering the same NAME replaces rather than duplicating.
    assert_eq!(d.len(), 1);
}

#[test]
fn registry_iter_yields_named_definitions() {
    let d = ToolRegistry::new()
        .with_tool(SampleTool)
        .with_tool(RunCommandTool);

    // The named iterator reports an exact length without consuming itself.
    let iter = d.iter();
    assert_eq!(iter.len(), 2);

    let mut names: Vec<&str> = iter.map(|(name, _def)| name).collect();
    names.sort_unstable();
    assert_eq!(names, vec!["run_command", "sample"]);

    // `&ToolRegistry` also implements IntoIterator with matching definitions.
    for (name, def) in &d {
        assert_eq!(name, def.name.as_str());
    }
}

#[test]
fn registry_with_tool_owned_chaining() {
    let d = ToolRegistry::new()
        .with_tool(SampleTool)
        .with_tool(RunCommandTool);
    assert_eq!(d.len(), 2);
    assert!(!d.is_empty());

    let defs = d.definitions();
    let mut names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect();
    names.sort_unstable();
    assert_eq!(names, vec!["run_command", "sample"]);
}

#[test]
fn registry_default_is_empty() {
    let d = ToolRegistry::default();
    assert!(d.is_empty());
    assert_eq!(d.len(), 0);
}

#[tokio::test]
async fn registry_replaces_on_duplicate_name() {
    struct AlternateSample;
    impl RustTool for AlternateSample {
        type Params = PathParams;
        const NAME: &'static str = "sample";
        const DESCRIPTION: &'static str = "Alternate sample";
        async fn call(
            &self,
            params: Self::Params,
            _ctx: &ToolContext,
        ) -> Result<ToolOutput, ToolError> {
            Ok(format!("alt: {}", params.path).into())
        }
    }

    let mut d = ToolRegistry::new();
    d.register(SampleTool);
    d.register(AlternateSample);
    assert_eq!(d.len(), 1);

    let result = d
        .dispatch("sample", serde_json::json!({"path": "x"}), &test_ctx())
        .await;
    assert_eq!(result.unwrap().content(), "alt: x");
}

#[tokio::test]
async fn registry_tool_returning_error() {
    struct FailingTool;
    impl RustTool for FailingTool {
        type Params = EmptyParams;
        const NAME: &'static str = "fail";
        const DESCRIPTION: &'static str = "Always fails";
        async fn call(
            &self,
            _params: Self::Params,
            _ctx: &ToolContext,
        ) -> Result<ToolOutput, ToolError> {
            Err(ToolError::new("intentional failure"))
        }
    }

    let mut d = ToolRegistry::new();
    d.register(FailingTool);
    let result = d.dispatch("fail", serde_json::json!({}), &test_ctx()).await;
    assert_eq!(result.unwrap_err(), ToolError::new("intentional failure"));
}

#[test]
fn registry_debug_shows_tool_names() {
    let mut d = ToolRegistry::new();
    d.register(SampleTool);
    let dbg = format!("{d:?}");
    assert!(dbg.contains("ToolRegistry"));
    assert!(dbg.contains("sample"));
    assert!(dbg.contains("tool_count: 1"));
}

// ── Async-specific tests ────────────────────────────────────────

/// A tool that actually awaits a tokio sleep, proving async dispatch works.
struct AsyncSleepTool;

impl RustTool for AsyncSleepTool {
    type Params = EmptyParams;
    const NAME: &'static str = "async_sleep";
    const DESCRIPTION: &'static str = "Sleeps briefly then returns.";

    async fn call(
        &self,
        _params: Self::Params,
        _ctx: &ToolContext,
    ) -> Result<ToolOutput, ToolError> {
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        Ok("slept".into())
    }
}

#[tokio::test]
async fn async_tool_with_tokio_sleep() {
    let mut d = ToolRegistry::new();
    d.register(AsyncSleepTool);
    let result = d
        .dispatch("async_sleep", serde_json::json!({}), &test_ctx())
        .await;
    assert_eq!(result.unwrap().content(), "slept");
}

/// A tool that reads a file using `tokio::fs`.
struct AsyncReadFileTool;

#[derive(Deserialize, schemars::JsonSchema)]
struct ReadFileParams {
    /// Path to the file to read.
    path: String,
}

impl RustTool for AsyncReadFileTool {
    type Params = ReadFileParams;
    const NAME: &'static str = "read_file";
    const DESCRIPTION: &'static str = "Reads a file asynchronously.";

    async fn call(
        &self,
        params: Self::Params,
        _ctx: &ToolContext,
    ) -> Result<ToolOutput, ToolError> {
        tokio::fs::read_to_string(&params.path)
            .await
            .map(ToolOutput::from)
            .map_err(|e| ToolError::new(format!("IO error: {e}")))
    }
}

#[tokio::test]
async fn async_tool_with_tokio_fs() {
    let tmp = tempfile::NamedTempFile::new().expect("create tempfile");
    std::fs::write(tmp.path(), "hello async").expect("write tempfile");

    let mut d = ToolRegistry::new();
    d.register(AsyncReadFileTool);

    let path_str = tmp.path().to_str().expect("path to str").to_owned();
    let result = d
        .dispatch(
            "read_file",
            serde_json::json!({"path": path_str}),
            &test_ctx(),
        )
        .await;
    assert_eq!(result.unwrap().content(), "hello async");
}

#[tokio::test]
async fn async_tool_tokio_fs_missing_file() {
    let mut d = ToolRegistry::new();
    d.register(AsyncReadFileTool);
    let result = d
        .dispatch(
            "read_file",
            serde_json::json!({"path": "/nonexistent/file.txt"}),
            &test_ctx(),
        )
        .await;
    let err = result.unwrap_err();
    assert!(
        err.message.contains("IO error"),
        "Expected IO error, got: {err}"
    );
}

/// A tool that uses a tokio channel to receive its result, proving
/// the full async machinery works end-to-end.
struct ChannelTool {
    tx: tokio::sync::mpsc::Sender<String>,
    rx: std::sync::Mutex<Option<tokio::sync::mpsc::Receiver<String>>>,
}

impl ChannelTool {
    fn new() -> Self {
        let (tx, rx) = tokio::sync::mpsc::channel(1);
        Self {
            tx,
            rx: std::sync::Mutex::new(Some(rx)),
        }
    }
}

impl RustTool for ChannelTool {
    type Params = EmptyParams;
    const NAME: &'static str = "channel_tool";
    const DESCRIPTION: &'static str = "Awaits a value from a channel.";

    async fn call(
        &self,
        _params: Self::Params,
        _ctx: &ToolContext,
    ) -> Result<ToolOutput, ToolError> {
        let mut rx = self
            .rx
            .lock()
            .unwrap()
            .take()
            .ok_or_else(|| ToolError::new("channel already consumed"))?;
        rx.recv()
            .await
            .map(ToolOutput::from)
            .ok_or_else(|| ToolError::new("channel closed"))
    }
}

#[tokio::test]
async fn async_tool_awaits_channel() {
    let tool = ChannelTool::new();
    let tx = tool.tx.clone();

    let mut d = ToolRegistry::new();
    d.register(tool);

    // Send the value from another task.
    let ctx = test_ctx();
    let dispatch_future = d.dispatch("channel_tool", serde_json::json!({}), &ctx);
    let send_future = async move {
        tx.send("from_channel".to_string()).await.unwrap();
    };

    let (result, ()) = tokio::join!(dispatch_future, send_future);
    assert_eq!(result.unwrap().content(), "from_channel");
}

// ── Concurrent dispatch tests ───────────────────────────────────

#[tokio::test]
async fn concurrent_dispatches_to_different_tools() {
    let mut d = ToolRegistry::new();
    d.register(SampleTool);
    d.register(AsyncSleepTool);
    d.register(RunCommandTool);

    let ctx = test_ctx();
    let (r1, r2, r3) = tokio::join!(
        d.dispatch("sample", serde_json::json!({"path": "a"}), &ctx),
        d.dispatch("async_sleep", serde_json::json!({}), &ctx),
        d.dispatch("run_command", serde_json::json!({"command": "ls"}), &ctx),
    );

    assert_eq!(r1.unwrap().content(), "a");
    assert_eq!(r2.unwrap().content(), "slept");
    assert_eq!(r3.unwrap().content(), "Ran: ls");
}

#[tokio::test]
async fn concurrent_dispatches_to_same_tool() {
    let mut d = ToolRegistry::new();
    d.register(SampleTool);

    let ctx = test_ctx();
    let futs: Vec<_> = (0..10)
        .map(|i| d.dispatch("sample", serde_json::json!({"path": format!("p{i}")}), &ctx))
        .collect();

    let results = futures::future::join_all(futs).await;
    for (i, r) in results.into_iter().enumerate() {
        assert_eq!(r.unwrap().content(), format!("p{i}"));
    }
}

// ── Schema / doc comment tests ──────────────────────────────────

#[derive(Deserialize, schemars::JsonSchema)]
struct DocumentedParams {
    /// The target hostname to connect to.
    hostname: String,
    /// Port number (1-65535).
    port: u16,
    /// Optional timeout in seconds.
    #[serde(default)]
    timeout: Option<f64>,
}

struct DocumentedTool;
impl RustTool for DocumentedTool {
    type Params = DocumentedParams;
    const NAME: &'static str = "connect";
    const DESCRIPTION: &'static str = "Connects to a remote host.";
    async fn call(&self, p: Self::Params, _ctx: &ToolContext) -> Result<ToolOutput, ToolError> {
        Ok(format!("{}:{}:{:?}", p.hostname, p.port, p.timeout).into())
    }
}

#[test]
fn schema_contains_field_descriptions() {
    let def = definition_of(&DocumentedTool).expect("schema");
    let schema = &def.parameter_schema;

    // Check the properties contain our fields.
    let props = schema["properties"].as_object().expect("properties object");
    assert!(props.contains_key("hostname"), "missing hostname");
    assert!(props.contains_key("port"), "missing port");
    assert!(props.contains_key("timeout"), "missing timeout");

    // Check the descriptions from doc comments made it through.
    let hostname_desc = props["hostname"]["description"]
        .as_str()
        .expect("hostname description");
    assert!(
        hostname_desc.contains("hostname"),
        "hostname description should mention 'hostname', got: {hostname_desc}"
    );

    let port_desc = props["port"]["description"]
        .as_str()
        .expect("port description");
    assert!(
        port_desc.contains("1-65535"),
        "port description should mention range, got: {port_desc}"
    );
}

#[test]
fn schema_required_vs_optional_fields() {
    let def = definition_of(&DocumentedTool).expect("schema");
    let schema = &def.parameter_schema;

    let required = schema["required"]
        .as_array()
        .expect("required should be an array");

    // hostname and port are required, timeout is Option → not required.
    assert!(
        required.iter().any(|v| v == "hostname"),
        "hostname required"
    );
    assert!(required.iter().any(|v| v == "port"), "port required");
    assert!(
        !required.iter().any(|v| v == "timeout"),
        "timeout should NOT be required"
    );
}

#[tokio::test]
async fn dispatch_with_optional_field_missing() {
    let mut d = ToolRegistry::new();
    d.register(DocumentedTool);

    // Dispatch without `timeout` (it has serde(default)).
    let result = d
        .dispatch(
            "connect",
            serde_json::json!({"hostname": "example.com", "port": 443}),
            &test_ctx(),
        )
        .await;
    assert_eq!(result.unwrap().content(), "example.com:443:None");
}

#[tokio::test]
async fn dispatch_with_optional_field_present() {
    let mut d = ToolRegistry::new();
    d.register(DocumentedTool);

    let result = d
        .dispatch(
            "connect",
            serde_json::json!({"hostname": "localhost", "port": 8080, "timeout": 30.0}),
            &test_ctx(),
        )
        .await;
    assert_eq!(result.unwrap().content(), "localhost:8080:Some(30.0)");
}

#[tokio::test]
async fn dispatch_with_extra_fields_ignored() {
    // serde's default behavior ignores unknown fields.
    let mut d = ToolRegistry::new();
    d.register(SampleTool);

    let result = d
        .dispatch(
            "sample",
            serde_json::json!({"path": "/tmp/x", "unknown_field": 42}),
            &test_ctx(),
        )
        .await;
    assert_eq!(result.unwrap().content(), "/tmp/x");
}

// ── BoxFuture / ErasedTool edge case tests ──────────────────────

#[tokio::test]
async fn erased_dispatch_preserves_borrow_lifetime() {
    // Ensures the BoxToolFuture lifetime is tied to &self correctly,
    // i.e. the registry can be borrowed immutably while the future runs.
    let mut d = ToolRegistry::new();
    d.register(AsyncSleepTool);
    d.register(SampleTool);

    // Dispatch two calls on the same registry reference.
    let r1 = d
        .dispatch("async_sleep", serde_json::json!({}), &test_ctx())
        .await;
    let r2 = d
        .dispatch("sample", serde_json::json!({"path": "test"}), &test_ctx())
        .await;

    assert_eq!(r1.unwrap().content(), "slept");
    assert_eq!(r2.unwrap().content(), "test");
}

#[tokio::test]
async fn dispatch_returns_meaningful_error_for_wrong_type() {
    let mut d = ToolRegistry::new();
    d.register(RunCommandTool);

    // `command` expects a String, pass an object instead.
    let result = d
        .dispatch(
            "run_command",
            serde_json::json!({"command": {"nested": "object"}}),
            &test_ctx(),
        )
        .await;
    let err = result.unwrap_err();
    assert!(
        err.message
            .contains("Failed to deserialize tool parameters"),
        "Error should mention deserialization failure, got: {err}"
    );
}

// ── Cache invalidation tests ────────────────────────────────────────

#[test]
fn cache_invalidation_on_register_remove_clear() {
    let mut reg = ToolRegistry::new();
    assert_eq!(reg.definitions_slice().len(), 0);

    reg.register(SampleTool);
    let slice1 = reg.definitions_slice();
    assert_eq!(slice1.len(), 1);
    assert_eq!(slice1[0].name, "sample");

    // Same slice returned from cache
    let slice1_cached = reg.definitions_slice();
    assert!(std::sync::Arc::ptr_eq(&slice1, &slice1_cached));

    // register invalidates cache
    reg.register(EmptyParamTool);
    let slice2 = reg.definitions_slice();
    assert_eq!(slice2.len(), 2);
    assert!(!std::sync::Arc::ptr_eq(&slice1, &slice2));

    // remove invalidates cache
    assert!(reg.remove("sample"));
    let slice3 = reg.definitions_slice();
    assert_eq!(slice3.len(), 1);
    assert_eq!(slice3[0].name, "empty");

    // clear invalidates cache
    reg.clear();
    let slice4 = reg.definitions_slice();
    assert_eq!(slice4.len(), 0);
}

#[test]
fn cached_definitions_poison_recovery() {
    let mut reg = ToolRegistry::new();
    reg.register(SampleTool);
    let slice = reg.definitions_slice();
    assert_eq!(slice.len(), 1);

    // Intentionally poison the cached_definitions lock in another thread
    let clone_lock = std::sync::Arc::new(crate::compat::RwLock::new(Some(slice)));
    let thread_lock = std::sync::Arc::clone(&clone_lock);
    let handle = std::thread::spawn(move || {
        let guard = crate::compat::write_lock(&thread_lock).expect("lock acquisition");
        std::hint::black_box(&guard);
        panic!("intentional poison for lock recovery test");
    });
    let join_result = handle.join();
    assert!(join_result.is_err());

    // Recovering read and write locks succeeds without panic
    let recovered_read = crate::compat::read_lock_recover(&clone_lock);
    assert!(recovered_read.is_some());
    drop(recovered_read);

    let mut recovered_write = crate::compat::write_lock_recover(&clone_lock);
    *recovered_write = None;
    drop(recovered_write);
}

// ── Batch dispatch & ToolEffect tests ───────────────────────────────

struct DestructiveTool;

impl RustTool for DestructiveTool {
    type Params = EmptyParams;
    const NAME: &'static str = "wipe_data";
    const DESCRIPTION: &'static str = "Destructive data wiper";
    const EFFECT: ToolEffect = ToolEffect::Destructive;

    async fn call(
        &self,
        _params: Self::Params,
        _ctx: &ToolContext,
    ) -> Result<ToolOutput, ToolError> {
        Ok("wiped".into())
    }
}

#[tokio::test]
async fn batch_dispatch_parallel_and_sequential_modes() {
    let mut reg = ToolRegistry::new();
    reg.register(SampleTool);
    reg.register(EmptyParamTool);
    reg.register(DestructiveTool);

    let ctx = test_ctx();

    // 1. Parallel execution with non-destructive calls
    let calls = vec![
        BatchToolCall {
            id: "call_1".into(),
            name: "sample".into(),
            arguments: serde_json::json!({ "path": "path1" }),
        },
        BatchToolCall {
            id: "call_2".into(),
            name: "empty".into(),
            arguments: serde_json::json!({}),
        },
    ];

    let outcomes = reg
        .dispatch_batch(&calls, BatchExecutionMode::Parallel, &ctx)
        .await;
    assert_eq!(outcomes.len(), 2);
    assert_eq!(outcomes[0].id, "call_1");
    assert_eq!(outcomes[0].result.as_ref().unwrap().content(), "path1");
    assert_eq!(outcomes[1].id, "call_2");
    assert_eq!(outcomes[1].result.as_ref().unwrap().content(), "ok");

    // 2. Parallel execution rejecting destructive tool without confirmation
    let destructive_calls = vec![BatchToolCall {
        id: "call_wipe".into(),
        name: "wipe_data".into(),
        arguments: serde_json::json!({}),
    }];

    let err_outcomes = reg
        .dispatch_batch(&destructive_calls, BatchExecutionMode::Parallel, &ctx)
        .await;
    assert_eq!(err_outcomes.len(), 1);
    let outcome_err = err_outcomes[0].result.as_ref().unwrap_err();
    assert!(outcome_err.message.contains("Destructive effect"));

    // 3. Parallel execution confirming destructive tool via ParallelAllowDestructive
    let ok_outcomes = reg
        .dispatch_batch(
            &destructive_calls,
            BatchExecutionMode::ParallelAllowDestructive,
            &ctx,
        )
        .await;
    assert_eq!(ok_outcomes.len(), 1);
    assert_eq!(ok_outcomes[0].result.as_ref().unwrap().content(), "wiped");

    // 4. Parallel execution confirming destructive tool via AllowDestructiveBatch context extension
    let conf_ctx = test_ctx();
    conf_ctx
        .set_ext(AllowDestructiveBatch(true))
        .expect("set_ext succeeds");
    let ok_ctx_outcomes = reg
        .dispatch_batch(&destructive_calls, BatchExecutionMode::Parallel, &conf_ctx)
        .await;
    assert_eq!(ok_ctx_outcomes.len(), 1);
    assert_eq!(
        ok_ctx_outcomes[0].result.as_ref().unwrap().content(),
        "wiped"
    );

    // 5. SequentialStopOnError halts on first error
    let seq_calls = vec![
        BatchToolCall {
            id: "s1".into(),
            name: "sample".into(),
            arguments: serde_json::json!({ "path": "first" }),
        },
        BatchToolCall {
            id: "s2".into(),
            name: "non_existent_tool".into(),
            arguments: serde_json::json!({}),
        },
        BatchToolCall {
            id: "s3".into(),
            name: "sample".into(),
            arguments: serde_json::json!({ "path": "third" }),
        },
    ];

    let seq_outcomes = reg
        .dispatch_batch(&seq_calls, BatchExecutionMode::SequentialStopOnError, &ctx)
        .await;
    assert_eq!(seq_outcomes.len(), 2);
    assert!(seq_outcomes[0].result.is_ok());
    assert!(seq_outcomes[1].result.is_err());
}