monothink 0.4.7

/*
* This file is part of monothink.
*
* monothink is free software: you can redistribute it and/or modify it under the terms
* of the GNU General Public License as published by the Free Software Foundation,
* either version 3 of the License, or (at your option) any later version.
*
* monothink is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
* See the GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License along with Wavething.
* If not, see <https://www.gnu.org/licenses/>.

*/

pub mod find_package;
pub mod list_dir;
pub mod read_file;
pub mod retrieve_observation;
pub mod search_text;
pub mod web_search;

use crate::config::{self};
use rmcp::{
    RoleClient, ServiceExt,
    model::{ClientCapabilities, ClientInfo, Implementation, Tool as RmcpTool},
    service::{ClientInitializeError, RunningService},
    transport::StreamableHttpClientTransport,
};
use serde::{Deserialize, Serialize};
use std::{
    collections::BTreeMap,
    fs::{self},
    io::ErrorKind,
    path::{Path, PathBuf},
};

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
pub enum DirectoryEntryKind {
    File,
    Directory,
}

impl std::fmt::Display for DirectoryEntryKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DirectoryEntryKind::File => {
                write!(f, "File")
            }
            DirectoryEntryKind::Directory => {
                write!(f, "Directory")
            }
        }
    }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
struct DirectoryEntry {
    path: String,
    kind: DirectoryEntryKind,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct DirectoryListing {
    entries: Vec<DirectoryEntry>,
    total_entries: usize,
    truncated: bool,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
struct Location {
    path: PathBuf,
    line: usize,
    //column: u32,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
struct Range {
    path: PathBuf,
    start: usize,
    end: usize,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct FileContent {
    range: Range,
    content: String,
    information: String,
}

// path is absolute
fn list_dir(path: &Path, recursive: bool) -> Result<DirectoryListing, ToolError> {
    let mut result = DirectoryListing {
        entries: Vec::new(),
        total_entries: 0,
        truncated: false,
    };
    let mut backlog = vec![path.to_path_buf()];

    while let Some(dir) = backlog.pop() {
        if !dir.is_dir() {
            return Result::Err(ToolError::DirectoryEntryNotFound {
                path: path.to_string_lossy().to_string(),
                kind: DirectoryEntryKind::Directory,
                suggestion: "Provide a directory to the list_dir tool.".to_string(),
            });
        }

        let paths = match fs::read_dir(dir) {
            Ok(paths) => paths,
            Err(err) if err.kind() == ErrorKind::PermissionDenied => {
                // @todo Log warning and skip this directory
                continue;
            }
            Err(err) => return Err(io_error_to_tool_error(err, path)),
        };

        for entry in paths.flatten() {
            if entry.path().is_dir() {
                if recursive {
                    backlog.push(entry.path());
                }
                result.entries.push(DirectoryEntry {
                    path: entry.path().to_string_lossy().to_string(),
                    kind: DirectoryEntryKind::Directory,
                });
            } else {
                result.entries.push(DirectoryEntry {
                    path: entry.path().to_string_lossy().to_string(),
                    kind: DirectoryEntryKind::File,
                });
            }
        }
    }

    result.total_entries = result.entries.len();

    Ok(result)
}

fn allowed_canonical_path(path_string: &Path, base_dir: &Path) -> Result<PathBuf, ToolError> {
    // some LLMs take the description to verbatim and provide quotation marks for root
    let path_string = if path_string == "\"\"" {
        Path::new("")
    } else {
        path_string
    };

    let canonical_base = base_dir
        .canonicalize()
        .map_err(|err| io_error_to_tool_error(err, path_string))?;
    let path = base_dir
        .join(path_string)
        .canonicalize()
        .map_err(|err| io_error_to_tool_error(err, path_string))?;

    if path.starts_with(canonical_base) {
        Ok(path)
    } else {
        Err(ToolError::PathTraversal {
            requested: path_string.to_string_lossy().to_string(),
            base: base_dir.to_string_lossy().to_string(),
        })
    }
}

fn io_error_to_tool_error(err: std::io::Error, path: &Path) -> ToolError {
    let path_string = path.to_string_lossy().to_string();
    match err.kind() {
        ErrorKind::NotFound => ToolError::DirectoryEntryNotFound {
            path: path_string,
            kind: DirectoryEntryKind::File,
            suggestion: "Use list_dir tool to explore available files.".to_string(),
        },
        ErrorKind::PermissionDenied => ToolError::PermissionDenied { path: path_string },
        _ => ToolError::Io {
            path: Some(path_string),
            source: err,
        },
    }
}

#[derive(Debug)]
pub enum ToolError {
    /// File or directory does not exist
    DirectoryEntryNotFound {
        path: String,
        kind: DirectoryEntryKind,
        suggestion: String,
    },
    /// Invalid arguments provided to the tool
    InvalidArgs { argument: String, reason: String },
    /// Network error
    NetworkError { reason: String },
    /// A tool or resource is not available
    NotAvailable { name: String, reason: String },
    /// Entry is not a File
    NotAFile { path: String, suggestion: String },
    /// Path is outside the allowed base directory
    PathTraversal { requested: String, base: String },
    /// Permission denied accessing path
    PermissionDenied { path: String },
    /// Resource not found (e.g., observation ID)
    ResourceNotFound {
        resource_type: ResourceType,
        identifier: String,
    },
    /// Other IO error
    Io {
        path: Option<String>,
        source: std::io::Error,
    },
    /// error when communicating with a MCP server
    Mcp { kind: String, info: String }, // @todo better values
}

impl std::fmt::Display for ToolError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ToolError::DirectoryEntryNotFound {
                path,
                kind,
                suggestion,
            } => {
                write!(f, "Error: {} '{}' not found. {}", kind, path, suggestion)
            }
            ToolError::InvalidArgs { argument, reason } => {
                write!(
                    f,
                    "Error: Invalid value for argument '{}': {}",
                    argument, reason
                )
            }
            ToolError::NetworkError { reason } => {
                write!(f, "Error: '{}' when making a network request.", reason)
            }
            ToolError::NotAvailable { name, reason } => {
                write!(f, "Error: '{}' is not avaible. {}", name, reason)
            }
            ToolError::NotAFile { path, suggestion } => {
                write!(f, "Error: '{}' is not a file. {}", path, suggestion)
            }
            ToolError::PathTraversal { requested, base } => {
                write!(
                    f,
                    "Error: Path '{}' is outside the allowed directory '{}'. \
                     Only paths within the project root are accessible.",
                    requested, base
                )
            }
            ToolError::PermissionDenied { path } => {
                write!(
                    f,
                    "Error: Permission denied when accessing '{}'. \
                     This file/directory requires elevated permissions. \
                     You cannot access this resource. Continue with available data \
                     or ask the user to check permissions.",
                    path
                )
            }
            ToolError::ResourceNotFound {
                resource_type,
                identifier,
            } => {
                write!(
                    f,
                    "Error: {} with identifier '{}' not found. \
                     Check the conversation history for valid {} IDs.",
                    resource_type, identifier, resource_type
                )
            }
            ToolError::Io { path, source } => {
                if let Some(p) = path {
                    write!(f, "Error reading '{}': {}", p, source)
                } else {
                    write!(f, "Error: {}", source)
                }
            }
            ToolError::Mcp { kind, info } => {
                write!(f, "Error {} accessing MCP with info '{}'.", kind, info)
            }
        }
    }
}
impl std::error::Error for ToolError {}

#[derive(Debug, Clone)]
pub enum ResourceType {
    Observation,
    Path,
}

impl std::fmt::Display for ResourceType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ResourceType::Observation => {
                write!(f, "Observation")
            }
            ResourceType::Path => {
                write!(f, "Path")
            }
        }
    }
}

