trusty-mpm 1.8.2

//! Daemon-side implementation of the `disk_survey` MCP tool (#6927).
//!
//! Why: `StateBackend` (in `mcp_backend.rs`) has to service the Disk
//! dashboard's survey, and inlining its body there would push that file over
//! the 500-SLOC production cap — the same reason `mcp_console.rs` and
//! `mcp_project.rs` exist. This module is also where the tool's real-world
//! inputs are resolved: the operator's workspace root and keep-list, the live
//! session store, the delegation registry, `gh`, and the daemon's shared size
//! index. [`crate::disk::survey_run`] itself resolves none of them, which is
//! what keeps it hermetically testable.
//! What: one free async function, [`disk_survey`], returning DOC-73 §16.5's
//! tree as JSON.
//! Test: `crate::mcp::tests::dispatch_disk_survey_tool` drives the dispatch
//! path against the mock backend; `crate::disk::survey_tests` covers the survey
//! itself over scratch repositories.
//!
//! READ-ONLY. Nothing here removes, prunes, or writes anything — DOC-73 §16.6
//! item 4 owns the clear action, behind an explicit operator confirm.

use std::cell::RefCell;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};

use serde_json::Value;

use crate::core::trusty_tools_config::{self, TrustyToolsConfig};
use crate::daemon::disk_survey_cache::{Plan, SurveyKey};
use crate::daemon::state::DaemonState;
use crate::disk::size_index::DirSize;
use crate::disk::survey::GroupBy;
use crate::disk::survey_run::{self, DiskProbes, run};
use crate::session_manager::worktree_ownership::AgentWorktreeOwner;
use crate::session_manager::worktree_reclaim::{
    BranchPrState, PrIndex, pr_state_for_branch_within,
};
use crate::session_manager::worktree_reclaim_gh::GH_TIMEOUT;
use crate::session_manager::worktree_registry::ScannedWorktree;
use crate::session_manager::worktree_safety::inspect_dirt;

/// The largest `budget_seconds` this tool will honour (#7313).
///
/// Why 55 and not "whatever the caller asked for": every MCP caller reaches the
/// daemon through `tm serve --stdio`, whose reqwest client carries a 60-second
/// `REQUEST_TIMEOUT` (`crate::bin::tm::commands::serve_stdio`, a path this
/// module cannot import — the two numbers are bound instead by
/// `the_forwarding_timeout_leaves_headroom_over_the_disk_survey_budget_clamp`
/// in that module's own tests). A larger budget therefore cannot produce a
/// survey the caller will ever see: the bridge gives up at 60 s and returns a
/// transport error while the daemon keeps working, which is the #6929 failure
/// mode over again — a 502 carrying no survey where a truncated one was
/// available. 55 leaves five seconds for serialization and the round trip.
/// The console already clamps on its own side (`trusty-console`'s
/// `routes/disk.rs`); this is the same rule at the tool, for every other caller.
/// Test: `a_budget_past_the_bridge_timeout_is_clamped`.
pub const MAX_BUDGET_SECONDS: u64 = 55;

/// Apply [`MAX_BUDGET_SECONDS`], saying whether it bit (#7313).
///
/// What: the budget to use, and whether the caller's figure was reduced — which
/// the response reports as `budget_clamped`, so a caller that asked for 120 and
/// got a partial pass can tell the clamp from a genuinely slow fleet.
/// Test: `a_budget_past_the_bridge_timeout_is_clamped`,
/// `a_budget_inside_the_bridge_timeout_is_untouched`.
fn clamp_budget(requested: Option<u64>) -> (Option<u64>, bool) {
    match requested {
        Some(s) if s > MAX_BUDGET_SECONDS => (Some(MAX_BUDGET_SECONDS), true),
        other => (other, false),
    }
}

/// Read the caller's `group_by` argument (#7313).
///
/// Why an error rather than a silent fall-through on an unknown value: a typo
/// would otherwise return a complete, plausible survey with the roll-up the
/// caller asked for simply absent, and nothing in the payload saying why.
/// Test: `an_unknown_group_by_is_rejected`.
fn parse_group_by(group_by: Option<&str>) -> Result<GroupBy, String> {
    match group_by {
        None => Ok(GroupBy::None),
        Some("session") => Ok(GroupBy::Session),
        Some(other) => Err(format!(
            "disk_survey: unknown `group_by` value `{other}`; the only supported value is `session`"
        )),
    }
}

