kaynine-core 0.1.0

use crate::error::ToolError;
use crate::ids::{RunId, SessionId, ToolCallId, TurnId};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::path::PathBuf;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ToolSpec {
    pub name: String,
    pub description: String,
    pub parameters_schema: Value,
    pub concurrency: Concurrency,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum Concurrency {
    Sequential,
    ParallelSafe,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "capability")]
pub enum Capability {
    FileRead {
        path: PathBuf,
    },
    FileWrite {
        path: PathBuf,
    },
    ProcessSpawn {
        program: String,
        cwd: PathBuf,
    },
    NetworkAccess {
        origin: String,
    },
    SecretUse {
        credential: String,
    },
    ExternalMutation {
        namespace: String,
        action: String,
        resource: String,
    },
    Product {
        namespace: String,
        action: String,
        resource: Value,
    },
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct PreparedToolCall {
    pub call_id: ToolCallId,
    pub name: String,
    pub arguments: Value,
    pub capabilities: Vec<Capability>,
}

pub struct ToolContext {
    pub session_id: SessionId,
    pub run_id: RunId,
    pub turn_id: TurnId,
    pub call_id: ToolCallId,
}

pub struct ToolExecutionContext {
    pub cancel: CancellationToken,
    pub progress: mpsc::Sender<ToolProgress>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ToolProgress {
    pub call_id: ToolCallId,
    pub message: String,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolOutput {
    pub is_error: bool,
    pub text: String,
}

#[async_trait]
pub trait Tool: Send + Sync {
    fn spec(&self) -> ToolSpec;

    async fn prepare(
        &self,
        arguments: Value,
        context: &ToolContext,
    ) -> Result<PreparedToolCall, ToolError>;

    async fn execute(
        &self,
        call: PreparedToolCall,
        context: ToolExecutionContext,
    ) -> Result<ToolOutput, ToolError>;
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::ToolError;
    use crate::ids::{RunId, SessionId, ToolCallId, TurnId};
    use tokio_util::sync::CancellationToken;

    struct PassthroughTool;

    #[async_trait::async_trait]
    impl Tool for PassthroughTool {
        fn spec(&self) -> ToolSpec {
            ToolSpec {
                name: "passthrough".into(),
                description: String::new(),
                parameters_schema: serde_json::json!({"type": "object"}),
                concurrency: Concurrency::Sequential,
            }
        }

        async fn prepare(
            &self,
            arguments: Value,
            context: &ToolContext,
        ) -> Result<PreparedToolCall, ToolError> {
            Ok(PreparedToolCall {
                call_id: context.call_id.clone(),
                name: self.spec().name,
                arguments,
                capabilities: Vec::new(),
            })
        }

        async fn execute(
            &self,
            call: PreparedToolCall,
            _context: ToolExecutionContext,
        ) -> Result<ToolOutput, ToolError> {
            Ok(ToolOutput {
                is_error: false,
                text: call.arguments.to_string(),
            })
        }
    }

    #[tokio::test]
    async fn prepare_then_execute_roundtrip() {
        let tool = PassthroughTool;
        let ctx = ToolContext {
            session_id: SessionId::from("s"),
            run_id: RunId::from("r"),
            turn_id: TurnId::from("t"),
            call_id: ToolCallId::from("c1"),
        };
        let prepared = tool
            .prepare(serde_json::json!({"x": 1}), &ctx)
            .await
            .unwrap();
        assert_eq!(prepared.call_id, ToolCallId::from("c1"));
        assert_eq!(prepared.arguments, serde_json::json!({"x": 1}));
        let (tx, _rx) = tokio::sync::mpsc::channel(8);
        let out = tool
            .execute(
                prepared,
                ToolExecutionContext {
                    cancel: CancellationToken::new(),
                    progress: tx,
                },
            )
            .await
            .unwrap();
        assert_eq!(out.text, "{\"x\":1}");
    }
}