agent_engine/tools/

discovery.rs

//! Model-facing discovery and exact-activation builtins (Task 17, spec
//! §4.2, §7.2, §7.3, §7.7): `search_tools` and `activate_tools`.
//!
//! Both tools operate ONLY on the capability context the host injects into
//! [`crate::tools::ToolCapabilities::tool_activation`]:
//!
//! - a PASSIVE catalog snapshot (compact descriptors + digests; reading it
//!   starts no process, touches no network, exposes no schema);
//! - the RETAINED per-stream [`SharedSessionToolSet`] handle, so a
//!   confirmed activation mutates the same set the `ExecutionGate` and the
//!   next provider round consume;
//! - a host-supplied [`ActivationAuthority`], populated exclusively by host
//!   policy — never from model-authored JSON arguments.
//!
//! `search_tools` returns bounded compact descriptors with stable
//! [`ToolId`]s only — no full schemas, no factories, no source paths, no
//! credentials. `activate_tools` performs deterministic atomic bulk
//! activation of EXACT known ids through host grant issuance; source-wide,
//! provider-wide, wildcard, unknown, and malformed identities fail typed
//! with zero mutation. Confirmation authority is host-only: explicit host
//! auto-approve authorizes directly; otherwise the interactive host prompt
//! (when present) must approve the bounded exact id list with y/yes. Absent
//! prompt or any other answer, the request is denied before any grant, set,
//! or schema-generation mutation — model-authored JSON never confers
//! authority.

use serde_json::{json, Value};

use super::activation::{activate_model_initiated, ActivationAuthority, SharedSessionToolSet};
use super::catalog::{DiscoveryIndex, DiscoveryQuery, SearchLimits, ToolCatalog, ToolId};
use super::{Tool, ToolContext, ToolOrigin, ToolOutput};
use crate::{Result, RuntimeError};
use agent_core::BoundedText;

/// Default bounded result-count budget for `search_tools`.
pub const SEARCH_TOOLS_MAX_RESULTS: usize = 16;
/// Default bounded serialized byte budget for `search_tools` hits.
pub const SEARCH_TOOLS_MAX_RESULT_BYTES: usize = 8 * 1024;
/// Maximum number of ids in one `activate_tools` batch.
pub const ACTIVATE_TOOLS_MAX_BATCH: usize = 16;
/// Echo bound for hostile identity strings inside typed denials.
const ID_ECHO_MAX_BYTES: usize = 128;

/// Capability context for the discovery/activation builtins. Constructed by
/// the host per round from the round's catalog snapshot, the retained
/// session-set handle, and host confirmation policy.
#[derive(Clone)]
pub struct ActivationCapability {
    catalog: ToolCatalog,
    session_set: SharedSessionToolSet,
    authority: ActivationAuthority,
    /// Whether `activate_tools` may ask the host (the y/n confirm dialog)
    /// when `authority` is `Unauthorized`. `false` = `tools.activation_confirm
    /// = deny`: fail closed without ever raising a prompt.
    host_prompt_allowed: bool,
}

impl ActivationCapability {
    pub fn new(
        catalog: ToolCatalog,
        session_set: SharedSessionToolSet,
        authority: ActivationAuthority,
    ) -> Self {
        Self {
            catalog,
            session_set,
            authority,
            host_prompt_allowed: true,
        }
    }

    /// Set whether the host confirm prompt may be raised for an
    /// `Unauthorized` authority (`tools.activation_confirm`: `prompt` →
    /// true, `deny` → false). Default true.
    pub fn with_host_prompt(mut self, allowed: bool) -> Self {
        self.host_prompt_allowed = allowed;
        self
    }

    pub fn host_prompt_allowed(&self) -> bool {
        self.host_prompt_allowed
    }

    pub fn catalog(&self) -> &ToolCatalog {
        &self.catalog
    }

    pub fn session_set(&self) -> &SharedSessionToolSet {
        &self.session_set
    }

    pub fn authority(&self) -> ActivationAuthority {
        self.authority
    }

