tokensave 6.1.2

// Rust guideline compliant 2025-10-17
//! MCP server that reads JSON-RPC 2.0 messages from stdin and writes
//! responses to stdout.
//!
//! The server exposes code graph tools via the Model Context Protocol,
//! allowing AI assistants to query the code graph interactively.

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use serde_json::{json, Value};

use crate::errors::Result;
use crate::global_db::GlobalDb;
use crate::tokensave::TokenSave;

use super::tools::{explore_call_budget, get_tool_definitions_with_budget, handle_tool_call};
use super::transport::{ErrorCode, JsonRpcRequest, JsonRpcResponse};

/// Runtime statistics for the MCP server.
pub struct ServerStats {
    started_at: Instant,
    total_requests: AtomicU64,
    tool_calls: AtomicU64,
    errors: AtomicU64,
}

impl ServerStats {
    fn new() -> Self {
        Self {
            started_at: Instant::now(),
            total_requests: AtomicU64::new(0),
            tool_calls: AtomicU64::new(0),
            errors: AtomicU64::new(0),
        }
    }
}

/// Cache duration for version checks (15 minutes).
const VERSION_CHECK_INTERVAL: Duration = Duration::from_mins(15);

/// Hand-maintained schema documentation for the `tokensave://schema` resource.
/// Mirrors `src/db/migrations.rs::create_schema`. Update both together.
const SCHEMA_MARKDOWN: &str = r"# tokensave SQLite schema

The on-disk database lives at `.tokensave/tokensave.db` (per-branch variants
under multi-branch mode). All tables are plain SQLite; safe to query with any
client. WAL mode is used, so readers do not block writers.

## Tables

### `nodes` — every indexed symbol
- `id` TEXT PRIMARY KEY — content-hashed identifier (changes when symbol moves or renames)
- `kind` TEXT — e.g. `function`, `struct`, `trait`, `impl`, `method`, `module`, `file`
- `name` TEXT — local identifier
- `qualified_name` TEXT — language-style path (e.g. `crate::module::Type::method`)
- `file_path` TEXT — relative to the project root
- `start_line`, `end_line` INTEGER — 1-based inclusive line range of the symbol
- `start_column`, `end_column` INTEGER — 0-based column range
- `attrs_start_line` INTEGER — first line of leading doc-comments / attributes (or `start_line` if none)
- `signature` TEXT NULL — extracted source-level signature
- `docstring` TEXT NULL — leading doc-comment
- `visibility` TEXT — one of `public`, `pub_crate`, `pub_super`, `private`
- `is_async` INTEGER (0/1)
- `branches`, `loops`, `returns`, `max_nesting`, `unsafe_blocks`, `unchecked_calls`, `assertions` INTEGER — complexity metrics
- `updated_at` INTEGER — UNIX epoch seconds

Indexes: `kind`, `name`, `qualified_name`, `file_path`, `(file_path,start_line)`, `lower(name)`.

### `edges` — directed relationships between nodes
- `id` INTEGER PRIMARY KEY AUTOINCREMENT
- `source` TEXT — FK → `nodes.id` (CASCADE DELETE)
- `target` TEXT — FK → `nodes.id` (CASCADE DELETE)
- `kind` TEXT — one of `contains`, `calls`, `returns`, `type_of`, `uses`, `implements`, `extends`, `annotates`, `derives_macro`, `receives`
- `line` INTEGER NULL — source line of the relationship

Unique constraint: `(source, target, kind, COALESCE(line, -1))`. Indexes on `source`, `target`, `kind`, `(source,kind)`, `(target,kind)`.

### `files` — index bookkeeping
- `path` TEXT PRIMARY KEY
- `content_hash` TEXT — sha256 of file contents at index time
- `size` INTEGER — file size in bytes
- `modified_at`, `indexed_at` INTEGER — UNIX epoch seconds
- `node_count` INTEGER — number of nodes extracted from this file

### `unresolved_refs` — references the resolver could not bind
- `from_node_id` FK → `nodes.id`
- `reference_name` TEXT
- `reference_kind` TEXT
- `line`, `col` INTEGER
- `file_path` TEXT

### `vectors` — optional embeddings (semantic search backend)
- `node_id` PRIMARY KEY FK → `nodes.id`
- `embedding` BLOB
- `model` TEXT, `created_at` INTEGER

### `metadata` — key/value store
Common keys: `tokens_saved`, schema-version markers.

### `memory_decisions`, `memory_code_areas`
Hand-recorded notes from `tokensave_record_decision` / `tokensave_record_code_area`. FTS5 mirror tables exist for `nodes` (`nodes_fts`) and `memory_decisions` (`memory_decisions_fts`).

## Recipes

### Find every impl block of a trait
```sql
SELECT n.id, n.qualified_name, n.file_path, n.start_line
FROM nodes n
JOIN edges e ON e.source = n.id
WHERE e.kind = 'implements'
  AND e.target IN (SELECT id FROM nodes WHERE qualified_name = ?1);
```

### Top callers of a node
```sql
SELECT n.qualified_name, COUNT(*) AS call_count
FROM edges e
JOIN nodes n ON n.id = e.source
WHERE e.target = ?1 AND e.kind = 'calls'
GROUP BY n.qualified_name
ORDER BY call_count DESC
LIMIT 20;
```

### Files modified since last index
Compare `files.modified_at` against the live filesystem mtime — `tokensave_affected` does this with extra git plumbing.

### Largest functions by line span
```sql
SELECT qualified_name, file_path, end_line - start_line + 1 AS lines
FROM nodes
WHERE kind IN ('function', 'method')
ORDER BY lines DESC
LIMIT 20;
```

## Gotchas
- `nodes.id` is a content hash, so it changes when the symbol moves. For cross-run lookups use `qualified_name` (or `tokensave_by_qualified_name`).
- `edges.kind = 'calls'` may reference a *trait method* node rather than the resolved concrete impl — trait dispatch is not currently rewritten.
- `derives_macro` edges record `#[derive(...)]` usage but generated impls are not in the graph.
";

/// Build the per-file staleness banner inserted at the top of any tool
/// response that referenced files the in-line sync couldn't refresh.
///
/// The shape mimics codegraph's #428 banner: name each pending file with
/// its edit age (how long since the on-disk mtime), and direct the agent
/// to `Read` those specific files. The rest of the response is treated
/// as authoritative — distinct from the previous binary "STALE INDEX"
/// warning that asked the agent to distrust the whole answer.
fn format_per_file_staleness_banner(
    project_root: &std::path::Path,
    stale_files: &[String],
) -> String {
    let now_secs = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs() as i64;

    let mut lines = Vec::with_capacity(stale_files.len() + 2);
    lines.push(format!(
        "WARNING: {} file(s) referenced below were edited after the last sync. \
         Read these directly; the rest of this response reflects the current index:",
        stale_files.len()
    ));
    for path in stale_files {
        let age = file_mtime_secs(project_root, path).map_or(0, |m| now_secs.saturating_sub(m));
        lines.push(format!("  - {path} (edited {})", humanize_age(age)));
    }
    lines.push("Run `tokensave sync` to refresh the index.".to_string());
    lines.join("\n")
}