/// Survey every managed worktree for the console Disk view (#6927).
///
/// Why: see the module docs. The whole pass runs on the blocking pool because
/// every probe under it is a synchronous subprocess — `git status`,
/// `git worktree list`, `gh pr list` — and a byte walk; running that on a
/// runtime worker would stall the daemon's other routes.
/// What: resolves the workspace root and keep-list from the operator's config,
/// snapshots the live session claims, and hands `disk::survey_run::run` a
/// `gh`-backed pull-request probe plus a measurement closure over the daemon's
/// SHARED [`crate::disk::size_index::DirSizeIndex`] — never a second index, so
/// a poll every few seconds costs cache hits rather than walks, and the mutex
/// is held for each measurement alone rather than for the whole pass.
/// `budget_seconds` bounds the whole pass, byte walks included (#6929);
/// worktrees past it are listed as `review`, never omitted and never `stale`,
/// and a project or root past it reports no byte figure. It is clamped to
/// [`MAX_BUDGET_SECONDS`] (#7313) — see that constant for why a larger figure
/// cannot produce a survey any caller receives.
/// `group_by` (#7313) adds the per-session roll-up to the response.
/// #8985: a fleet too large for the clamp is answered from the last complete
/// pass, which an unbudgeted background pass keeps current — see
/// [`crate::daemon::disk_survey_cache`]. The response's `freshness`,
/// `age_seconds` and `background_pass` say which answer it is. Only a
/// budgeted call is answered from the cache; an omitted budget always runs
/// live.
/// Test: `crate::disk::survey_tests`, `disk_survey_cache_tests`,
/// `an_omitted_budget_always_runs_a_live_survey`, and
/// `dispatch_disk_survey_tool` for the dispatch wiring.
pub async fn disk_survey(
    state: &Arc<DaemonState>,
    project: Option<&str>,
    budget_seconds: Option<u64>,
    group_by: Option<&str>,
) -> Result<Value, String> {
    let group = parse_group_by(group_by)?;
    // #7313: clamped BEFORE the deadline is computed, so the clamp is what the
    // pass actually runs under rather than a figure reported beside a longer one.
    let (budget_seconds, budget_clamped) = clamp_budget(budget_seconds);
    let key = SurveyKey {
        project: project.map(str::to_string),
        group_by: group_by.map(str::to_string),
    };
    let cache = state.disk_survey_cache();
    // #8985 review: an omitted budget asks for a live, unbounded survey, so it
    // never consults the cache. Its complete pass is still stored, and clears
    // the key's budget bit, through `after_live`.
    let plan = match budget_seconds {
        Some(_) => cache.plan(&key, Instant::now()),
        None => Plan::Live,
    };
    let mut value = match plan {
        Plan::Serve { survey, refresh } => {
            if refresh {
                spawn_background_pass(state, key, group);
            }
            survey
        }
        Plan::Live => {
            let started = Instant::now();
            let deadline = budget_seconds.map(|s| started + Duration::from_secs(s));
            let live = survey_pass(state, project, deadline, group).await?;
            let (answer, refresh) = cache.after_live(&key, live, started, Instant::now());
            if refresh {
                spawn_background_pass(state, key, group);
            }
            answer
        }
    };
    // #7313: the clamp is a TRANSPORT fact, not a survey fact — the survey
    // ran a full pass under whatever deadline it was handed and has nothing
    // to say about what the caller originally asked for. Recorded here, on
    // the response, so it stays out of `DiskSurvey`'s shape.
    if let Value::Object(map) = &mut value {
        map.insert("budget_clamped".to_string(), Value::Bool(budget_clamped));
    }
    Ok(value)
}

/// Run one unbudgeted pass for `key` off the caller's path (#8985).
///
/// Why: the caller holds the refresh slot [`DiskSurveyCache::plan`] or
/// [`DiskSurveyCache::after_live`] gave it, and must hand it back whatever the
/// pass does — [`DiskSurveyCache::after_background`] always frees it.
/// Test: `a_budgeted_call_is_answered_from_the_background_pass`.
///
/// [`DiskSurveyCache::plan`]: crate::daemon::disk_survey_cache::DiskSurveyCache::plan
/// [`DiskSurveyCache::after_live`]: crate::daemon::disk_survey_cache::DiskSurveyCache::after_live
/// [`DiskSurveyCache::after_background`]: crate::daemon::disk_survey_cache::DiskSurveyCache::after_background
fn spawn_background_pass(state: &Arc<DaemonState>, key: SurveyKey, group: GroupBy) {
    let pass_state = Arc::clone(state);
    let project = key.project.clone();
    let pass = async move { survey_pass(&pass_state, project.as_deref(), None, group).await };
    drop(settle_in_background(state, key, pass));
}