    /// Narrow MCP grant invalidation (Task 19): revoke EXACTLY one MCP
    /// `ToolId` from the RETAINED session set, and only when the caller's
    /// session identity matches the retained set's session. Core tools and
    /// sibling activations are untouched; the catalog is never mutated.
    /// Returns true iff an activation was actually removed.
    pub fn revoke_exact_mcp_grant(
        &self,
        session: &crate::tools::activation::SessionId,
        server: &str,
        server_tool_name: &str,
    ) -> bool {
        let mut set = self
            .session_set
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if set.session() != session {
            return false;
        }
        set.revoke_exact(&ToolId::mcp(server, server_tool_name))
            .is_ok()
    }

    /// Narrow extension grant invalidation (Task 20): revoke EXACTLY one
    /// extension `ToolId` from the RETAINED session set, and only when the
    /// caller's session identity matches the retained set's session. Core
    /// tools and sibling activations are untouched; the catalog is never
    /// mutated. Returns true iff an activation was actually removed.
    pub fn revoke_exact_extension_grant(
        &self,
        session: &crate::tools::activation::SessionId,
        plugin_id: &str,
        tool_name: &str,
    ) -> bool {
        let mut set = self
            .session_set
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if set.session() != session {
            return false;
        }
        set.revoke_exact(&ToolId::extension(plugin_id, tool_name))
            .is_ok()
    }
}

impl std::fmt::Debug for ActivationCapability {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ActivationCapability")
            .field("authority", &self.authority)
            .finish_non_exhaustive()
    }
}

fn require_capability(ctx: &ToolContext) -> Result<ActivationCapability> {
    ctx.capabilities.tool_activation.clone().ok_or_else(|| {
        RuntimeError::Tool("tool discovery/activation is not available in this context".to_string())
    })
}

/// Interactive host confirmation for model-initiated activation (M1).
///
/// Returns `ModelConfirmed` ONLY when the host's interactive prompt exists
/// and the host answers an explicit `y`/`yes` (case-insensitive, trimmed)
/// to the bounded exact id list. Every other outcome — no prompt handle, a
/// closed prompt channel, cancellation (`None`), or any other answer —
/// stays `Unauthorized` (fail-closed). The id list shown to the host is
/// exactly the parsed batch: each id passed `ToolId::parse` (bounded,
/// canonical) and the batch is capped at [`ACTIVATE_TOOLS_MAX_BATCH`].
async fn confirm_activation_with_host(
    prompt: Option<&crate::tools::SecretPromptHandle>,
    tool_ids: &[ToolId],
) -> ActivationAuthority {
    let Some(prompt) = prompt else {
        return ActivationAuthority::Unauthorized;
    };
    let mut ids: Vec<&str> = tool_ids.iter().map(ToolId::as_str).collect();
    ids.sort_unstable();
    let question = format!(
        "The model requests session-scoped activation of {} exact tool id(s):\n{}\n\
         Type y/yes to allow exactly these tools; anything else denies.",
        ids.len(),
        ids.join("\n")
    );
    match prompt
        .prompt(
            crate::session::CONFIRM_ACTIVATION_PROMPT_TITLE.to_string(),
            question,
        )
        .await
    {
        Some(answer) => {
            let normalized = answer.trim().to_ascii_lowercase();
            if normalized == "y" || normalized == "yes" {
                ActivationAuthority::ModelConfirmed
            } else {
                ActivationAuthority::Unauthorized
            }
        }
        None => ActivationAuthority::Unauthorized,
    }
}

fn require_query(params: &Value) -> Result<DiscoveryQuery> {
    let raw = params["query"]
        .as_str()
        .ok_or_else(|| RuntimeError::Tool("Missing 'query' parameter (string)".to_string()))?;
    DiscoveryQuery::parse(raw)
        .map_err(|err| RuntimeError::Tool(format!("invalid discovery query: {err}")))
}

// ── search_tools ────────────────────────────────────────────────────────────

