archy/core/archipelago/src/assistant/mod.rs

//! D-02: "one assistant, many front doors." A shared assistant service —
//! one curated tool registry, one backend selector, one place the model key
//! lives — used today by AIUI chat (`CallerScope::LocalOperator`) and, by
//! design, extensible to mesh/LoRa callers (`CallerScope::Mesh`) and later
//! Pine voice without recreating a second, divergent security model.
//!
//! This is the tracer slice for Phase 13 (D-01, D-02, D-06): a typed
//! question reaches exactly one curated, read-only tool
//! (`tools::system_disk_status_tool`) via the Claude backend, dispatched
//! through the SAME `handle_system_disk_status` RPC handler every other
//! authenticated caller uses. See `13-01-PLAN.md` for the full spine.

pub mod backends;
pub mod confirm;
pub mod grants;
pub mod loop_;
pub mod tools;

use std::collections::{BTreeSet, HashMap, HashSet};
use std::path::Path;
use std::sync::{Arc, Mutex};

use anyhow::Result;
use serde::{Deserialize, Serialize};

use crate::api::rpc::RpcHandler;

/// D-16's ten permission categories. All default-closed on a fresh node —
/// nothing is shared with the model until deliberately granted. Serialized
/// with `rename_all = "kebab-case"` so the wire form (`"ai-local"`, etc.)
/// matches `neode-ui/src/stores/aiPermissions.ts`'s category ids
/// one-for-one — 13-05's acceptance criterion diffs the two lists by hand,
/// but the serde rename is what keeps them from drifting silently again.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum PermissionCategory {
    Apps,
    System,
    Network,
    Wallet,
    Files,
    Media,
    Search,
    AiLocal,
    Notes,
    Bitcoin,
}

impl PermissionCategory {
    /// All ten categories, for grants-get/list-tools enumeration and for
    /// tests that need to assert over the whole set.
    pub const ALL: [PermissionCategory; 10] = [
        PermissionCategory::Apps,
        PermissionCategory::System,
        PermissionCategory::Network,
        PermissionCategory::Wallet,
        PermissionCategory::Files,
        PermissionCategory::Media,
        PermissionCategory::Search,
        PermissionCategory::AiLocal,
        PermissionCategory::Notes,
        PermissionCategory::Bitcoin,
    ];
}

/// D-02's promoted primary noun: a caller identity carrying the permission
/// scope its tool calls resolve authority through. "A mesh peer" and "the
/// local operator in AIUI" are two variants of it; Pine voice will be a
/// third (not built in this phase — no `Voice` variant exists yet, by
/// design, until that phase actually needs one).
#[derive(Debug, Clone)]
pub enum CallerScope {
    /// A mesh/LoRa peer. Not exercised by this plan (mesh's existing
    /// `!ai` path is Q&A-only, per `mesh/listener/assist.rs`'s own doc
    /// comment) — the variant exists so the shape is right when a future
    /// plan wires mesh callers into the shared loop. `authorized` is where
    /// that future plan threads the existing per-caller
    /// `trusted_only`/`allowed_contacts`/`denied_askers` resolution
    /// (`api/rpc/mesh/assistant.rs`) through: a mesh peer's authority is
    /// never wider than the operator's own persisted grants, only ever a
    /// subset of them (all-or-nothing today; a future plan may narrow this
    /// to a per-peer category subset without changing this variant's
    /// shape).
    Mesh { peer_id: String, authorized: bool },
    /// The authenticated operator using AIUI, identified by their neode-ui
    /// session. This is the only variant this tracer's `assistant.chat`
    /// RPC constructs.
    LocalOperator { session_id: String },
}

impl CallerScope {
    /// The sole source of tool authority `execute_tool` reads. No
    /// `execute_tool` branch may read a caller-specific field directly
    /// instead of going through this — that would reintroduce the
    /// mesh-only assumption D-02 exists to retire.
    ///
    /// Both variants resolve through the SAME persisted `Grants` store
    /// (D-16's default-closed categories, per `data_dir`) — never a
    /// hardcoded default, and never two divergent sources of authority.
    pub async fn granted_categories(&self, data_dir: &Path) -> BTreeSet<PermissionCategory> {
        let persisted = grants::Grants::load(data_dir).await;
        match self {
            CallerScope::LocalOperator { .. } => persisted.categories().clone(),
            // A mesh peer can never exceed what the operator opened: its
            // authority is the persisted grants intersected with whether
            // this peer is itself authorized to use the assistant at all
            // (today a single boolean; a future plan may narrow this to a
            // per-peer category subset without changing this call site).
            CallerScope::Mesh { authorized, .. } => {
                if *authorized {
                    persisted.categories().clone()
                } else {
                    BTreeSet::new()
                }
            }
        }
    }
}

