runifold-tool 0.5.4

use std::{future::Future, marker::PhantomData};

use runifold_core::{CapabilityId, EffectClass, RiskLevel};
use schemars::{JsonSchema, schema_for};
use serde::{Serialize, de::DeserializeOwned};

use crate::{Tool, ToolContext, ToolDescriptor, ToolError, ToolErrorKind, ToolFuture, ToolOutput};

/// A typed asynchronous Rust function exposed through the canonical Tool
/// boundary.
pub struct FunctionTool<Input, Output, Handler> {
    descriptor: ToolDescriptor,
    handler: Handler,
    types: PhantomData<fn(Input) -> Output>,
}

impl<Input, Output, Handler> FunctionTool<Input, Output, Handler>
where
    Input: JsonSchema,
    Output: JsonSchema,
{
    /// Creates a typed Tool with generated input and output JSON Schemas.
    ///
    /// The default effect is [`EffectClass::Pure`] and the default risk is
    /// [`RiskLevel::Low`]. Callers must explicitly override these values for
    /// functions that read or modify external state.
    pub fn new(name: impl Into<String>, description: impl Into<String>, handler: Handler) -> Self {
        Self {
            descriptor: ToolDescriptor {
                id: CapabilityId::new(),
                name: name.into(),
                version: "1".into(),
                description: description.into(),
                input_schema: schema_for!(Input).to_value(),
                output_schema: schema_for!(Output).to_value(),
                effect: EffectClass::Pure,
                risk: RiskLevel::Low,
                metadata: std::collections::BTreeMap::new(),
            },
            handler,
            types: PhantomData,
        }
    }
}

impl<Input, Handler> FunctionTool<Input, ToolOutput, Handler>
where
    Input: JsonSchema,
{
    /// Creates a typed Tool whose handler returns canonical rich content.
    ///
    /// Unlike [`Self::new`], this constructor preserves the returned
    /// [`ToolOutput`] instead of serializing it into JSON text. The default
    /// output schema is permissive because rich presentation content and
    /// optional structured content are validated by the canonical Tool
    /// boundary rather than one generated Rust output type.
    pub fn new_rich(
        name: impl Into<String>,
        description: impl Into<String>,
        handler: Handler,
    ) -> Self {
        Self {
            descriptor: ToolDescriptor {
                id: CapabilityId::new(),
                name: name.into(),
                version: "1".into(),
                description: description.into(),
                input_schema: schema_for!(Input).to_value(),
                output_schema: serde_json::json!({}),
                effect: EffectClass::Pure,
                risk: RiskLevel::Low,
                metadata: std::collections::BTreeMap::new(),
            },
            handler,
            types: PhantomData,
        }
    }
}

impl<Input, Output, Handler> FunctionTool<Input, Output, Handler> {
    /// Replaces the stable capability identity.
    #[must_use]
    pub const fn capability_id(mut self, id: CapabilityId) -> Self {
        self.descriptor.id = id;
        self
    }

    /// Sets the semantic Tool contract version.
    #[must_use]
    pub fn version(mut self, version: impl Into<String>) -> Self {
        self.descriptor.version = version.into();
        self
    }

    /// Declares external-effect behavior.
    #[must_use]
    pub const fn effect(mut self, effect: EffectClass) -> Self {
        self.descriptor.effect = effect;
        self
    }

    /// Declares policy risk.
    #[must_use]
    pub const fn risk(mut self, risk: RiskLevel) -> Self {
        self.descriptor.risk = risk;
        self
    }

    /// Adds host-only namespaced metadata.
    #[must_use]
    pub fn metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
        self.descriptor.metadata.insert(key.into(), value);
        self
    }

    /// Replaces the successful output schema.
    ///
    /// This is primarily useful for rich Tools that attach typed
    /// `structured_content` alongside media. The registry compiles and
    /// enforces the schema before exposing a successful result to an Agent.
    #[must_use]
    pub fn output_schema(mut self, schema: serde_json::Value) -> Self {
        self.descriptor.output_schema = schema;
        self
    }
}

impl<Input, Output, Handler> std::fmt::Debug for FunctionTool<Input, Output, Handler> {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("FunctionTool")
            .field("descriptor", &self.descriptor)
            .finish_non_exhaustive()
    }
}

impl<Input, Output, Handler, HandlerFuture> Tool for FunctionTool<Input, Output, Handler>
where
    Input: DeserializeOwned + JsonSchema + Send + 'static,
    Output: JsonSchema + Serialize + Send + 'static,
    Handler: Fn(Input, ToolContext) -> HandlerFuture + Send + Sync,
    HandlerFuture: Future<Output = Result<Output, ToolError>> + Send + 'static,
{
    fn descriptor(&self) -> &ToolDescriptor {
        &self.descriptor
    }

    fn invoke(
        &self,
        input: serde_json::Value,
        context: ToolContext,
    ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
        let input = match decode_input(input) {
            Ok(input) => input,
            Err(error) => {
                return Box::pin(async move { Err(error) });
            }
        };
        let future = (self.handler)(input, context);
        Box::pin(async move {
            let output = future.await?;
            let value = serde_json::to_value(output).map_err(|error| {
                ToolError::local(
                    ToolErrorKind::InvalidOutput,
                    format!("typed Tool output cannot be serialized: {error}"),
                )
            })?;
            Ok(ToolOutput::model_visible(value))
        })
    }
}

