ObsiGate/backend/tools/registry.py

"""Tool registry — declarative registration and uniform execution.

Tools are registered with the :func:`tool` decorator. Each tool declares its
input model (used both for JSON Schema exposure and argument validation), its
risk level, and the fronts (in-app / MCP) it is exposed to.

Execution goes through :func:`call_tool`, which centralizes argument
validation, vault permission checks, confirmation gating, and audit logging.
"""

from __future__ import annotations

import logging
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

from pydantic import BaseModel, ValidationError

from backend.tools.audit import log_tool_call
from backend.tools.context import (
    ToolConfirmationRequired,
    ToolContext,
    ToolError,
    ToolNotFoundError,
    ToolPermissionError,
    ToolRisk,
    ToolScope,
    ToolValidationError,
)
from backend.tools.schemas import ToolResult

logger = logging.getLogger("obsigate.tools.registry")


@dataclass(frozen=True)
class ToolSpec:
    """Static description of a registered tool."""

    name: str
    description: str
    input_model: type[BaseModel]
    handler: Callable[[ToolContext, Any], Any]
    risk: ToolRisk = ToolRisk.READ
    scopes: tuple[ToolScope, ...] = (ToolScope.IN_APP, ToolScope.MCP)
    requires_vault: bool = False

    @property
    def requires_confirmation(self) -> bool:
        """Mutating tools always require an explicit confirmation."""
        return self.risk != ToolRisk.READ

    def parameters_schema(self) -> dict[str, Any]:
        """JSON Schema of the tool arguments (for LLM function calling)."""
        schema = self.input_model.model_json_schema()
        schema.pop("title", None)
        return schema

    def openai_schema(self) -> dict[str, Any]:
        """OpenAI-compatible function-calling schema."""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters_schema(),
            },
        }


_REGISTRY: dict[str, ToolSpec] = {}


def tool(
    *,
    name: str,
    description: str,
    input_model: type[BaseModel],
    risk: ToolRisk = ToolRisk.READ,
    scopes: tuple[ToolScope, ...] = (ToolScope.IN_APP, ToolScope.MCP),
    requires_vault: bool = False,
) -> Callable[[Callable[[ToolContext, Any], Any]], Callable[[ToolContext, Any], Any]]:
    """Register a tool. Returns the original handler unchanged."""

    def decorator(func: Callable[[ToolContext, Any], Any]) -> Callable[[ToolContext, Any], Any]:
        if name in _REGISTRY:
            raise ValueError(f"Duplicate tool name: {name}")
        _REGISTRY[name] = ToolSpec(
            name=name,
            description=description,
            input_model=input_model,
            handler=func,
            risk=risk,
            scopes=scopes,
            requires_vault=requires_vault,
        )
        return func

    return decorator


def get_tool(name: str) -> ToolSpec | None:
    """Return the spec for *name*, or ``None``."""
    return _REGISTRY.get(name)


def list_tools(*, scope: ToolScope | None = None) -> list[ToolSpec]:
    """Return registered tools, optionally filtered by scope."""
    specs = list(_REGISTRY.values())
    if scope is not None:
        specs = [s for s in specs if scope in s.scopes]
    return specs


def get_tool_schemas(*, scope: ToolScope | None = None) -> list[dict[str, Any]]:
    """Return OpenAI-compatible schemas for registered tools."""
    return [spec.openai_schema() for spec in list_tools(scope=scope)]


def _audit(ctx: ToolContext, spec: ToolSpec, arguments: dict[str, Any], *, ok: bool, error: str | None = None) -> None:
    if not ctx.audit_enabled:
        return
    try:
        log_tool_call(
            username=ctx.username,
            mode=ctx.mode.value,
            tool=spec.name,
            arguments=arguments,
            ok=ok,
            vault=arguments.get("vault"),
            ip=ctx.ip,
            error=error,
        )
    except Exception as e:
        logger.debug(f"Tool audit failed for '{spec.name}': {e}")


def call_tool(
    name: str,
    ctx: ToolContext,
    arguments: dict[str, Any] | None = None,
    *,
    confirm: bool = False,
) -> ToolResult:
    """Validate, authorize, execute and audit a tool call.

    Args:
        name: Registered tool name.
        ctx: Caller context (identity, mode, confirmation state).
        arguments: Raw arguments to validate against the tool input model.
        confirm: Explicit one-shot confirmation (two-step ``apply`` phase).

    Raises:
        ToolNotFoundError: Unknown tool.
        ToolValidationError: Arguments fail validation.
        ToolPermissionError: Vault access denied or path escapes the vault.
        ToolConfirmationRequired: Mutating tool invoked without confirmation.
        ToolError: Any other execution error.
    """
    spec = _REGISTRY.get(name)
    if spec is None:
        raise ToolNotFoundError(f"Unknown tool: {name}", details={"tool": name})

    arguments = arguments or {}

    try:
        params = spec.input_model.model_validate(arguments)
    except ValidationError as e:
        error = ToolValidationError(
            f"Invalid arguments for '{name}'",
            details={"errors": e.errors(include_url=False)},
        )
        _audit(ctx, spec, arguments, ok=False, error=error.code)
        raise error from e

    vault = getattr(params, "vault", None)
    if spec.requires_vault and vault:
        try:
            ctx.require_vault_access(vault)
        except ToolPermissionError as e:
            _audit(ctx, spec, arguments, ok=False, error=e.code)
            raise

    if spec.requires_confirmation and not (confirm or ctx.confirmed):
        raise ToolConfirmationRequired(spec.name, arguments)

    started = time.perf_counter()
    try:
        data = spec.handler(ctx, params)
    except ToolError as e:
        _audit(ctx, spec, arguments, ok=False, error=e.code)
        raise
    except Exception as e:
        logger.error(f"Tool '{name}' failed: {e}")
        exec_error = ToolError(f"Tool '{name}' failed: {e}", code="tool_execution_error")
        _audit(ctx, spec, arguments, ok=False, error=exec_error.code)
        raise exec_error from e

    duration_ms = round((time.perf_counter() - started) * 1000, 2)
    logger.debug(f"Tool '{name}' ok in {duration_ms}ms")
    _audit(ctx, spec, arguments, ok=True)
    return ToolResult(ok=True, data=data)