"""Tool registry: the MCP-style boundary between the agent and the world. Every tool declares an input schema, an output schema, permissions (read_only, default True), a timeout, and error behavior. Every call — success or failure — is appended to an audit trail with the arguments, a SHA-256 hash of the result, and a UTC timestamp. Write tools additionally require an explicit policy-authorization flag per call. If the `mcp` package is installed, tools/mcp_server.py exposes this same registry as a real MCP server; otherwise this registry itself is the MCP-style boundary. """ from __future__ import annotations import concurrent.futures import hashlib import json import time import uuid from datetime import datetime, timezone from typing import Any, Callable, Dict, List, Optional, Type from pydantic import BaseModel, Field, ValidationError from sensorflow.retro import store class ToolPermissionError(PermissionError): pass class ToolSpec(BaseModel): name: str description: str input_schema: Dict[str, Any] output_schema: Dict[str, Any] read_only: bool = True timeout_s: float = 10.0 error_behavior: str = "raise" # raise | return_error class AuditRecord(BaseModel): call_id: str analysis_id: str tool: str timestamp: str args: Dict[str, Any] status: str # ok | error | timeout | denied result_hash: Optional[str] # sha256 of canonical result JSON elapsed_ms: float error: Optional[str] = None authorized_write: bool = False class _RegisteredTool(BaseModel): spec: ToolSpec fn: Callable[..., BaseModel] input_model: Type[BaseModel] output_model: Type[BaseModel] model_config = {"arbitrary_types_allowed": True} def _hash_result(result: Any) -> str: blob = json.dumps(result, sort_keys=True, default=str) return hashlib.sha256(blob.encode()).hexdigest() class ToolCallResult(BaseModel): ok: bool tool: str call_id: str result: Optional[Dict[str, Any]] = None error: Optional[str] = None elapsed_ms: float = 0.0 class ToolRegistry: """Audited, permissioned, schema-validated tool dispatch.""" def __init__(self, analysis_id: Optional[str] = None, persist_audit: bool = True): self._tools: Dict[str, _RegisteredTool] = {} self.analysis_id = analysis_id or f"adhoc-{uuid.uuid4().hex[:8]}" self.persist_audit = persist_audit self.audit_log: List[AuditRecord] = [] # ------------------------------------------------------------ registration def register(self, name: str, fn: Callable[..., BaseModel], input_model: Type[BaseModel], output_model: Type[BaseModel], description: str, read_only: bool = True, timeout_s: float = 10.0, error_behavior: str = "raise") -> None: if name in self._tools: raise ValueError(f"tool {name!r} already registered") if error_behavior not in ("raise", "return_error"): raise ValueError("error_behavior must be 'raise' or 'return_error'") spec = ToolSpec(name=name, description=description, input_schema=input_model.model_json_schema(), output_schema=output_model.model_json_schema(), read_only=read_only, timeout_s=timeout_s, error_behavior=error_behavior) self._tools[name] = _RegisteredTool(spec=spec, fn=fn, input_model=input_model, output_model=output_model) def specs(self) -> List[ToolSpec]: return [t.spec for t in self._tools.values()] def get_spec(self, name: str) -> ToolSpec: return self._tools[name].spec # ------------------------------------------------------------------ audit def _audit(self, record: AuditRecord) -> None: self.audit_log.append(record) if self.persist_audit: store.append_audit(self.analysis_id, record.model_dump()) # ------------------------------------------------------------------- call def call(self, name: str, args: Dict[str, Any], policy_authorization: bool = False) -> ToolCallResult: call_id = uuid.uuid4().hex[:12] t0 = time.perf_counter() ts = datetime.now(timezone.utc).isoformat() def audit(status: str, result_hash: Optional[str] = None, error: Optional[str] = None) -> None: self._audit(AuditRecord( call_id=call_id, analysis_id=self.analysis_id, tool=name, timestamp=ts, args=args, status=status, result_hash=result_hash, elapsed_ms=round((time.perf_counter() - t0) * 1000, 2), error=error, authorized_write=policy_authorization)) if name not in self._tools: audit("error", error="unknown tool") raise KeyError(f"unknown tool {name!r}") tool = self._tools[name] if not tool.spec.read_only and not policy_authorization: audit("denied", error="write tool requires policy_authorization=True") raise ToolPermissionError( f"tool {name!r} is a WRITE tool and requires an explicit " "policy_authorization=True flag on the call") try: validated = tool.input_model(**args) except ValidationError as exc: audit("error", error=f"input validation failed: {exc}") if tool.spec.error_behavior == "return_error": return ToolCallResult(ok=False, tool=name, call_id=call_id, error=f"input validation failed: {exc}", elapsed_ms=(time.perf_counter() - t0) * 1000) raise try: with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: future = pool.submit(tool.fn, validated) output = future.result(timeout=tool.spec.timeout_s) except concurrent.futures.TimeoutError: audit("timeout", error=f"timed out after {tool.spec.timeout_s}s") if tool.spec.error_behavior == "return_error": return ToolCallResult(ok=False, tool=name, call_id=call_id, error=f"timeout after {tool.spec.timeout_s}s", elapsed_ms=(time.perf_counter() - t0) * 1000) raise TimeoutError(f"tool {name!r} timed out after {tool.spec.timeout_s}s") except Exception as exc: audit("error", error=f"{exc.__class__.__name__}: {exc}") if tool.spec.error_behavior == "return_error": return ToolCallResult(ok=False, tool=name, call_id=call_id, error=f"{exc.__class__.__name__}: {exc}", elapsed_ms=(time.perf_counter() - t0) * 1000) raise if not isinstance(output, tool.output_model): output = tool.output_model(**(output if isinstance(output, dict) else output.model_dump())) result_dict = output.model_dump() audit("ok", result_hash=_hash_result(result_dict)) return ToolCallResult(ok=True, tool=name, call_id=call_id, result=result_dict, elapsed_ms=round((time.perf_counter() - t0) * 1000, 2))