impl<Input, Handler, HandlerFuture> Tool for FunctionTool<Input, ToolOutput, Handler>
where
    Input: DeserializeOwned + JsonSchema + Send + 'static,
    Handler: Fn(Input, ToolContext) -> HandlerFuture + Send + Sync,
    HandlerFuture: Future<Output = Result<ToolOutput, ToolError>> + Send + 'static,
{
    fn descriptor(&self) -> &ToolDescriptor {
        &self.descriptor
    }

    fn invoke(
        &self,
        input: serde_json::Value,
        context: ToolContext,
    ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
        let input = match decode_input(input) {
            Ok(input) => input,
            Err(error) => {
                return Box::pin(async move { Err(error) });
            }
        };
        let future = (self.handler)(input, context);
        Box::pin(future)
    }
}

fn decode_input<Input: DeserializeOwned>(input: serde_json::Value) -> Result<Input, ToolError> {
    serde_json::from_value(input).map_err(|error| {
        ToolError::local(
            ToolErrorKind::InvalidInput,
            format!("typed Tool input is invalid: {error}"),
        )
    })
}

#[cfg(test)]
mod tests {
    use std::sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    };

    use runifold_core::{Budget, BudgetTracker, CapabilitySet, RunContext};
    use runifold_model::{ContentPart, MediaSource};
    use schemars::JsonSchema;
    use serde::{Deserialize, Serialize};
    use serde_json::json;

    use super::FunctionTool;
    use crate::{Tool, ToolError, ToolErrorKind, ToolOutput, ToolRegistry};

    #[derive(Deserialize, JsonSchema)]
    struct AddInput {
        left: i64,
        right: i64,
    }

    #[derive(JsonSchema, Serialize)]
    struct AddOutput {
        sum: i64,
    }

    #[test]
    fn typed_function_generates_schemas_and_runs_through_registry() {
        let calls = Arc::new(AtomicUsize::new(0));
        let observed = calls.clone();
        let tool = Arc::new(FunctionTool::new(
            "add",
            "adds two integers",
            move |input: AddInput, _context| {
                let observed = observed.clone();
                async move {
                    observed.fetch_add(1, Ordering::SeqCst);
                    Ok(AddOutput {
                        sum: input.left + input.right,
                    })
                }
            },
        ));
        let descriptor = tool.descriptor();
        assert_eq!(
            descriptor.input_schema["required"],
            json!(["left", "right"])
        );
        assert_eq!(descriptor.output_schema["required"], json!(["sum"]));

        let mut capabilities = CapabilitySet::new();
        capabilities.grant(descriptor.capability());
        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
        let mut registry = ToolRegistry::new();
        registry.register(tool).unwrap();

        let output = futures_executor::block_on(registry.invoke(
            "add",
            json!({"left": 2, "right": 3}),
            &run,
        ))
        .unwrap();

        assert_eq!(output.structured_content, Some(json!({"sum": 5})));
        assert_eq!(calls.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn invalid_typed_input_never_calls_handler() {
        let calls = Arc::new(AtomicUsize::new(0));
        let observed = calls.clone();
        let tool = Arc::new(FunctionTool::new(
            "add",
            "adds two integers",
            move |_input: AddInput, _context| {
                let observed = observed.clone();
                async move {
                    observed.fetch_add(1, Ordering::SeqCst);
                    Ok(AddOutput { sum: 0 })
                }
            },
        ));
        let mut capabilities = CapabilitySet::new();
        capabilities.grant(tool.descriptor().capability());
        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
        let mut registry = ToolRegistry::new();
        registry.register(tool).unwrap();

        let error = futures_executor::block_on(registry.invoke("add", json!({"left": 2}), &run))
            .unwrap_err();

        assert_eq!(error.kind, ToolErrorKind::InvalidInput);
        assert_eq!(calls.load(Ordering::SeqCst), 0);
    }

    #[test]
    fn rich_function_preserves_image_and_structured_content() {
        let tool = Arc::new(
            FunctionTool::new_rich(
                "kline",
                "returns a K-line chart",
                |input: AddInput, _context| async move {
                    Ok::<_, ToolError>(
                        ToolOutput::rich(vec![
                            ContentPart::text("K-line chart"),
                            ContentPart::Image {
                                source: MediaSource::Url {
                                    url: "https://example.com/kline.png".into(),
                                    media_type: Some("image/png".into()),
                                },
                            },
                        ])
                        .with_structured_content(json!({
                            "left": input.left,
                            "right": input.right,
                        })),
                    )
                },
            )
            .output_schema(json!({
                "type": "object",
                "required": ["left", "right"],
                "properties": {
                    "left": { "type": "integer" },
                    "right": { "type": "integer" }
                }
            })),
        );
        let mut capabilities = CapabilitySet::new();
        capabilities.grant(tool.descriptor().capability());
        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
        let mut registry = ToolRegistry::new();
        registry.register(tool).unwrap();

        let output = futures_executor::block_on(registry.invoke(
            "kline",
            json!({"left": 20, "right": 22}),
            &run,
        ))
        .unwrap();

        assert_eq!(
            output.structured_content,
            Some(json!({"left": 20, "right": 22}))
        );
        assert!(matches!(
            &output.content[1],
            ContentPart::Image {
                source: MediaSource::Url { media_type, .. }
            } if media_type.as_deref() == Some("image/png")
        ));
    }
}