/// Read the on-disk mtime (UNIX seconds) for `relative_path` joined onto
/// `project_root`. Returns `None` when the file is missing or stat fails.
fn file_mtime_secs(project_root: &std::path::Path, relative_path: &str) -> Option<i64> {
    let abs = project_root.join(relative_path);
    let meta = std::fs::metadata(&abs).ok()?;
    let modified = meta.modified().ok()?;
    let secs = modified
        .duration_since(std::time::UNIX_EPOCH)
        .ok()?
        .as_secs() as i64;
    Some(secs)
}

/// Render a duration in seconds as a compact phrase: `"5s ago"`,
/// `"3m ago"`, `"2h ago"`, `"4d ago"`. Used in the staleness banner so
/// the agent can judge how stale "still stale" actually is.
fn humanize_age(secs: i64) -> String {
    if secs < 60 {
        format!("{}s ago", secs.max(0))
    } else if secs < 3_600 {
        format!("{}m ago", secs / 60)
    } else if secs < 86_400 {
        format!("{}h ago", secs / 3_600)
    } else {
        format!("{}d ago", secs / 86_400)
    }
}

/// Cached result of a latest-version check against GitHub releases.
struct VersionCheckState {
    latest: Option<String>,
    checked_at: Option<Instant>,
}

/// The MCP server wrapping a `TokenSave` instance.
// Lock ordering: file_token_map -> tool_call_counts (never nested)
pub struct McpServer {
    cg: TokenSave,
    stats: ServerStats,
    tool_call_counts: std::sync::Mutex<HashMap<String, u64>>,
    /// Approximate token count per indexed file (`file_path` -> tokens).
    file_token_map: std::sync::Mutex<HashMap<String, u64>>,
    /// Running total of tokens saved by serving from the graph.
    tokens_saved: AtomicU64,
    /// Tokens already flushed to the worldwide counter this session.
    last_flushed_tokens: AtomicU64,
    /// UNIX timestamp of last worldwide flush (0 = never).
    last_flush_at: AtomicI64,
    /// User-level database tracking all projects (best-effort).
    global_db: Option<GlobalDb>,
    /// Cached latest-version check result.
    version_cache: std::sync::Mutex<VersionCheckState>,
    /// Pending JSON-RPC notifications to send before the next response.
    pending_notifications: std::sync::Mutex<Vec<Value>>,
    /// When the MCP server was started from a subdirectory of the project root,
    /// this holds the relative path prefix (e.g. `"src/mcp"`). Listing tools
    /// use it as the default path filter. `None` when cwd == project root.
    scope_prefix: Option<String>,
    /// Set to `true` after `shutdown` runs once; makes shutdown idempotent so
    /// callers can invoke it explicitly after `run` returns without re-running
    /// persistence logic.
    shutdown_done: AtomicBool,
    /// When true, every `tools/call` response gains a `_meta.duration_us`
    /// field measuring the handler's pure execution time. Toggled by
    /// `tokensave serve --timings`. Off by default to keep responses clean.
    timings_enabled: AtomicBool,
    /// UNIX timestamp (secs) of the most recent staleness check started by
    /// the server. Read-modify-update via `compare_exchange` in
    /// [`maybe_sync_if_stale`](Self::maybe_sync_if_stale) so concurrent
    /// tool calls don't pile on the same walk.
    last_staleness_check_at: AtomicI64,
    /// Cached worktree-vs-index mismatch detection for this session. `None`
    /// when no mismatch exists (the common case) or detection was skipped
    /// (not a git repo / git missing). Computed once at startup so we
    /// spawn at most one pair of `git rev-parse` per session no matter how
    /// many tool calls fire. See [`crate::worktree`] and #312.
    worktree_mismatch: Option<crate::worktree::WorktreeIndexMismatch>,
    /// Flipped to `true` once [`Self::run_startup_catch_up_sync`] finishes
    /// (#414). Production code never reads this; tests poll it via
    /// [`Self::wait_for_startup_catch_up`] so they can race-free assert on
    /// the index state after the detached catch-up task completes.
    startup_catch_up_done: AtomicBool,
}

impl McpServer {
    /// Creates a new MCP server backed by the given code graph.
    ///
    /// Index freshness is maintained by a lazy staleness check
    /// ([`maybe_sync_if_stale`](Self::maybe_sync_if_stale)) invoked at the
    /// start of every `tools/call` and gated by a 30 s cooldown — there
    /// is no background watcher task. This replaces the
    /// `notify-debouncer-full` watcher removed in v6.x (#80), which was
    /// the source of severe CPU and memory pressure on large monorepos
    /// where nested ignored directories (`apps/*/node_modules`,
    /// `packages/*/target`) drove unbounded event traffic and `FileId`
    /// cache growth.
    pub async fn new(cg: TokenSave, scope_prefix: Option<String>) -> Arc<Self> {
        let file_token_map = cg.get_file_token_map().await.unwrap_or_default();
        let persisted = cg.get_tokens_saved().await.unwrap_or(0);
        let global_db = GlobalDb::open().await;
        // Register this project in the global DB with its current tokens
        if let Some(ref gdb) = global_db {
            gdb.upsert(cg.project_root(), persisted).await;
        }

        // Detect borrowed-worktree index once at startup so every read
        // tool can cheaply prefix a heads-up. Two git rev-parse spawns
        // worst case (#312). spawn_blocking because the underlying
        // `Command::output()` can sit on slow disks.
        let worktree_mismatch = {
            let project_root = cg.project_root().to_path_buf();
            tokio::task::spawn_blocking(move || {
                let cwd = std::env::current_dir().ok()?;
                crate::worktree::detect_worktree_index_mismatch(&cwd, &project_root)
            })
            .await
            .ok()
            .flatten()
        };

        let server = Arc::new(Self {
            cg,
            stats: ServerStats::new(),
            tool_call_counts: std::sync::Mutex::new(HashMap::new()),
            file_token_map: std::sync::Mutex::new(file_token_map),
            tokens_saved: AtomicU64::new(persisted),
            last_flushed_tokens: AtomicU64::new(persisted),
            last_flush_at: AtomicI64::new(0),
            global_db,
            version_cache: std::sync::Mutex::new(VersionCheckState {
                latest: None,
                checked_at: None,
            }),
            pending_notifications: std::sync::Mutex::new(Vec::new()),
            scope_prefix,
            shutdown_done: AtomicBool::new(false),
            timings_enabled: AtomicBool::new(false),
            last_staleness_check_at: AtomicI64::new(0),
            worktree_mismatch,
            startup_catch_up_done: AtomicBool::new(false),
        });

        // Catch-up sync (#414): pick up changes made while the server
        // was down — terminal `git pull`, IDE edits before the agent
        // launched, files touched by another tool. Detached + weak so
        // it never extends the server's lifetime; non-blocking so MCP
        // `initialize` doesn't wait on the walk.
        {
            let weak = Arc::downgrade(&server);
            tokio::spawn(async move {
                if let Some(s) = weak.upgrade() {
                    s.run_startup_catch_up_sync().await;
                }
            });
        }

        server
    }

