harn-serve 0.8.84

use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::path::{Path, PathBuf};

use harn_parser::{Attribute, Node, TypeExpr};

use crate::limits::{limits_and_budget_from_attributes, BudgetSpec, RouteLimits};
use crate::DispatchError;

#[derive(Clone, Debug, PartialEq)]
pub struct ExportedParam {
    pub name: String,
    pub type_expr: Option<TypeExpr>,
    pub input_schema: serde_json::Value,
    pub has_default: bool,
    pub rest: bool,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ExportedCallableKind {
    Function,
    Pipeline,
}

#[derive(Clone, Debug)]
pub struct ExportedFunction {
    pub name: String,
    pub kind: ExportedCallableKind,
    pub params: Vec<ExportedParam>,
    pub return_type: Option<TypeExpr>,
    pub input_schema: serde_json::Value,
    pub output_schema: Option<serde_json::Value>,
    /// Scopes the caller's credential must carry to invoke this function.
    /// Populated from `@scopes("...", "...")` attribute literals on the
    /// declaration; empty when no attribute is present, meaning the route
    /// is unrestricted beyond whatever scopes the auth method enforces
    /// globally.
    pub required_scopes: BTreeSet<String>,
    /// Rate / backpressure ceilings declared via `@limits(...)`. `None`
    /// when the route is unbounded — the dispatch path short-circuits
    /// cheaply when both `limits` and `budget` are absent.
    pub limits: Option<RouteLimits>,
    /// Per-dispatch resource budget declared via `@budget(...)` (LLM
    /// cost / token / pg query / MCP call ceilings). `None` when no
    /// budget caps were declared.
    pub budget: Option<BudgetSpec>,
    /// HTTP route this function answers when hosted by `harn serve site`.
    /// Populated from a `@route("METHOD", "/path")` attribute, or
    /// inferred from a `handler_*` naming convention when the attribute is
    /// absent. `None` for functions that are dispatch-only (API/A2A/MCP)
    /// and not meant to be reached over a bare HTTP path.
    pub route: Option<RouteSpec>,
}

/// An HTTP method + path a `.harn` handler answers under `harn serve
/// site`. Declared with `@route("GET", "/users/{id}")` or inferred from
/// the `handler_*` naming convention.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RouteSpec {
    /// Uppercased HTTP method (`GET`, `POST`, …), or `*` to answer every
    /// method on the path — the handler inspects `req.method` itself.
    pub method: String,
    /// axum-style path with `{param}` captures, always rooted at `/`.
    pub path: String,
}

/// A `HARN-SRV-*` diagnostic raised while building the export catalog.
///
/// These flag the malformed `@route(...)` / `@scopes(...)` attribute
/// forms that the collector would otherwise drop silently — leaving a
/// handler mis-routed, unmounted, or less scope-restricted than the
/// author intended. They are surfaced by the serve adapters at startup
/// (see [`emit_export_diagnostics`]) rather than aborting catalog
/// construction, so one bad attribute doesn't take down a script whose
/// other handlers are fine.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ExportDiagnostic {
    /// Stable code so log scanners and editors can key on the condition.
    pub code: &'static str,
    /// 1-based source line of the offending attribute (0 when unknown).
    pub line: usize,
    pub message: String,
}

/// `@route` carries an argument that is not a string literal, so the
/// method/path positions are ambiguous and the handler is not mounted.
pub const ROUTE_ARG_NOT_STRING: &str = "HARN-SRV-001";
/// `@route` has the wrong number of arguments — it takes a path, or a
/// method and a path. The handler is not mounted.
pub const ROUTE_BAD_ARITY: &str = "HARN-SRV-002";
/// `@scopes` carries an argument that is not a string literal; that
/// scope requirement is dropped, leaving the route less restricted.
pub const SCOPES_ARG_NOT_STRING: &str = "HARN-SRV-003";

impl std::fmt::Display for ExportDiagnostic {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.line > 0 {
            write!(f, "{}: {} (line {})", self.code, self.message, self.line)
        } else {
            write!(f, "{}: {}", self.code, self.message)
        }
    }
}

