archy/core/archipelago/src/assistant/tools.rs

//! D-06: a curated, hand-written tool registry. Never derived from
//! `api::rpc::dispatcher`'s method table — every capability the chat has is
//! a deliberate decision recorded here, and the model never sees the full
//! RPC surface. No `schemars` — that crate is absent from `Cargo.toml` and
//! from 13-RESEARCH.md's Package Legitimacy Audit, so `parameters` below is
//! a hand-written JSON Schema object literal instead.

use std::collections::{BTreeSet, HashMap};

use anyhow::{Context, Result};
use serde::Deserialize;
use serde_json::{json, Value};

use super::PermissionCategory;

/// The backend-agnostic in/out of a tool invocation — the same shape
/// regardless of which adapter (Ollama/Claude/Routstr) produced it.
#[derive(Debug, Clone)]
pub struct ToolCall {
    pub id: String,
    pub name: String,
    pub arguments: Value,
}

#[derive(Debug, Clone)]
pub struct ToolResult {
    pub call_id: String,
    pub content: String,
    pub is_error: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Role {
    System,
    User,
    Assistant,
    Tool,
}

#[derive(Debug, Clone)]
pub struct ChatMessage {
    pub role: Role,
    /// Plain text, or (a future plan's) D-10-wrapped untrusted content.
    pub text: Option<String>,
    /// Assistant-authored tool calls made THIS turn (role: Assistant).
    pub tool_calls: Vec<ToolCall>,
    /// Tool results fed back THIS turn (role: Tool).
    pub tool_results: Vec<ToolResult>,
}

/// D-06: one curated, hand-written tool. Never generated from the RPC
/// dispatcher — the curated set IS the D-09 authority boundary.
#[derive(Clone)]
pub struct ToolDef {
    pub name: &'static str,
    pub description: &'static str,
    /// JSON Schema `{"type":"object","properties":{...},"required":[...]}`,
    /// hand-written and pinned adjacent to the args struct it must never
    /// drift from — see `disk_status_schema_round_trips_required_keys`.
    pub parameters: Value,
    pub category: PermissionCategory,
    /// D-07: true => confirm gate, no exceptions. There are no destructive
    /// tools in this tracer's registry; `execute_tool` refuses this branch
    /// with a not-yet-implemented error until 13-08 fills it in.
    pub destructive: bool,
}

/// Args for `system_disk_status` — takes no parameters.
#[derive(Debug, Deserialize)]
pub struct SystemDiskStatusArgs {}

impl ToolDef {
    /// Deserialize + validate model-produced arguments before ANY
    /// execution. Never coerce, never guess, never panic on a mismatch —
    /// refuse and let the caller turn the error into a tool result the
    /// model can recover from.
    ///
    /// This tracer's registry has exactly one tool, so this is a direct
    /// deserialize; a future plan adding a second tool dispatches by
    /// `self.name` here before deserializing into that tool's own args type.
    pub fn validate(&self, raw: &Value) -> Result<SystemDiskStatusArgs> {
        serde_json::from_value(raw.clone())
            .context("tool arguments did not match the declared schema")
    }
}

/// `system_disk_status` — category `System`, read-only. Reports free and
/// total disk space on this node via the same `system.disk-status` handler
/// every other authenticated caller uses.
pub fn system_disk_status_tool() -> ToolDef {
    ToolDef {
        name: "system_disk_status",
        description: "Report free and total disk space on this Archipelago node.",
        parameters: json!({
            "type": "object",
            "properties": {},
            "required": [],
        }),
        category: PermissionCategory::System,
        destructive: false,
    }
}

/// D-06's curated allowlist, name-indexed.
pub struct ToolRegistry {
    tools: HashMap<&'static str, ToolDef>,
}

impl ToolRegistry {
    pub fn get(&self, name: &str) -> Option<&ToolDef> {
        self.tools.get(name)
    }

    /// The subset of the registry visible to a caller with `grants`. D-16:
    /// an unconfigured node's system prompt should advertise close to zero
    /// tools — the model should never even see a tool it can't use.
    pub fn visible_to(&self, grants: &BTreeSet<PermissionCategory>) -> Vec<ToolDef> {
        self.tools
            .values()
            .filter(|t| grants.contains(&t.category))
            .cloned()
            .collect()
    }
}

/// The curated D-06 registry. This tracer registers exactly one tool.
pub fn registry() -> ToolRegistry {
    let mut tools = HashMap::new();
    let tool = system_disk_status_tool();
    tools.insert(tool.name, tool);
    ToolRegistry { tools }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The schema sent to the model and the struct used to deserialize its
    /// output must never silently drift apart. Round-trip the schema's
    /// declared `required` keys through the args struct.
    #[test]
    fn disk_status_schema_round_trips_required_keys() {
        let tool = system_disk_status_tool();
        let required = tool
            .parameters
            .get("required")
            .and_then(|r| r.as_array())
            .cloned()
            .unwrap_or_default();

        let mut obj = serde_json::Map::new();
        for key in &required {
            if let Some(k) = key.as_str() {
                obj.insert(k.to_string(), Value::Null);
            }
        }
        let value = Value::Object(obj);
        let parsed: Result<SystemDiskStatusArgs, _> = serde_json::from_value(value);
        assert!(parsed.is_ok(), "schema/args struct drift: {:?}", parsed.err());
    }

    #[test]
    fn registry_visible_to_respects_grants() {
        let reg = registry();
        let mut grants = BTreeSet::new();
        assert!(reg.visible_to(&grants).is_empty());
        grants.insert(PermissionCategory::System);
        assert_eq!(reg.visible_to(&grants).len(), 1);
    }
}