    /// Returns the active scope prefix, if the server was launched from a subdirectory.
    pub fn scope_prefix(&self) -> Option<&str> {
        self.scope_prefix.as_deref()
    }

    /// Enables or disables per-call timing reporting. When enabled, every
    /// `tools/call` response gains a `_meta.duration_us` field with the
    /// handler's pure execution time in microseconds. Useful for profiling
    /// where time is spent inside the index vs. on the JSON-RPC/stdio
    /// transport. Safe to flip at any time — the next call observes the
    /// new setting.
    pub fn set_timings_enabled(&self, enabled: bool) {
        self.timings_enabled
            .store(enabled, std::sync::atomic::Ordering::Relaxed);
    }

    /// Returns whether timing reporting is currently enabled.
    pub fn timings_enabled(&self) -> bool {
        self.timings_enabled
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Test-only accessor for the backing `TokenSave`. Exposed so
    /// integration tests can drive the staleness pipeline directly,
    /// bypassing the 30 s cooldown in
    /// [`maybe_sync_if_stale`](Self::maybe_sync_if_stale).
    #[doc(hidden)]
    pub fn cg(&self) -> &TokenSave {
        &self.cg
    }

    /// Adds the approximate token count for the given file paths to the
    /// running saved-tokens counter and persists it to the database.
    /// Returns the delta (tokens saved by this call).
    async fn accumulate_tokens_saved(&self, file_paths: &[String]) -> u64 {
        if file_paths.is_empty() {
            return 0;
        }
        debug_assert!(
            file_paths.iter().all(|p| !p.is_empty()),
            "accumulate_tokens_saved received empty file path"
        );
        let delta = {
            let Ok(map) = self.file_token_map.lock() else {
                return 0;
            };
            let mut total: u64 = 0;
            for path in file_paths {
                if let Some(&tokens) = map.get(path.as_str()) {
                    total += tokens;
                }
            }
            total
        };
        if delta > 0 {
            let new_total = self.tokens_saved.fetch_add(delta, Ordering::Relaxed) + delta;
            // Persist to DB (best-effort, don't block on failure)
            let _ = self.cg.set_tokens_saved(new_total).await;
            // Also increment the resettable local counter
            let _ = self.cg.add_local_counter(delta).await;
            // Best-effort update to global DB
            if let Some(ref gdb) = self.global_db {
                gdb.upsert(self.cg.project_root(), new_total).await;
            }
        }
        delta
    }

    /// Re-read the file-to-token-count map from the DB and swap it into the
    /// cached `file_token_map`. Called after each lazy sync triggered by
    /// [`maybe_sync_if_stale`](Self::maybe_sync_if_stale) so the accounting
    /// tracks newly indexed / removed files.
    pub async fn refresh_file_token_map(&self) {
        // best-effort; leave stale map in place if the DB read fails
        let Ok(fresh) = self.cg.get_file_token_map().await else {
            return;
        };
        if let Ok(mut guard) = self.file_token_map.lock() {
            *guard = fresh;
        }
    }

    /// Catch-up sync run once at startup (#414). Bypasses the 30 s
    /// cooldown in [`Self::maybe_sync_if_stale`] so changes made while
    /// the server was down — a terminal `git pull`, IDE edits before
    /// the agent launched, files touched by another tool — are
    /// reconciled by the time the first MCP tool call arrives. The
    /// staleness-check stamp is updated on the way out so the first
    /// tool call doesn't re-walk the tree.
    ///
    /// The completion flag is flipped on every exit path (including
    /// errors) so [`Self::wait_for_startup_catch_up`] never hangs.
    pub async fn run_startup_catch_up_sync(&self) {
        let stale = self.cg.find_stale_files().await;
        if !stale.is_empty() {
            if let Err(e) = self.cg.sync_if_stale_silent(&stale).await {
                eprintln!("[tokensave] startup catch-up sync failed: {e}");
                self.startup_catch_up_done.store(true, Ordering::Release);
                return;
            }
        }
        self.refresh_file_token_map().await;
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs() as i64;
        self.last_staleness_check_at.store(now, Ordering::Release);
        self.startup_catch_up_done.store(true, Ordering::Release);
    }

    /// Returns `true` once the detached
    /// [`Self::run_startup_catch_up_sync`] task has finished (success
    /// or error). Production code never needs this — the MCP loop runs
    /// regardless of catch-up state — but tests poll it to avoid
    /// racing the catch-up task against later DB assertions.
    pub fn startup_catch_up_done(&self) -> bool {
        self.startup_catch_up_done.load(Ordering::Acquire)
    }

    /// Polls [`Self::startup_catch_up_done`] with a 25 ms interval up
    /// to `timeout`, returning `true` if catch-up completed within the
    /// budget. Tests use this to make the otherwise-detached #414
    /// task observable.
    pub async fn wait_for_startup_catch_up(&self, timeout: std::time::Duration) -> bool {
        let deadline = tokio::time::Instant::now() + timeout;
        while !self.startup_catch_up_done() {
            if tokio::time::Instant::now() >= deadline {
                return false;
            }
            tokio::time::sleep(std::time::Duration::from_millis(25)).await;
        }
        true
    }

    /// Walk the project tree, sync any stale files, and refresh the
    /// file-to-token-count map — but only if at least 30 s have passed
    /// since the last successful sync. The cooldown is the gate: while
    /// it holds, this returns immediately, so dropping it into every
    /// `tools/call` handler is cheap.
    ///
    /// Concurrent callers are serialized via
    /// [`Self::last_staleness_check_at`]: the first caller stamps `now`
    /// into the field with `compare_exchange`; later callers within the
    /// same window see the stamp and bail. If the actual sync work
    /// fails, the stamp still advances — failure to walk the tree
    /// should not cause every subsequent tool call to retry.
    pub async fn maybe_sync_if_stale(&self) {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs() as i64;
        let last_sync = self.cg.last_sync_timestamp().await;
        if now.saturating_sub(last_sync) < 30 {
            return;
        }

        let previous = self.last_staleness_check_at.load(Ordering::Acquire);
        if now.saturating_sub(previous) < 30 {
            return;
        }
        if self
            .last_staleness_check_at
            .compare_exchange(previous, now, Ordering::AcqRel, Ordering::Acquire)
            .is_err()
        {
            return;
        }

        let stale = self.cg.find_stale_files().await;
        if !stale.is_empty() {
            if let Err(e) = self.cg.sync_if_stale_silent(&stale).await {
                eprintln!("[tokensave] lazy sync failed: {e}");
                return;
            }
        }
        // Always refresh: a sibling MCP peer may have synced the DB
        // between our cooldown windows, in which case `stale` is empty
        // here but our in-memory `file_token_map` is still pre-sync.
        self.refresh_file_token_map().await;
    }

    /// Internal: snapshot of the current `file_token_map`. Exposed for
    /// integration tests only; not part of the stable public API.
    #[doc(hidden)]
    pub fn file_token_map_snapshot(&self) -> HashMap<String, u64> {
        self.file_token_map
            .lock()
            .map(|g| g.clone())
            .unwrap_or_default()
    }

    /// Flushes pending tokens to the worldwide counter if at least 30 seconds
    /// have elapsed since the last flush. Best-effort, never blocks for long.
    async fn maybe_flush_worldwide(&self) {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs() as i64;
        let last = self.last_flush_at.load(Ordering::Relaxed);
        if now - last < 30 {
            return;
        }
        // Mark as attempted immediately to prevent re-entry.
        self.last_flush_at.store(now, Ordering::Relaxed);

        let current = self.tokens_saved.load(Ordering::Relaxed);
        let last_flushed = self.last_flushed_tokens.load(Ordering::Relaxed);
        if current <= last_flushed {
            return;
        }
        let delta = current - last_flushed;

        let success = tokio::task::spawn_blocking(move || {
            let mut config = crate::user_config::UserConfig::load();
            config.pending_upload += delta;
            if config.upload_enabled && crate::cloud::flush_pending(config.pending_upload).is_some()
            {
                config.pending_upload = 0;
                config.last_upload_at = std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs() as i64;
                config.save();
                return true;
            }
            config.save();
            false
        })
        .await
        .unwrap_or(false);

        if success {
            self.last_flushed_tokens.store(current, Ordering::Relaxed);
        }
    }

    /// Returns a version-update warning if a newer release is available.
    /// Results are cached for `VERSION_CHECK_INTERVAL` (15 minutes).
    async fn check_version_update(&self) -> Option<String> {
        let current = env!("CARGO_PKG_VERSION");

        // Fast path: serve from cache if still fresh.
        {
            let cache = self.version_cache.lock().ok()?;
            if let Some(checked_at) = cache.checked_at {
                if checked_at.elapsed() < VERSION_CHECK_INTERVAL {
                    let latest = cache.latest.as_deref()?;
                    return if crate::cloud::is_newer_minor_version(current, latest) {
                        Some(format!(
                            "⚠️ tokensave v{current} is installed, but v{latest} is available. \
                             Run `tokensave upgrade` to update."
                        ))
                    } else {
                        None
                    };
                }
            }
        }

        // Cache miss or expired – fetch from GitHub (best-effort, 1 s timeout).
        let latest = tokio::task::spawn_blocking(crate::cloud::fetch_latest_version)
            .await
            .ok()
            .flatten();

        // Update cache regardless of fetch outcome so we don't retry immediately.
        if let Ok(mut cache) = self.version_cache.lock() {
            cache.latest.clone_from(&latest);
            cache.checked_at = Some(Instant::now());
        }

        let latest = latest?;
        if crate::cloud::is_newer_minor_version(current, &latest) {
            Some(format!(
                "⚠️ tokensave v{current} is installed, but v{latest} is available. \
                 Run `tokensave upgrade` to update."
            ))
        } else {
            None
        }
    }

    /// Process a single raw JSON-RPC line and write the response.
    /// Used to replay a peeked `initialize` message that was consumed before
    /// the server's main loop started.
    pub async fn handle_and_write(
        &self,
        line: &str,
        transport: &mut impl super::transport::McpTransport,
    ) {
        let parsed: std::result::Result<super::transport::JsonRpcRequest, _> =
            serde_json::from_str(line);
        let response = match parsed {
            Ok(request) => self.handle_request(&request).await,
            Err(e) => Some(super::transport::JsonRpcResponse::error(
                Value::Null,
                super::transport::ErrorCode::ParseError,
                format!("failed to parse JSON-RPC request: {e}"),
            )),
        };
        if let Some(resp) = response {
            let json_str = serde_json::to_string(&resp).unwrap_or_default();
            let _ = transport.write_line(&json_str).await;
            let _ = transport.flush().await;
        }
    }

    /// Runs the server, reading JSON-RPC requests from stdin and writing
    /// responses to stdout. Runs until stdin is closed or a shutdown signal
    /// (SIGINT/SIGTERM) is received, then performs graceful cleanup.
    pub async fn run(&self, transport: &mut impl super::transport::McpTransport) -> Result<()> {
        debug_assert!(
            self.stats.total_requests.load(Ordering::Relaxed) == 0,
            "server run() called on an already-used server"
        );

        loop {
            let line: String = {
                #[cfg(unix)]
                {
                    #[allow(clippy::expect_used)]
                    let mut sigterm =
                        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
                            .expect("failed to register SIGTERM handler");
                    tokio::select! {
                        result = transport.read_line() => {
                            match result {
                                Ok(Some(line)) => line,
                                _ => break,
                            }
                        }
                        _ = tokio::signal::ctrl_c() => break,
                        _ = sigterm.recv() => break,
                    }
                }
                #[cfg(not(unix))]
                {
                    tokio::select! {
                        result = transport.read_line() => {
                            match result {
                                Ok(Some(line)) => line,
                                _ => break,
                            }
                        }
                        _ = tokio::signal::ctrl_c() => break,
                    }
                }
            };

            let line = line.trim().to_string();
            if line.is_empty() {
                continue;
            }

            // Parse the incoming JSON
            let parsed: std::result::Result<JsonRpcRequest, _> = serde_json::from_str(&line);

            let response = match parsed {
                Ok(request) => self.handle_request(&request).await,
                Err(e) => Some(JsonRpcResponse::error(
                    Value::Null,
                    ErrorCode::ParseError,
                    format!("failed to parse JSON-RPC request: {e}"),
                )),
            };

            // Drain and write any pending notifications (e.g., version warnings).
            {
                let notifications: Vec<Value> = self
                    .pending_notifications
                    .lock()
                    .map(|mut p| p.drain(..).collect())
                    .unwrap_or_default();
                for notification in notifications {
                    if let Ok(s) = serde_json::to_string(&notification) {
                        let _ = transport.write_line(&format!("{s}\n")).await;
                        let _ = transport.flush().await;
                    }
                }
            }

            // Write response (if any) as a single line to stdout
            if let Some(resp) = response {
                let json_line = match serde_json::to_string(&resp) {
                    Ok(s) => s,
                    Err(e) => {
                        eprintln!("failed to serialize response: {e}");
                        continue;
                    }
                };
                let output = format!("{json_line}\n");
                if let Err(e) = transport.write_line(&output).await {
                    eprintln!("failed to write response: {e}");
                    break;
                }
                if let Err(e) = transport.flush().await {
                    eprintln!("failed to flush stdout: {e}");
                    break;
                }
            }
        }

        self.shutdown().await;
        Ok(())
    }

    /// Persists the tokens-saved counter, flushes pending tokens to the
    /// worldwide counter, checkpoints the WAL, and logs a session summary.
    ///
    /// Idempotent — safe to call multiple times. `run` invokes it once when
    /// its main loop exits; callers (e.g. `main.rs`, tests) may invoke it
    /// explicitly afterwards without re-running the persistence logic.
    pub async fn shutdown(&self) {
        // Idempotency guard: only run the persistence path once.
        if self.shutdown_done.swap(true, Ordering::SeqCst) {
            return;
        }

        let uptime = self.stats.started_at.elapsed();
        let tool_calls = self.stats.tool_calls.load(Ordering::Relaxed);
        let tokens_saved = self.tokens_saved.load(Ordering::Relaxed);

        // Persist final tokens-saved value
        if let Err(e) = self.cg.set_tokens_saved(tokens_saved).await {
            eprintln!("[tokensave] warning: failed to persist tokens_saved on shutdown: {e}");
        }

        // Update global DB with final count and checkpoint it
        if let Some(ref gdb) = self.global_db {
            gdb.upsert(self.cg.project_root(), tokens_saved).await;
            gdb.checkpoint().await;
        }

        // Flush remaining delta to worldwide counter (what periodic flushes missed)
        let last_flushed = self.last_flushed_tokens.load(Ordering::Relaxed);
        if tokens_saved > last_flushed {
            let delta = tokens_saved - last_flushed;
            let mut config = crate::user_config::UserConfig::load();
            config.pending_upload += delta;
            if config.upload_enabled {
                if let Some(_total) = crate::cloud::flush_pending(config.pending_upload) {
                    config.pending_upload = 0;
                    let now = std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .unwrap_or_default()
                        .as_secs() as i64;
                    config.last_upload_at = now;
                }
            }
            config.save();
        }

        // Checkpoint WAL to merge it into the main database file
        if let Err(e) = self.cg.checkpoint().await {
            eprintln!("[tokensave] warning: failed to checkpoint WAL on shutdown: {e}");
        }

        eprintln!(
            "[tokensave] shutdown: {} tool calls, ~{} tokens saved, uptime {}s",
            tool_calls,
            tokens_saved,
            uptime.as_secs()
        );
    }

    /// Dispatches a parsed JSON-RPC request to the appropriate handler.
    ///
    /// Returns `None` for notifications (requests without an `id`).
    pub(crate) async fn handle_request(&self, request: &JsonRpcRequest) -> Option<JsonRpcResponse> {
        debug_assert!(
            !request.method.is_empty(),
            "handle_request called with empty method"
        );
        self.stats.total_requests.fetch_add(1, Ordering::Relaxed);
        let id = request.id.clone();

        let result = match request.method.as_str() {
            "initialize" => Some(Self::handle_initialize(id)),
            "initialized" => {
                // Notification - no response required
                None
            }
            "notifications/initialized" => {
                // Alternative notification path - no response required
                None
            }
            "tools/list" => Some(self.handle_tools_list(id).await),
            "tools/call" => Some(self.handle_tools_call(id, request.params.as_ref()).await),
            "resources/list" => Some(Self::handle_resources_list(id)),
            "resources/read" => Some(
                self.handle_resources_read(id, request.params.as_ref())
                    .await,
            ),
            "ping" | "logging/setLevel" => Some(JsonRpcResponse::success(id, json!({}))),
            _ => Some(JsonRpcResponse::error(
                id,
                ErrorCode::MethodNotFound,
                format!("method not found: {}", request.method),
            )),
        };

        // Track errors
        if let Some(ref resp) = result {
            if resp.error.is_some() {
                self.stats.errors.fetch_add(1, Ordering::Relaxed);
            }
        }

        result
    }

    /// Handles the `initialize` method, returning server capabilities.
    fn handle_initialize(id: Value) -> JsonRpcResponse {
        JsonRpcResponse::success(
            id,
            json!({
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {},
                    "resources": {},
                    "logging": {}
                },
                "serverInfo": {
                    "name": "tokensave",
                    "version": env!("CARGO_PKG_VERSION")
                },
                "instructions": "tokensave is a code-graph MCP server. \
                    Start with tokensave_context for any code exploration task \
                    — it returns relevant symbols, relationships, and code \
                    snippets for a natural-language query. Use tokensave_search \
                    to find specific symbols by name. Discovery and analysis \
                    tools are read-only and safe to call in parallel. Edit \
                    and session-memory tools can mutate local project state \
                    and declare readOnlyHint=false. \
                    When a tool result contains a `tokensave_metrics:` line, \
                    report the savings to the user (e.g. 'TokenSave\\'d ~N tokens')."
            }),
        )
    }