/// Print catalog diagnostics to stderr at server startup, matching the
/// `[harn] …` banner the adapters already emit. Standalone serve
/// commands call this so authors see malformed attributes immediately;
/// embedders that build a router directly can read
/// [`ExportCatalog::diagnostics`] and render them in their own UI.
pub fn emit_export_diagnostics(diagnostics: &[ExportDiagnostic]) {
    for diagnostic in diagnostics {
        eprintln!("[harn] warning: {diagnostic}");
    }
}

#[derive(Clone, Debug)]
pub struct ExportCatalog {
    pub script_path: PathBuf,
    pub functions: BTreeMap<String, ExportedFunction>,
    /// Non-fatal `HARN-SRV-*` diagnostics gathered while collecting the
    /// route/scope attributes. Empty for a well-formed script.
    pub diagnostics: Vec<ExportDiagnostic>,
}

impl ExportCatalog {
    pub fn from_path(path: &Path) -> Result<Self, DispatchError> {
        let source = fs::read_to_string(path).map_err(|error| {
            DispatchError::Io(format!("failed to read {}: {error}", path.display()))
        })?;
        let program = harn_parser::parse_source(&source).map_err(|error| {
            DispatchError::Validation(format!("failed to parse {}: {error}", path.display()))
        })?;

        let mut functions = BTreeMap::new();
        let mut diagnostics = Vec::new();
        for node in &program {
            let (attrs, inner) = harn_parser::peel_attributes(node);
            let Node::FnDecl {
                name,
                params,
                return_type,
                is_pub,
                ..
            } = &inner.node
            else {
                continue;
            };
            if !*is_pub {
                continue;
            }

            let (limits, budget) = limits_and_budget_from_attributes(attrs);
            functions.insert(
                name.clone(),
                ExportedFunction {
                    name: name.clone(),
                    kind: ExportedCallableKind::Function,
                    params: exported_params(params),
                    return_type: return_type.clone(),
                    input_schema: harn_vm::json_schema_for_typed_params(params),
                    output_schema: return_type
                        .as_ref()
                        .and_then(harn_vm::json_schema_for_type_expr),
                    required_scopes: scopes_from_attributes(attrs, name, &mut diagnostics),
                    limits,
                    budget,
                    route: route_from_attributes(attrs, name, &mut diagnostics),
                },
            );
        }

        let has_public_exports = !functions.is_empty();
        for node in &program {
            let (attrs, inner) = harn_parser::peel_attributes(node);
            let Node::Pipeline {
                name,
                params,
                return_type,
                is_pub,
                ..
            } = &inner.node
            else {
                continue;
            };
            if has_public_exports && !*is_pub {
                continue;
            }
            let required_scopes = scopes_from_attributes(attrs, name, &mut diagnostics);
            let (limits, budget) = limits_and_budget_from_attributes(attrs);
            functions
                .entry(name.clone())
                .or_insert_with(|| ExportedFunction {
                    name: name.clone(),
                    kind: ExportedCallableKind::Pipeline,
                    params: pipeline_exported_params(params),
                    return_type: return_type.clone(),
                    input_schema: pipeline_input_schema(params),
                    output_schema: return_type
                        .as_ref()
                        .and_then(harn_vm::json_schema_for_type_expr),
                    required_scopes,
                    limits,
                    budget,
                    // Pipelines are dispatch-only; they never carry an
                    // HTTP route. Only `pub fn` handlers participate in
                    // `harn serve site`.
                    route: None,
                });
        }

        Ok(Self {
            script_path: path.to_path_buf(),
            functions,
            diagnostics,
        })
    }

    pub fn function(&self, name: &str) -> Option<&ExportedFunction> {
        self.functions.get(name)
    }

    /// Non-fatal `HARN-SRV-*` diagnostics gathered while collecting the
    /// route/scope attributes. Empty for a well-formed script.
    pub fn diagnostics(&self) -> &[ExportDiagnostic] {
        &self.diagnostics
    }
}

