lean-ctx 3.9.19

use super::super::{
    CallToolRequestParams, CallToolResult, CrpMode, ErrorData, LeanCtxServer, elicitation, helpers,
    is_shell_tool_name, permission_inheritance, post_process,
};
use super::dispatch_and_post_process;
use crate::core::ocla::response_cache::{
    CachedResponse, ResponseCache, ResponseCacheKey, global_response_cache,
};
use rmcp::model::{ContentBlock, Meta};
use serde::Deserialize;
use serde_json::{Map, Value};
use std::time::{Duration, Instant};

const CACHEABLE_TOOLS: [&str; 3] = ["ctx_search", "ctx_tree", "ctx_glob"];

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct CachedCallToolResult {
    content: Vec<ContentBlock>,
    #[serde(default)]
    structured_content: Option<Value>,
    #[serde(default)]
    is_error: Option<bool>,
    #[serde(default, rename = "_meta")]
    meta: Option<Meta>,
}

pub(super) fn response_cache_key(
    tool_name: &str,
    arguments: Option<&Map<String, Value>>,
    project_root: &str,
) -> Option<ResponseCacheKey> {
    CACHEABLE_TOOLS.contains(&tool_name).then(|| {
        let mut input = Vec::with_capacity(project_root.len() + 1);
        input.extend_from_slice(project_root.as_bytes());
        input.push(0);
        input.extend_from_slice(
            &serde_json::to_vec(&arguments).expect("JSON arguments must serialize"),
        );
        let digest = blake3::hash(&input);
        let mut hash_bytes = [0; 8];
        hash_bytes.copy_from_slice(&digest.as_bytes()[..8]);
        let arguments_hash = u64::from_be_bytes(hash_bytes);
        ResponseCacheKey::new(tool_name, arguments_hash, 0.0, 0)
    })
}

pub(super) fn cached_call_result(
    cache: &ResponseCache,
    key: &ResponseCacheKey,
) -> Option<CallToolResult> {
    let response = cache.get(key);
    crate::core::telemetry::global_metrics().record_cache(response.is_some());
    response.and_then(|cached| {
        serde_json::from_slice::<CachedCallToolResult>(&cached.body)
            .ok()
            .map(|cached| {
                let mut result = CallToolResult::success(cached.content);
                result.structured_content = cached.structured_content;
                result.is_error = cached.is_error;
                result.meta = cached.meta;
                result
            })
    })
}

pub(super) fn cache_call_result(
    cache: &ResponseCache,
    key: ResponseCacheKey,
    result: &CallToolResult,
) {
    if result.is_error == Some(true) {
        return;
    }
    let Ok(body) = serde_json::to_vec(result) else {
        return;
    };
    let tokens = crate::core::tokens::count_tokens(&String::from_utf8_lossy(&body))
        .try_into()
        .unwrap_or(u64::MAX);
    cache.put(
        key,
        CachedResponse {
            body,
            status: 200,
            tokens,
            created_at: Instant::now(),
            ttl: Duration::ZERO,
        },
    );
}