    /// Handles the `tools/list` method, returning all available tool definitions.
    async fn handle_tools_list(&self, id: Value) -> JsonRpcResponse {
        let node_count = self.cg.get_stats().await.map_or(0, |s| s.node_count);
        let budget = explore_call_budget(node_count);
        let tools = get_tool_definitions_with_budget(node_count, budget);
        JsonRpcResponse::success(id, json!({ "tools": tools }))
    }

    /// Handles the `resources/list` method, returning available resources.
    fn handle_resources_list(id: Value) -> JsonRpcResponse {
        JsonRpcResponse::success(
            id,
            json!({
                "resources": [
                    {
                        "uri": "tokensave://status",
                        "name": "Graph Status",
                        "description": "Code graph statistics: node/edge/file counts, languages, DB size, and index freshness.",
                        "mimeType": "application/json"
                    },
                    {
                        "uri": "tokensave://files",
                        "name": "File List",
                        "description": "All indexed project files grouped by directory with symbol counts.",
                        "mimeType": "text/plain"
                    },
                    {
                        "uri": "tokensave://overview",
                        "name": "Project Overview",
                        "description": "High-level project summary: language distribution, largest modules, and top entry points.",
                        "mimeType": "text/plain"
                    },
                    {
                        "uri": "tokensave://branches",
                        "name": "Tracked Branches",
                        "description": "List of tracked branches with DB sizes, parent branch, and last sync time. Empty if multi-branch is not active.",
                        "mimeType": "application/json"
                    },
                    {
                        "uri": "tokensave://schema",
                        "name": "SQLite Schema",
                        "description": "Documentation for the .tokensave/tokensave.db schema: tables, columns, indexes, and common query recipes. Use when MCP tools don't cover your query and you need to drop down to raw SQL.",
                        "mimeType": "text/markdown"
                    }
                ]
            }),
        )
    }