fn exported_params(params: &[harn_parser::TypedParam]) -> Vec<ExportedParam> {
    params
        .iter()
        .map(|param| ExportedParam {
            name: param.name.clone(),
            type_expr: param.type_expr.clone(),
            input_schema: param
                .type_expr
                .as_ref()
                .and_then(harn_vm::json_schema_for_type_expr)
                .unwrap_or_else(|| serde_json::json!({})),
            has_default: param.default_value.is_some(),
            rest: param.rest,
        })
        .collect()
}

fn pipeline_exported_params(params: &[String]) -> Vec<ExportedParam> {
    params
        .iter()
        .map(|name| ExportedParam {
            name: name.clone(),
            type_expr: None,
            input_schema: serde_json::json!({}),
            has_default: false,
            rest: false,
        })
        .collect()
}

/// Collect scope literals from any `@scopes(...)` attributes on a
/// declaration. Both positional and named arguments are accepted (named
/// args are useful for ergonomics like `@scopes(read: "personas:read")`
/// in callers that prefer key-value form); only string literals
/// contribute. Multiple `@scopes` attributes on the same declaration
/// union into one set.
fn scopes_from_attributes(
    attrs: &[Attribute],
    fn_name: &str,
    diagnostics: &mut Vec<ExportDiagnostic>,
) -> BTreeSet<String> {
    let mut set = BTreeSet::new();
    for attr in attrs {
        if attr.name != "scopes" {
            continue;
        }
        for arg in &attr.args {
            match &arg.value.node {
                Node::StringLiteral(value) | Node::RawStringLiteral(value) => {
                    set.insert(value.clone());
                }
                // A non-string scope is silently dropped by the
                // collector, which would leave the route *less*
                // restricted than the author wrote — worth a loud warning.
                _ => diagnostics.push(ExportDiagnostic {
                    code: SCOPES_ARG_NOT_STRING,
                    line: arg.span.line,
                    message: format!(
                        "`@scopes` on `{fn_name}` requires string-literal arguments; \
                         dropping a non-string scope leaves the route less restricted"
                    ),
                }),
            }
        }
    }
    set
}

/// Resolve the HTTP route a `pub fn` answers under `harn serve site`.
///
/// Two ways to declare one, in priority order:
///
/// 1. An explicit `@route("METHOD", "/path")` attribute. The first
///    positional string is the method (case-insensitive; `"*"` or
///    `"ANY"` matches every method), the second is the path. A
///    single-argument form `@route("/path")` defaults the method to
///    `GET`. Paths are normalized to start with `/`.
/// 2. The `handler_<name>` naming convention. `pub fn handler_health()`
///    is mounted at `GET|POST /health`; a bare `pub fn handler()` mounts
///    at the site root `/`. This keeps the zero-config path the issue
///    calls for ("mounts every exported `pub fn handler_*` at `/<name>`")
///    while letting authors opt into precise routing with the attribute.
///
/// A present-but-malformed `@route` does not fall back to the naming
/// convention: it records a `HARN-SRV-*` diagnostic and returns `None`,
/// so the author sees the mistake instead of a silently different route.
///
/// Returns `None` for any other `pub fn`, so a script can export helper
/// functions (reachable via the API/A2A/MCP dispatch adapters) without
/// every one of them grabbing an HTTP path.
fn route_from_attributes(
    attrs: &[Attribute],
    fn_name: &str,
    diagnostics: &mut Vec<ExportDiagnostic>,
) -> Option<RouteSpec> {
    // An explicit (even if malformed) `@route` overrides the naming
    // convention: an author who wrote one expects that path, not a
    // surprise fallback to `/<name>`. A malformed one yields `None` plus
    // a diagnostic, leaving the handler unmounted until they fix it.
    if attrs.iter().any(|attr| attr.name == "route") {
        return explicit_route_attribute(attrs, fn_name, diagnostics);
    }
    handler_convention_route(fn_name)
}