impl LeanCtxServer {
    pub(crate) async fn call_tool_guarded(
        &self,
        request: CallToolRequestParams,
    ) -> Result<CallToolResult, ErrorData> {
        self.check_idle_expiry().await;
        self.resolve_roots_once().await;
        elicitation::increment_call();

        let original_name = request.name.as_ref().to_string();
        let (resolved_name, resolved_args) = if original_name == "ctx" {
            let sub = request
                .arguments
                .as_ref()
                .and_then(|a| a.get("tool"))
                .and_then(|v| v.as_str())
                .map(std::string::ToString::to_string)
                .ok_or_else(|| {
                    ErrorData::invalid_params("'tool' is required for ctx meta-tool", None)
                })?;
            let tool_name = if sub.starts_with("ctx_") {
                sub
            } else {
                format!("ctx_{sub}")
            };
            let mut args = request.arguments.unwrap_or_default();
            args.remove("tool");
            (tool_name, Some(args))
        } else {
            (original_name, request.arguments)
        };
        let name = resolved_name.as_str();
        let args = resolved_args.as_ref();
        let query = args
            .and_then(|values| values.get("query").and_then(serde_json::Value::as_str))
            .or_else(|| {
                args.and_then(|values| values.get("task").and_then(serde_json::Value::as_str))
            })
            .unwrap_or(name);
        let session_id = self.session.read().await.id.clone();
        let agent_id = match self.agent_id.read().await.clone() {
            Some(agent_id) => agent_id,
            None => self
                .presence_agent_id
                .read()
                .await
                .clone()
                .unwrap_or_else(|| "mcp-agent".to_owned()),
        };
        let config = crate::core::config::Config::load_arc();
        let decision_context = config
            .decision_loop
            .enabled
            .then(|| {
                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                    crate::core::decision_loop_runtime::DecisionLoopRuntime::get_or_init()
                        .on_tool_start(name, query, &session_id, &agent_id)
                }))
                .map_err(|_| tracing::warn!(tool = name, "decision loop start panicked"))
                .ok()
            })
            .flatten();
        *self.task_envelope.write().await = crate::core::task_spine::TaskSpine::current();

        if let Some(denied) = Self::guard_role_and_policy(name) {
            finish_decision_loop(decision_context.as_ref(), args, &denied);
            return Ok(denied);
        }

        // ctx_call is a meta-dispatcher: the egress DLP and permission-
        // inheritance gates below must inspect the INNER tool + arguments, or
        // the universal invoker becomes a policy bypass (#1008 security pass).
        // Role/rate/workflow gates for the inner tool already run inside the
        // dispatch layer; these two ran only on the wrapper name before.
        let inner_call: Option<(String, Option<serde_json::Map<String, serde_json::Value>>)> =
            if name == "ctx_call" {
                helpers::get_str(args, "name").map(|inner_name| {
                    let inner_args = args
                        .and_then(|m| m.get("arguments"))
                        .and_then(serde_json::Value::as_object)
                        .cloned();
                    (inner_name, inner_args)
                })
            } else {
                None
            };
        let (guard_name, guard_args): (&str, Option<&serde_json::Map<_, _>>) = match &inner_call {
            Some((n, a)) => (n.as_str(), a.as_ref()),
            None => (name, args),
        };

        if let Some(blocked) = Self::guard_egress(guard_name, guard_args) {
            finish_decision_loop(decision_context.as_ref(), args, &blocked);
            return Ok(blocked);
        }

        if let Some(blocked) = self.guard_workflow(name).await {
            finish_decision_loop(decision_context.as_ref(), args, &blocked);
            return Ok(blocked);
        }

        // #794: cost cap guard — block tool calls when session cost exceeds the
        // configured limit. ctx_session is exempt so the agent can inspect
        // budget status and override the cap.
        if name != "ctx_session"
            && let Some(cap_msg) =
                crate::core::budget_tracker::BudgetTracker::global().cost_cap_message()
        {
            let result = CallToolResult::error(vec![ContentBlock::text(cap_msg)]);
            finish_decision_loop(decision_context.as_ref(), args, &result);
            return Ok(result);
        }

        // #990: determine machine-readability *before* the once-per-session
        // decorations below. A machine-readable invocation (e.g. ctx_outline
        // format=json) must reach the client byte-exact and parseable, so every
        // prose decoration and terse compression is suppressed and the pure
        // pre-decoration body is restored at the end (see the `machine_readable`
        // guard near the end of this function). Computing it here — not after
        // dispatch — means such a call also never *consumes* a latched
        // once-per-session flag (auto-context briefing, rules tip) whose prose
        // we would then discard, so those surface on the next human-facing call.
        //
        // `ctx_call` is a meta-dispatcher: the contract belongs to its *inner*
        // tool + inner arguments, not to ctx_call itself. Unwrap one level so
        // JSON reached via the lazy `ctx_call` path (the default advertised
        // surface, where ctx_outline is not a top-level tool) is just as
        // byte-exact as a direct call. This also covers JSON error envelopes
        // from the early rate-limit path, which the first-call auto-context
        // briefing would otherwise corrupt.
        let (mr_name, mr_args): (
            Option<String>,
            Option<&serde_json::Map<String, serde_json::Value>>,
        ) = if name == "ctx_call" {
            (
                helpers::get_str(args, "name"),
                args.and_then(|m| m.get("arguments"))
                    .and_then(serde_json::Value::as_object),
            )
        } else {
            (Some(name.to_string()), args)
        };
        let machine_readable = mr_name
            .as_deref()
            .and_then(|n| self.registry.as_ref().and_then(|r| r.get_arc(n)))
            .is_some_and(|tool| tool.produces_machine_readable(mr_args));

        // Skip the session wake-up briefing for machine-readable calls: the
        // pre-hook latches `session_initialized` via compare-exchange, so calling
        // it here would burn the once-per-session slot for a briefing we then
        // throw away. Deferring keeps the briefing intact for the next call.
        let auto_context = if machine_readable {
            None
        } else {
            let task = {
                let session = self.session.read().await;
                session.task.as_ref().map(|t| t.description.clone())
            };
            let project_root = {
                let session = self.session.read().await;
                session.project_root.clone()
            };
            let cache_timeout =
                tokio::time::timeout(std::time::Duration::from_secs(5), self.cache.write()).await;
            if let Ok(mut cache) = cache_timeout {
                crate::tools::autonomy::session_lifecycle_pre_hook(
                    &self.autonomy,
                    name,
                    &mut cache,
                    task.as_deref(),
                    project_root.as_deref(),
                    CrpMode::effective(),
                )
            } else {
                tracing::warn!("pre-dispatch: cache write-lock timeout (5s), skipping autonomy");
                None
            }
        };

        let args_fp = args
            .map(|a| {
                crate::core::loop_detection::LoopDetector::fingerprint(&serde_json::Value::Object(
                    a.clone(),
                ))
            })
            .unwrap_or_default();
        let throttle_result = {
            let fp = &args_fp;
            let detector_timeout = tokio::time::timeout(
                std::time::Duration::from_secs(3),
                self.loop_detector.write(),
            )
            .await;
            if let Ok(mut detector) = detector_timeout {
                let is_search = crate::core::loop_detection::LoopDetector::is_search_tool(name);
                let is_search_shell = name == "ctx_shell" && {
                    let cmd = args
                        .as_ref()
                        .and_then(|a| a.get("command"))
                        .and_then(|v| v.as_str())
                        .unwrap_or("");
                    crate::core::loop_detection::LoopDetector::is_search_shell_command(cmd)
                };

                if is_search || is_search_shell {
                    let search_pattern = args.and_then(|a| {
                        a.get("pattern")
                            .or_else(|| a.get("query"))
                            .and_then(|v| v.as_str())
                    });
                    let shell_pattern = if is_search_shell {
                        args.and_then(|a| a.get("command"))
                            .and_then(|v| v.as_str())
                            .and_then(helpers::extract_search_pattern_from_command)
                    } else {
                        None
                    };
                    let pat = search_pattern.or(shell_pattern.as_deref());
                    detector.record_search(name, fp, pat)
                } else {
                    detector.record_call(name, fp)
                }
            } else {
                tracing::warn!("pre-dispatch: loop_detector write-lock timeout (3s), skipping");
                crate::core::loop_detection::ThrottleResult::default()
            }
        };

        if throttle_result.level == crate::core::loop_detection::ThrottleLevel::Blocked {
            let msg = throttle_result.message.unwrap_or_default();
            let result = CallToolResult::success(vec![ContentBlock::text(msg)]);
            finish_decision_loop(decision_context.as_ref(), args, &result);
            return Ok(result);
        }

        let throttle_warning =
            if throttle_result.level == crate::core::loop_detection::ThrottleLevel::Reduced {
                throttle_result.message.clone()
            } else {
                None
            };

        let minimal = config.minimal_overhead_effective();

        // IDE permission inheritance: when enabled, mirror the host IDE's
        // bash/read/edit/grep permission rules onto the matching lean-ctx tool so
        // e.g. `ctx_shell` honors a `rm *: ask` rule instead of bypassing it.
        // Gated on the cheap effective() check so the default (off) pays no lock
        // cost on the hot path. Checks the ctx_call-unwrapped inner tool (#1008)
        // so the invoker cannot side-step an IDE deny.
        if config.permission_inheritance_effective()
            == crate::core::config::PermissionInheritance::On
        {
            let client_name = self.client_name.read().await.clone();
            let project_root = self.session.read().await.project_root.clone();
            let perm = permission_inheritance::check(
                &client_name,
                guard_name,
                guard_args,
                project_root.as_deref(),
                &config,
            );
            if let Some(blocked) = permission_inheritance::into_call_tool_result(&perm) {
                tracing::warn!(tool = guard_name, "held back by IDE permission inheritance");
                finish_decision_loop(decision_context.as_ref(), args, &blocked);
                return Ok(blocked);
            }
        }

        if let Some(msg) = post_process::budget_exhausted_message(name) {
            tracing::warn!(tool = name, "{msg}");
            let result = CallToolResult::success(vec![ContentBlock::text(msg)]);
            finish_decision_loop(decision_context.as_ref(), args, &result);
            return Ok(result);
        }

        if is_shell_tool_name(name) {
            crate::core::budget_tracker::BudgetTracker::global().record_shell();
        }

        let project_root = self
            .session
            .read()
            .await
            .project_root
            .clone()
            .unwrap_or_default();
        let cache_key = response_cache_key(name, args, &project_root);
        if let Some(cached) = cache_key
            .as_ref()
            .and_then(|key| cached_call_result(global_response_cache(), key))
        {
            finish_decision_loop(decision_context.as_ref(), args, &cached);
            return Ok(cached);
        }

        let result = dispatch_and_post_process(
            self,
            name,
            args,
            minimal,
            config,
            machine_readable,
            auto_context,
            throttle_warning,
            args_fp,
            decision_context,
        )
        .await;
        if let (Some(key), Ok(response)) = (cache_key, &result) {
            cache_call_result(global_response_cache(), key, response);
        }
        result
    }
}