    /// Handles the `resources/read` method, returning resource contents.
    async fn handle_resources_read(&self, id: Value, params: Option<&Value>) -> JsonRpcResponse {
        let uri = params.and_then(|p| p.get("uri")).and_then(|v| v.as_str());

        let Some(uri) = uri else {
            return JsonRpcResponse::error(
                id,
                ErrorCode::InvalidParams,
                "missing 'uri' in resources/read params".to_string(),
            );
        };

        match uri {
            "tokensave://status" => self.read_resource_status(id).await,
            "tokensave://files" => self.read_resource_files(id).await,
            "tokensave://overview" => self.read_resource_overview(id).await,
            "tokensave://branches" => self.read_resource_branches(id),
            "tokensave://schema" => Self::read_resource_schema(id),
            _ => JsonRpcResponse::error(
                id,
                ErrorCode::InvalidParams,
                format!("unknown resource URI: {uri}"),
            ),
        }
    }

    /// Returns the `SQLite` schema documentation as a markdown resource.
    /// Sourced from `src/db/migrations.rs::create_schema` — keep in sync.
    fn read_resource_schema(id: Value) -> JsonRpcResponse {
        JsonRpcResponse::success(
            id,
            json!({
                "contents": [{
                    "uri": "tokensave://schema",
                    "mimeType": "text/markdown",
                    "text": SCHEMA_MARKDOWN
                }]
            }),
        )
    }