// @todo also take sender to send system messages
pub async fn get_mcp_tools(
    tools: &BTreeMap<String, config::McpTool>,
) -> Vec<(RunningService<RoleClient, ClientInfo>, Vec<RmcpTool>)> {
    let mut result = Vec::new();

    for (name, mcp_tool) in tools {
        if let Ok((peer, tools)) = connect_and_fetch_tools(mcp_tool).await {
            result.push((peer, tools));
            // @todo send system message
        } else {
            // @todo send system message about failure
        }
    }

    result
}

async fn connect_and_fetch_tools(
    mcp_tool: &config::McpTool,
) -> Result<(RunningService<RoleClient, ClientInfo>, Vec<RmcpTool>), anyhow::Error> {
    let client = create_client(mcp_tool).await?;
    let tools = client.list_tools(Default::default()).await?.tools;
    Ok((client, tools))
}

async fn create_client(
    mcp_tool: &config::McpTool,
) -> Result<RunningService<RoleClient, ClientInfo>, Box<ClientInitializeError>> {
    let transport = StreamableHttpClientTransport::from_uri(mcp_tool.url.clone());
    let client_info_impl = Implementation::new(
        env!("CARGO_PKG_NAME").to_string(),
        env!("CARGO_PKG_VERSION").to_string(),
    );
    let client_info = ClientInfo::new(ClientCapabilities::default(), client_info_impl);

    client_info.serve(transport).await.map_err(|e| {
        // @todo tracing::error!("client error: {e:?}");
        Box::new(e)
    })
    // @note: This is for rmcp 3.x which rig does not yet support
    // it is left here because I expect rig to migrate soon
    // client_info.serve_with_lifecycle(
    //     transport,
    //     rmcp::ClientLifecycleMode::Auto {
    //         preferred_versions: vec![rmcp::model::ProtocolVersion::V_2026_07_28],
    //         legacy_version: Some(rmcp::model::ProtocolVersion::V_2025_11_25),
    //     },
    // )
    // .await
}