/// Run `pass` in the background and hand its result, and the refresh slot,
/// back to the cache (#8985).
///
/// What: `pass` runs in its own task, so a panic arrives here as a join
/// error and is recorded as a failed pass — the slot is freed either way.
/// The returned handle lets a test wait for the hand-back.
/// Test: `a_failed_or_panicking_background_pass_hands_the_slot_back`.
fn settle_in_background<F>(
    state: &Arc<DaemonState>,
    key: SurveyKey,
    pass: F,
) -> tokio::task::JoinHandle<()>
where
    F: std::future::Future<Output = Result<Value, String>> + Send + 'static,
{
    let state = Arc::clone(state);
    tokio::spawn(async move {
        let started = Instant::now();
        // #8985 review: a panic inside `pass` must not strand the slot.
        let result = tokio::spawn(pass)
            .await
            .unwrap_or_else(|e| Err(format!("disk_survey: background pass panicked: {e}")));
        state
            .disk_survey_cache()
            .after_background(&key, result, started);
    })
}

/// One survey pass under `deadline`, serialized (#6927, #8985).
///
/// What: the body `disk_survey` ran before #8985. With a deadline, each
/// measurement waits for the shared index no longer than the time left, so a
/// background pass holding the index for a walk cannot push a live call past
/// its clamp.
async fn survey_pass(
    state: &Arc<DaemonState>,
    project: Option<&str>,
    deadline: Option<Instant>,
    group_by: GroupBy,
) -> Result<Value, String> {
    let config = TrustyToolsConfig::load();
    let repos_root = trusty_tools_config::workspace_root(&config);
    // #6927: read FALLIBLY, separately from the lenient `load` above. A config
    // that will not parse yields a keep-list that keeps everything and says so
    // in `keep_list.error`, rather than an empty one that protects nothing.
    let keep_list = trusty_tools_config::load_disk_keep_list();
    // The claim set is a SNAPSHOT, which is correct here and would not be on a
    // destructive path: this tool only displays, and the delete path re-reads
    // liveness per candidate immediately before each removal (#2919).
    let manager = state.session_manager().await;
    // #7232: through the crate's single claim producer, so this report and the
    // reclaim sweep agree about which claims are tombstones. The console is not
    // a managed session, so it names no caller — every LIVE claim it sees is
    // foreign, which is the conservative reading.
    let claims = manager.workspace_claims(None).await;
    let index = state.disk_size_index();
    let state_for_agents = Arc::clone(state);
    let project = project.map(str::to_string);
    // #7357: this tool is the entry point, so it resolves the adopted anchors
    // under the daemon's own framework root; `survey_run::run` takes them.
    let adopted = crate::project::adopted_anchors_under(state.framework_root());

    tokio::task::spawn_blocking(move || {
        let agent_state = |owner: &AgentWorktreeOwner| {
            crate::daemon::services::agent_worktree_reap::delegation_state_for_agent(
                &state_for_agents,
                &owner.agent_id,
            )
        };
        // One `gh` index per repository, not one per worktree — the same
        // amortization `survey_with_index` applies, for the same reason.
        let indexes: RefCell<BTreeMap<PathBuf, PrIndex>> = RefCell::new(BTreeMap::new());
        // #6929: the survey's remaining time is the ceiling on both `gh` calls
        // this probe can make, so a worktree admitted late cannot spend twenty
        // seconds of `GH_TIMEOUT` past the deadline the console is holding.
        let pr_state = |scanned: &ScannedWorktree, left: Option<Duration>| -> BranchPrState {
            pr_for(&indexes, scanned, left)
        };
        // #6927 review: the lock spans ONE measurement, not the pass. Held
        // across `run` it would also cover every `git status` and `gh pr list`
        // that `classify` shells out to, so a second `disk_survey` — or the
        // #6926 background refresher — would queue behind minutes of network
        // work for an index it only wanted to read.
        // #6929: `budget` is whatever the survey deadline has left. Passing it
        // through is what stops one cold walk from spending the index's fixed
        // 30-second ceiling and overrunning the survey the console is waiting
        // on.
        let measure = |path: &Path, budget: Option<Duration>| -> Option<DirSize> {
            measure_within_budget(&index, path, budget)
        };
        let probes = DiskProbes {
            pr_state: &pr_state,
            claims: &claims,
            agent_state: &agent_state,
            dirt: &inspect_dirt,
            measure: &measure,
            now: &survey_run::SYSTEM_CLOCK,
        };
        let survey = run(
            &repos_root,
            &keep_list,
            &probes,
            deadline,
            project.as_deref(),
            group_by,
            &adopted,
        );
        serde_json::to_value(survey).map_err(|e| format!("disk_survey: serialize error: {e}"))
    })
    .await
    .map_err(|e| format!("disk_survey: the survey pass panicked: {e}"))?
}

