lakehouse/crates/gateway/src/tools/service.rs

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)
}