    /// Returns graph statistics as a JSON resource.
    async fn read_resource_status(&self, id: Value) -> JsonRpcResponse {
        match self.cg.get_stats().await {
            Ok(stats) => {
                let text = serde_json::to_string_pretty(&stats).unwrap_or_default();
                JsonRpcResponse::success(
                    id,
                    json!({
                        "contents": [{
                            "uri": "tokensave://status",
                            "mimeType": "application/json",
                            "text": text
                        }]
                    }),
                )
            }
            Err(e) => JsonRpcResponse::error(
                id,
                ErrorCode::InternalError,
                format!("failed to read graph stats: {e}"),
            ),
        }
    }

    /// Returns the file list as a text resource (grouped by directory).
    async fn read_resource_files(&self, id: Value) -> JsonRpcResponse {
        match self.cg.get_all_files().await {
            Ok(mut files) => {
                files.sort_by(|a, b| a.path.cmp(&b.path));
                let mut groups: std::collections::BTreeMap<String, Vec<String>> =
                    std::collections::BTreeMap::new();
                for f in &files {
                    let dir = f.path.rfind('/').map_or(".", |i| &f.path[..i]).to_string();
                    #[allow(clippy::map_unwrap_or)]
                    let name = f
                        .path
                        .rfind('/')
                        .map(|i| &f.path[i + 1..])
                        .unwrap_or(&f.path);
                    groups
                        .entry(dir)
                        .or_default()
                        .push(format!("{} ({} symbols)", name, f.node_count));
                }
                let mut lines = Vec::new();
                lines.push(format!("{} indexed files", files.len()));
                for (dir, entries) in &groups {
                    lines.push(format!("\n{}/ ({} files)", dir, entries.len()));
                    for entry in entries {
                        lines.push(format!("  {entry}"));
                    }
                }
                let text = lines.join("\n");
                JsonRpcResponse::success(
                    id,
                    json!({
                        "contents": [{
                            "uri": "tokensave://files",
                            "mimeType": "text/plain",
                            "text": text
                        }]
                    }),
                )
            }
            Err(e) => JsonRpcResponse::error(
                id,
                ErrorCode::InternalError,
                format!("failed to read file list: {e}"),
            ),
        }
    }

    /// Returns a high-level project overview as a text resource.
    async fn read_resource_overview(&self, id: Value) -> JsonRpcResponse {
        let stats = match self.cg.get_stats().await {
            Ok(s) => s,
            Err(e) => {
                return JsonRpcResponse::error(
                    id,
                    ErrorCode::InternalError,
                    format!("failed to read graph stats: {e}"),
                );
            }
        };

        let mut lines = Vec::new();
        lines.push(format!("Project: {}", self.cg.project_root().display()));
        lines.push(format!(
            "Graph: {} nodes, {} edges, {} files",
            stats.node_count, stats.edge_count, stats.file_count
        ));

        // Language distribution
        if !stats.files_by_language.is_empty() {
            lines.push("\nLanguages:".to_string());
            let mut langs: Vec<_> = stats.files_by_language.iter().collect();
            langs.sort_by(|a, b| b.1.cmp(a.1));
            for (lang, count) in &langs {
                lines.push(format!("  {lang} ({count} files)"));
            }
        }

        // Node kind distribution (top 10)
        if !stats.nodes_by_kind.is_empty() {
            lines.push("\nSymbol kinds:".to_string());
            let mut kinds: Vec<_> = stats.nodes_by_kind.iter().collect();
            kinds.sort_by(|a, b| b.1.cmp(a.1));
            for (kind, count) in kinds.iter().take(10) {
                lines.push(format!("  {kind} ({count})"));
            }
        }

        let text = lines.join("\n");
        JsonRpcResponse::success(
            id,
            json!({
                "contents": [{
                    "uri": "tokensave://overview",
                    "mimeType": "text/plain",
                    "text": text
                }]
            }),
        )
    }

