scryer-mcp 0.3.0

use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::Instant;
use tiktoken_rs::CoreBPE;
use tracing::{debug, error};

use scryer_db::{ScryerDb, ToolInvocationMetric};

pub const MAX_PAYLOAD_TOKENS: usize = 1000;

/// Token counting, payload budget enforcement and local usage telemetry, backed by `tiktoken-rs` and
/// the project database.
///
/// Telemetry records measured quantities only (tool name, response tokens, truncation, latency,
/// session); never arguments or responses. See ADR 0015.
#[derive(Clone)]
pub struct TokenSavingsMiddleware {
    db: ScryerDb,
    bpe: Arc<CoreBPE>,
}

impl TokenSavingsMiddleware {
    /// Initialize the middleware using `cl100k_base` tokenizer.
    pub fn new(db: ScryerDb) -> anyhow::Result<Self> {
        let bpe = tiktoken_rs::cl100k_base()?;
        Ok(Self {
            db,
            bpe: Arc::new(bpe),
        })
    }

    /// Access the underlying tokenizer.
    pub fn bpe(&self) -> &CoreBPE {
        &self.bpe
    }

    /// Count exact tokens for a text string using BPE.
    pub fn count_tokens(&self, text: &str) -> usize {
        self.bpe.encode_with_special_tokens(text).len()
    }

    /// Enforce the payload token limit.
    /// - If `no_truncate` is true, bypasses truncation regardless of token count.
    /// - If `max_tokens` is provided and > 0, enforces that limit; if 0, truncation is disabled.
    /// - If `max_tokens` is None, defaults to `MAX_PAYLOAD_TOKENS` (1,000).
    pub fn enforce_limit(
        &self,
        payload: &str,
        max_tokens: Option<usize>,
        no_truncate: bool,
    ) -> (String, usize, bool) {
        let count = self.count_tokens(payload);
        if no_truncate {
            return (payload.to_string(), count, false);
        }

        let budget = max_tokens.unwrap_or(MAX_PAYLOAD_TOKENS);
        if budget == 0 || count <= budget {
            return (payload.to_string(), count, false);
        }

        let budget_str = if budget == 1000 {
            "1,000".to_string()
        } else {
            budget.to_string()
        };
        let more = budget * 2;
        let notice = format!(
            "\n\n... [Output truncated to enforce hard {budget_str} token limit. Call again with 'max_tokens: {more}' or 'no_truncate: true', or narrow with offset, limit or filters.]"
        );
        let notice_tokens = self.count_tokens(&notice);
        let token_budget = budget.saturating_sub(notice_tokens);

        // Longest prefix that fits, found by bisection on char boundaries, then backed up to a
        // natural break: the last newline for line-oriented text, else the last comma, so a
        // compact (single-line) JSON payload is not cut mid-token.
        let bounds: Vec<usize> = payload
            .char_indices()
            .map(|(i, _)| i)
            .chain(std::iter::once(payload.len()))
            .collect();
        let (mut lo, mut hi) = (0, bounds.len() - 1);
        while lo < hi {
            let mid = (lo + hi).div_ceil(2);
            if self.count_tokens(&payload[..bounds[mid]]) <= token_budget {
                lo = mid;
            } else {
                hi = mid - 1;
            }
        }
        let fit = &payload[..bounds[lo]];
        let break_at = if fit.contains('\n') {
            fit.rfind('\n')
        } else {
            fit.rfind(',').map(|i| i + 1)
        };
        let mut truncated = break_at.map_or(fit, |i| &fit[..i]).trim_end().to_string();

        truncated.push_str(&notice);
        let actual = self.count_tokens(&truncated);
        (truncated, actual, true)
    }

    /// Backward-compatible limit enforcement using default 1,000 token limit.
    pub fn enforce_default_limit(&self, payload: &str) -> (String, usize, bool) {
        self.enforce_limit(payload, None, false)
    }

    /// Record a tool invocation and asynchronously persist it to `ToolInvocationMetric`.
    ///
    /// Only sizes and flags are stored: `payload_tokens` is the response as sent, `truncated`
    /// whether the agent was told there is more. The columns that held the retired
    /// "naive baseline" estimate are written as 0.
    pub fn record_invocation(
        &self,
        session_id: String,
        project_id: Option<u64>,
        tool_name: String,
        payload_tokens: u64,
        truncated: bool,
        duration: Instant,
    ) {
        if !recording_enabled() {
            return;
        }
        let duration_ms = duration.elapsed().as_millis() as u32;
        let db = self.db.clone();
        let now = scryer_db::time::now_rfc3339();

        tokio::spawn(async move {
            let mut guard = db.lock().await;
            let result = ToolInvocationMetric::create()
                .session_id(session_id)
                .project_id(project_id)
                .tool_name(tool_name)
                .payload_tokens(payload_tokens)
                .estimated_naive_tokens(0)
                .tokens_saved(0)
                .truncated(truncated)
                .execution_duration_ms(duration_ms)
                .created_at(now)
                .exec(&mut *guard)
                .await;

            if let Err(e) = result {
                error!("Failed to persist ToolInvocationMetric: {e}");
            } else {
                debug!(payload_tokens, truncated, "Persisted ToolInvocationMetric");
            }
        });
    }