/// Bundles what `execute_tool` needs regardless of which backend produced
/// the tool call: the curated registry, the caller's resolved authority,
/// and a handle back to the SAME `RpcHandler` every other authenticated
/// caller dispatches through — never an AI-only backdoor.
///
/// Also carries the AI-SPEC §4b.1 "≤ 2 consecutive validation failures per
/// tool name" counter. Construct via [`ToolExecCtx::new`] — the counter
/// field is private so every call site shares the same reset/note logic
/// rather than reimplementing it.
pub struct ToolExecCtx {
    pub registry: tools::ToolRegistry,
    pub caller: CallerScope,
    pub handler: Arc<RpcHandler>,
    /// D-07/D-11: the confirm gate every destructive tool suspends on.
    /// Defaults to the process-wide gate (`confirm::global()`) so the loop
    /// and the `assistant.confirm-tool`/`assistant.pending` RPC handlers
    /// share one pending queue; tests inject a fresh gate per test via
    /// [`ToolExecCtx::with_confirm_gate`] for isolation.
    pub confirm: Arc<confirm::ConfirmGate>,
    validation_failures: Mutex<HashMap<String, u32>>,
    /// 13-08 on-device UAT (T-13-50): actions the human declined this turn,
    /// keyed by the same canonical `(tool_name, validated_args)` identity
    /// the nonce binds. A declined action must not re-prompt within the
    /// turn — the model retrying after "the user declined" would otherwise
    /// mint a fresh confirmation and re-open the dialog until the human
    /// gives in, which is the habituation failure, mechanized.
    declined_actions: Mutex<HashSet<String>>,
}

impl ToolExecCtx {
    pub fn new(
        registry: tools::ToolRegistry,
        caller: CallerScope,
        handler: Arc<RpcHandler>,
    ) -> Self {
        Self::with_confirm_gate(registry, caller, handler, confirm::global())
    }

    pub fn with_confirm_gate(
        registry: tools::ToolRegistry,
        caller: CallerScope,
        handler: Arc<RpcHandler>,
        confirm: Arc<confirm::ConfirmGate>,
    ) -> Self {
        Self {
            registry,
            caller,
            handler,
            confirm,
            validation_failures: Mutex::new(HashMap::new()),
            declined_actions: Mutex::new(HashSet::new()),
        }
    }

    /// The human declined this exact action; remember it for the rest of
    /// the turn so it can never re-prompt.
    pub(crate) fn note_declined(&self, action_key: String) {
        self.declined_actions
            .lock()
            .expect("declined_actions mutex poisoned")
            .insert(action_key);
    }

    /// Whether this exact action was already declined this turn.
    pub(crate) fn was_declined(&self, action_key: &str) -> bool {
        self.declined_actions
            .lock()
            .expect("declined_actions mutex poisoned")
            .contains(action_key)
    }

    /// A tool name's argument validation just succeeded — its consecutive
    /// failure streak resets.
    pub(crate) fn reset_validation_failures(&self, tool_name: &str) {
        self.validation_failures
            .lock()
            .expect("validation_failures mutex poisoned")
            .remove(tool_name);
    }

    /// A tool name's argument validation just failed. Returns `true` if
    /// this failure is the third (or later) consecutive one for this tool
    /// name — the signal `run_loop` uses to abort the turn with an apology
    /// rather than spinning (D-05, AI-SPEC §4b.1).
    pub(crate) fn note_validation_failure(&self, tool_name: &str) -> bool {
        let mut map = self
            .validation_failures
            .lock()
            .expect("validation_failures mutex poisoned");
        let count = map.entry(tool_name.to_string()).or_insert(0);
        *count += 1;
        *count > 2
    }