    fn read_resource_branches(&self, id: Value) -> JsonRpcResponse {
        let tokensave_dir = crate::config::get_tokensave_dir(self.cg.project_root());
        let current = self.cg.active_branch();

        let branches: Vec<Value> = match crate::branch_meta::load_branch_meta(&tokensave_dir) {
            Some(meta) => meta
                .branches
                .iter()
                .map(|(name, entry)| {
                    let db_path = tokensave_dir.join(&entry.db_file);
                    let size_bytes = db_path.metadata().map_or(0, |m| m.len());
                    json!({
                        "name": name,
                        "db_file": entry.db_file,
                        "parent": entry.parent,
                        "size_bytes": size_bytes,
                        "last_synced_at": entry.last_synced_at,
                        "is_current": current == Some(name.as_str()),
                        "is_default": name == &meta.default_branch,
                    })
                })
                .collect(),
            None => vec![],
        };

        let output = json!({
            "branch_count": branches.len(),
            "branches": branches,
        });
        let text = serde_json::to_string_pretty(&output).unwrap_or_default();
        JsonRpcResponse::success(
            id,
            json!({
                "contents": [{
                    "uri": "tokensave://branches",
                    "mimeType": "application/json",
                    "text": text
                }]
            }),
        )
    }

    /// Handles the `tools/call` method, dispatching to the appropriate tool handler.
    async fn handle_tools_call(&self, id: Value, params: Option<&Value>) -> JsonRpcResponse {
        debug_assert!(
            !id.is_null(),
            "handle_tools_call called with null request id"
        );
        let Some(params) = params else {
            return JsonRpcResponse::error(
                id,
                ErrorCode::InvalidParams,
                "missing params for tools/call".to_string(),
            );
        };

        let Some(tool_name) = params.get("name").and_then(|v| v.as_str()) else {
            return JsonRpcResponse::error(
                id,
                ErrorCode::InvalidParams,
                "missing 'name' in tools/call params".to_string(),
            );
        };

        let arguments = params.get("arguments").cloned().unwrap_or(json!({}));

        // Notification-free freshness: walk the tree and resync any stale
        // files, gated by a 30 s cooldown. Replaces the embedded watcher
        // (see McpServer::new). No-op on the hot path most of the time.
        self.maybe_sync_if_stale().await;

        self.stats.tool_calls.fetch_add(1, Ordering::Relaxed);
        eprintln!("[tokensave] tool call: {tool_name}");
        if let Ok(mut counts) = self.tool_call_counts.lock() {
            *counts.entry(tool_name.to_string()).or_insert(0) += 1;
        }

        let server_stats = if tool_name == "tokensave_status" {
            Some(self.server_stats_json().await)
        } else {
            None
        };

        let timings_enabled = self.timings_enabled();
        let handler_start = if timings_enabled {
            Some(std::time::Instant::now())
        } else {
            None
        };
        let dispatch_outcome = handle_tool_call(
            &self.cg,
            tool_name,
            arguments,
            server_stats,
            self.scope_prefix(),
        )
        .await;
        let handler_elapsed_us = handler_start.map(|t| t.elapsed().as_micros() as u64);
        match dispatch_outcome {
            Ok(mut result) => {
                if let Some(us) = handler_elapsed_us {
                    let obj = result.value.as_object_mut();
                    if let Some(map) = obj {
                        let meta = map.entry("_meta").or_insert_with(|| json!({}));
                        if let Some(meta_obj) = meta.as_object_mut() {
                            meta_obj.insert("duration_us".to_string(), json!(us));
                        }
                    }
                }
                let raw_file_tokens = self.accumulate_tokens_saved(&result.touched_files).await;
                crate::monitor::write_entry(
                    self.cg.project_root(),
                    "tokensave",
                    tool_name,
                    raw_file_tokens,
                    raw_file_tokens,
                );
                self.maybe_flush_worldwide().await;

                // Estimate approximate token count of the graph response.
                let response_tokens: u64 = result
                    .value
                    .get("content")
                    .and_then(|c| c.as_array())
                    .map_or(0, |arr| {
                        let total_chars: usize = arr
                            .iter()
                            .filter_map(|item| item.get("text").and_then(|t| t.as_str()))
                            .map(str::len)
                            .sum();
                        (total_chars / 4) as u64
                    });

                // Append per-call token savings to the response content.
                if raw_file_tokens > 0 {
                    if let Some(content) = result
                        .value
                        .get_mut("content")
                        .and_then(|c| c.as_array_mut())
                    {
                        content.push(json!({"type": "text", "text": format!(
                            "\ntokensave_metrics: before={raw_file_tokens} after={response_tokens}"
                        )}));
                    }
                }

                // Persist to the cross-project savings ledger (best-effort, non-blocking).
                {
                    let project_path_str = self.cg.project_root().to_string_lossy().to_string();
                    let tool_name_owned = tool_name.to_string();
                    let ts = crate::tokensave::current_timestamp();
                    tokio::spawn(async move {
                        if let Some(gdb) = crate::global_db::GlobalDb::open().await {
                            gdb.record_savings(
                                &project_path_str,
                                &tool_name_owned,
                                raw_file_tokens,
                                response_tokens,
                                ts,
                            )
                            .await;
                        }
                    });
                }

                // Prepend version-update warning + queue logging notification.
                if let Some(warning) = self.check_version_update().await {
                    if let Some(content) = result
                        .value
                        .get_mut("content")
                        .and_then(|c| c.as_array_mut())
                    {
                        content.insert(0, json!({"type": "text", "text": &warning}));
                    }
                    if let Ok(mut pending) = self.pending_notifications.lock() {
                        pending.push(json!({
                            "jsonrpc": "2.0",
                            "method": "notifications/message",
                            "params": {
                                "level": "warning",
                                "logger": "tokensave",
                                "data": warning
                            }
                        }));
                    }
                }

                // Per-file staleness banner (#428 design): files this response
                // referenced that are still pending after the in-line sync
                // attempt get a focused banner naming them with edit ages,
                // telling the agent to Read THOSE files directly while
                // treating the rest of the response as authoritative.
                // Replaces the previous all-or-nothing "STALE INDEX"
                // warning that made agents distrust the entire answer.
                if !result.touched_files.is_empty() {
                    let stale_files = self.cg.check_file_staleness(&result.touched_files).await;
                    if !stale_files.is_empty() {
                        let still_stale = match self.cg.sync_if_stale(&stale_files).await {
                            Ok(false) => false,        // sync completed; files now fresh
                            Ok(true) | Err(_) => true, // still stale (lock contention / sync error)
                        };
                        if still_stale {
                            let banner = format_per_file_staleness_banner(
                                self.cg.project_root(),
                                &stale_files,
                            );
                            // Machine-readable marker. Same shape as before
                            // so existing scrapers keep working.
                            let stale_json = serde_json::to_string(&stale_files)
                                .unwrap_or_else(|_| "[]".to_string());
                            let marker = format!("\ntokensave_graph_stale: {stale_json}");
                            debug_assert!(
                                result.value.is_object(),
                                "tool result must be a JSON object so graph_stale can be attached"
                            );
                            if let Some(obj) = result.value.as_object_mut() {
                                obj.insert("graph_stale".to_string(), json!(stale_files));
                            }
                            if let Some(content) = result
                                .value
                                .get_mut("content")
                                .and_then(|c| c.as_array_mut())
                            {
                                content.insert(0, json!({"type": "text", "text": &banner}));
                                content.push(json!({"type": "text", "text": marker}));
                            }
                        }
                    }
                }

                // Warn if serving from a fallback (ancestor) branch DB.
                if let Some(warning) = self.cg.fallback_warning() {
                    let warning = format!("WARNING: {warning}");
                    if let Some(content) = result
                        .value
                        .get_mut("content")
                        .and_then(|c| c.as_array_mut())
                    {
                        content.insert(0, json!({"type": "text", "text": &warning}));
                    }
                }

                // Check overall index age (warn if older than 1 hour).
                // Uses `last_sync_timestamp` (sync execution time) not the
                // max file `indexed_at` — a no-change sync still updates the
                // sync metadata even though no file gets a fresh `indexed_at`,
                // so a per-file fallback fires the warning forever on quiet
                // repos (#86).
                {
                    let last_time = self.cg.last_sync_timestamp().await;
                    let now = std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .unwrap_or_default()
                        .as_secs() as i64;
                    let age_secs = now - last_time;
                    if last_time > 0 && age_secs > 3600 {
                        let hours = age_secs / 3600;
                        let mins = (age_secs % 3600) / 60;
                        let warning = if hours >= 24 {
                            format!(
                                "WARNING: Index last synced {}d {}h ago. Run `tokensave sync` to update.",
                                hours / 24, hours % 24
                            )
                        } else {
                            format!(
                                "WARNING: Index last synced {hours}h {mins}m ago. Run `tokensave sync` to update."
                            )
                        };
                        if let Some(content) = result
                            .value
                            .get_mut("content")
                            .and_then(|c| c.as_array_mut())
                        {
                            content.insert(0, json!({"type": "text", "text": &warning}));
                        }
                    }
                }

                // Borrowed-worktree heads-up (#312). Inserted LAST so it
                // appears FIRST in the response — the index serving the
                // wrong branch is the most serious of these warnings to
                // surface to the agent.
                if let Some(ref m) = self.worktree_mismatch {
                    let notice = crate::worktree::worktree_mismatch_notice(m);
                    if let Some(content) = result
                        .value
                        .get_mut("content")
                        .and_then(|c| c.as_array_mut())
                    {
                        content.insert(0, json!({"type": "text", "text": notice}));
                    }
                }

                JsonRpcResponse::success(id, result.value)
            }
            Err(e) => JsonRpcResponse::error(
                id,
                ErrorCode::InternalError,
                format!("tool execution failed: {e}"),
            ),
        }
    }