/// Bounded, deterministic, descriptor-only capability search over the
/// passive catalog snapshot. Pure and local: no factory invocation, no
/// process start, no network, no schema exposure, no grant.
pub struct SearchToolsTool;

#[async_trait::async_trait]
impl Tool for SearchToolsTool {
    fn effect(&self) -> crate::tools::catalog::ToolEffect {
        crate::tools::catalog::ToolEffect::ReadOnly
    }

    fn name(&self) -> &str {
        "search_tools"
    }

    fn description(&self) -> &str {
        "Search locally known tools by keyword. Returns bounded compact \
         descriptors with stable tool ids (never full schemas). Use \
         activate_tools with an exact returned id to request activation."
    }

    fn origin(&self) -> ToolOrigin {
        ToolOrigin::Builtin
    }

    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Case-insensitive keyword matched against tool ids, summaries, and tags."
                }
            },
            "required": ["query"]
        })
    }

    async fn execute(&self, params: Value, ctx: ToolContext) -> Result<String> {
        let capability = require_capability(&ctx)?;
        let query = require_query(&params)?;
        let limits = SearchLimits::new(SEARCH_TOOLS_MAX_RESULTS, SEARCH_TOOLS_MAX_RESULT_BYTES)
            .expect("static search_tools budgets are within caps");
        let index = DiscoveryIndex::build(capability.catalog())
            .map_err(|err| RuntimeError::Tool(format!("tool discovery unavailable: {err}")))?;
        let results = index.search(&query, &limits);
        let body = json!({
            "generation": results.generation(),
            "truncated": results.truncated(),
            "tools": results.hits(),
        });
        serde_json::to_string(&body)
            .map_err(|err| RuntimeError::Tool(format!("failed to serialize search results: {err}")))
    }
}

// ── activate_tools ──────────────────────────────────────────────────────────

/// Model-initiated deterministic exact activation. Every requested id must
/// be an exact known cataloged `ToolId`; the whole batch validates first
/// (host grant issuance: exact id, current generation/digest, source trust)
/// and applies atomically in stable id order with exactly one session
/// schema-generation advance. Requires host-supplied confirmation
/// authority; a model-authored JSON flag can never substitute for it.
pub struct ActivateToolsTool;

#[async_trait::async_trait]
impl Tool for ActivateToolsTool {
    fn name(&self) -> &str {
        "activate_tools"
    }

    fn description(&self) -> &str {
        "Request session-scoped activation of exact tool ids returned by \
         search_tools. Activation is subject to the host's \
         tools.activation_confirm policy (auto: granted; prompt: the user is \
         asked y/n; deny: always refused); only the exact requested tools are \
         activated for this session."
    }

    fn origin(&self) -> ToolOrigin {
        ToolOrigin::Builtin
    }

    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "tools": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Exact stable tool ids (namespace:name) from search_tools. No wildcards."
                }
            },
            "required": ["tools"]
        })
    }

    async fn execute(&self, params: Value, ctx: ToolContext) -> Result<String> {
        self.activate(params, ctx).await.map(|(summary, _)| summary)
    }

    /// The text summary first (compaction, previews and non-Anthropic
    /// providers read it), then one `tool_reference` per activated tool,
    /// keyed by canonical tool id. On the Anthropic transport with deferred
    /// loading, the wire layer resolves each id to its wire name and sends
    /// only the references, which makes the deferred definitions callable
    /// without changing the cached prompt prefix. Everywhere else the
    /// references are dropped and the text stands alone.
    async fn execute_rich(&self, params: Value, ctx: ToolContext) -> Result<ToolOutput> {
        let (summary, ids) = self.activate(params, ctx).await?;
        let mut blocks = Vec::with_capacity(ids.len() + 1);
        blocks.push(json!({"type": "text", "text": summary}));
        blocks.extend(
            ids.iter()
                .map(|id| json!({"type": "tool_reference", "tool_name": id})),
        );
        Ok(ToolOutput::Blocks { blocks, summary })
    }
}