    /// Aggregate measured usage across one session or all sessions.
    pub async fn get_metrics(&self, session_id: Option<&str>) -> anyhow::Result<UsageSummary> {
        let mut guard = self.db.lock().await;
        let metrics = match session_id {
            Some(sid) => {
                ToolInvocationMetric::filter(ToolInvocationMetric::fields().session_id().eq(sid))
                    .exec(&mut *guard)
                    .await?
            }
            None => ToolInvocationMetric::all().exec(&mut *guard).await?,
        };
        Ok(summarize(
            metrics
                .into_iter()
                .filter(|m| m.tool_name != METRICS_TOOL)
                .map(|m| UsageRow {
                    id: m.id,
                    session_id: m.session_id,
                    tool_name: m.tool_name,
                    payload_tokens: m.payload_tokens,
                    truncated: m.truncated,
                    duration_ms: m.execution_duration_ms as u64,
                }),
        ))
    }
}

/// Telemetry is on unless `SCRYER_TELEMETRY` is `off`, `0` or `false`. Rows stay in the user's own
/// database; nothing is sent anywhere.
pub fn recording_enabled() -> bool {
    !matches!(
        std::env::var("SCRYER_TELEMETRY").ok().as_deref(),
        Some("off" | "0" | "false")
    )
}

/// The tool that reads these rows; its own calls are not usage worth reporting.
const METRICS_TOOL: &str = "get_token_savings_metrics";

/// One recorded call, reduced to what the summary needs.
#[derive(Debug, Clone)]
pub struct UsageRow {
    /// Insertion order: calls of one session are ordered by it.
    pub id: u64,
    pub session_id: String,
    pub tool_name: String,
    pub payload_tokens: u64,
    pub truncated: bool,
    pub duration_ms: u64,
}

/// Usage of one tool.
#[derive(
    Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize, schemars::JsonSchema,
)]
pub struct ToolUsage {
    pub calls: u64,
    /// Total tokens of the responses sent for this tool.
    pub response_tokens: u64,
    pub mean_response_tokens: f64,
    /// Calls whose answer told the agent there is more (cut at the budget, or a page with a
    /// continuation).
    pub truncated: u64,
    /// Calls followed by another Scryer call in the same session.
    pub followed_up: u64,
    /// `followed_up / calls`, in 0..=1.
    pub follow_up_rate: f64,
    pub mean_duration_ms: f64,
}

/// Measured usage; nothing here is an estimate of savings.
///
/// A call is *followed up* when the same session makes another Scryer call after it. The server
/// cannot see what the agent does outside Scryer (a `grep` after a Scryer answer), so this is a
/// lower bound on how often an answer was not enough. The last call of a session has no next call
/// and counts as not followed up.
#[derive(
    Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize, schemars::JsonSchema,
)]
pub struct UsageSummary {
    pub total_calls: u64,
    pub total_response_tokens: u64,
    pub total_truncated: u64,
    pub total_followed_up: u64,
    pub follow_up_rate: f64,
    pub mean_duration_ms: f64,
    pub by_tool: BTreeMap<String, ToolUsage>,
}

/// Aggregate recorded calls. Calls of one session are taken in `id` order for the follow-up count.
pub fn summarize(rows: impl IntoIterator<Item = UsageRow>) -> UsageSummary {
    let mut rows: Vec<UsageRow> = rows.into_iter().collect();
    rows.sort_by(|a, b| (&a.session_id, a.id).cmp(&(&b.session_id, b.id)));

    let mut summary = UsageSummary::default();
    let mut duration_total = 0u64;
    let mut tool_durations: BTreeMap<String, u64> = BTreeMap::new();
    for (i, row) in rows.iter().enumerate() {
        let followed = rows
            .get(i + 1)
            .is_some_and(|next| next.session_id == row.session_id);
        let tool = summary.by_tool.entry(row.tool_name.clone()).or_default();
        tool.calls += 1;
        tool.response_tokens += row.payload_tokens;
        tool.truncated += u64::from(row.truncated);
        tool.followed_up += u64::from(followed);
        *tool_durations.entry(row.tool_name.clone()).or_default() += row.duration_ms;

        summary.total_calls += 1;
        summary.total_response_tokens += row.payload_tokens;
        summary.total_truncated += u64::from(row.truncated);
        summary.total_followed_up += u64::from(followed);
        duration_total += row.duration_ms;
    }
    let ratio = |n: u64, d: u64| if d > 0 { n as f64 / d as f64 } else { 0.0 };
    summary.follow_up_rate = ratio(summary.total_followed_up, summary.total_calls);
    summary.mean_duration_ms = ratio(duration_total, summary.total_calls);
    for (name, tool) in summary.by_tool.iter_mut() {
        tool.mean_response_tokens = ratio(tool.response_tokens, tool.calls);
        tool.follow_up_rate = ratio(tool.followed_up, tool.calls);
        tool.mean_duration_ms = ratio(tool_durations[name], tool.calls);
    }
    summary
}