    /// Returns the current server runtime statistics as a JSON value.
    pub async fn server_stats_json(&self) -> Value {
        let uptime = self.stats.started_at.elapsed();
        let tool_counts: Value = self
            .tool_call_counts
            .lock()
            .map(|counts| json!(*counts))
            .unwrap_or(json!({}));

        let mut stats = json!({
            "uptime_secs": uptime.as_secs(),
            "total_requests": self.stats.total_requests.load(Ordering::Relaxed),
            "tool_calls": self.stats.tool_calls.load(Ordering::Relaxed),
            "errors": self.stats.errors.load(Ordering::Relaxed),
            "tool_call_counts": tool_counts,
            "approx_tokens_saved": self.tokens_saved.load(Ordering::Relaxed),
        });

        if let Some(ref gdb) = self.global_db {
            if let Some(global_total) = gdb.global_tokens_saved().await {
                let local = self.tokens_saved.load(Ordering::Relaxed);
                stats["global_tokens_saved"] = json!(global_total.saturating_sub(local));
            }
        }

        // Surface the verbose worktree-mismatch warning when present, so
        // `tokensave_status` is the one tool whose output is loud about
        // serving a borrowed index (#312).
        if let Some(ref m) = self.worktree_mismatch {
            stats["worktree_mismatch"] = json!({
                "worktree_root": m.worktree_root.display().to_string(),
                "index_root": m.index_root.display().to_string(),
                "warning": crate::worktree::worktree_mismatch_warning(m),
            });
        }

        stats
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod staleness_banner_tests {
    use super::{format_per_file_staleness_banner, humanize_age};
    use std::fs;
    use tempfile::tempdir;

    #[test]
    fn humanize_age_picks_right_unit() {
        assert_eq!(humanize_age(0), "0s ago");
        assert_eq!(humanize_age(45), "45s ago");
        assert_eq!(humanize_age(125), "2m ago");
        assert_eq!(humanize_age(3_700), "1h ago");
        assert_eq!(humanize_age(90_000), "1d ago");
    }

    #[test]
    fn banner_lists_stale_files_with_age() {
        let tmp = tempdir().unwrap();
        let root = tmp.path();
        fs::create_dir_all(root.join("src")).unwrap();
        fs::write(root.join("src/a.rs"), "fn a() {}").unwrap();
        fs::write(root.join("src/b.rs"), "fn b() {}").unwrap();

        let stale = vec!["src/a.rs".to_string(), "src/b.rs".to_string()];
        let banner = format_per_file_staleness_banner(root, &stale);
        assert!(banner.contains("2 file(s) referenced below were edited"));
        assert!(banner.contains("src/a.rs ("));
        assert!(banner.contains("src/b.rs ("));
        assert!(banner.contains("ago)"));
        assert!(banner.contains("tokensave sync"));
        // Critical UX shift: should NOT say "STALE INDEX" — the whole
        // point of #428 is to scope the warning, not blanket-distrust
        // the entire response.
        assert!(!banner.contains("STALE INDEX"));
    }

    #[test]
    fn banner_handles_missing_file_gracefully() {
        let tmp = tempdir().unwrap();
        let stale = vec!["does/not/exist.rs".to_string()];
        let banner = format_per_file_staleness_banner(tmp.path(), &stale);
        // Missing files still get listed (e.g. file deleted between
        // sync and tool response). Age falls back to 0s.
        assert!(banner.contains("does/not/exist.rs"));
    }
}