    /// Whether any tool name has crossed the consecutive-failure threshold
    /// this turn. Checked by `run_loop` after each batch of tool calls.
    pub(crate) fn should_abort(&self) -> bool {
        self.validation_failures
            .lock()
            .expect("validation_failures mutex poisoned")
            .values()
            .any(|&c| c > 2)
    }
}

/// D-16/AI-SPEC §4b.3: one static, phase-authored persona and confirm-gate
/// statement, never assembled from prior model output and never editable
/// by AIUI. Appends **only** the currently-granted-category tools' names
/// and descriptions — an ungranted tool's name never appears in this
/// string, which is the prompt-side half of the two-layer defense (the
/// `execute_tool` grant re-check in `loop_.rs` is the other, and the gate
/// that actually matters — see `ungranted_tool_absent_from_system_prompt`
/// and `settings_tool_respects_category_grant`).
const SYSTEM_PROMPT_PREAMBLE: &str = "You are the Archipelago node's operator-control assistant. \
Only use the tools explicitly listed below for this turn — never invent a tool name or call one \
that is not listed here, even if it sounds like something this node could plausibly do. Every \
write requires a human confirmation you cannot bypass, skip, or pre-approve on the user's \
behalf; the node itself presents that confirmation to the operator in a trusted dialog the \
moment you call the tool. So when the user asks for something a listed tool does, call the tool \
directly — never ask for permission or confirmation in your text first. A text pre-ask is worse \
than redundant: it stalls the action behind a reply you cannot act on, and it trains the \
operator to rubber-stamp. If the user asks for something outside the tools listed below \
(including anything touching keys, seeds, wallet spends, federation trust, or a factory reset), \
refuse plainly and, if there is a real path in neode-ui's Settings screen for it, name that path \
instead of fabricating a tool call.";

pub fn build_system_prompt(visible_tools: &[tools::ToolDef]) -> String {
    let mut prompt = String::from(SYSTEM_PROMPT_PREAMBLE);
    if visible_tools.is_empty() {
        prompt.push_str(
            "\n\nNo permission categories are granted on this node right now, so no tools are \
            available this turn. Say so plainly if asked to do something — do not guess, and do \
            not pretend a category is open.",
        );
    } else {
        prompt.push_str("\n\nTools available this turn:\n");
        for tool in visible_tools {
            prompt.push_str(&format!("- {}: {}\n", tool.name, tool.description));
        }
    }
    prompt
}