fn finish_decision_loop(
    context: Option<&crate::core::decision_loop_runtime::TaskContext>,
    args: Option<&Map<String, Value>>,
    result: &CallToolResult,
) {
    let Some(context) = context else {
        return;
    };
    let input_tokens = args
        .and_then(|args| serde_json::to_string(args).ok())
        .map_or(0, |input| (input.len() / 4) as u64);
    let output_tokens = format!("{result:?}").len() as u64 / 4;
    if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
        crate::core::decision_loop_runtime::DecisionLoopRuntime::get_or_init().on_tool_end(
            context,
            input_tokens,
            output_tokens,
            "mcp-tool",
            result.is_error != Some(true),
        );
    }))
    .is_err()
    {
        tracing::warn!("decision loop end panicked");
    }
}

#[cfg(test)]
mod tests {
    use super::{
        CallToolResult, ContentBlock, Duration, Meta, ResponseCache, Value, cache_call_result,
        cached_call_result, response_cache_key,
    };

    #[test]
    fn cached_result_preserves_meta() {
        let cache = ResponseCache::new(8, Duration::from_mins(1));
        let key = response_cache_key("ctx_search", None, "/project").expect("cacheable tool");
        let mut result = CallToolResult::success(vec![ContentBlock::text("cached")]);
        let mut meta = Meta::new();
        meta.0
            .insert("cache_hint".to_owned(), Value::String("stable".to_owned()));
        result.meta = Some(meta);

        cache_call_result(&cache, key.clone(), &result);

        let cached = cached_call_result(&cache, &key).expect("response should be cached");
        assert_eq!(cached.meta, result.meta);
    }

    #[test]
    fn ctx_read_is_not_response_cached() {
        assert!(response_cache_key("ctx_read", None, "/project").is_none());
    }

    #[test]
    fn remaining_response_cache_tools_are_cacheable() {
        for tool_name in ["ctx_search", "ctx_tree", "ctx_glob"] {
            assert!(response_cache_key(tool_name, None, "/project").is_some());
        }
    }
}