bitrouter 1.0.0-alpha.34

//! `PolicyHook` — a `language_model::PreRequestHook` enforcing per-API-key
//! policy: model allow/deny, spend ceiling, expiry, tool-access rules,
//! request-rate limits.
//!
//! The caller's `policy_id` is read from the auth module's `PipelineContext`
//! metadata (not by importing the auth crate's event type). Spend and rate
//! are read from the sibling [`crate::metering::MeteringStore`] — a direct
//! concrete-type call, no SDK trait in between.

use std::sync::Arc;

use async_trait::async_trait;

use bitrouter_sdk::PluginId;
use bitrouter_sdk::Result;
use bitrouter_sdk::language_model::{DenyReason, HookDecision, PipelineContext, PreRequestHook};

use crate::metering::{MeteringStore, TimeWindow};
use crate::policy::store::PolicyStore;

/// Enforces per-API-key policy at Stage 1.
pub struct PolicyHook {
    store: Arc<PolicyStore>,
    /// Optional — without a [`MeteringStore`], spend / rate ceilings cannot
    /// be enforced (model allow/deny, expiry, and tool-access checks still
    /// are).
    metering: Option<MeteringStore>,
}

impl PolicyHook {
    /// Build a `PolicyHook` over a policy store, optionally with a metering
    /// store for spend + rate enforcement.
    pub fn new(store: Arc<PolicyStore>, metering: Option<MeteringStore>) -> Self {
        Self { store, metering }
    }

    /// Read the caller's `policy_id` from the auth module's context metadata.
    fn policy_id(ctx: &PipelineContext) -> Option<String> {
        ctx.get_metadata(&PluginId::new("bitrouter-auth"))
            .and_then(|m| m.get("policy_id"))
            .and_then(|v| v.as_str())
            .map(|s| s.to_string())
    }
}

impl PolicyHook {
    async fn evaluate(&self, ctx: &PipelineContext, model: &str) -> Result<HookDecision> {
        let policy_id = Self::policy_id(ctx);
        // No policy bound → no constraints (the combination is permissive).
        let ids: Vec<&str> = policy_id.as_deref().into_iter().collect();
        let effective = self.store.effective_for(&ids);

        // 1. model allow / deny
        if let Err(violation) = effective.check_model(model) {
            return Ok(HookDecision::Deny(DenyReason::Forbidden(
                violation.to_string(),
            )));
        }

        // 2. hard expiry
        if let Err(violation) = effective.check_expiry(chrono::Utc::now()) {
            return Ok(HookDecision::Deny(DenyReason::Forbidden(
                violation.to_string(),
            )));
        }

        // 3. tool-access rules — checked against the request's declared tools
        if effective.has_tool_restriction() {
            let tools = ctx
                .generation_prompt()
                .into_iter()
                .flat_map(|prompt| &prompt.tools)
                .map(|tool| tool.name());
            if let Err(violation) = effective.check_tools(tools) {
                return Ok(HookDecision::Deny(DenyReason::Forbidden(
                    violation.to_string(),
                )));
            }
        }

        // 4. spend ceiling — only enforceable with a MeteringStore
        if effective.max_spend_micro_usd.is_some()
            && let Some(metering) = &self.metering
        {
            let Some(spent) = metering
                .get_enforceable_spend(ctx.caller().api_key_id(), TimeWindow::ThisMonth)
                .await?
            else {
                return Ok(HookDecision::Deny(DenyReason::Forbidden(
                    "spend ceiling cannot be evaluated: prior charge evidence is incomplete"
                        .to_string(),
                )));
            };
            if let Err(violation) = effective.check_spend(spent) {
                return Ok(HookDecision::Deny(DenyReason::Forbidden(
                    violation.to_string(),
                )));
            }
        }

        // 5. request-rate ceiling — also reads the MeteringStore. A rate
        //    violation maps to 429 (RateLimited) rather than 403, with a
        //    Retry-After hint.
        if effective.max_requests_per_minute.is_some()
            && let Some(metering) = &self.metering
        {
            let rate = metering.get_rate(ctx.caller().api_key_id()).await?;
            let observed = rate.requests_per_minute.round().max(0.0) as u32;
            if let Err(violation) = effective.check_rate(observed) {
                tracing::debug!(%violation, "policy rate limit hit");
                return Ok(HookDecision::Deny(DenyReason::RateLimited {
                    retry_after: Some(60),
                }));
            }
        }

        Ok(HookDecision::Allow)
    }
}

struct FrozenPolicyBinding {
    model: String,
    policy_id: Option<String>,
}

#[async_trait]
impl PreRequestHook for PolicyHook {
    async fn check(&self, ctx: &mut PipelineContext) -> Result<HookDecision> {
        let binding = FrozenPolicyBinding {
            model: ctx.model().into(),
            policy_id: Self::policy_id(ctx),
        };
        let result = self.evaluate(ctx, &binding.model).await?;
        if matches!(&result, HookDecision::Allow) {
            ctx.insert_extension(Arc::new(binding));
        }
        Ok(result)
    }

    async fn revalidate_context(&self, ctx: &PipelineContext) -> Result<HookDecision> {
        let binding = ctx.extension::<FrozenPolicyBinding>().ok_or_else(|| {
            bitrouter_sdk::BitrouterError::bad_request(
                "policy revalidation lost its ingress binding",
            )
        })?;
        if binding.policy_id != Self::policy_id(ctx) {
            return Err(bitrouter_sdk::BitrouterError::bad_request(
                "policy binding changed after admission",
            ));
        }
        // Stage 1 checks the normalized selector, not the later policy-selected
        // model. Preserve that contract while rechecking live policy and spend.
        self.evaluate(ctx, &binding.model).await
    }
}