1use 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
14pub struct ToolRegistry {
19 pub builtins: Arc<BuiltinTools>,
21 pub mcp: Arc<Mutex<ToolExecutor>>,
24 pub mcp_tool_defs: Vec<Tool>,
26 pub builtin_names: HashSet<String>,
28}
29
30impl ToolRegistry {
31 pub async fn build(workdir: PathBuf, config: &Config) -> Self {
33 let ctx = ToolContext::new(workdir);
34 let builtins = Arc::new(BuiltinTools::new(ctx));
35 let builtin_names: HashSet<String> = builtins.names().into_iter().collect();
36
37 let mut mcp_executor = ToolExecutor::new();
38 let mut mcp_tool_defs: Vec<Tool> = Vec::new();
39
40 if !config.mcp_servers.is_empty() {
41 let mut discovery = ToolDiscovery::new();
42 let oauth = leviath_mcp::OAuthClient::new();
43 let store_path = leviath_mcp::AuthStore::default_path();
44 let now = unix_now_secs();
45 let credentials = credential_store_or_warn(crate::credentials::store_for(
51 config.security.credential_store,
52 ));
53 for server_cfg in &config.mcp_servers {
54 let auth_header = match resolve_bearer(
59 &oauth,
60 &server_cfg.name,
61 store_path.as_deref(),
62 now,
63 credentials.as_deref(),
64 )
65 .await
66 {
67 Ok(header) => header,
68 Err(e) => {
69 tracing::warn!(server = %server_cfg.name, error = %e, "MCP auth unavailable - skipping");
70 continue;
71 }
72 };
73 let auth_was_resolved = auth_header.is_some();
76 match discovery
77 .discover_from_config_with_auth(
78 server_cfg,
79 auth_header,
80 &config.security.allow_env_vars,
81 )
82 .await
83 {
84 Ok((_tool_metas, mut client)) => {
85 if auth_was_resolved && let Some(path) = store_path.clone() {
89 client.set_refresher(std::sync::Arc::new(
90 leviath_mcp::StoredTokenRefresher::new(
91 server_cfg.name.clone(),
92 path,
93 ),
94 ));
95 }
96 let mut reserved: HashSet<String> = builtin_names.clone();
101 reserved.extend(mcp_tool_defs.iter().map(|t| t.name.clone()));
102 let advertised = mcp_executor.add_client_advertised(
103 server_cfg.name.clone(),
104 client,
105 &reserved,
106 );
107 for meta in advertised {
108 mcp_tool_defs.push(Tool {
109 name: meta.name,
110 description: meta.description,
111 parameters: meta.schema,
112 });
113 }
114 tracing::info!(server = %server_cfg.name, "Connected MCP server");
115 }
116 Err(e) => {
117 let span = tracing::warn_span!(
118 "mcp_server_connect_failed",
119 server = tracing::field::Empty,
120 error = tracing::field::Empty
121 );
122 let _enter = span.enter();
123 span.record("server", tracing::field::display(&server_cfg.name));
124 span.record("error", tracing::field::display(&e));
125 tracing::warn!("Failed to connect MCP server - skipping");
126 }
127 }
128 }
129 }
130
131 Self {
132 builtins,
133 mcp: Arc::new(Mutex::new(mcp_executor)),
134 mcp_tool_defs,
135 builtin_names,
136 }
137 }
138
139 pub fn all_tool_defs(&self) -> Vec<Tool> {
141 let mut tools = self.builtins.tool_defs();
142 tools.extend(BuiltinTools::subagent_tool_defs());
143 tools.extend_from_slice(&self.mcp_tool_defs);
144 tools
145 }
146
147 pub async fn shutdown(&self) {
149 let mut mcp = self.mcp.lock().await;
150 let _ = mcp.shutdown_all().await;
155 }
156}
157
158pub(crate) fn unix_now_secs() -> u64 {
162 std::time::SystemTime::now()
163 .duration_since(std::time::UNIX_EPOCH)
164 .map(|d| d.as_secs())
165 .unwrap_or(0)
166}
167
168pub(crate) fn credential_store_or_warn(
181 resolved: crate::credentials::Resolved,
182) -> Option<Box<dyn leviath_core::CredentialStore>> {
183 match resolved {
184 Ok(store) => store,
185 Err(e) => {
186 tracing::warn!("{e}. MCP servers needing OAuth will appear logged out.");
187 None
188 }
189 }
190}
191
192pub(crate) async fn resolve_bearer(
193 oauth: &leviath_mcp::OAuthClient,
194 server_name: &str,
195 store_path: Option<&std::path::Path>,
196 now: u64,
197 credentials: Option<&dyn leviath_core::CredentialStore>,
198) -> anyhow::Result<Option<(String, String)>> {
199 match store_path {
200 Some(path) => {
201 oauth
202 .authorization_header_with(server_name, path, now, credentials)
203 .await
204 }
205 None => Ok(None),
206 }
207}
208
209pub fn default_tool_policy(tool_name: &str, is_builtin: bool) -> ToolPolicy {
213 match leviath_tools::canonical_tool_name(tool_name) {
216 "read_file" | "read_files" | "list_dir" => ToolPolicy::Allow,
217 "context_write" | "context_append" | "context_read" | "context_delete" | "context_list"
222 | "todo_add" | "todo_done" | "todo_note" => ToolPolicy::Allow,
226 "write_file" | "edit_file" | "shell" => ToolPolicy::Ask,
227 "spawn_agent" | "check_agent" | "wait_for_agent" | "send_to_agent" | "kill_agent" => {
244 ToolPolicy::Allow
245 }
246 "ask_user_text" | "ask_user_choice" | "ask_user_confirm" | "edit_document" => {
250 ToolPolicy::Allow
251 }
252 _ => {
253 let _ = is_builtin;
255 ToolPolicy::Ask
256 }
257 }
258}
259
260fn restrictiveness(p: ToolPolicy) -> u8 {
262 match p {
263 ToolPolicy::Allow => 0,
264 ToolPolicy::Ask => 1,
265 ToolPolicy::Deny => 2,
266 }
267}
268
269fn stricter(a: ToolPolicy, b: ToolPolicy) -> ToolPolicy {
271 if restrictiveness(b) > restrictiveness(a) {
272 b
273 } else {
274 a
275 }
276}
277
278pub fn clamp_by_effect(
301 tool_name: &str,
302 arguments: &serde_json::Value,
303 policy: ToolPolicy,
304 write_policy: &dyn Fn() -> ToolPolicy,
305) -> ToolPolicy {
306 if leviath_tools::canonical_tool_name(tool_name) != "shell" {
307 return policy;
308 }
309 let Some(command) = arguments.get("command").and_then(|v| v.as_str()) else {
310 return policy;
311 };
312 if crate::shell_keys::writes_a_file(command) {
313 return stricter(policy, write_policy());
314 }
315 policy
316}
317
318pub fn escaping_write_refusal(
341 tool_name: &str,
342 arguments: &serde_json::Value,
343 workdir: &std::path::Path,
344) -> Option<String> {
345 if leviath_tools::canonical_tool_name(tool_name) != "shell" {
346 return None;
347 }
348 let command = arguments.get("command").and_then(|v| v.as_str())?;
349 let escaping = crate::shell_keys::write_target_paths(command)
350 .into_iter()
351 .find(|target| {
352 let joined = match std::path::Path::new(target).is_absolute() {
353 true => std::path::PathBuf::from(target),
354 false => workdir.join(target),
355 };
356 !leviath_core::resolves_within(&joined, workdir)
357 })?;
358 Some(format!(
359 "[denied] Shell redirect writes to '{escaping}', which is outside the working directory \
360 ({}). The `write_file` tool refuses the same path; a redirect is not a way around it. \
361 Write inside the workspace instead.",
362 workdir.display()
363 ))
364}
365
366pub fn declared_write_bytes(tool_name: &str, arguments: &serde_json::Value) -> Option<u64> {
375 let field = match leviath_tools::canonical_tool_name(tool_name) {
376 "write_file" => "content",
377 "edit_file" => "new_str",
378 _ => return None,
379 };
380 arguments
381 .get(field)
382 .and_then(|v| v.as_str())
383 .map(|s| s.len() as u64)
384}
385
386pub fn write_budget_refusal(
401 tool_name: &str,
402 arguments: &serde_json::Value,
403 workdir: &std::path::Path,
404 budget: &crate::daemon::tool_service::WriteBudget,
405) -> Option<String> {
406 let declared = declared_write_bytes(tool_name, arguments);
407 let writes_something = declared.is_some()
408 || (leviath_tools::canonical_tool_name(tool_name) == "shell"
409 && arguments
410 .get("command")
411 .and_then(|v| v.as_str())
412 .is_some_and(crate::shell_keys::writes_a_file));
413 if !writes_something {
414 return None;
415 }
416 budget.check(workdir, declared.unwrap_or(0)).refusal()
417}
418
419pub fn measured_write_bytes(
437 tool_name: &str,
438 arguments: &serde_json::Value,
439 workdir: &std::path::Path,
440) -> u64 {
441 if declared_write_bytes(tool_name, arguments).is_some() {
442 return 0;
443 }
444 if leviath_tools::canonical_tool_name(tool_name) != "shell" {
445 return 0;
446 }
447 let Some(command) = arguments.get("command").and_then(|v| v.as_str()) else {
448 return 0;
449 };
450 crate::shell_keys::write_target_paths(command)
451 .into_iter()
452 .map(|target| {
453 let path = match std::path::Path::new(&target).is_absolute() {
454 true => std::path::PathBuf::from(&target),
455 false => workdir.join(&target),
456 };
457 std::fs::metadata(&path).map(|m| m.len()).unwrap_or(0)
458 })
459 .sum()
460}
461
462const BLUEPRINT_LOOSENABLE: &[&str] = &["web_search", "web_fetch"];
489
490fn blueprint_loosenable(tool_name: &str) -> bool {
492 leviath_tools::tool_name_spellings(tool_name).any(|n| BLUEPRINT_LOOSENABLE.contains(&n))
493}
494
495pub fn resolve_policy(
521 tool_name: &str,
522 is_builtin: bool,
523 launch_overrides: &HashMap<String, ToolPolicy>,
524 stage_permissions: &HashMap<String, String>,
525 agent_permissions: &HashMap<String, String>,
526 global_permissions: &HashMap<String, ToolPolicy>,
527 blueprint_may_loosen: bool,
528) -> ToolPolicy {
529 let ceiling = by_any_spelling(global_permissions, tool_name).copied();
530
531 let blueprint = by_any_spelling(stage_permissions, tool_name)
533 .or_else(|| by_any_spelling(agent_permissions, tool_name))
534 .map(|s| parse_policy_str(s));
535
536 let configured = match (blueprint, ceiling) {
537 (Some(b), Some(c)) => stricter(b, c),
538 (Some(b), None) => {
542 let default = default_tool_policy(tool_name, is_builtin);
543 match blueprint_may_loosen || blueprint_loosenable(tool_name) {
544 true => b,
545 false => stricter(b, default),
546 }
547 }
548 (None, Some(c)) => c,
549 (None, None) => default_tool_policy(tool_name, is_builtin),
550 };
551
552 if configured == ToolPolicy::Deny {
554 return ToolPolicy::Deny;
555 }
556
557 by_any_spelling(launch_overrides, tool_name)
558 .or_else(|| launch_overrides.get("*"))
559 .copied()
560 .unwrap_or(configured)
561}
562
563pub fn session_approval_keys(tool_name: &str, arguments: &serde_json::Value) -> Vec<String> {
577 if leviath_tools::canonical_tool_name(tool_name) != "shell" {
578 return vec![tool_name.to_string()];
579 }
580 let Some(command) = arguments.get("command").and_then(|v| v.as_str()) else {
581 return Vec::new();
582 };
583 crate::shell_keys::command_keys(command)
584}
585
586fn by_any_spelling<'a, V>(map: &'a HashMap<String, V>, tool_name: &str) -> Option<&'a V> {
597 leviath_tools::tool_name_spellings(tool_name).find_map(|name| map.get(name))
598}
599
600fn parse_policy_str(s: &str) -> ToolPolicy {
609 match s.to_lowercase().as_str() {
610 "allow" => ToolPolicy::Allow,
611 "deny" => ToolPolicy::Deny,
612 _ => ToolPolicy::Ask,
613 }
614}
615
616#[cfg(test)]
617mod mcp_registry_tests {
618 use super::*;
619 use crate::test_support::with_tracing;
620 use leviath_mcp::MCPServerConfig;
621
622 const STUB_INIT_AND_LIST: &str = r#"
628import sys, json
629
630def respond(id, result):
631 msg = json.dumps({"jsonrpc": "2.0", "id": id, "result": result})
632 sys.stdout.write(msg + "\n")
633 sys.stdout.flush()
634
635for line in sys.stdin:
636 line = line.strip()
637 if not line:
638 continue
639 req = json.loads(line)
640 method = req.get("method", "")
641 id_ = req.get("id")
642 if method == "initialize":
643 respond(id_, {"capabilities": {"tools": {"listChanged": True}}, "protocolVersion": "2024-11-05"})
644 elif method == "notifications/initialized":
645 pass
646 elif method == "tools/list":
647 respond(id_, {"tools": [{"name": "echo", "description": "echo tool", "inputSchema": {}}]})
648 elif method == "tools/call":
649 args = req.get("params", {}).get("arguments", {})
650 if args.get("fail"):
651 respond(id_, {"content": [{"type": "text", "text": "it broke"}], "isError": True})
652 else:
653 respond(id_, {"content": [{"type": "text", "text": "echoed!"}], "isError": False})
654 else:
655 respond(id_, {"error": {"code": -32601, "message": "method not found"}})
656"#;
657
658 fn config_with_mcp_server(command: &str, args: Vec<&str>) -> Config {
659 Config {
660 mcp_servers: vec![MCPServerConfig::stdio(
661 "stub-server",
662 command,
663 args.into_iter().map(String::from).collect(),
664 )],
665 ..Config::default()
666 }
667 }
668
669 async fn with_temp_home<F, Fut, T>(body: F) -> T
673 where
674 F: FnOnce() -> Fut,
675 Fut: std::future::Future<Output = T>,
676 {
677 let dir = tempfile::tempdir().unwrap();
678 temp_env::async_with_vars(
679 [("LEVIATH_HOME", Some(dir.path().to_str().unwrap()))],
680 body(),
681 )
682 .await
683 }
684
685 #[tokio::test]
686 async fn build_connects_mcp_server_and_registers_its_tools() {
687 with_tracing(|| {});
688 let registry = with_temp_home(|| async {
689 let config = config_with_mcp_server("python3", vec!["-c", STUB_INIT_AND_LIST]);
690 ToolRegistry::build(std::env::temp_dir(), &config).await
691 })
692 .await;
693
694 assert_eq!(registry.mcp_tool_defs.len(), 1);
695 assert_eq!(registry.mcp_tool_defs[0].name, "echo");
696
697 registry.shutdown().await;
698 }
699
700 #[tokio::test]
701 async fn build_advertises_two_servers_and_namespaces_a_collision() {
702 with_tracing(|| {});
707 let registry = with_temp_home(|| async {
708 let config = Config {
709 mcp_servers: vec![
710 MCPServerConfig::stdio(
711 "alpha",
712 "python3",
713 vec!["-c".to_string(), STUB_INIT_AND_LIST.to_string()],
714 ),
715 MCPServerConfig::stdio(
716 "beta",
717 "python3",
718 vec!["-c".to_string(), STUB_INIT_AND_LIST.to_string()],
719 ),
720 ],
721 ..Config::default()
722 };
723 ToolRegistry::build(std::env::temp_dir(), &config).await
724 })
725 .await;
726
727 let names: Vec<&str> = registry
728 .mcp_tool_defs
729 .iter()
730 .map(|t| t.name.as_str())
731 .collect();
732 assert!(names.contains(&"echo"), "names: {names:?}");
734 assert!(names.contains(&"beta__echo"), "names: {names:?}");
735 registry.shutdown().await;
736 }
737
738 async fn mock_http_mcp_server() -> String {
741 use axum::response::IntoResponse;
742 use axum::routing::post;
743 use axum::{Json, Router};
744 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
745 let base = format!("http://{}", listener.local_addr().unwrap());
746 let app = Router::new().route(
747 "/mcp",
748 post(|body: String| async move {
751 let req: serde_json::Value = serde_json::from_str(&body).unwrap();
752 let id = req.get("id").cloned().unwrap_or(serde_json::json!(1));
753 let result = match req.get("method").and_then(|m| m.as_str()) {
754 Some("initialize") => {
755 serde_json::json!({"capabilities": {}, "protocolVersion": "2024-11-05"})
756 }
757 Some("tools/list") => {
758 serde_json::json!({"tools": [{"name": "remote_tool", "inputSchema": {}}]})
759 }
760 _ => serde_json::json!({}),
761 };
762 (
763 [(axum::http::header::CONTENT_TYPE, "application/json")],
764 Json(serde_json::json!({"jsonrpc": "2.0", "id": id, "result": result}))
765 .into_response()
766 .into_body(),
767 )
768 .into_response()
769 }),
770 );
771 tokio::spawn(std::future::IntoFuture::into_future(axum::serve(
772 listener, app,
773 )));
774 base
775 }
776
777 #[tokio::test]
778 async fn build_attaches_a_refresher_to_an_authenticated_http_server() {
779 with_tracing(|| {});
782 let base = mock_http_mcp_server().await;
783 let registry = with_temp_home(|| async {
784 let mut store = leviath_mcp::AuthStore::default();
786 store.set(
787 "remote",
788 leviath_mcp::ServerAuth {
789 access_token: "live-token".to_string(),
790 expires_at: u64::MAX,
791 ..Default::default()
792 },
793 );
794 store
795 .save(&leviath_mcp::AuthStore::default_path().unwrap())
796 .unwrap();
797
798 let config = Config {
799 mcp_servers: vec![MCPServerConfig::http("remote", format!("{base}/mcp"))],
800 ..Config::default()
801 };
802 ToolRegistry::build(std::env::temp_dir(), &config).await
803 })
804 .await;
805
806 assert_eq!(registry.mcp_tool_defs.len(), 1);
807 assert_eq!(registry.mcp_tool_defs[0].name, "remote_tool");
808 registry.shutdown().await;
809 }
810
811 #[tokio::test]
812 async fn build_skips_mcp_server_that_fails_to_connect() {
813 with_tracing(|| {});
817 let registry = with_temp_home(|| async {
818 let config = config_with_mcp_server("definitely-not-a-real-binary-xyz", vec![]);
819 ToolRegistry::build(std::env::temp_dir(), &config).await
820 })
821 .await;
822
823 assert!(registry.mcp_tool_defs.is_empty());
824 }
825
826 #[tokio::test]
827 async fn build_skips_http_server_whose_token_cannot_be_refreshed() {
828 with_tracing(|| {});
833 let registry = with_temp_home(|| async {
834 let mut store = leviath_mcp::AuthStore::default();
836 store.set(
837 "remote",
838 leviath_mcp::ServerAuth {
839 token_endpoint: "http://127.0.0.1:1/token".to_string(),
840 access_token: "expired".to_string(),
841 refresh_token: Some("good".to_string()),
842 expires_at: 1,
843 ..Default::default()
844 },
845 );
846 store
847 .save(&leviath_mcp::AuthStore::default_path().unwrap())
848 .unwrap();
849
850 let config = Config {
851 mcp_servers: vec![MCPServerConfig::http("remote", "http://127.0.0.1:1/mcp")],
852 ..Config::default()
853 };
854 ToolRegistry::build(std::env::temp_dir(), &config).await
855 })
856 .await;
857 assert!(registry.mcp_tool_defs.is_empty());
858 }
859
860 #[test]
863 fn an_unreachable_credential_store_warns_rather_than_failing_tool_setup() {
864 assert!(
865 credential_store_or_warn(Err("no keychain here".to_string())).is_none(),
866 "an unreachable store yields no credentials"
867 );
868 assert!(
869 credential_store_or_warn(Ok(None)).is_none(),
870 "and so does the file backend"
871 );
872 assert!(
873 credential_store_or_warn(Ok(Some(Box::new(leviath_core::MemoryStore::new()))))
874 .is_some()
875 );
876 }
877
878 #[tokio::test]
879 async fn resolve_bearer_without_a_store_is_none() {
880 let oauth = leviath_mcp::OAuthClient::new();
881 let header = resolve_bearer(&oauth, "srv", None, 0, None).await.unwrap();
882 assert!(header.is_none());
883 }
884
885 #[tokio::test]
886 async fn shutdown_with_no_servers_is_a_noop() {
887 let config = Config::default();
888 let registry = ToolRegistry::build(std::env::temp_dir(), &config).await;
889 registry.shutdown().await; }
891}
892
893#[cfg(test)]
894mod policy_tests {
895 use super::*;
896
897 fn shell_call(command: &str) -> serde_json::Value {
900 serde_json::json!({ "command": command })
901 }
902
903 fn deny() -> ToolPolicy {
907 ToolPolicy::Deny
908 }
909
910 fn allow() -> ToolPolicy {
911 ToolPolicy::Allow
912 }
913
914 #[test]
921 fn a_redirect_outside_the_workdir_is_refused() {
922 let dir = tempfile::tempdir().expect("tempdir");
923 for command in [
924 "echo pwn > /root/.bashrc",
925 "cat notes.md > ../escaped.txt",
926 "echo x >> ../../etc/hosts",
927 ] {
928 let refusal = escaping_write_refusal("shell", &shell_call(command), dir.path());
929 assert!(
930 refusal.is_some(),
931 "{command:?} writes outside the workdir and must be refused"
932 );
933 let refusal = refusal.expect("just asserted");
934 assert!(
935 refusal.contains("outside the working directory"),
936 "{refusal}"
937 );
938 }
939 }
940
941 #[test]
944 fn a_redirect_inside_the_workdir_is_allowed() {
945 let dir = tempfile::tempdir().expect("tempdir");
946 for command in [
947 "echo x > out.txt",
948 "echo x > sub/dir/out.txt",
949 "cat a >> ./notes.md",
950 "ninja > /dev/null 2>&1",
952 "cat a 2>/dev/null",
953 "ls",
955 ] {
956 assert_eq!(
957 escaping_write_refusal("shell", &shell_call(command), dir.path()),
958 None,
959 "{command:?} stays inside and must not be refused"
960 );
961 }
962 }
963
964 #[test]
971 fn an_absolute_path_into_the_workdir_is_allowed() {
972 let dir = tempfile::tempdir().expect("tempdir");
973 let inside = dir.path().join("out.txt");
974 let command = format!("echo x > {}", inside.display());
975 assert_eq!(
976 escaping_write_refusal("shell", &shell_call(&command), dir.path()),
977 None
978 );
979 }
980
981 #[test]
984 fn the_refusal_covers_the_alias_and_ignores_other_tools() {
985 let dir = tempfile::tempdir().expect("tempdir");
986 assert!(
987 escaping_write_refusal("bash", &shell_call("echo x > /root/.bashrc"), dir.path(),)
988 .is_some()
989 );
990 assert_eq!(
991 escaping_write_refusal(
992 "write_file",
993 &serde_json::json!({ "path": "/root/.bashrc", "content": "x" }),
994 dir.path(),
995 ),
996 None
997 );
998 assert_eq!(
1000 escaping_write_refusal("shell", &serde_json::json!({}), dir.path()),
1001 None
1002 );
1003 }
1004
1005 #[test]
1010 fn a_declaring_tool_reports_its_size_and_a_shell_call_does_not() {
1011 assert_eq!(
1012 declared_write_bytes("write_file", &serde_json::json!({"content": "abcd"})),
1013 Some(4)
1014 );
1015 assert_eq!(
1016 declared_write_bytes("edit_file", &serde_json::json!({"new_str": "abc"})),
1017 Some(3)
1018 );
1019 assert_eq!(
1023 declared_write_bytes("shell", &shell_call("echo x > out.txt")),
1024 None
1025 );
1026 assert_eq!(
1028 declared_write_bytes("write_file", &serde_json::json!({})),
1029 None
1030 );
1031 }
1032
1033 #[test]
1037 fn measuring_charges_the_shell_and_not_the_declaring_tools() {
1038 let dir = tempfile::tempdir().expect("tempdir");
1039 std::fs::write(dir.path().join("out.txt"), "0123456789").expect("write");
1040
1041 assert_eq!(
1042 measured_write_bytes("shell", &shell_call("echo x > out.txt"), dir.path()),
1043 10
1044 );
1045 assert_eq!(
1047 measured_write_bytes(
1048 "write_file",
1049 &serde_json::json!({"path": "a", "content": "abcd"}),
1050 dir.path()
1051 ),
1052 0
1053 );
1054 assert_eq!(
1057 measured_write_bytes("shell", &shell_call("echo x > absent.txt"), dir.path()),
1058 0
1059 );
1060 assert_eq!(
1062 measured_write_bytes("shell", &shell_call("ls -la"), dir.path()),
1063 0
1064 );
1065 let absolute = dir.path().join("out.txt");
1069 let command = format!("echo x > {}", absolute.display());
1070 assert_eq!(
1071 measured_write_bytes("shell", &shell_call(&command), dir.path()),
1072 10
1073 );
1074 assert_eq!(
1076 measured_write_bytes("shell", &serde_json::json!({}), dir.path()),
1077 0
1078 );
1079 assert_eq!(
1081 measured_write_bytes("read_file", &serde_json::json!({"path": "a"}), dir.path()),
1082 0
1083 );
1084 }
1085
1086 #[test]
1089 fn a_denied_write_tool_denies_a_shell_redirect() {
1090 assert_eq!(
1091 clamp_by_effect(
1092 "shell",
1093 &shell_call("echo pwn > /root/.bashrc"),
1094 ToolPolicy::Allow,
1095 &deny,
1096 ),
1097 ToolPolicy::Deny
1098 );
1099 assert_eq!(
1101 clamp_by_effect(
1102 "bash",
1103 &shell_call("echo pwn >> ~/.profile"),
1104 ToolPolicy::Allow,
1105 &deny,
1106 ),
1107 ToolPolicy::Deny
1108 );
1109 }
1110
1111 #[test]
1114 fn the_clamp_never_loosens_a_shell_call() {
1115 assert_eq!(
1116 clamp_by_effect(
1117 "shell",
1118 &shell_call("echo x > out"),
1119 ToolPolicy::Ask,
1120 &allow,
1121 ),
1122 ToolPolicy::Ask
1123 );
1124 assert_eq!(
1125 clamp_by_effect(
1126 "shell",
1127 &shell_call("echo x > out"),
1128 ToolPolicy::Deny,
1129 &allow,
1130 ),
1131 ToolPolicy::Deny
1132 );
1133 }
1134
1135 #[test]
1138 fn a_call_that_writes_nothing_is_untouched() {
1139 for command in ["ls -la", "cat a 2>/dev/null", "grep x f", "sort < in"] {
1140 assert_eq!(
1141 clamp_by_effect("shell", &shell_call(command), ToolPolicy::Allow, &deny,),
1142 ToolPolicy::Allow,
1143 "{command:?} writes nothing"
1144 );
1145 }
1146 }
1147
1148 #[test]
1152 fn the_write_policy_is_resolved_only_when_a_call_actually_writes() {
1153 let calls = std::cell::Cell::new(0);
1154 let resolve = || {
1155 calls.set(calls.get() + 1);
1156 ToolPolicy::Deny
1157 };
1158
1159 clamp_by_effect(
1160 "read_file",
1161 &serde_json::json!({"path": "a"}),
1162 ToolPolicy::Allow,
1163 &resolve,
1164 );
1165 clamp_by_effect("shell", &shell_call("ls -la"), ToolPolicy::Allow, &resolve);
1166 clamp_by_effect(
1167 "shell",
1168 &shell_call("cat a 2>/dev/null"),
1169 ToolPolicy::Allow,
1170 &resolve,
1171 );
1172 assert_eq!(
1173 calls.get(),
1174 0,
1175 "nothing here writes, so nothing should resolve"
1176 );
1177
1178 clamp_by_effect(
1179 "shell",
1180 &shell_call("echo x > f"),
1181 ToolPolicy::Allow,
1182 &resolve,
1183 );
1184 assert_eq!(calls.get(), 1, "a writing call resolves it exactly once");
1185 }
1186
1187 #[test]
1190 fn a_non_shell_tool_and_a_malformed_call_are_untouched() {
1191 assert_eq!(
1192 clamp_by_effect(
1193 "read_file",
1194 &serde_json::json!({ "path": "a > b" }),
1195 ToolPolicy::Allow,
1196 &deny,
1197 ),
1198 ToolPolicy::Allow
1199 );
1200 assert_eq!(
1201 clamp_by_effect(
1202 "shell",
1203 &serde_json::json!({ "not_a_command": 1 }),
1204 ToolPolicy::Allow,
1205 &deny,
1206 ),
1207 ToolPolicy::Allow
1208 );
1209 }
1210
1211 fn blueprint_says(tool: &str, policy: &str, may_loosen: bool) -> ToolPolicy {
1215 let mut agent = HashMap::new();
1216 agent.insert(tool.to_string(), policy.to_string());
1217 resolve_policy(
1218 tool,
1219 true,
1220 &HashMap::new(),
1221 &HashMap::new(),
1222 &agent,
1223 &HashMap::new(),
1224 may_loosen,
1225 )
1226 }
1227
1228 #[test]
1233 fn a_blueprint_cannot_loosen_a_tool_the_user_never_configured() {
1234 for tool in ["shell", "write_file", "edit_file"] {
1235 assert_eq!(
1236 blueprint_says(tool, "allow", false),
1237 ToolPolicy::Ask,
1238 "{tool} must fall back to its built-in default"
1239 );
1240 }
1241 assert_eq!(
1245 blueprint_says("spawn_agent", "allow", false),
1246 ToolPolicy::Allow
1247 );
1248 }
1249
1250 #[test]
1253 fn a_blueprint_may_still_tighten_anything() {
1254 assert_eq!(blueprint_says("shell", "deny", false), ToolPolicy::Deny);
1255 assert_eq!(blueprint_says("read_file", "ask", false), ToolPolicy::Ask);
1256 }
1257
1258 #[test]
1261 fn a_blueprint_may_preapprove_the_read_only_web_tools() {
1262 assert_eq!(
1263 blueprint_says("web_fetch", "allow", false),
1264 ToolPolicy::Allow
1265 );
1266 assert_eq!(
1267 blueprint_says("web_search", "allow", false),
1268 ToolPolicy::Allow
1269 );
1270 }
1271
1272 #[test]
1274 fn an_opted_in_blueprint_may_loosen_anything() {
1275 assert_eq!(blueprint_says("shell", "allow", true), ToolPolicy::Allow);
1276 }
1277
1278 #[test]
1281 fn a_configured_ceiling_still_bounds_a_blueprint() {
1282 let mut agent = HashMap::new();
1283 agent.insert("shell".to_string(), "allow".to_string());
1284 let mut global = HashMap::new();
1285 global.insert("shell".to_string(), ToolPolicy::Allow);
1286 assert_eq!(
1287 resolve_policy(
1288 "shell",
1289 true,
1290 &HashMap::new(),
1291 &HashMap::new(),
1292 &agent,
1293 &global,
1294 false,
1295 ),
1296 ToolPolicy::Allow,
1297 "naming the tool in the user's own config is the per-agent grant"
1298 );
1299
1300 global.insert("shell".to_string(), ToolPolicy::Deny);
1301 assert_eq!(
1302 resolve_policy(
1303 "shell",
1304 true,
1305 &HashMap::new(),
1306 &HashMap::new(),
1307 &agent,
1308 &global,
1309 true,
1310 ),
1311 ToolPolicy::Deny,
1312 "a configured deny is terminal even for an opted-in blueprint"
1313 );
1314 }
1315
1316 #[test]
1321 fn the_bundled_agents_resolve_unchanged() {
1322 for agent in crate::bundled::BUNDLED_AGENTS {
1323 let (_, manifest) = agent
1324 .files
1325 .iter()
1326 .find(|(rel, _)| rel.ends_with("agent.leviath"))
1327 .expect("every bundled agent ships a manifest");
1328 let bp = leviath_core::manifest::parse_manifest(manifest)
1329 .expect("every bundled agent's manifest parses");
1330 let perms = bp.agent_tool_permissions();
1331 for (tool, declared) in &perms {
1332 let clamped = resolve_policy(
1333 tool,
1334 true,
1335 &HashMap::new(),
1336 &HashMap::new(),
1337 &perms,
1338 &HashMap::new(),
1339 false,
1340 );
1341 let unclamped = resolve_policy(
1342 tool,
1343 true,
1344 &HashMap::new(),
1345 &HashMap::new(),
1346 &perms,
1347 &HashMap::new(),
1348 true,
1349 );
1350 assert_eq!(
1351 clamped, unclamped,
1352 "{}'s {tool} = {declared:?} changed meaning under the allowlist",
1353 agent.name
1354 );
1355 }
1356 }
1357 }
1358
1359 #[test]
1360 fn test_default_policy_read_file() {
1361 assert_eq!(default_tool_policy("read_file", true), ToolPolicy::Allow);
1362 assert_eq!(default_tool_policy("list_dir", true), ToolPolicy::Allow);
1363 }
1364
1365 #[test]
1366 fn test_default_policy_write_tools() {
1367 assert_eq!(default_tool_policy("write_file", true), ToolPolicy::Ask);
1368 assert_eq!(default_tool_policy("edit_file", true), ToolPolicy::Ask);
1369 assert_eq!(default_tool_policy("bash", true), ToolPolicy::Ask);
1370 }
1371
1372 #[test]
1373 fn test_default_policy_ask_user_tools_allow_by_default() {
1374 assert_eq!(
1377 default_tool_policy("ask_user_text", true),
1378 ToolPolicy::Allow
1379 );
1380 assert_eq!(
1381 default_tool_policy("ask_user_choice", true),
1382 ToolPolicy::Allow
1383 );
1384 assert_eq!(
1385 default_tool_policy("ask_user_confirm", true),
1386 ToolPolicy::Allow
1387 );
1388 assert_eq!(
1389 default_tool_policy("edit_document", true),
1390 ToolPolicy::Allow
1391 );
1392 }
1393
1394 #[test]
1395 fn test_resolve_policy_launch_override_wins() {
1396 let mut launch = HashMap::new();
1397 launch.insert("bash".to_string(), ToolPolicy::Allow);
1398 let policy = resolve_policy(
1399 "bash",
1400 true,
1401 &launch,
1402 &HashMap::new(),
1403 &HashMap::new(),
1404 &HashMap::new(),
1405 false,
1406 );
1407 assert_eq!(policy, ToolPolicy::Allow);
1408 }
1409
1410 #[test]
1411 fn test_resolve_policy_yolo_wins() {
1412 let mut launch = HashMap::new();
1413 launch.insert("*".to_string(), ToolPolicy::Allow);
1414 let policy = resolve_policy(
1415 "bash",
1416 true,
1417 &launch,
1418 &HashMap::new(),
1419 &HashMap::new(),
1420 &HashMap::new(),
1421 false,
1422 );
1423 assert_eq!(policy, ToolPolicy::Allow);
1424 }
1425
1426 #[test]
1428 fn test_resolve_policy_stage_may_tighten_global() {
1429 let mut stage = HashMap::new();
1430 stage.insert("bash".to_string(), "deny".to_string());
1431 let mut global = HashMap::new();
1432 global.insert("bash".to_string(), ToolPolicy::Allow);
1433 let policy = resolve_policy(
1434 "bash",
1435 true,
1436 &HashMap::new(),
1437 &stage,
1438 &HashMap::new(),
1439 &global,
1440 false,
1441 );
1442 assert_eq!(policy, ToolPolicy::Deny);
1443 }
1444
1445 #[test]
1451 fn test_resolve_policy_stage_cannot_loosen_global() {
1452 let mut stage = HashMap::new();
1453 stage.insert("bash".to_string(), "allow".to_string());
1454 let mut global = HashMap::new();
1455 global.insert("bash".to_string(), ToolPolicy::Deny);
1456 let policy = resolve_policy(
1457 "bash",
1458 true,
1459 &HashMap::new(),
1460 &stage,
1461 &HashMap::new(),
1462 &global,
1463 false,
1464 );
1465 assert_eq!(policy, ToolPolicy::Deny);
1466 }
1467
1468 #[test]
1472 fn test_resolve_policy_blueprint_free_when_user_silent() {
1473 let mut agent = HashMap::new();
1474 agent.insert("web_fetch".to_string(), "allow".to_string());
1475 let policy = resolve_policy(
1476 "web_fetch",
1477 false,
1478 &HashMap::new(),
1479 &HashMap::new(),
1480 &agent,
1481 &HashMap::new(),
1482 false,
1483 );
1484 assert_eq!(policy, ToolPolicy::Allow);
1485 }
1486
1487 #[test]
1492 fn test_yolo_does_not_override_configured_deny() {
1493 let mut launch = HashMap::new();
1494 launch.insert("*".to_string(), ToolPolicy::Allow);
1495 let mut global = HashMap::new();
1496 global.insert("bash".to_string(), ToolPolicy::Deny);
1497 let policy = resolve_policy(
1498 "bash",
1499 true,
1500 &launch,
1501 &HashMap::new(),
1502 &HashMap::new(),
1503 &global,
1504 false,
1505 );
1506 assert_eq!(policy, ToolPolicy::Deny);
1507 }
1508
1509 #[test]
1511 fn test_named_allow_does_not_override_configured_deny() {
1512 let mut launch = HashMap::new();
1513 launch.insert("bash".to_string(), ToolPolicy::Allow);
1514 let mut global = HashMap::new();
1515 global.insert("bash".to_string(), ToolPolicy::Deny);
1516 let policy = resolve_policy(
1517 "bash",
1518 true,
1519 &launch,
1520 &HashMap::new(),
1521 &HashMap::new(),
1522 &global,
1523 false,
1524 );
1525 assert_eq!(policy, ToolPolicy::Deny);
1526 }
1527
1528 #[test]
1531 fn test_yolo_does_not_override_blueprint_deny() {
1532 let mut launch = HashMap::new();
1533 launch.insert("*".to_string(), ToolPolicy::Allow);
1534 let mut agent = HashMap::new();
1535 agent.insert("bash".to_string(), "deny".to_string());
1536 let policy = resolve_policy(
1537 "bash",
1538 true,
1539 &launch,
1540 &HashMap::new(),
1541 &agent,
1542 &HashMap::new(),
1543 false,
1544 );
1545 assert_eq!(policy, ToolPolicy::Deny);
1546 }
1547
1548 #[test]
1550 fn test_yolo_still_collapses_ask_to_allow() {
1551 let mut launch = HashMap::new();
1552 launch.insert("*".to_string(), ToolPolicy::Allow);
1553 let mut global = HashMap::new();
1554 global.insert("bash".to_string(), ToolPolicy::Ask);
1555 let policy = resolve_policy(
1556 "bash",
1557 true,
1558 &launch,
1559 &HashMap::new(),
1560 &HashMap::new(),
1561 &global,
1562 false,
1563 );
1564 assert_eq!(policy, ToolPolicy::Allow);
1565 }
1566
1567 #[test]
1568 fn test_resolve_policy_falls_through_to_default() {
1569 let policy = resolve_policy(
1570 "bash",
1571 true,
1572 &HashMap::new(),
1573 &HashMap::new(),
1574 &HashMap::new(),
1575 &HashMap::new(),
1576 false,
1577 );
1578 assert_eq!(policy, ToolPolicy::Ask);
1579 }
1580
1581 #[test]
1584 fn test_default_policy_unknown_tools() {
1585 assert_eq!(default_tool_policy("unknown_tool", false), ToolPolicy::Ask);
1586 assert_eq!(default_tool_policy("mcp_tool", false), ToolPolicy::Ask);
1587 assert_eq!(default_tool_policy("custom_thing", true), ToolPolicy::Ask);
1588 }
1589
1590 #[test]
1594 fn test_resolve_policy_agent_cannot_loosen_global() {
1595 let mut agent = HashMap::new();
1596 agent.insert("bash".to_string(), "allow".to_string());
1597 let mut global = HashMap::new();
1598 global.insert("bash".to_string(), ToolPolicy::Deny);
1599 let policy = resolve_policy(
1600 "bash",
1601 true,
1602 &HashMap::new(),
1603 &HashMap::new(),
1604 &agent,
1605 &global,
1606 false,
1607 );
1608 assert_eq!(policy, ToolPolicy::Deny);
1609 }
1610
1611 #[test]
1614 fn test_resolve_policy_global_ask_bounds_blueprint_allow() {
1615 let mut agent = HashMap::new();
1616 agent.insert("write_file".to_string(), "allow".to_string());
1617 let mut global = HashMap::new();
1618 global.insert("write_file".to_string(), ToolPolicy::Ask);
1619 let policy = resolve_policy(
1620 "write_file",
1621 true,
1622 &HashMap::new(),
1623 &HashMap::new(),
1624 &agent,
1625 &global,
1626 false,
1627 );
1628 assert_eq!(policy, ToolPolicy::Ask);
1629 }
1630
1631 #[test]
1632 fn test_resolve_policy_launch_override_specific_beats_wildcard() {
1633 let mut launch = HashMap::new();
1634 launch.insert("bash".to_string(), ToolPolicy::Deny);
1635 launch.insert("*".to_string(), ToolPolicy::Allow);
1636 let policy = resolve_policy(
1637 "bash",
1638 true,
1639 &launch,
1640 &HashMap::new(),
1641 &HashMap::new(),
1642 &HashMap::new(),
1643 false,
1644 );
1645 assert_eq!(policy, ToolPolicy::Deny);
1647 }
1648
1649 #[test]
1650 fn test_resolve_policy_global_overrides_default() {
1651 let mut global = HashMap::new();
1652 global.insert("read_file".to_string(), ToolPolicy::Deny);
1653 let policy = resolve_policy(
1654 "read_file",
1655 true,
1656 &HashMap::new(),
1657 &HashMap::new(),
1658 &HashMap::new(),
1659 &global,
1660 false,
1661 );
1662 assert_eq!(policy, ToolPolicy::Deny);
1663 }
1664
1665 #[test]
1666 fn test_resolve_policy_stage_deny() {
1667 let mut stage = HashMap::new();
1668 stage.insert("bash".to_string(), "deny".to_string());
1669 let policy = resolve_policy(
1670 "bash",
1671 true,
1672 &HashMap::new(),
1673 &stage,
1674 &HashMap::new(),
1675 &HashMap::new(),
1676 false,
1677 );
1678 assert_eq!(policy, ToolPolicy::Deny);
1679 }
1680
1681 #[test]
1682 fn test_resolve_policy_stage_ask() {
1683 let mut stage = HashMap::new();
1684 stage.insert("read_file".to_string(), "ask".to_string());
1685 let policy = resolve_policy(
1686 "read_file",
1687 true,
1688 &HashMap::new(),
1689 &stage,
1690 &HashMap::new(),
1691 &HashMap::new(),
1692 false,
1693 );
1694 assert_eq!(policy, ToolPolicy::Ask);
1695 }
1696
1697 #[test]
1698 fn test_resolve_policy_unknown_stage_string_defaults_to_ask() {
1699 let mut stage = HashMap::new();
1700 stage.insert("bash".to_string(), "unknown_policy".to_string());
1701 let policy = resolve_policy(
1702 "bash",
1703 true,
1704 &HashMap::new(),
1705 &stage,
1706 &HashMap::new(),
1707 &HashMap::new(),
1708 false,
1709 );
1710 assert_eq!(policy, ToolPolicy::Ask);
1711 }
1712
1713 #[test]
1716 fn test_parse_policy_str_values() {
1717 assert_eq!(parse_policy_str("allow"), ToolPolicy::Allow);
1718 assert_eq!(parse_policy_str("Allow"), ToolPolicy::Allow);
1719 assert_eq!(parse_policy_str("ALLOW"), ToolPolicy::Allow);
1720 assert_eq!(parse_policy_str("deny"), ToolPolicy::Deny);
1721 assert_eq!(parse_policy_str("Deny"), ToolPolicy::Deny);
1722 assert_eq!(parse_policy_str("ask"), ToolPolicy::Ask);
1723 assert_eq!(parse_policy_str("Ask"), ToolPolicy::Ask);
1724 assert_eq!(parse_policy_str("anything_else"), ToolPolicy::Ask);
1725 assert_eq!(parse_policy_str(""), ToolPolicy::Ask);
1726 }
1727
1728 #[tokio::test]
1731 async fn test_tool_registry_build_no_mcp() {
1732 let config = Config::default();
1733 let workdir = std::env::current_dir().unwrap();
1734 let registry = ToolRegistry::build(workdir, &config).await;
1735
1736 assert!(!registry.builtin_names.is_empty());
1738 assert!(registry.mcp_tool_defs.is_empty());
1740 }
1741
1742 #[tokio::test]
1743 async fn test_tool_registry_all_tool_defs() {
1744 let config = Config::default();
1745 let workdir = std::env::current_dir().unwrap();
1746 let registry = ToolRegistry::build(workdir, &config).await;
1747
1748 let all_defs = registry.all_tool_defs();
1749 assert!(!all_defs.is_empty());
1750
1751 let names: Vec<&str> = all_defs.iter().map(|t| t.name.as_str()).collect();
1753 assert!(names.contains(&"read_file"));
1754 }
1755
1756 #[tokio::test]
1757 async fn test_tool_registry_builtin_names_consistent() {
1758 let config = Config::default();
1759 let workdir = std::env::current_dir().unwrap();
1760 let registry = ToolRegistry::build(workdir, &config).await;
1761
1762 let names_from_builtins: HashSet<String> = registry.builtins.names().into_iter().collect();
1764 assert_eq!(registry.builtin_names, names_from_builtins);
1765 }
1766
1767 fn shell_args(command: &str) -> serde_json::Value {
1772 serde_json::json!({ "command": command })
1773 }
1774
1775 #[test]
1778 fn the_bash_alias_is_scoped_like_shell() {
1779 assert_eq!(
1780 session_approval_keys("bash", &shell_args("ls -la")),
1781 ["shell:ls"]
1782 );
1783 }
1784
1785 #[test]
1788 fn other_tools_are_keyed_by_name() {
1789 assert_eq!(
1790 session_approval_keys("read_file", &serde_json::json!({ "path": "a" })),
1791 ["read_file"]
1792 );
1793 }
1794
1795 #[test]
1798 fn a_shell_call_without_a_command_is_not_grantable() {
1799 assert!(session_approval_keys("shell", &serde_json::json!({})).is_empty());
1800 }
1801
1802 #[test]
1804 fn test_resolve_policy_launch_overrides_stage_ask() {
1805 let mut launch = HashMap::new();
1806 launch.insert("bash".to_string(), ToolPolicy::Allow);
1807 let mut stage = HashMap::new();
1808 stage.insert("bash".to_string(), "ask".to_string());
1809 let policy = resolve_policy(
1810 "bash",
1811 true,
1812 &launch,
1813 &stage,
1814 &HashMap::new(),
1815 &HashMap::new(),
1816 false,
1817 );
1818 assert_eq!(policy, ToolPolicy::Allow);
1819 }
1820
1821 #[test]
1824 fn test_resolve_policy_launch_cannot_override_stage_deny() {
1825 let mut launch = HashMap::new();
1826 launch.insert("bash".to_string(), ToolPolicy::Allow);
1827 let mut stage = HashMap::new();
1828 stage.insert("bash".to_string(), "deny".to_string());
1829 let policy = resolve_policy(
1830 "bash",
1831 true,
1832 &launch,
1833 &stage,
1834 &HashMap::new(),
1835 &HashMap::new(),
1836 false,
1837 );
1838 assert_eq!(policy, ToolPolicy::Deny);
1839 }
1840
1841 #[test]
1842 fn test_resolve_policy_stage_overrides_agent() {
1843 let mut stage = HashMap::new();
1844 stage.insert("bash".to_string(), "deny".to_string());
1845 let mut agent = HashMap::new();
1846 agent.insert("bash".to_string(), "allow".to_string());
1847 let policy = resolve_policy(
1848 "bash",
1849 true,
1850 &HashMap::new(),
1851 &stage,
1852 &agent,
1853 &HashMap::new(),
1854 false,
1855 );
1856 assert_eq!(policy, ToolPolicy::Deny);
1857 }
1858
1859 #[test]
1860 fn test_resolve_policy_agent_overrides_global() {
1861 let mut agent = HashMap::new();
1862 agent.insert("write_file".to_string(), "deny".to_string());
1863 let mut global = HashMap::new();
1864 global.insert("write_file".to_string(), ToolPolicy::Allow);
1865 let policy = resolve_policy(
1866 "write_file",
1867 true,
1868 &HashMap::new(),
1869 &HashMap::new(),
1870 &agent,
1871 &global,
1872 false,
1873 );
1874 assert_eq!(policy, ToolPolicy::Deny);
1875 }
1876
1877 #[test]
1878 fn test_resolve_policy_wildcard_launch_with_missing_specific() {
1879 let mut launch = HashMap::new();
1880 launch.insert("*".to_string(), ToolPolicy::Allow);
1881 let policy = resolve_policy(
1883 "unknown_tool",
1884 false,
1885 &launch,
1886 &HashMap::new(),
1887 &HashMap::new(),
1888 &HashMap::new(),
1889 false,
1890 );
1891 assert_eq!(policy, ToolPolicy::Allow);
1892 }
1893
1894 #[test]
1895 fn test_resolve_policy_mcp_tool_defaults_to_ask() {
1896 let policy = resolve_policy(
1897 "mcp_custom_tool",
1898 false,
1899 &HashMap::new(),
1900 &HashMap::new(),
1901 &HashMap::new(),
1902 &HashMap::new(),
1903 false,
1904 );
1905 assert_eq!(policy, ToolPolicy::Ask);
1906 }
1907
1908 #[test]
1909 fn test_resolve_policy_read_file_default_is_allow() {
1910 let policy = resolve_policy(
1911 "read_file",
1912 true,
1913 &HashMap::new(),
1914 &HashMap::new(),
1915 &HashMap::new(),
1916 &HashMap::new(),
1917 false,
1918 );
1919 assert_eq!(policy, ToolPolicy::Allow);
1920 }
1921
1922 #[test]
1923 fn test_resolve_policy_list_dir_default_is_allow() {
1924 let policy = resolve_policy(
1925 "list_dir",
1926 true,
1927 &HashMap::new(),
1928 &HashMap::new(),
1929 &HashMap::new(),
1930 &HashMap::new(),
1931 false,
1932 );
1933 assert_eq!(policy, ToolPolicy::Allow);
1934 }
1935
1936 #[test]
1937 fn test_resolve_policy_write_file_default_is_ask() {
1938 let policy = resolve_policy(
1939 "write_file",
1940 true,
1941 &HashMap::new(),
1942 &HashMap::new(),
1943 &HashMap::new(),
1944 &HashMap::new(),
1945 false,
1946 );
1947 assert_eq!(policy, ToolPolicy::Ask);
1948 }
1949
1950 #[test]
1951 fn test_resolve_policy_edit_file_default_is_ask() {
1952 let policy = resolve_policy(
1953 "edit_file",
1954 true,
1955 &HashMap::new(),
1956 &HashMap::new(),
1957 &HashMap::new(),
1958 &HashMap::new(),
1959 false,
1960 );
1961 assert_eq!(policy, ToolPolicy::Ask);
1962 }
1963
1964 #[tokio::test]
1965 async fn test_tool_registry_shutdown_no_panic() {
1966 let config = Config::default();
1967 let workdir = std::env::current_dir().unwrap();
1968 let registry = ToolRegistry::build(workdir, &config).await;
1969 registry.shutdown().await;
1970 }
1971
1972 #[tokio::test]
1973 async fn test_tool_registry_all_defs_includes_subagent() {
1974 let config = Config::default();
1975 let workdir = std::env::current_dir().unwrap();
1976 let registry = ToolRegistry::build(workdir, &config).await;
1977 let all_defs = registry.all_tool_defs();
1978 let names: Vec<&str> = all_defs.iter().map(|t| t.name.as_str()).collect();
1979 assert!(names.contains(&"spawn_agent"));
1981 }
1982
1983 #[test]
1986 fn test_default_policy_search_is_ask() {
1987 assert_eq!(default_tool_policy("search", true), ToolPolicy::Ask);
1988 }
1989
1990 #[test]
1991 fn test_default_policy_glob_is_ask() {
1992 assert_eq!(default_tool_policy("glob", true), ToolPolicy::Ask);
1993 }
1994
1995 #[test]
1996 fn test_default_policy_http_request_is_ask() {
1997 assert_eq!(default_tool_policy("http_request", true), ToolPolicy::Ask);
1998 }
1999
2000 #[test]
2001 fn test_default_policy_read_file_not_builtin_still_allow() {
2002 assert_eq!(default_tool_policy("read_file", false), ToolPolicy::Allow);
2004 }
2005
2006 #[test]
2007 fn test_default_policy_list_dir_not_builtin_still_allow() {
2008 assert_eq!(default_tool_policy("list_dir", false), ToolPolicy::Allow);
2009 }
2010
2011 #[test]
2018 fn a_permission_written_as_an_alias_reaches_the_tool() {
2019 let policy = |layer: &str, spelling: &str, called: &str| {
2020 let mut launch = HashMap::new();
2021 let mut stage = HashMap::new();
2022 let mut agent = HashMap::new();
2023 let mut global = HashMap::new();
2024 match layer {
2025 "launch" => {
2026 launch.insert(spelling.to_string(), ToolPolicy::Allow);
2027 }
2028 "stage" => {
2029 stage.insert(spelling.to_string(), "deny".to_string());
2030 }
2031 "agent" => {
2032 agent.insert(spelling.to_string(), "deny".to_string());
2033 }
2034 _ => {
2035 global.insert(spelling.to_string(), ToolPolicy::Deny);
2036 }
2037 }
2038 resolve_policy(called, true, &launch, &stage, &agent, &global, false)
2039 };
2040
2041 assert_eq!(policy("launch", "bash", "shell"), ToolPolicy::Allow);
2043 assert_eq!(policy("stage", "bash", "shell"), ToolPolicy::Deny);
2044 assert_eq!(policy("agent", "bash", "shell"), ToolPolicy::Deny);
2045 assert_eq!(policy("global", "bash", "shell"), ToolPolicy::Deny);
2046
2047 assert_eq!(policy("launch", "shell", "bash"), ToolPolicy::Allow);
2049 assert_eq!(policy("global", "shell", "bash"), ToolPolicy::Deny);
2050
2051 assert_eq!(policy("global", "read_file", "write_file"), ToolPolicy::Ask);
2053 }
2054
2055 #[test]
2058 fn the_default_shell_policy_covers_both_spellings() {
2059 assert_eq!(default_tool_policy("shell", true), ToolPolicy::Ask);
2060 assert_eq!(default_tool_policy("bash", true), ToolPolicy::Ask);
2061 }
2062
2063 #[test]
2068 fn the_context_tools_do_not_prompt() {
2069 for tool in [
2070 "context_write",
2071 "context_append",
2072 "context_read",
2073 "context_delete",
2074 "context_list",
2075 "read_files",
2076 ] {
2077 assert_eq!(
2078 default_tool_policy(tool, true),
2079 ToolPolicy::Allow,
2080 "{tool} must not raise a prompt by default"
2081 );
2082 }
2083 }
2084
2085 #[test]
2088 fn test_resolve_policy_agent_deny() {
2089 let mut agent = HashMap::new();
2090 agent.insert("read_file".to_string(), "deny".to_string());
2091 let policy = resolve_policy(
2092 "read_file",
2093 true,
2094 &HashMap::new(),
2095 &HashMap::new(),
2096 &agent,
2097 &HashMap::new(),
2098 false,
2099 );
2100 assert_eq!(policy, ToolPolicy::Deny);
2101 }
2102
2103 #[test]
2106 fn test_resolve_policy_agent_unknown_string_defaults_to_ask() {
2107 let mut agent = HashMap::new();
2108 agent.insert("bash".to_string(), "foobar".to_string());
2109 let policy = resolve_policy(
2110 "bash",
2111 true,
2112 &HashMap::new(),
2113 &HashMap::new(),
2114 &agent,
2115 &HashMap::new(),
2116 false,
2117 );
2118 assert_eq!(policy, ToolPolicy::Ask);
2119 }
2120
2121 #[test]
2124 fn test_resolve_policy_global_allow_overrides_default_ask() {
2125 let mut global = HashMap::new();
2126 global.insert("bash".to_string(), ToolPolicy::Allow);
2127 let policy = resolve_policy(
2128 "bash",
2129 true,
2130 &HashMap::new(),
2131 &HashMap::new(),
2132 &HashMap::new(),
2133 &global,
2134 false,
2135 );
2136 assert_eq!(policy, ToolPolicy::Allow);
2137 }
2138
2139 #[tokio::test]
2140 async fn test_tool_registry_all_defs_includes_all_subagent_tools() {
2141 let config = Config::default();
2142 let workdir = std::env::current_dir().unwrap();
2143 let registry = ToolRegistry::build(workdir, &config).await;
2144 let all_defs = registry.all_tool_defs();
2145 let names: Vec<&str> = all_defs.iter().map(|t| t.name.as_str()).collect();
2146
2147 for expected in &[
2148 "spawn_agent",
2149 "check_agent",
2150 "wait_for_agent",
2151 "send_to_agent",
2152 "kill_agent",
2153 ] {
2154 assert!(names.contains(expected));
2155 }
2156 }
2157
2158 #[tokio::test]
2161 async fn test_tool_registry_builtin_names_has_expected_tools() {
2162 let config = Config::default();
2163 let workdir = std::env::current_dir().unwrap();
2164 let registry = ToolRegistry::build(workdir, &config).await;
2165
2166 for name in &["read_file", "list_dir"] {
2168 assert!(registry.builtin_names.contains(*name));
2169 }
2170
2171 assert!(!registry.builtin_names.contains("spawn_agent"));
2173 }
2174
2175 #[tokio::test]
2178 async fn test_tool_registry_all_defs_no_mcp_when_none_configured() {
2179 let config = Config::default();
2180 let workdir = std::env::current_dir().unwrap();
2181 let registry = ToolRegistry::build(workdir, &config).await;
2182 assert!(registry.mcp_tool_defs.is_empty());
2183
2184 let all_defs = registry.all_tool_defs();
2186 let builtin_count = registry.builtins.tool_defs().len();
2187 let subagent_count = leviath_tools::BuiltinTools::subagent_tool_defs().len();
2188 assert_eq!(all_defs.len(), builtin_count + subagent_count);
2189 }
2190
2191 #[test]
2197 fn test_resolve_policy_full_chain_deny_is_terminal() {
2198 let mut launch = HashMap::new();
2199 launch.insert("bash".to_string(), ToolPolicy::Allow);
2200 let mut stage = HashMap::new();
2201 stage.insert("bash".to_string(), "deny".to_string());
2202 let mut agent = HashMap::new();
2203 agent.insert("bash".to_string(), "deny".to_string());
2204 let mut global = HashMap::new();
2205 global.insert("bash".to_string(), ToolPolicy::Deny);
2206
2207 let policy = resolve_policy("bash", true, &launch, &stage, &agent, &global, false);
2208 assert_eq!(policy, ToolPolicy::Deny);
2209 }
2210
2211 #[test]
2214 fn test_resolve_policy_full_chain_launch_relaxes_ask() {
2215 let mut launch = HashMap::new();
2216 launch.insert("bash".to_string(), ToolPolicy::Allow);
2217 let mut stage = HashMap::new();
2218 stage.insert("bash".to_string(), "ask".to_string());
2219 let mut agent = HashMap::new();
2220 agent.insert("bash".to_string(), "ask".to_string());
2221 let mut global = HashMap::new();
2222 global.insert("bash".to_string(), ToolPolicy::Ask);
2223
2224 let policy = resolve_policy("bash", true, &launch, &stage, &agent, &global, false);
2225 assert_eq!(policy, ToolPolicy::Allow);
2226 }
2227
2228 #[tokio::test]
2232 async fn test_tool_registry_build_with_failing_mcp_server() {
2233 use leviath_mcp::MCPServerConfig;
2234
2235 let bad_server = MCPServerConfig::stdio(
2236 "bad-server",
2237 "/nonexistent/binary/that/does/not/exist",
2238 vec![],
2239 );
2240 let config = Config {
2241 mcp_servers: vec![bad_server],
2242 ..Config::default()
2243 };
2244
2245 let workdir = std::env::current_dir().unwrap();
2246 let registry = ToolRegistry::build(workdir, &config).await;
2248
2249 assert!(registry.mcp_tool_defs.is_empty());
2251 assert!(!registry.builtin_names.is_empty());
2253 }
2254
2255 const ALLOWED_WITHOUT_ASKING: &[&str] = &[
2285 "read_file",
2286 "read_files",
2287 "list_dir",
2288 "context_write",
2289 "context_append",
2290 "context_read",
2291 "context_delete",
2292 "context_list",
2293 "todo_add",
2299 "todo_done",
2300 "todo_note",
2301 "spawn_agent",
2302 "check_agent",
2303 "wait_for_agent",
2304 "send_to_agent",
2305 "kill_agent",
2306 "ask_user_text",
2307 "ask_user_choice",
2308 "ask_user_confirm",
2309 "edit_document",
2310 ];
2311
2312 fn advertised_tool_names() -> Vec<String> {
2319 let dir = tempfile::tempdir().expect("tempdir");
2320 let builtins = leviath_tools::BuiltinTools::new(leviath_tools::ToolContext::new(
2321 dir.path().to_path_buf(),
2322 ));
2323 let mut names: Vec<String> = builtins
2324 .tool_defs()
2325 .into_iter()
2326 .map(|t| t.name)
2327 .chain(
2328 leviath_tools::BuiltinTools::subagent_tool_defs()
2329 .into_iter()
2330 .map(|t| t.name),
2331 )
2332 .collect();
2333 names.sort();
2334 names.dedup();
2335 names
2336 }
2337
2338 #[test]
2341 fn the_tool_discovery_finds_the_real_toolset() {
2342 let names = advertised_tool_names();
2343 let count = names.len();
2346 assert!(count >= 15, "only {count} tools discovered: {names:?}");
2347 for expected in ["read_file", "write_file", "shell"] {
2348 assert!(names.contains(&expected.to_string()), "{expected} missing");
2349 }
2350 }
2351
2352 #[test]
2358 fn only_reviewed_tools_run_without_asking() {
2359 for name in advertised_tool_names() {
2360 let policy = default_tool_policy(&name, true);
2361 let reviewed = ALLOWED_WITHOUT_ASKING.contains(&name.as_str());
2362 match reviewed {
2363 true => assert_eq!(
2364 policy,
2365 ToolPolicy::Allow,
2366 "{name} is on the reviewed unprompted list but does not resolve to Allow"
2367 ),
2368 false => assert_ne!(
2369 policy,
2370 ToolPolicy::Allow,
2371 "{name} runs with no prompt and is not on the reviewed list"
2372 ),
2373 }
2374 }
2375 }
2376
2377 #[test]
2383 fn an_unknown_tool_asks_rather_than_running() {
2384 for name in [
2385 "definitely_not_a_real_tool",
2386 "mcp__someserver__delete_everything",
2387 "",
2388 ] {
2389 assert_eq!(default_tool_policy(name, false), ToolPolicy::Ask, "{name}");
2390 assert_eq!(default_tool_policy(name, true), ToolPolicy::Ask, "{name}");
2391 }
2392 }
2393
2394 #[test]
2402 fn the_effectful_tools_ask_under_every_spelling() {
2403 for name in ["write_file", "edit_file", "shell"] {
2404 for spelling in leviath_tools::tool_name_spellings(name) {
2405 assert_eq!(
2406 default_tool_policy(spelling, true),
2407 ToolPolicy::Ask,
2408 "{spelling:?} (a spelling of {name}) does not ask"
2409 );
2410 }
2411 }
2412 }
2413}