fn explicit_route_attribute(
    attrs: &[Attribute],
    fn_name: &str,
    diagnostics: &mut Vec<ExportDiagnostic>,
) -> Option<RouteSpec> {
    let attr = attrs.iter().find(|attr| attr.name == "route")?;
    let literals: Vec<&str> = attr
        .args
        .iter()
        .filter_map(|arg| match &arg.value.node {
            Node::StringLiteral(value) | Node::RawStringLiteral(value) => Some(value.as_str()),
            _ => None,
        })
        .collect();

    // Any non-string argument makes the method/path positions ambiguous
    // (e.g. `@route("GET", some_var)` would otherwise collapse to the
    // single-arg form and mis-mount at `/GET`), so refuse to guess.
    if literals.len() != attr.args.len() {
        diagnostics.push(ExportDiagnostic {
            code: ROUTE_ARG_NOT_STRING,
            line: attr.span.line,
            message: format!(
                "`@route` on `{fn_name}` requires string-literal arguments \
                 (`@route(\"/path\")` or `@route(\"METHOD\", \"/path\")`); handler not mounted"
            ),
        });
        return None;
    }

    match literals.as_slice() {
        // `@route("/path")` — method defaults to GET.
        [path] => Some(RouteSpec {
            method: "GET".to_string(),
            path: normalize_route_path(path),
        }),
        // `@route("METHOD", "/path")` — explicit method.
        [method, path] => Some(RouteSpec {
            method: normalize_route_method(method),
            path: normalize_route_path(path),
        }),
        // Zero args (`@route()`) or three-plus: the method/path pair is
        // under- or over-specified, so the route is undefined.
        _ => {
            diagnostics.push(ExportDiagnostic {
                code: ROUTE_BAD_ARITY,
                line: attr.span.line,
                message: format!(
                    "`@route` on `{fn_name}` takes a path or a method and a path \
                     (`@route(\"/path\")` or `@route(\"METHOD\", \"/path\")`), \
                     found {} arguments; handler not mounted",
                    literals.len()
                ),
            });
            None
        }
    }
}

fn handler_convention_route(fn_name: &str) -> Option<RouteSpec> {
    let path = match fn_name {
        "handler" => "/".to_string(),
        other => {
            let suffix = other.strip_prefix("handler_")?;
            if suffix.is_empty() {
                return None;
            }
            format!("/{suffix}")
        }
    };
    // Convention handlers answer both GET and POST so a script can serve
    // a read and a form-style write from one function without an explicit
    // attribute; the handler discriminates on `req.method`.
    Some(RouteSpec {
        method: "*".to_string(),
        path,
    })
}

fn normalize_route_method(method: &str) -> String {
    let upper = method.trim().to_ascii_uppercase();
    if upper == "ANY" || upper.is_empty() {
        "*".to_string()
    } else {
        upper
    }
}

fn normalize_route_path(path: &str) -> String {
    let trimmed = path.trim();
    if trimmed.starts_with('/') {
        trimmed.to_string()
    } else {
        format!("/{trimmed}")
    }
}