/// Measure `path` through the shared index, waiting for it no longer than
/// `budget` (#8985).
///
/// Why: a background pass can hold the index for a whole walk; a live pass
/// that queued behind it would run past its clamp.
/// What: no budget takes the lock unconditionally. With one, the lock is
/// awaited for at most `budget`; a timeout is `None` (no byte figure), and a
/// won lock measures with what the wait did not spend.
/// Test: `a_held_index_costs_a_budgeted_measurement_no_more_than_its_budget`.
fn measure_within_budget(
    index: &parking_lot::Mutex<crate::disk::size_index::DirSizeIndex>,
    path: &Path,
    budget: Option<Duration>,
) -> Option<DirSize> {
    let Some(budget) = budget else {
        return survey_run::measure(&mut index.lock(), path, None);
    };
    let asked = Instant::now();
    let mut index = index.try_lock_for(budget)?;
    let left = budget.saturating_sub(asked.elapsed());
    survey_run::measure(&mut index, path, Some(left))
}

/// The pull-request state of one scanned worktree's branch.
///
/// Why: named rather than inlined so the truncation fallback is stated once.
/// The bulk index cannot reach a branch older than its window, and on this
/// repository that is nearly every worktree the dashboard is about — so an
/// unresolved branch is retried with a targeted lookup, exactly as the reclaim
/// survey does. Without it the view would render `review` for worktrees whose
/// pull requests merged months ago.
///
/// What `left` does (#6929): it is the survey's remaining wall clock, and it
/// bounds BOTH calls rather than each of them separately. `GH_TIMEOUT` is a
/// fixed ten seconds, and this function can spend it twice, so a worktree the
/// loop admitted with a millisecond left used to run twenty seconds past the
/// survey's deadline — past the console's thirty-second transport too, which
/// is why the Disk view answered 502 with no survey rather than a truncated
/// one. The fallback is skipped outright when the first call consumed the
/// budget: the bulk index's own answer (`Unknown` or `LookupFailed`) stands,
/// and both block, so skipping it can only be conservative.
/// Test: `a_budgeted_survey_answers_within_its_budget` covers the contract in
/// `survey_run`; this wiring is exercised by `dispatch_disk_survey_tool`.
fn pr_for(
    indexes: &RefCell<BTreeMap<PathBuf, PrIndex>>,
    scanned: &ScannedWorktree,
    left: Option<Duration>,
) -> BranchPrState {
    // A local deadline, so the SECOND call is bounded by what the first one
    // left rather than by the same figure over again.
    let until = left.map(|l| Instant::now() + l);
    let mut cache = indexes.borrow_mut();
    let index = cache
        .entry(scanned.registry_root.clone())
        .or_insert_with(|| PrIndex::from_gh_within(&scanned.registry_root, gh_budget(until)));
    let pr = index.state_for(scanned.branch.as_deref());
    if matches!(
        pr,
        BranchPrState::Unknown | BranchPrState::LookupFailed { .. }
    ) && !index.is_complete()
        && let Some(branch) = scanned.branch.as_deref()
    {
        let budget = gh_budget(until);
        if !budget.is_zero() {
            return pr_state_for_branch_within(&scanned.registry_root, branch, budget);
        }
    }
    pr
}

/// How long a `gh` call may run, given the survey's own deadline.
///
/// What: the time left, capped at [`GH_TIMEOUT`]; [`GH_TIMEOUT`] when the
/// survey has no deadline; and `Duration::ZERO` when it has none left, which
/// the caller reads as "do not start this call".
/// Test: `gh_budget_never_exceeds_the_fixed_ceiling`,
/// `gh_budget_is_zero_once_the_survey_deadline_has_passed`.
fn gh_budget(until: Option<Instant>) -> Duration {
    match until {
        None => GH_TIMEOUT,
        Some(d) => d
            .checked_duration_since(Instant::now())
            .unwrap_or(Duration::ZERO)
            .min(GH_TIMEOUT),
    }
}

