from __future__ import annotations import copy import hashlib import json import tempfile import unittest from pathlib import Path from .policy import ( CompositionError, PROVIDER_CREDENTIAL_DEFINITIONS, PROVIDER_PROFILE_NAMES, RAW_CREDENTIAL_VARIABLES, assert_single_lm_studio_registration, compose_profile, isolated_settings, load_profile, ) ROOT = Path(__file__).resolve().parents[4] PROFILES = ROOT / "profiles" / "nono" LM_STUDIO = ROOT / "extensions" / "lm-studio" / "index.ts" PI_INSTALL = "/home/jb/.bun/install/global/node_modules/@earendil-works/pi-coding-agent" JITI_INSTALL = "/home/jb/.bun/install/global/node_modules/jiti" GLOBAL_INSTRUCTIONS = "/home/jb/.pi/agent/AGENTS.md" class PolicyContractTests(unittest.TestCase): def test_cores_are_fail_closed_and_separate_state(self) -> None: for context in ("personal", "work"): with self.subTest(context=context): profile = load_profile(PROFILES / f"pi-core-{context}.json") self.assertEqual(profile["extends"], "default") self.assertEqual(profile["workdir"], {"access": "none"}) self.assertEqual(profile["network"], {"block": True}) self.assertEqual( set(profile["environment"]["deny_vars"]), RAW_CREDENTIAL_VARIABLES, ) self.assertNotIn("SSH_AUTH_SOCK", profile["environment"]["allow_vars"]) self.assertIn("system_write_linux", profile["groups"]["exclude"]) self.assertIn("user_tools", profile["groups"]["exclude"]) grants = json.dumps(profile["filesystem"], sort_keys=True) self.assertIn(f"pi-sbx/{context}", grants) other = "work" if context == "personal" else "personal" self.assertNotIn(f"pi-sbx/{other}", grants) self.assertNotIn("$WORKDIR", grants) self.assertNotIn('"/tmp"', grants) self.assertNotIn("$TMPDIR", grants) self.assertEqual( set(profile["filesystem"]["allow"]), { f"$XDG_CACHE_HOME/pi-sbx/{context}", f"$XDG_STATE_HOME/pi-sbx/{context}/agent", f"$XDG_STATE_HOME/pi-sbx/{context}/sessions", }, ) self.assertEqual( profile["filesystem"]["read"], [PI_INSTALL, JITI_INSTALL], ) self.assertEqual( set(profile["filesystem"]["read_file"]), {GLOBAL_INSTRUCTIONS, "/usr/bin/node"}, ) self.assertEqual( set(profile["filesystem"]["allow_file"]), {"/dev/null", "/dev/stderr", "/dev/stdout", "/dev/tty"}, ) for capability in ("write", "write_file", "unix_socket"): self.assertEqual(profile["filesystem"].get(capability, []), []) def test_provider_profiles_select_one_route(self) -> None: for provider, profile_name in PROVIDER_PROFILE_NAMES.items(): with self.subTest(provider=provider): profile = load_profile(PROFILES / f"{profile_name}.json") network = profile["network"] self.assertFalse(network["block"]) if provider == "lm-studio": self.assertEqual(set(network), {"block", "connect_port"}) self.assertEqual(network["connect_port"], [1234]) self.assertNotIn("credentials", network) else: self.assertEqual( set(network), {"block", "allow_domain", "credentials", "custom_credentials"}, ) self.assertEqual(network["allow_domain"], []) self.assertEqual(network["credentials"], [provider]) self.assertEqual( network["custom_credentials"], {provider: PROVIDER_CREDENTIAL_DEFINITIONS[provider]}, ) self.assertNotIn("connect_port", network) self.assertEqual(profile["workdir"], {"access": "none"}) self.assertIn("system_read_linux_core", profile["groups"]["exclude"]) self.assertIn("system_read_macos", profile["groups"]["exclude"]) self.assertTrue(profile["environment"]["allow_vars"]) self.assertEqual( set(profile["environment"]["deny_vars"]), RAW_CREDENTIAL_VARIABLES, ) self.assertEqual(profile["open_urls"]["allow_origins"], []) self.assertFalse(profile["allow_launch_services"]) def test_trusted_composition_replaces_block_with_one_provider(self) -> None: core = load_profile(PROFILES / "pi-core-personal.json") provider = load_profile(PROFILES / "provider-lm-studio.json") with tempfile.TemporaryDirectory() as directory: temporary_directory = Path(directory) generated = compose_profile( core, provider, workdir_access="read", temporary_directory=temporary_directory, lm_studio_extension=LM_STUDIO, ) self.assertEqual(generated["workdir"], {"access": "read"}) self.assertIn(str(temporary_directory), generated["filesystem"]["allow"]) self.assertEqual( generated["environment"]["set_vars"]["TMPDIR"], str(temporary_directory), ) self.assertEqual(generated["network"]["connect_port"], [1234]) self.assertFalse(generated["network"]["block"]) self.assertNotIn("command_policies", generated) self.assertIn(str(LM_STUDIO), generated["filesystem"]["read"]) with self.assertRaisesRegex(CompositionError, "cannot load"): compose_profile( core, load_profile(PROFILES / "provider-openai.json"), workdir_access="read", temporary_directory=temporary_directory, lm_studio_extension=LM_STUDIO, ) with self.assertRaisesRegex(CompositionError, "private owned"): compose_profile( core, provider, workdir_access="read", temporary_directory=Path("/tmp"), lm_studio_extension=LM_STUDIO, ) mutated_core = copy.deepcopy(core) mutated_core["filesystem"]["allow"].append("/tmp") with tempfile.TemporaryDirectory() as directory: with self.assertRaisesRegex(CompositionError, "differs from reviewed"): compose_profile( mutated_core, provider, workdir_access="read", temporary_directory=Path(directory), lm_studio_extension=LM_STUDIO, ) def test_generated_settings_have_no_privileged_resources(self) -> None: normal = isolated_settings("openai") self.assertEqual(normal["packages"], []) self.assertEqual(normal["extensions"], []) self.assertEqual(normal["defaultTools"], ["read", "edit", "write"]) self.assertNotIn("bash", normal["defaultTools"]) self.assertEqual(normal["defaultProjectTrust"], "never") local = isolated_settings("lm-studio", lm_studio_extension=LM_STUDIO) self.assertEqual(local["extensions"], [str(LM_STUDIO)]) with tempfile.TemporaryDirectory() as directory: unreviewed = Path(directory) / "lm-studio.ts" unreviewed.write_text("unreviewed\n", encoding="utf-8") with self.assertRaisesRegex(CompositionError, "reviewed lock source"): isolated_settings("lm-studio", lm_studio_extension=unreviewed) serialized = json.dumps(local).lower() for forbidden in ("pi-workflow", "update", "background", "notify"): self.assertNotIn(forbidden, serialized) def test_lm_studio_registration_is_unique_and_uses_exact_ip(self) -> None: lock = json.loads((ROOT / "config" / "pi-package-lock.json").read_text(encoding="utf-8")) entry = next(item for item in lock["extensions"] if item["name"] == "lm-studio") self.assertEqual(entry["source"], "extensions/lm-studio/index.ts") self.assertEqual(entry["status"], "blocked") self.assertEqual( entry["reviewed"]["sha256"], hashlib.sha256(LM_STUDIO.read_bytes()).hexdigest(), ) assert_single_lm_studio_registration([LM_STUDIO]) contents = LM_STUDIO.read_text(encoding="utf-8") self.assertIn('"http://127.0.0.1:1234/v1"', contents) self.assertNotIn("http://localhost:1234", contents) with tempfile.TemporaryDirectory() as directory: duplicate = Path(directory) / "duplicate.ts" duplicate.write_text( 'pi.registerProvider("lm-studio", {});\n', encoding="utf-8", ) with self.assertRaisesRegex(CompositionError, "exactly one"): assert_single_lm_studio_registration([LM_STUDIO, duplicate]) duplicate.write_text( 'pi.registerProvider("lm-studio", {});\n' 'pi.registerProvider("lm-studio", {});\n', encoding="utf-8", ) with self.assertRaisesRegex(CompositionError, "found 2"): assert_single_lm_studio_registration([duplicate]) duplicate.write_text( 'const baseUrl = "http://localhost:1234/v1";\n' 'pi.registerProvider("lm-studio", {});\n', encoding="utf-8", ) with self.assertRaisesRegex(CompositionError, "legacy localhost"): assert_single_lm_studio_registration([duplicate]) if __name__ == "__main__": unittest.main()