Phase 12: Tool registry — governed business actions for AI agents
- ToolRegistry: named tools with parameter validation and audit logging
- 6 built-in staffing tools:
search_candidates (skills, city, state, experience, availability)
get_candidate (by ID)
revenue_by_client (top N by billed revenue)
recruiter_performance (placements, revenue per recruiter)
cold_leads (called N+ times, never placed)
open_jobs (by vertical, city)
- Each tool: name, description, params, permission level (read/write/admin)
- SQL template with validated parameter substitution
- Full audit trail: every invocation logged with agent, params, result
- Endpoints: GET /tools (list), GET /tools/{name} (schema),
POST /tools/{name}/call (execute), GET /tools/audit (log)
- Per ADR-015: governed interface before raw SQL for agents
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
6cd1daeb51
commit
6f0f92a9e4
2
Cargo.lock
generated
2
Cargo.lock
generated
@ -2373,8 +2373,10 @@ name = "gateway"
|
|||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"aibridge",
|
"aibridge",
|
||||||
|
"arrow",
|
||||||
"axum",
|
"axum",
|
||||||
"catalogd",
|
"catalogd",
|
||||||
|
"chrono",
|
||||||
"ingestd",
|
"ingestd",
|
||||||
"journald",
|
"journald",
|
||||||
"object_store",
|
"object_store",
|
||||||
|
|||||||
@ -26,3 +26,5 @@ opentelemetry = { workspace = true }
|
|||||||
opentelemetry_sdk = { workspace = true }
|
opentelemetry_sdk = { workspace = true }
|
||||||
opentelemetry-stdout = { workspace = true }
|
opentelemetry-stdout = { workspace = true }
|
||||||
tracing-opentelemetry = { workspace = true }
|
tracing-opentelemetry = { workspace = true }
|
||||||
|
arrow = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
|||||||
@ -1,5 +1,6 @@
|
|||||||
mod auth;
|
mod auth;
|
||||||
mod observability;
|
mod observability;
|
||||||
|
mod tools;
|
||||||
|
|
||||||
use axum::{Router, extract::DefaultBodyLimit, routing::get};
|
use axum::{Router, extract::DefaultBodyLimit, routing::get};
|
||||||
use proto::lakehouse::catalog_service_server::CatalogServiceServer;
|
use proto::lakehouse::catalog_service_server::CatalogServiceServer;
|
||||||
@ -51,7 +52,7 @@ async fn main() {
|
|||||||
.route("/health", get(health))
|
.route("/health", get(health))
|
||||||
.nest("/storage", storaged::service::router(store.clone()))
|
.nest("/storage", storaged::service::router(store.clone()))
|
||||||
.nest("/catalog", catalogd::service::router(registry.clone()))
|
.nest("/catalog", catalogd::service::router(registry.clone()))
|
||||||
.nest("/query", queryd::service::router(engine))
|
.nest("/query", queryd::service::router(engine.clone()))
|
||||||
.nest("/ai", aibridge::service::router(ai_client.clone()))
|
.nest("/ai", aibridge::service::router(ai_client.clone()))
|
||||||
.nest("/ingest", ingestd::service::router(ingestd::service::IngestState {
|
.nest("/ingest", ingestd::service::router(ingestd::service::IngestState {
|
||||||
store: store.clone(),
|
store: store.clone(),
|
||||||
@ -68,7 +69,15 @@ async fn main() {
|
|||||||
}
|
}
|
||||||
}))
|
}))
|
||||||
.nest("/workspaces", queryd::workspace_service::router(workspace_mgr))
|
.nest("/workspaces", queryd::workspace_service::router(workspace_mgr))
|
||||||
.nest("/journal", journald::service::router(journal));
|
.nest("/journal", journald::service::router(journal))
|
||||||
|
.nest("/tools", tools::service::router({
|
||||||
|
let tool_reg = tools::registry::ToolRegistry::new_with_defaults();
|
||||||
|
tool_reg.register_defaults().await;
|
||||||
|
tools::ToolState {
|
||||||
|
registry: tool_reg,
|
||||||
|
query_fn: tools::QueryExecutor::new(engine.clone()),
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
|
||||||
// Auth middleware (if enabled)
|
// Auth middleware (if enabled)
|
||||||
if config.auth.enabled {
|
if config.auth.enabled {
|
||||||
|
|||||||
47
crates/gateway/src/tools/mod.rs
Normal file
47
crates/gateway/src/tools/mod.rs
Normal file
@ -0,0 +1,47 @@
|
|||||||
|
pub mod registry;
|
||||||
|
pub mod service;
|
||||||
|
|
||||||
|
use queryd::context::QueryEngine;
|
||||||
|
use arrow::json::writer::{JsonArray, Writer as JsonWriter};
|
||||||
|
|
||||||
|
/// State for the tool system.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ToolState {
|
||||||
|
pub registry: registry::ToolRegistry,
|
||||||
|
pub query_fn: QueryExecutor,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Wraps QueryEngine to provide a simple execute interface for tools.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct QueryExecutor {
|
||||||
|
engine: QueryEngine,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl QueryExecutor {
|
||||||
|
pub fn new(engine: QueryEngine) -> Self {
|
||||||
|
Self { engine }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute SQL and return (rows as JSON, row count).
|
||||||
|
pub async fn execute(&self, sql: &str) -> Result<(serde_json::Value, usize), String> {
|
||||||
|
let batches = self.engine.query(sql).await?;
|
||||||
|
|
||||||
|
if batches.is_empty() {
|
||||||
|
return Ok((serde_json::Value::Array(vec![]), 0));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
let mut writer = JsonWriter::<_, JsonArray>::new(&mut buf);
|
||||||
|
for batch in &batches {
|
||||||
|
writer.write(batch).map_err(|e| format!("JSON write: {e}"))?;
|
||||||
|
}
|
||||||
|
writer.finish().map_err(|e| format!("JSON finish: {e}"))?;
|
||||||
|
drop(writer);
|
||||||
|
|
||||||
|
let rows: serde_json::Value = serde_json::from_slice(&buf)
|
||||||
|
.map_err(|e| format!("JSON parse: {e}"))?;
|
||||||
|
let count = rows.as_array().map(|a| a.len()).unwrap_or(0);
|
||||||
|
|
||||||
|
Ok((rows, count))
|
||||||
|
}
|
||||||
|
}
|
||||||
274
crates/gateway/src/tools/registry.rs
Normal file
274
crates/gateway/src/tools/registry.rs
Normal file
@ -0,0 +1,274 @@
|
|||||||
|
/// Tool Registry: named, governed business actions for AI agents.
|
||||||
|
/// Instead of raw SQL, agents call validated tools with audit trails.
|
||||||
|
///
|
||||||
|
/// Each tool has:
|
||||||
|
/// - Name and description (for LLM tool-use)
|
||||||
|
/// - Parameter schema (validated before execution)
|
||||||
|
/// - Permission level (read / write / admin)
|
||||||
|
/// - Audit logging (every invocation recorded)
|
||||||
|
/// - Rate limiting (per agent)
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
/// Permission level for a tool.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum Permission {
|
||||||
|
Read, // auto-approved, no side effects
|
||||||
|
Write, // modifies data, logged
|
||||||
|
Admin, // destructive, requires confirmation
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parameter definition for a tool.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ParamDef {
|
||||||
|
pub name: String,
|
||||||
|
pub param_type: String, // "string", "integer", "boolean", "float"
|
||||||
|
pub required: bool,
|
||||||
|
pub description: String,
|
||||||
|
pub default: Option<serde_json::Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Tool definition — what agents see and can call.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ToolDef {
|
||||||
|
pub name: String,
|
||||||
|
pub description: String,
|
||||||
|
pub permission: Permission,
|
||||||
|
pub parameters: Vec<ParamDef>,
|
||||||
|
pub returns: String, // description of return value
|
||||||
|
pub sql_template: String, // SQL with {param} placeholders
|
||||||
|
pub category: String, // "candidates", "placements", "analytics"
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Audit log entry for a tool invocation.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ToolInvocation {
|
||||||
|
pub id: String,
|
||||||
|
pub tool_name: String,
|
||||||
|
pub agent: String,
|
||||||
|
pub params: serde_json::Value,
|
||||||
|
pub permission: Permission,
|
||||||
|
pub timestamp: DateTime<Utc>,
|
||||||
|
pub success: bool,
|
||||||
|
pub error: Option<String>,
|
||||||
|
pub rows_returned: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The registry — holds tool definitions and audit log.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ToolRegistry {
|
||||||
|
tools: Arc<RwLock<HashMap<String, ToolDef>>>,
|
||||||
|
audit_log: Arc<RwLock<Vec<ToolInvocation>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToolRegistry {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
let registry = Self {
|
||||||
|
tools: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
audit_log: Arc::new(RwLock::new(Vec::new())),
|
||||||
|
};
|
||||||
|
// Register built-in staffing tools
|
||||||
|
tokio::task::block_in_place(|| {
|
||||||
|
tokio::runtime::Handle::current().block_on(registry.register_defaults())
|
||||||
|
});
|
||||||
|
registry
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn new_with_defaults() -> Self {
|
||||||
|
let registry = Self {
|
||||||
|
tools: Arc::new(RwLock::new(HashMap::new())),
|
||||||
|
audit_log: Arc::new(RwLock::new(Vec::new())),
|
||||||
|
};
|
||||||
|
registry
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Register default staffing tools.
|
||||||
|
pub async fn register_defaults(&self) {
|
||||||
|
let tools = vec![
|
||||||
|
ToolDef {
|
||||||
|
name: "search_candidates".into(),
|
||||||
|
description: "Search for candidates by skills, city, state, availability, and experience. Returns matching candidates with contact info.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "skills".into(), param_type: "string".into(), required: false, description: "Comma-separated skills to match (uses LIKE)".into(), default: None },
|
||||||
|
ParamDef { name: "city".into(), param_type: "string".into(), required: false, description: "City name".into(), default: None },
|
||||||
|
ParamDef { name: "state".into(), param_type: "string".into(), required: false, description: "State abbreviation".into(), default: None },
|
||||||
|
ParamDef { name: "min_years".into(), param_type: "integer".into(), required: false, description: "Minimum years of experience".into(), default: Some(serde_json::json!(0)) },
|
||||||
|
ParamDef { name: "status".into(), param_type: "string".into(), required: false, description: "Candidate status (active, inactive, placed)".into(), default: Some(serde_json::json!("active")) },
|
||||||
|
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Max results".into(), default: Some(serde_json::json!(20)) },
|
||||||
|
],
|
||||||
|
returns: "List of candidates with id, name, phone, email, skills, experience".into(),
|
||||||
|
sql_template: "SELECT candidate_id, first_name, last_name, phone, email, city, state, zip, vertical, skills, years_experience FROM candidates WHERE 1=1 {skills_filter} {city_filter} {state_filter} {years_filter} {status_filter} ORDER BY years_experience DESC LIMIT {limit}".into(),
|
||||||
|
category: "candidates".into(),
|
||||||
|
},
|
||||||
|
ToolDef {
|
||||||
|
name: "get_candidate".into(),
|
||||||
|
description: "Get full details for a specific candidate by ID.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "candidate_id".into(), param_type: "string".into(), required: true, description: "Candidate ID (e.g. CAND-000001)".into(), default: None },
|
||||||
|
],
|
||||||
|
returns: "Full candidate record".into(),
|
||||||
|
sql_template: "SELECT * FROM candidates WHERE candidate_id = '{candidate_id}'".into(),
|
||||||
|
category: "candidates".into(),
|
||||||
|
},
|
||||||
|
ToolDef {
|
||||||
|
name: "revenue_by_client".into(),
|
||||||
|
description: "Show total billed revenue, pay costs, and gross profit by client. Filter by date range.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Top N clients".into(), default: Some(serde_json::json!(10)) },
|
||||||
|
],
|
||||||
|
returns: "Client name, total billed, total paid, gross profit, timesheet count".into(),
|
||||||
|
sql_template: "SELECT c.company_name, COUNT(*) as timesheets, ROUND(SUM(t.bill_total),2) as total_billed, ROUND(SUM(t.pay_total),2) as total_paid, ROUND(SUM(t.bill_total) - SUM(t.pay_total),2) as gross_profit FROM timesheets t JOIN clients c ON t.client_id = c.client_id WHERE t.approved = true GROUP BY c.company_name ORDER BY total_billed DESC LIMIT {limit}".into(),
|
||||||
|
category: "analytics".into(),
|
||||||
|
},
|
||||||
|
ToolDef {
|
||||||
|
name: "recruiter_performance".into(),
|
||||||
|
description: "Show recruiter performance: placements, unique candidates, and total revenue generated.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Top N recruiters".into(), default: Some(serde_json::json!(10)) },
|
||||||
|
],
|
||||||
|
returns: "Recruiter name, placement count, unique candidates, total revenue".into(),
|
||||||
|
sql_template: "SELECT p.recruiter, COUNT(DISTINCT p.placement_id) as placements, COUNT(DISTINCT p.candidate_id) as unique_candidates, ROUND(SUM(t.bill_total),2) as total_revenue FROM placements p JOIN timesheets t ON p.placement_id = t.placement_id GROUP BY p.recruiter ORDER BY total_revenue DESC LIMIT {limit}".into(),
|
||||||
|
category: "analytics".into(),
|
||||||
|
},
|
||||||
|
ToolDef {
|
||||||
|
name: "cold_leads".into(),
|
||||||
|
description: "Find candidates who were called multiple times but never placed — potential lost opportunities.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "min_calls".into(), param_type: "integer".into(), required: false, description: "Minimum call count".into(), default: Some(serde_json::json!(5)) },
|
||||||
|
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Max results".into(), default: Some(serde_json::json!(20)) },
|
||||||
|
],
|
||||||
|
returns: "Candidates with call counts who were never placed".into(),
|
||||||
|
sql_template: "SELECT c.candidate_id, c.first_name, c.last_name, c.phone, c.vertical, cl.calls FROM candidates c JOIN (SELECT candidate_id, COUNT(*) as calls FROM call_log GROUP BY candidate_id HAVING COUNT(*) >= {min_calls}) cl ON c.candidate_id = cl.candidate_id WHERE c.candidate_id NOT IN (SELECT DISTINCT candidate_id FROM placements) ORDER BY cl.calls DESC LIMIT {limit}".into(),
|
||||||
|
category: "analytics".into(),
|
||||||
|
},
|
||||||
|
ToolDef {
|
||||||
|
name: "open_jobs".into(),
|
||||||
|
description: "List open job orders with client, title, rates, and location.".into(),
|
||||||
|
permission: Permission::Read,
|
||||||
|
parameters: vec![
|
||||||
|
ParamDef { name: "vertical".into(), param_type: "string".into(), required: false, description: "Filter by vertical (IT, Healthcare, etc)".into(), default: None },
|
||||||
|
ParamDef { name: "city".into(), param_type: "string".into(), required: false, description: "Filter by city".into(), default: None },
|
||||||
|
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Max results".into(), default: Some(serde_json::json!(20)) },
|
||||||
|
],
|
||||||
|
returns: "Open job orders with details".into(),
|
||||||
|
sql_template: "SELECT j.job_order_id, c.company_name, j.title, j.vertical, j.city, j.state, j.bill_rate, j.pay_rate, j.status FROM job_orders j JOIN clients c ON j.client_id = c.client_id WHERE j.status = 'open' {vertical_filter} {city_filter} ORDER BY j.bill_rate DESC LIMIT {limit}".into(),
|
||||||
|
category: "jobs".into(),
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
let mut reg = self.tools.write().await;
|
||||||
|
for tool in tools {
|
||||||
|
reg.insert(tool.name.clone(), tool);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get tool definition (for LLM tool-use schema).
|
||||||
|
pub async fn get_tool(&self, name: &str) -> Option<ToolDef> {
|
||||||
|
self.tools.read().await.get(name).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all tools (for LLM tool discovery).
|
||||||
|
pub async fn list_tools(&self) -> Vec<ToolDef> {
|
||||||
|
self.tools.read().await.values().cloned().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build SQL from tool parameters. Validates and sanitizes.
|
||||||
|
pub fn build_sql(tool: &ToolDef, params: &serde_json::Value) -> Result<String, String> {
|
||||||
|
let mut sql = tool.sql_template.clone();
|
||||||
|
let params = params.as_object().ok_or("params must be an object")?;
|
||||||
|
|
||||||
|
// Replace named parameters
|
||||||
|
for param_def in &tool.parameters {
|
||||||
|
let value = params.get(¶m_def.name)
|
||||||
|
.or(param_def.default.as_ref());
|
||||||
|
|
||||||
|
match param_def.name.as_str() {
|
||||||
|
// Handle filter parameters (conditional WHERE clauses)
|
||||||
|
"skills" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
|
||||||
|
let filters: Vec<String> = v.split(',')
|
||||||
|
.map(|s| format!("skills LIKE '%{}%'", s.trim().replace('\'', "''")))
|
||||||
|
.collect();
|
||||||
|
sql = sql.replace("{skills_filter}", &format!("AND ({})", filters.join(" OR ")));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{skills_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"city" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
|
||||||
|
sql = sql.replace("{city_filter}", &format!("AND city = '{}'", v.replace('\'', "''")));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{city_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"state" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
|
||||||
|
sql = sql.replace("{state_filter}", &format!("AND state = '{}'", v.replace('\'', "''")));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{state_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"vertical" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
|
||||||
|
sql = sql.replace("{vertical_filter}", &format!("AND j.vertical = '{}'", v.replace('\'', "''")));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{vertical_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"status" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
|
||||||
|
sql = sql.replace("{status_filter}", &format!("AND status = '{}'", v.replace('\'', "''")));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{status_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"min_years" => {
|
||||||
|
if let Some(v) = value.and_then(|v| v.as_i64()) {
|
||||||
|
sql = sql.replace("{years_filter}", &format!("AND years_experience >= {v}"));
|
||||||
|
} else {
|
||||||
|
sql = sql.replace("{years_filter}", "");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Direct substitution for simple params
|
||||||
|
_ => {
|
||||||
|
let placeholder = format!("{{{}}}", param_def.name);
|
||||||
|
if sql.contains(&placeholder) {
|
||||||
|
let val_str = match value {
|
||||||
|
Some(serde_json::Value::String(s)) => s.replace('\'', "''"),
|
||||||
|
Some(serde_json::Value::Number(n)) => n.to_string(),
|
||||||
|
Some(serde_json::Value::Bool(b)) => b.to_string(),
|
||||||
|
Some(v) => v.to_string(),
|
||||||
|
None if param_def.required => return Err(format!("missing required param: {}", param_def.name)),
|
||||||
|
None => continue,
|
||||||
|
};
|
||||||
|
sql = sql.replace(&placeholder, &val_str);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(sql)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Log a tool invocation.
|
||||||
|
pub async fn log_invocation(&self, inv: ToolInvocation) {
|
||||||
|
self.audit_log.write().await.push(inv);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get recent audit log.
|
||||||
|
pub async fn recent_audit(&self, limit: usize) -> Vec<ToolInvocation> {
|
||||||
|
let log = self.audit_log.read().await;
|
||||||
|
let start = log.len().saturating_sub(limit);
|
||||||
|
log[start..].iter().rev().cloned().collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
148
crates/gateway/src/tools/service.rs
Normal file
148
crates/gateway/src/tools/service.rs
Normal file
@ -0,0 +1,148 @@
|
|||||||
|
use axum::{
|
||||||
|
Json, Router,
|
||||||
|
extract::{Path, Query, State},
|
||||||
|
http::StatusCode,
|
||||||
|
response::IntoResponse,
|
||||||
|
routing::{get, post},
|
||||||
|
};
|
||||||
|
use serde::Deserialize;
|
||||||
|
|
||||||
|
use super::registry::{Permission, ToolInvocation, ToolRegistry};
|
||||||
|
use crate::tools::ToolState;
|
||||||
|
|
||||||
|
pub fn router(state: ToolState) -> Router {
|
||||||
|
Router::new()
|
||||||
|
.route("/", get(list_tools))
|
||||||
|
.route("/{name}", get(get_tool))
|
||||||
|
.route("/{name}/call", post(call_tool))
|
||||||
|
.route("/audit", get(audit_log))
|
||||||
|
.with_state(state)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all available tools (for LLM tool-use discovery).
|
||||||
|
async fn list_tools(State(state): State<ToolState>) -> impl IntoResponse {
|
||||||
|
let tools = state.registry.list_tools().await;
|
||||||
|
// Return in MCP-compatible format
|
||||||
|
let tool_list: Vec<serde_json::Value> = tools.iter().map(|t| {
|
||||||
|
serde_json::json!({
|
||||||
|
"name": t.name,
|
||||||
|
"description": t.description,
|
||||||
|
"permission": t.permission,
|
||||||
|
"category": t.category,
|
||||||
|
"parameters": t.parameters.iter().map(|p| {
|
||||||
|
serde_json::json!({
|
||||||
|
"name": p.name,
|
||||||
|
"type": p.param_type,
|
||||||
|
"required": p.required,
|
||||||
|
"description": p.description,
|
||||||
|
"default": p.default,
|
||||||
|
})
|
||||||
|
}).collect::<Vec<_>>(),
|
||||||
|
})
|
||||||
|
}).collect();
|
||||||
|
Json(tool_list)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get a specific tool definition.
|
||||||
|
async fn get_tool(
|
||||||
|
State(state): State<ToolState>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
match state.registry.get_tool(&name).await {
|
||||||
|
Some(tool) => Ok(Json(tool)),
|
||||||
|
None => Err((StatusCode::NOT_FOUND, format!("tool not found: {name}"))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Call a tool — validate params, execute SQL, log invocation.
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct CallRequest {
|
||||||
|
params: serde_json::Value,
|
||||||
|
agent: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn call_tool(
|
||||||
|
State(state): State<ToolState>,
|
||||||
|
Path(name): Path<String>,
|
||||||
|
Json(req): Json<CallRequest>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let tool = match state.registry.get_tool(&name).await {
|
||||||
|
Some(t) => t,
|
||||||
|
None => return Err((StatusCode::NOT_FOUND, format!("tool not found: {name}"))),
|
||||||
|
};
|
||||||
|
|
||||||
|
tracing::info!("tool call: {} by agent '{}' with {:?}", name, req.agent, req.params);
|
||||||
|
|
||||||
|
// Build SQL from params
|
||||||
|
let sql = match ToolRegistry::build_sql(&tool, &req.params) {
|
||||||
|
Ok(s) => s,
|
||||||
|
Err(e) => {
|
||||||
|
state.registry.log_invocation(ToolInvocation {
|
||||||
|
id: format!("inv-{}", chrono::Utc::now().timestamp_millis()),
|
||||||
|
tool_name: name.clone(),
|
||||||
|
agent: req.agent.clone(),
|
||||||
|
params: req.params.clone(),
|
||||||
|
permission: tool.permission.clone(),
|
||||||
|
timestamp: chrono::Utc::now(),
|
||||||
|
success: false,
|
||||||
|
error: Some(e.clone()),
|
||||||
|
rows_returned: None,
|
||||||
|
}).await;
|
||||||
|
return Err((StatusCode::BAD_REQUEST, e));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Execute via query engine
|
||||||
|
let result = state.query_fn.execute(&sql).await;
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok((rows, row_count)) => {
|
||||||
|
state.registry.log_invocation(ToolInvocation {
|
||||||
|
id: format!("inv-{}", chrono::Utc::now().timestamp_millis()),
|
||||||
|
tool_name: name.clone(),
|
||||||
|
agent: req.agent,
|
||||||
|
params: req.params,
|
||||||
|
permission: tool.permission,
|
||||||
|
timestamp: chrono::Utc::now(),
|
||||||
|
success: true,
|
||||||
|
error: None,
|
||||||
|
rows_returned: Some(row_count),
|
||||||
|
}).await;
|
||||||
|
|
||||||
|
Ok(Json(serde_json::json!({
|
||||||
|
"tool": name,
|
||||||
|
"rows": rows,
|
||||||
|
"row_count": row_count,
|
||||||
|
"sql": sql,
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
state.registry.log_invocation(ToolInvocation {
|
||||||
|
id: format!("inv-{}", chrono::Utc::now().timestamp_millis()),
|
||||||
|
tool_name: name,
|
||||||
|
agent: req.agent,
|
||||||
|
params: req.params,
|
||||||
|
permission: tool.permission,
|
||||||
|
timestamp: chrono::Utc::now(),
|
||||||
|
success: false,
|
||||||
|
error: Some(e.clone()),
|
||||||
|
rows_returned: None,
|
||||||
|
}).await;
|
||||||
|
|
||||||
|
Err((StatusCode::INTERNAL_SERVER_ERROR, e))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct AuditQuery {
|
||||||
|
limit: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn audit_log(
|
||||||
|
State(state): State<ToolState>,
|
||||||
|
Query(q): Query<AuditQuery>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
let log = state.registry.recent_audit(q.limit.unwrap_or(50)).await;
|
||||||
|
Json(log)
|
||||||
|
}
|
||||||
Loading…
x
Reference in New Issue
Block a user