fn pipeline_input_schema(params: &[String]) -> serde_json::Value {
    serde_json::json!({
        "type": "object",
        "properties": params
            .iter()
            .map(|name| (name.clone(), serde_json::json!({})))
            .collect::<serde_json::Map<_, _>>(),
        "required": params,
    })
}

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

    #[test]
    fn export_catalog_only_includes_public_functions() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("server.harn");
        std::fs::write(
            &path,
            r#"
fn hidden() { return "nope" }
pub fn greet(name: string, excited: bool = false) -> string {
  if excited { return "hi!" }
  return name
}
"#,
        )
        .expect("write script");

        let catalog = ExportCatalog::from_path(&path).expect("catalog");
        assert!(catalog.function("hidden").is_none());
        let greet = catalog.function("greet").expect("greet export");
        assert_eq!(greet.params.len(), 2);
        assert_eq!(greet.input_schema["type"], "object");
        assert_eq!(
            greet.output_schema.as_ref().expect("output")["type"],
            "string"
        );
    }

    #[test]
    fn export_catalog_captures_scopes_attribute_from_function_decl() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("server.harn");
        std::fs::write(
            &path,
            r#"
@scopes("personas:read", "sessions:write")
pub fn list_sessions() -> string {
  return "ok"
}

pub fn ping() -> string {
  return "pong"
}
"#,
        )
        .expect("write script");

        let catalog = ExportCatalog::from_path(&path).expect("catalog");
        let list = catalog.function("list_sessions").expect("list_sessions");
        assert_eq!(
            list.required_scopes,
            BTreeSet::from(["personas:read".to_string(), "sessions:write".to_string()])
        );
        let ping = catalog.function("ping").expect("ping");
        assert!(ping.required_scopes.is_empty());
    }

    #[test]
    fn export_catalog_parses_limits_and_budget_attributes() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("server.harn");
        std::fs::write(
            &path,
            r#"
@limits(
    per_tenant: "100/min",
    per_route: "5000/min",
    burst: 50,
    algorithm: "sliding_window",
    in_flight_max: 20,
)
@budget(llm_cost_usd: 0.50, mcp_calls: 20)
pub fn create() -> string { return "ok" }

pub fn ping() -> string { return "pong" }
"#,
        )
        .expect("write script");

        let catalog = ExportCatalog::from_path(&path).expect("catalog");
        let create = catalog.function("create").expect("create export");
        let limits = create.limits.as_ref().expect("limits parsed");
        assert_eq!(limits.per_tenant.unwrap().count, 100);
        assert_eq!(limits.per_route.unwrap().count, 5_000);
        assert_eq!(limits.burst, Some(50));
        assert_eq!(limits.algorithm, crate::limits::Algorithm::SlidingWindow);
        assert_eq!(limits.in_flight_max, Some(20));
        let budget = create.budget.as_ref().expect("budget parsed");
        assert_eq!(budget.llm_cost_usd, Some(0.50));
        assert_eq!(budget.mcp_calls, Some(20));

        // Routes without the attributes get None — the dispatch path
        // short-circuits without consulting the registry.
        let ping = catalog.function("ping").expect("ping export");
        assert!(ping.limits.is_none());
        assert!(ping.budget.is_none());
    }

    #[test]
    fn route_attribute_parses_method_and_path() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("server.harn");
        std::fs::write(
            &path,
            r#"
@route("POST", "/users/{id}")
pub fn update_user(req: dict) -> dict { return req }

@route("/health")
pub fn liveness(req: dict) -> dict { return req }

@route("any", "metrics")
pub fn metrics(req: dict) -> dict { return req }

pub fn helper(req: dict) -> dict { return req }
"#,
        )
        .expect("write script");

        let catalog = ExportCatalog::from_path(&path).expect("catalog");
        let update = catalog.function("update_user").expect("update_user");
        assert_eq!(
            update.route,
            Some(RouteSpec {
                method: "POST".to_string(),
                path: "/users/{id}".to_string()
            })
        );
        // Single-arg form defaults to GET.
        let liveness = catalog.function("liveness").expect("liveness");
        assert_eq!(
            liveness.route,
            Some(RouteSpec {
                method: "GET".to_string(),
                path: "/health".to_string()
            })
        );
        // `any` lowercases to the `*` wildcard; a path missing its leading
        // slash is normalized.
        let metrics = catalog.function("metrics").expect("metrics");
        assert_eq!(
            metrics.route,
            Some(RouteSpec {
                method: "*".to_string(),
                path: "/metrics".to_string()
            })
        );
        // A plain `pub fn` with no attribute and no `handler_` prefix is
        // dispatch-only — it gets no HTTP route.
        let helper = catalog.function("helper").expect("helper");
        assert_eq!(helper.route, None);
    }

    #[test]
    fn handler_naming_convention_infers_route() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("server.harn");
        std::fs::write(
            &path,
            r"
pub fn handler(req: dict) -> dict { return req }
pub fn handler_echo(req: dict) -> dict { return req }
",
        )
        .expect("write script");

        let catalog = ExportCatalog::from_path(&path).expect("catalog");
        // Bare `handler` mounts at the site root.
        assert_eq!(
            catalog.function("handler").expect("handler").route,
            Some(RouteSpec {
                method: "*".to_string(),
                path: "/".to_string()
            })
        );
        // `handler_echo` mounts at `/echo`, answering every method.
        assert_eq!(
            catalog
                .function("handler_echo")
                .expect("handler_echo")
                .route,
            Some(RouteSpec {
                method: "*".to_string(),
                path: "/echo".to_string()
            })
        );
    }

    #[test]
    fn export_catalog_falls_back_to_legacy_pipelines_without_public_exports() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("server.harn");
        std::fs::write(
            &path,
            r"
pipeline default(task) {
  __io_println(task)
}
",
        )
        .expect("write script");

        let catalog = ExportCatalog::from_path(&path).expect("catalog");
        let default = catalog.function("default").expect("default pipeline");
        assert_eq!(default.kind, ExportedCallableKind::Pipeline);
        assert_eq!(default.params[0].name, "task");
    }

    /// Build a catalog from inline source, asserting it parses cleanly.
    fn catalog_from_source(source: &str) -> ExportCatalog {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("server.harn");
        std::fs::write(&path, source).expect("write script");
        ExportCatalog::from_path(&path).expect("catalog")
    }

    #[test]
    fn well_formed_attributes_emit_no_diagnostics() {
        let catalog = catalog_from_source(
            r#"
@scopes("personas:read")
@route("POST", "/users/{id}")
pub fn update_user(req: dict) -> dict { return req }

@route("/health")
pub fn liveness(req: dict) -> dict { return req }
"#,
        );
        assert!(
            catalog.diagnostics().is_empty(),
            "unexpected diagnostics: {:?}",
            catalog.diagnostics()
        );
    }

    #[test]
    fn route_with_non_string_arg_is_diagnosed_and_unmounted() {
        // The second arg is an identifier, not a string literal. Left
        // unchecked the collector would treat this as `@route("GET")` and
        // mis-mount the handler at `/GET`.
        let catalog = catalog_from_source(
            r#"
pub fn make_path(req: dict) -> string { return "/x" }

@route("GET", make_path)
pub fn handler_users(req: dict) -> dict { return req }
"#,
        );
        let handler = catalog.function("handler_users").expect("handler_users");
        assert_eq!(
            handler.route, None,
            "a malformed @route must not fall back to the handler_ convention route"
        );
        let codes: Vec<&str> = catalog.diagnostics().iter().map(|d| d.code).collect();
        assert_eq!(codes, vec![ROUTE_ARG_NOT_STRING]);
    }

    #[test]
    fn route_with_zero_args_is_diagnosed_and_unmounted() {
        let catalog = catalog_from_source(
            r"
@route()
pub fn handler_status(req: dict) -> dict { return req }
",
        );
        let handler = catalog.function("handler_status").expect("handler_status");
        assert_eq!(handler.route, None);
        let codes: Vec<&str> = catalog.diagnostics().iter().map(|d| d.code).collect();
        assert_eq!(codes, vec![ROUTE_BAD_ARITY]);
    }

    #[test]
    fn route_with_too_many_args_is_diagnosed_and_unmounted() {
        let catalog = catalog_from_source(
            r#"
@route("GET", "/x", "/y")
pub fn handler_overspecified(req: dict) -> dict { return req }
"#,
        );
        let handler = catalog
            .function("handler_overspecified")
            .expect("handler_overspecified");
        assert_eq!(handler.route, None);
        let codes: Vec<&str> = catalog.diagnostics().iter().map(|d| d.code).collect();
        assert_eq!(codes, vec![ROUTE_BAD_ARITY]);
    }

    #[test]
    fn scopes_with_non_string_arg_is_diagnosed_but_keeps_valid_scopes() {
        let catalog = catalog_from_source(
            r#"
pub fn make_scope(req: dict) -> string { return "sessions:write" }

@scopes("personas:read", make_scope)
pub fn list_sessions() -> string { return "ok" }
"#,
        );
        let list = catalog.function("list_sessions").expect("list_sessions");
        // The valid literal is still enforced; only the bad arg is dropped.
        assert_eq!(
            list.required_scopes,
            BTreeSet::from(["personas:read".to_string()])
        );
        let diagnostic = catalog
            .diagnostics()
            .iter()
            .find(|d| d.code == SCOPES_ARG_NOT_STRING)
            .expect("scopes diagnostic");
        assert!(diagnostic.message.contains("list_sessions"));
    }
}