impl ActivateToolsTool {
    /// Authorize and apply one activation batch. Returns the JSON summary
    /// text and the sorted canonical ids that were requested and are now
    /// active for the session.
    async fn activate(&self, params: Value, ctx: ToolContext) -> Result<(String, Vec<String>)> {
        let capability = require_capability(&ctx)?;
        let requested = params["tools"].as_array().ok_or_else(|| {
            RuntimeError::Tool("Missing 'tools' parameter (array of tool id strings)".to_string())
        })?;
        if requested.is_empty() {
            return Err(RuntimeError::Tool(
                "no tool ids requested: 'tools' must name at least one exact id".to_string(),
            ));
        }
        if requested.len() > ACTIVATE_TOOLS_MAX_BATCH {
            return Err(RuntimeError::Tool(format!(
                "too many tool ids requested: {} exceeds limit {ACTIVATE_TOOLS_MAX_BATCH}",
                requested.len()
            )));
        }
        // Parse-at-boundary: every entry must be an exact canonical ToolId.
        // Source-wide ("builtin"), wildcard ("ns:*"), empty, and oversized
        // spellings fail typed here, before any authority/grant work.
        let mut tool_ids = Vec::with_capacity(requested.len());
        for entry in requested {
            let raw = entry
                .as_str()
                .ok_or_else(|| RuntimeError::Tool("tool ids must be strings".to_string()))?;
            let id = ToolId::parse(raw).map_err(|err| {
                RuntimeError::Tool(format!(
                    "invalid tool id {:?}: {err} (exact namespace:name ids only; \
                     wildcards and source-wide requests are not accepted)",
                    BoundedText::new(raw, ID_ECHO_MAX_BYTES).text
                ))
            })?;
            tool_ids.push(id);
        }
        // Host confirmation policy (M1): authority is host-only and exact.
        // `tools.activation_confirm = auto` (default) or explicit server
        // auto-approve authorizes directly; `prompt`: when an
        // interactive host prompt exists, the HOST is asked to approve this
        // bounded exact id list and only an explicit y/yes authorizes.
        // Absent prompt, denial, or cancellation stays fail-closed. The
        // model-authored JSON arguments are never consulted for authority.
        // The prompt await happens strictly BEFORE the session-set lock is
        // taken; no lock is held across an await.
        // `tools.activation_confirm = deny` never raises the prompt.
        let authority = match capability.authority() {
            ActivationAuthority::ModelConfirmed => ActivationAuthority::ModelConfirmed,
            ActivationAuthority::Unauthorized if !capability.host_prompt_allowed() => {
                ActivationAuthority::Unauthorized
            }
            ActivationAuthority::Unauthorized => {
                confirm_activation_with_host(ctx.capabilities.secret_prompt.as_ref(), &tool_ids)
                    .await
            }
        };
        // Authority check + host grant issuance + atomic bulk activation.
        // The write guard is short-lived and never held across an await; the
        // resulting schema generation is captured inside the SAME critical
        // section (M3) so the reported value is exactly the completed
        // batch's generation, never a later concurrent one.
        let (activated, schema_generation) = {
            let mut set = capability
                .session_set()
                .write()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            let applied =
                activate_model_initiated(authority, &mut set, capability.catalog(), &tool_ids)
                    .map_err(|err| RuntimeError::Tool(format!("tool activation denied: {err}")))?;
            (applied, set.schema_generation())
        };
        let ids: Vec<&str> = {
            let mut sorted: Vec<&ToolId> = tool_ids.iter().collect();
            sorted.sort();
            sorted.into_iter().map(ToolId::as_str).collect()
        };
        let body = json!({
            "activated": ids,
            "count": activated,
            "schema_generation": schema_generation,
        });
        let summary = serde_json::to_string(&body).map_err(|err| {
            RuntimeError::Tool(format!("failed to serialize activation result: {err}"))
        })?;
        let ids = ids.into_iter().map(str::to_string).collect();
        Ok((summary, ids))
    }
}