/// Entry point: run one chat turn for `caller` through the shared loop.
/// Builds the visible-tool set from the caller's granted categories only
/// (D-16 — the model should never even see a tool it can't use), selects a
/// backend per D-04's chain (Ollama first, Claude fallback), and runs it to
/// a final answer. 13-10 Task 2 adds D-08 history persistence on top of
/// this.
pub async fn chat(
    handler: Arc<RpcHandler>,
    caller: CallerScope,
    user_text: String,
) -> Result<String> {
    let registry = tools::registry();
    let grants = caller.granted_categories(handler.data_dir()).await;
    let visible_tools = registry.visible_to(&grants);

    let (backend, backend_id) = backends::select_backend(handler.data_dir()).await;
    tracing::info!(backend = %backend_id, "assistant.chat: backend selected for this turn");

    let system_prompt = build_system_prompt(&visible_tools);

    let history = vec![tools::ChatMessage {
        role: tools::Role::User,
        text: Some(user_text),
        tool_calls: vec![],
        tool_results: vec![],
    }];

    let ctx = ToolExecCtx::new(registry, caller, handler);

    loop_::run_loop(
        backend.as_ref(),
        &system_prompt,
        &visible_tools,
        history,
        &ctx,
    )
    .await
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::api::rpc::RpcHandler;

    /// A minimal but real `RpcHandler` for tests, matching
    /// `loop_::tests::test_rpc_handler` — a fresh temp `data_dir`, no
    /// orchestrator.
    async fn test_rpc_handler() -> (Arc<RpcHandler>, tempfile::TempDir) {
        let tmp = tempfile::tempdir().expect("tempdir");
        let mut config = crate::config::Config::default();
        config.data_dir = tmp.path().to_path_buf();
        let state_manager = Arc::new(crate::state::StateManager::new());
        let metrics_store = Arc::new(crate::monitoring::MetricsStore::new());
        let session_store =
            crate::session::SessionStore::new_for_tests(tmp.path().join("sessions.json"));
        let handler = RpcHandler::new(
            config,
            state_manager,
            metrics_store,
            session_store,
            None,
            None,
        )
        .await
        .expect("RpcHandler::new");
        (Arc::new(handler), tmp)
    }

    /// A minimal installed-and-running app entry, seeded into the scanner
    /// snapshot so the app-lifecycle tools' business-rule id resolution
    /// (`container-list`) finds it — the same cached-state path the real
    /// node serves, no orchestrator or podman needed.
    fn installed_entry(app_id: &str) -> crate::data_model::PackageDataEntry {
        use crate::data_model::{
            Description, Manifest, PackageDataEntry, PackageState, StaticFiles,
        };
        PackageDataEntry {
            state: PackageState::Running,
            health: None,
            exit_code: None,
            static_files: StaticFiles {
                license: String::new(),
                instructions: String::new(),
                icon: String::new(),
            },
            manifest: Manifest {
                id: app_id.to_string(),
                title: app_id.to_string(),
                version: String::new(),
                description: Description {
                    short: String::new(),
                    long: String::new(),
                },
                release_notes: String::new(),
                license: String::new(),
                wrapper_repo: String::new(),
                upstream_repo: String::new(),
                support_site: String::new(),
                marketing_site: String::new(),
                donation_url: None,
                author: None,
                website: None,
                interfaces: None,
                tier: None,
            },
            installed: None,
            install_progress: None,
            uninstall_stage: None,
            available_update: None,
        }
    }

    /// Like `test_rpc_handler`, but with one installed, running app seeded
    /// into the state snapshot so destructive app-lifecycle calls pass
    /// business-rule validation and actually reach the confirm gate.
    async fn test_rpc_handler_with_app(app_id: &str) -> (Arc<RpcHandler>, tempfile::TempDir) {
        let tmp = tempfile::tempdir().expect("tempdir");
        let mut config = crate::config::Config::default();
        config.data_dir = tmp.path().to_path_buf();
        let state_manager = Arc::new(crate::state::StateManager::new());

        let mut data = crate::data_model::DataModel::new();
        data.server_info.status_info.containers_scanned = true;
        data.package_data
            .insert(app_id.to_string(), installed_entry(app_id));
        state_manager.update_data(data).await;

        let metrics_store = Arc::new(crate::monitoring::MetricsStore::new());
        let session_store =
            crate::session::SessionStore::new_for_tests(tmp.path().join("sessions.json"));
        let handler = RpcHandler::new(
            config,
            state_manager,
            metrics_store,
            session_store,
            None,
            None,
        )
        .await
        .expect("RpcHandler::new");
        (Arc::new(handler), tmp)
    }

    async fn grant(handler: &Arc<RpcHandler>, category: PermissionCategory) {
        let mut g = grants::Grants::load(handler.data_dir()).await;
        g.set(category, true);
        g.save(handler.data_dir()).await.expect("save grants");
    }

    async fn wait_pending(gate: &confirm::ConfirmGate) -> confirm::PendingSnapshot {
        for _ in 0..500 {
            if let Some(snap) = gate.peek() {
                return snap;
            }
            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
        }
        panic!(
            "no pending confirmation appeared — the destructive branch did not suspend on the gate"
        );
    }

    /// S-01 / D-07: a destructive tool call suspends before any execution
    /// and produces a pending confirmation whose node-authored description
    /// names the resource; nothing runs until a human resolves it, and a
    /// decline returns an error result without executing.
    #[tokio::test]
    async fn destructive_tool_requires_confirm() {
        let (handler, _tmp) = test_rpc_handler_with_app("demo-app").await;
        grant(&handler, PermissionCategory::Apps).await;

        // Part 1 — the choke point directly: execute_tool suspends on the
        // gate, and a decline returns a declined error result (dispatch is
        // never reached — nothing was changed).
        let gate = Arc::new(confirm::ConfirmGate::new());
        let ctx = Arc::new(ToolExecCtx::with_confirm_gate(
            tools::registry(),
            CallerScope::LocalOperator {
                session_id: "s".to_string(),
            },
            handler.clone(),
            gate.clone(),
        ));
        let call = tools::ToolCall {
            id: "call-1".to_string(),
            name: "app_restart".to_string(),
            arguments: serde_json::json!({ "app_id": "demo-app" }),
        };
        let (ctx_task, call_task) = (ctx.clone(), call.clone());
        let suspended =
            tokio::spawn(async move { loop_::execute_tool(&call_task, &ctx_task).await });

        let snap = wait_pending(&gate).await;
        assert!(
            snap.description.contains("demo-app"),
            "the node-authored description must name the resource: {}",
            snap.description
        );
        assert!(
            !suspended.is_finished(),
            "execute_tool must stay suspended until the human decides"
        );

        gate.resolve(&snap.req_id, &snap.nonce, false)
            .expect("decline resolves");
        let result = suspended.await.expect("join");
        assert!(result.is_error, "a declined action must not execute");
        assert!(
            result.content.to_lowercase().contains("declined"),
            "the tool result must say the user declined: {}",
            result.content
        );

        // Part 2 — through the real loop: a scripted backend proposes the
        // destructive call, the loop suspends on the gate, and only after
        // the human answers does the turn finish.
        let gate2 = Arc::new(confirm::ConfirmGate::new());
        let ctx2 = ToolExecCtx::with_confirm_gate(
            tools::registry(),
            CallerScope::LocalOperator {
                session_id: "s".to_string(),
            },
            handler.clone(),
            gate2.clone(),
        );
        let backend = crate::assistant::backends::scripted::ScriptedBackend::new(vec![
            crate::assistant::backends::BackendTurn::ToolCalls(vec![call.clone()]),
            crate::assistant::backends::BackendTurn::Text(
                "Understood — I did not restart it.".to_string(),
            ),
        ]);
        let tools_list = vec![tools::app_restart_tool()];
        let loop_task = tokio::spawn(async move {
            loop_::run_loop(&backend, "sys", &tools_list, vec![], &ctx2).await
        });

        let snap2 = wait_pending(&gate2).await;
        assert!(
            !loop_task.is_finished(),
            "run_loop must stay suspended while the confirmation is pending"
        );
        gate2
            .resolve(&snap2.req_id, &snap2.nonce, false)
            .expect("decline resolves");
        let answer = loop_task.await.expect("join").expect("run_loop");
        assert_eq!(answer, "Understood — I did not restart it.");
    }

    /// 13-08 UAT / T-13-50: once the human declines an action, the same
    /// action re-requested in the same turn is refused immediately — no
    /// fresh pending confirmation is minted, so the dialog cannot re-open
    /// until the human gives in.
    #[tokio::test]
    async fn declined_action_never_reprompts_same_turn() {
        let (handler, _tmp) = test_rpc_handler_with_app("demo-app").await;
        grant(&handler, PermissionCategory::Apps).await;

        let gate = Arc::new(confirm::ConfirmGate::new());
        let ctx = Arc::new(ToolExecCtx::with_confirm_gate(
            tools::registry(),
            CallerScope::LocalOperator {
                session_id: "s".to_string(),
            },
            handler.clone(),
            gate.clone(),
        ));
        let call = tools::ToolCall {
            id: "call-1".to_string(),
            name: "app_restart".to_string(),
            arguments: serde_json::json!({ "app_id": "demo-app" }),
        };

        // First ask: suspend, then the human declines.
        let (ctx_task, call_task) = (ctx.clone(), call.clone());
        let suspended =
            tokio::spawn(async move { loop_::execute_tool(&call_task, &ctx_task).await });
        let snap = wait_pending(&gate).await;
        gate.resolve(&snap.req_id, &snap.nonce, false)
            .expect("decline resolves");
        let first = suspended.await.expect("join");
        assert!(first.is_error);

        // Second ask, same action, same turn: refused without suspending
        // and without minting a new pending confirmation.
        let mut second_call = call.clone();
        second_call.id = "call-2".to_string();
        let second = loop_::execute_tool(&second_call, &ctx).await;
        assert!(
            second.is_error,
            "a re-asked declined action must be an error"
        );
        assert!(
            second.content.to_lowercase().contains("already declined"),
            "the refusal must tell the model it was already declined: {}",
            second.content
        );
        assert!(
            gate.peek().is_none(),
            "no new pending confirmation may be minted for a declined action"
        );
    }

    /// S-06 / D-16: a fresh node's `LocalOperator` resolves to no granted
    /// categories at all.
    #[tokio::test]
    async fn fresh_node_grants_are_empty() {
        let (handler, _tmp) = test_rpc_handler().await;
        let caller = CallerScope::LocalOperator {
            session_id: "s".to_string(),
        };
        let granted = caller.granted_categories(handler.data_dir()).await;
        assert!(
            granted.is_empty(),
            "fresh node must grant nothing: {granted:?}"
        );
    }

    /// The system prompt built for a caller must never mention a tool
    /// whose category is not currently granted — the prompt filter is
    /// defense in depth, never the gate (S-05's gate is
    /// `settings_tool_respects_category_grant` in tools.rs), but it must
    /// still hold.
    #[test]
    fn ungranted_tool_absent_from_system_prompt() {
        let reg = tools::registry();
        let mut grants = BTreeSet::new();
        grants.insert(PermissionCategory::System);
        let visible = reg.visible_to(&grants);
        let prompt = build_system_prompt(&visible);

        for tool in reg.all() {
            if tool.category == PermissionCategory::System {
                continue;
            }
            assert!(
                !prompt.contains(tool.name),
                "ungranted tool {} leaked into the system prompt",
                tool.name
            );
        }
        // Sanity: at least one granted tool IS present, so this isn't
        // trivially passing because the prompt is empty.
        assert!(
            reg.visible_to(&grants)
                .iter()
                .any(|t| prompt.contains(t.name)),
            "expected at least one granted-category tool name in the prompt"
        );
    }

    /// D-16: revoking a category takes effect on the very next
    /// `granted_categories` resolution — not only on the next session /
    /// process restart.
    #[tokio::test]
    async fn grant_revocation_takes_effect_next_turn() {
        let (handler, _tmp) = test_rpc_handler().await;
        let caller = CallerScope::LocalOperator {
            session_id: "s".to_string(),
        };

        let mut g = grants::Grants::load(handler.data_dir()).await;
        g.set(PermissionCategory::System, true);
        g.save(handler.data_dir()).await.expect("save");

        let granted = caller.granted_categories(handler.data_dir()).await;
        assert!(granted.contains(&PermissionCategory::System));

        let mut g = grants::Grants::load(handler.data_dir()).await;
        g.set(PermissionCategory::System, false);
        g.save(handler.data_dir()).await.expect("save");

        // Same caller, same process, no restart — the very next resolution
        // must reflect the revocation.
        let granted = caller.granted_categories(handler.data_dir()).await;
        assert!(
            !granted.contains(&PermissionCategory::System),
            "revocation must take effect on the next resolution, not only on restart"
        );
    }

    /// D-02's promoted-primary contract: both `CallerScope` variants
    /// resolve authority through the SAME method, and a mesh peer's
    /// authority is never wider than the operator's own persisted grants.
    #[tokio::test]
    async fn every_caller_variant_resolves_authority_through_caller_scope() {
        let (handler, _tmp) = test_rpc_handler().await;

        let mut g = grants::Grants::load(handler.data_dir()).await;
        g.set(PermissionCategory::Bitcoin, true);
        g.save(handler.data_dir()).await.expect("save");

        let operator = CallerScope::LocalOperator {
            session_id: "s".to_string(),
        };
        let operator_grants = operator.granted_categories(handler.data_dir()).await;
        assert!(operator_grants.contains(&PermissionCategory::Bitcoin));

        let authorized_peer = CallerScope::Mesh {
            peer_id: "peer-1".to_string(),
            authorized: true,
        };
        let peer_grants = authorized_peer.granted_categories(handler.data_dir()).await;
        assert_eq!(
            peer_grants, operator_grants,
            "an authorized mesh peer's ceiling is exactly the operator's persisted grants"
        );

        let unauthorized_peer = CallerScope::Mesh {
            peer_id: "peer-2".to_string(),
            authorized: false,
        };
        let peer_grants = unauthorized_peer
            .granted_categories(handler.data_dir())
            .await;
        assert!(
            peer_grants.is_empty(),
            "an unauthorized mesh peer must resolve to no authority regardless of what the operator granted"
        );
    }
}