#[cfg(test)]
#[path = "mcp_disk_cache_tests.rs"]
mod mcp_disk_cache_tests;

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

    #[test]
    fn gh_budget_never_exceeds_the_fixed_ceiling() {
        let generous = Instant::now() + GH_TIMEOUT * 10;
        assert_eq!(gh_budget(Some(generous)), GH_TIMEOUT);
        assert_eq!(gh_budget(None), GH_TIMEOUT);
    }

    #[test]
    fn gh_budget_is_zero_once_the_survey_deadline_has_passed() {
        let spent = Instant::now() - Duration::from_secs(1);
        assert_eq!(gh_budget(Some(spent)), Duration::ZERO);
    }

    /// A budget past the stdio bridge's timeout is cut down to the clamp.
    ///
    /// Why this is a bug and not a preference: the daemon honoured a
    /// `budget_seconds: 120` faithfully and kept surveying, while the bridge
    /// carrying the answer gave up at 60 seconds and returned a transport
    /// error. The caller saw no survey at all where a 55-second truncated one
    /// was available — the #6929 failure mode, reached through the argument
    /// rather than through a slow fleet.
    ///
    /// Fails before the change: `clamp_budget` did not exist and 120 was passed
    /// through to the deadline verbatim.
    #[test]
    fn a_budget_past_the_bridge_timeout_is_clamped() {
        assert_eq!(clamp_budget(Some(120)), (Some(MAX_BUDGET_SECONDS), true));
        assert_eq!(
            Duration::from_secs(MAX_BUDGET_SECONDS),
            Duration::from_secs(55),
            "the clamp the schema's `maximum` advertises"
        );
    }

    #[test]
    fn a_budget_inside_the_bridge_timeout_is_untouched() {
        assert_eq!(clamp_budget(Some(30)), (Some(30), false));
        assert_eq!(
            clamp_budget(Some(MAX_BUDGET_SECONDS)),
            (Some(MAX_BUDGET_SECONDS), false),
            "the ceiling itself is not a clamp"
        );
        assert_eq!(clamp_budget(None), (None, false));
    }

    #[test]
    fn group_by_session_is_the_one_accepted_value() {
        assert_eq!(parse_group_by(None), Ok(GroupBy::None));
        assert_eq!(parse_group_by(Some("session")), Ok(GroupBy::Session));
    }

    /// An unknown `group_by` is an error, never a silent full survey.
    ///
    /// Why: a typo would otherwise return a complete, plausible payload with
    /// the roll-up simply missing and nothing saying why.
    #[test]
    fn an_unknown_group_by_is_rejected() {
        let err = parse_group_by(Some("project")).expect_err("unknown values must not be accepted");
        assert!(err.contains("project"), "{err}");
        assert!(
            err.contains("session"),
            "the error must name what IS accepted: {err}"
        );
    }

    /// #8985 error arm: an index held by another pass yields no byte figure
    /// once the budget is spent, instead of blocking the live call.
    #[test]
    fn a_held_index_costs_a_budgeted_measurement_no_more_than_its_budget() {
        let index = parking_lot::Mutex::new(crate::disk::size_index::DirSizeIndex::new());
        let dir = tempfile::tempdir().expect("temp dir");
        let held = index.lock();
        let started = Instant::now();
        let got = measure_within_budget(&index, dir.path(), Some(Duration::from_millis(100)));
        assert!(got.is_none(), "a held index must yield no figure: {got:?}");
        assert!(
            started.elapsed() < Duration::from_secs(5),
            "waited {:?}",
            started.elapsed()
        );
        drop(held);
        let free = measure_within_budget(&index, dir.path(), Some(Duration::from_secs(5)));
        assert!(free.is_some(), "a free index measures");
    }

    #[test]
    fn gh_budget_hands_back_the_time_that_is_actually_left() {
        let tight = Instant::now() + Duration::from_millis(200);
        let budget = gh_budget(Some(tight));
        assert!(!budget.is_zero(), "{budget:?}");
        assert!(budget <= Duration::from_millis(200), "{budget:?}");
    }
}