diff --git a/cycode/cli/apps/ai_guardrails/ides/__init__.py b/cycode/cli/apps/ai_guardrails/ides/__init__.py index bc48448a..05a0b171 100644 --- a/cycode/cli/apps/ai_guardrails/ides/__init__.py +++ b/cycode/cli/apps/ai_guardrails/ides/__init__.py @@ -35,23 +35,22 @@ def get_ide(name: str) -> IDE: return ide -def collect_all_session_contexts() -> tuple[dict[str, dict], dict]: +def collect_all_session_contexts() -> tuple[list[dict], dict]: """Sweep every registered IDE's session context, regardless of which IDE triggered the hook. - Returns ``(config_files_by_ide, plugins)``: the global MCP config file of each IDE that has - one (keyed by IDE name), and the enabled plugins merged across IDEs (first registered IDE - wins on a duplicate plugin key - plugins are IDE-agnostic marketplace artifacts). + Returns ``(config_files, plugins)``: the global MCP config files of every IDE, and the + enabled plugins merged across IDEs (first registered IDE wins on a duplicate plugin key - + plugins are IDE-agnostic marketplace artifacts). """ - config_files_by_ide: dict[str, dict] = {} + config_files: list[dict] = [] plugins: dict = {} for ide in IDES.values(): - global_config_file, enabled_plugins = ide.get_session_context() - if global_config_file: - config_files_by_ide[ide.name] = global_config_file + global_config_files, enabled_plugins = ide.get_session_context() + config_files.extend(global_config_files) for plugin_key, plugin in (enabled_plugins or {}).items(): plugins.setdefault(plugin_key, plugin) - return config_files_by_ide, plugins + return config_files, plugins def collect_all_skills() -> list[dict]: diff --git a/cycode/cli/apps/ai_guardrails/ides/_plugin_utils.py b/cycode/cli/apps/ai_guardrails/ides/_plugin_utils.py index 6f7d8917..d10e7e56 100644 --- a/cycode/cli/apps/ai_guardrails/ides/_plugin_utils.py +++ b/cycode/cli/apps/ai_guardrails/ides/_plugin_utils.py @@ -42,17 +42,17 @@ def load_plugin_json(path: Path) -> Optional[dict]: return None -def build_global_config_file(path: Path, mcp_servers: Optional[dict]) -> Optional[dict]: +def build_global_config_files(path: Path, mcp_servers: Optional[dict]) -> list[dict]: """Wrap a global (non-plugin) MCP config into the session-context file shape. - Returns ``{"path": , "content": <{"mcpServers": ...} JSON>}`` when - there are servers, else ``None``. ``content`` is normalized to the canonical + Returns ``[{"path": , "content": <{"mcpServers": ...} JSON>}]`` when + there are servers, else ``[]``. ``content`` is normalized to the canonical ``{"mcpServers": {...}}`` shape, dropping everything else in the source file. """ - servers = mcp_servers or {} + servers = mcp_servers if isinstance(mcp_servers, dict) else {} if not servers: - return None - return {'path': str(path), 'content': json.dumps({'mcpServers': servers})} + return [] + return [{'path': str(path), 'content': json.dumps({'mcpServers': servers})}] def walk_enabled_plugins( diff --git a/cycode/cli/apps/ai_guardrails/ides/base.py b/cycode/cli/apps/ai_guardrails/ides/base.py index 331ffe41..4bf0c1ac 100644 --- a/cycode/cli/apps/ai_guardrails/ides/base.py +++ b/cycode/cli/apps/ai_guardrails/ides/base.py @@ -176,19 +176,19 @@ def get_user_email(self) -> Optional[str]: """ return None - def get_session_context(self) -> tuple[Optional[dict], dict]: - """Return ``(global_config_file, enabled_plugins)`` for session-context reporting. + def get_session_context(self) -> tuple[list[dict], dict]: + """Return ``(global_config_files, enabled_plugins)`` for session-context reporting. - ``global_config_file`` is the IDE's global (non-plugin) MCP config as - ``{"path": , "content": }``, - or ``None`` when there is no global MCP config. ``enabled_plugins`` maps each - enabled plugin key to its metadata (including its own ``mcp_config_file`` - content and ``mcp_config_file_path``). + ``global_config_files`` are the IDE's global (non-plugin) MCP configs, each as + ``{"path": , "content": }``; + empty when there is none. ``enabled_plugins`` maps each enabled plugin key to + its metadata (including its own ``mcp_config_file`` content and + ``mcp_config_file_path``). - Default: ``(None, {})`` (no plugin system, no discoverable MCP config). + Default: ``([], {})`` (no plugin system, no discoverable MCP config). Override to surface MCP/plugin inventory. """ - return None, {} + return [], {} def get_skills(self) -> list[dict]: """Return the IDE's user-scope skills as ``[{"path", "content"}]``. @@ -198,7 +198,7 @@ def get_skills(self) -> list[dict]: read these files off endpoints can only ever return raw content. Kept separate from ``get_session_context`` rather than folded into its - ``global_config_file`` slot, which is normalized to an MCP server map. + ``global_config_files``, which are normalized to an MCP server map. Default: ``[]`` (the IDE has no skill system). Override to surface skills. """ diff --git a/cycode/cli/apps/ai_guardrails/ides/claude_code.py b/cycode/cli/apps/ai_guardrails/ides/claude_code.py index 74fed9f4..b95dd593 100644 --- a/cycode/cli/apps/ai_guardrails/ides/claude_code.py +++ b/cycode/cli/apps/ai_guardrails/ides/claude_code.py @@ -1,6 +1,7 @@ """Claude Code IDE integration for AI guardrails.""" import json +import re from collections.abc import Iterator from copy import deepcopy from pathlib import Path @@ -8,7 +9,7 @@ from cycode.cli.apps.ai_guardrails.consts import CYCODE_SCAN_PROMPT_COMMAND, CYCODE_SESSION_START_COMMAND from cycode.cli.apps.ai_guardrails.ides._plugin_utils import ( - build_global_config_file, + build_global_config_files, load_plugin_json, resolve_cached_plugin_dir, walk_enabled_plugins, @@ -253,6 +254,41 @@ def _locate(plugin_name: str, marketplace_name: str) -> Optional[Path]: ) +def _sanitize_mcp_name(name: str) -> str: + """Mirrors how Claude Code sanitizes a server name inside a tool name.""" + return re.sub(r'[^a-z0-9_-]', '_', name.lower()) + + +def _local_mcp_servers(config: Optional[dict], cwd: Optional[str]) -> Iterator[tuple[str, str]]: + """Yield ``(name as Claude Code namespaces it, config key)`` for each locally configured MCP server.""" + sources = [] + if config: + sources.append(get_mcp_servers(config)) + if cwd: + sources.append(((config.get('projects') or {}).get(cwd) or {}).get('mcpServers')) + if cwd: + sources.append((load_plugin_json(Path(cwd) / '.mcp.json') or {}).get('mcpServers')) + for servers in sources: + if isinstance(servers, dict): + for name in servers: + yield name, name + + settings = load_claude_settings() + for plugin_key, entry in (resolve_plugins(settings) if settings else {}).items(): + plugin_name = plugin_key.split('@', 1)[0] + for server in entry.get('mcp_server_names') or []: + yield f'plugin_{plugin_name}_{server}', server + + +def _resolve_mcp_server_name(server: str, config: Optional[dict], cwd: Optional[str]) -> str: + """Map the server segment of ``mcp____`` back to its config key; unknown names pass through.""" + wanted = _sanitize_mcp_name(server) + for namespaced, name in _local_mcp_servers(config, cwd): + if _sanitize_mcp_name(namespaced) == wanted: + return name + return server + + # --- IDE integration ---------------------------------------------------------- @@ -339,6 +375,8 @@ def parse_hook_payload(self, raw_payload: dict) -> AIHookPayload: config = load_claude_config() ide_user_email = _email_from_config(config) if config else None + if mcp_server_name: + mcp_server_name = _resolve_mcp_server_name(mcp_server_name, config, raw_payload.get('cwd')) return AIHookPayload( event_name=canonical_event, @@ -396,14 +434,14 @@ def get_user_email(self) -> Optional[str]: config = load_claude_config() return _email_from_config(config) if config else None - def get_session_context(self) -> tuple[Optional[dict], dict]: + def get_session_context(self) -> tuple[list[dict], dict]: config = load_claude_config() - global_config_file = build_global_config_file(_CLAUDE_CONFIG_PATH, get_mcp_servers(config)) if config else None + global_config_files = build_global_config_files(_CLAUDE_CONFIG_PATH, get_mcp_servers(config)) if config else [] settings = load_claude_settings() enriched_plugins = resolve_plugins(settings) if settings else {} - return global_config_file, enriched_plugins + return global_config_files, enriched_plugins def get_skills(self) -> list[dict]: return walk_skill_dirs(_claude_skills_dir()) diff --git a/cycode/cli/apps/ai_guardrails/ides/codex.py b/cycode/cli/apps/ai_guardrails/ides/codex.py index df0024b9..b3918d57 100644 --- a/cycode/cli/apps/ai_guardrails/ides/codex.py +++ b/cycode/cli/apps/ai_guardrails/ides/codex.py @@ -2,7 +2,9 @@ import json import os +import re import sys +from collections.abc import Iterator from pathlib import Path from typing import ClassVar, Optional @@ -15,7 +17,7 @@ from cycode.cli.apps.ai_guardrails.consts import CYCODE_SCAN_PROMPT_COMMAND, CYCODE_SESSION_START_COMMAND from cycode.cli.apps.ai_guardrails.ides._plugin_utils import ( - build_global_config_file, + build_global_config_files, load_plugin_json, resolve_cached_plugin_dir, walk_enabled_plugins, @@ -158,6 +160,42 @@ def _resolve_codex_plugins(config: dict) -> dict: ) +def _sanitize_mcp_name(name: str) -> str: + """Mirrors how Codex sanitizes a server name into a tool namespace.""" + return re.sub(r'[^A-Za-z0-9_]', '_', name) or '_' + + +def _local_mcp_server_names() -> Iterator[str]: + """Yield the config key of each locally configured Codex MCP server, plugin servers included.""" + config = _load_codex_config() + if not config: + return + servers = config.get('mcp_servers') + if isinstance(servers, dict): + yield from servers + for entry in _resolve_codex_plugins(config).values(): + yield from entry.get('mcp_server_names') or [] + + +def _split_mcp_tool_name(tool_name: str) -> tuple[Optional[str], Optional[str]]: + """Split ``mcp____`` into ``(server config key, tool)``; unknown namespaces split on ``__``.""" + rest = tool_name[len('mcp__') :] + names = list(_local_mcp_server_names()) + # Codex appends `_<12 hex>` to a namespace on collision or length overflow. + for suffix in ('__', r'_+[0-9a-f]{12}__'): + best: Optional[tuple[int, str]] = None + for name in names: + namespace = _sanitize_mcp_name(name).rstrip('_') + match = namespace and re.match(re.escape(namespace) + suffix, rest) + if match and (best is None or match.end() > best[0]): + best = (match.end(), name) + if best: + return best[1], rest[best[0] :] + + parts = tool_name.split('__') + return (parts[1] if len(parts) >= 2 else None), (parts[2] if len(parts) >= 3 else None) + + def _enable_codex_hooks_feature(scope: str, repo_path: Optional[Path] = None) -> tuple[bool, str]: """Set ``[features] hooks = true`` in Codex's ``config.toml``. @@ -247,12 +285,8 @@ def parse_hook_payload(self, raw_payload: dict) -> AIHookPayload: mcp_server_name = None mcp_tool_name = None mcp_arguments = None - if tool_name.startswith('mcp__'): - parts = tool_name.split('__') - if len(parts) >= 2: - mcp_server_name = parts[1] - if len(parts) >= 3: - mcp_tool_name = parts[2] + if canonical_event == AiHookEventType.MCP_EXECUTION: + mcp_server_name, mcp_tool_name = _split_mcp_tool_name(tool_name) mcp_arguments = tool_input return AIHookPayload( @@ -302,17 +336,17 @@ def build_session_payload(self, raw_payload: dict) -> AIHookPayload: def get_user_email(self) -> Optional[str]: return _email_from_auth() - def get_session_context(self) -> tuple[Optional[dict], dict]: + def get_session_context(self) -> tuple[list[dict], dict]: config = _load_codex_config() if not config: - return None, {} + return [], {} # Codex stores MCP servers under `[mcp_servers.]`; the global config # file becomes its own session-context file. Plugins (via # `[plugins."@"]`) carry their own config files. config_path = _codex_config_toml_path('user') - global_config_file = build_global_config_file(config_path, config.get('mcp_servers')) + global_config_files = build_global_config_files(config_path, config.get('mcp_servers')) enriched_plugins = _resolve_codex_plugins(config) - return global_config_file, enriched_plugins + return global_config_files, enriched_plugins def get_skills(self) -> list[dict]: return walk_skill_dirs(_codex_skills_dir()) diff --git a/cycode/cli/apps/ai_guardrails/ides/copilot.py b/cycode/cli/apps/ai_guardrails/ides/copilot.py index 16c8178b..77118ed5 100644 --- a/cycode/cli/apps/ai_guardrails/ides/copilot.py +++ b/cycode/cli/apps/ai_guardrails/ides/copilot.py @@ -33,7 +33,7 @@ from cycode.cli.apps.ai_guardrails.consts import CYCODE_SCAN_PROMPT_COMMAND, CYCODE_SESSION_START_COMMAND from cycode.cli.apps.ai_guardrails.ides._plugin_utils import ( - build_global_config_file, + build_global_config_files, load_plugin_json, walk_enabled_plugins, ) @@ -283,6 +283,15 @@ def _collect_installed_plugins() -> dict: # --- MCP tool-name splitting ------------------------------------------------------ +def _global_mcp_server_maps() -> list[tuple[Path, object]]: + """``(path, servers)`` of each runtime's global MCP config: VS Code's ``servers``, the agent's ``mcpServers``.""" + agent_config_path = _copilot_home() / _AGENT_MCP_CONFIG_FILENAME + return [ + (_vscode_mcp_config_path(), (_load_vscode_mcp_config() or {}).get('servers')), + (agent_config_path, (_load_jsonc(agent_config_path) or {}).get('mcpServers')), + ] + + def _known_mcp_server_names() -> list[str]: """Config-declared MCP server names, across both runtimes' config files. @@ -293,15 +302,7 @@ def _known_mcp_server_names() -> list[str]: Best-effort inventory: servers contributed by extensions, ``chat.mcp.discovery`` imports, dev containers, or non-default profiles are not discoverable from disk. """ - config = _load_vscode_mcp_config() - servers = (config or {}).get('servers') - names = list(servers.keys()) if isinstance(servers, dict) else [] - - agent_config = _load_jsonc(_copilot_home() / _AGENT_MCP_CONFIG_FILENAME) or {} - agent_servers = agent_config.get('mcpServers') - if isinstance(agent_servers, dict): - names.extend(agent_servers.keys()) - + names = [name for _, servers in _global_mcp_server_maps() if isinstance(servers, dict) for name in servers] for plugin in _collect_installed_plugins().values(): names.extend(plugin.get('mcp_server_names') or []) return names @@ -485,14 +486,11 @@ def build_session_payload(self, raw_payload: dict) -> AIHookPayload: source=raw_payload.get('source'), ) - def get_session_context(self) -> tuple[Optional[dict], dict]: - # VS Code's mcp.json uses `servers` as its top-level key; normalized to the - # canonical mcpServers shape by build_global_config_file. - config = _load_vscode_mcp_config() - global_config_file = ( - build_global_config_file(_vscode_mcp_config_path(), config.get('servers')) if config else None - ) - return global_config_file, _collect_installed_plugins() + def get_session_context(self) -> tuple[list[dict], dict]: + global_config_files = [ + file for path, servers in _global_mcp_server_maps() for file in build_global_config_files(path, servers) + ] + return global_config_files, _collect_installed_plugins() def get_skills(self) -> list[dict]: return walk_skill_dirs(_copilot_skills_dir()) diff --git a/cycode/cli/apps/ai_guardrails/ides/cursor.py b/cycode/cli/apps/ai_guardrails/ides/cursor.py index 1f41a483..14ac0d0b 100644 --- a/cycode/cli/apps/ai_guardrails/ides/cursor.py +++ b/cycode/cli/apps/ai_guardrails/ides/cursor.py @@ -6,7 +6,7 @@ from typing import ClassVar, Optional from cycode.cli.apps.ai_guardrails.consts import CYCODE_SCAN_PROMPT_COMMAND, CYCODE_SESSION_START_COMMAND -from cycode.cli.apps.ai_guardrails.ides._plugin_utils import build_global_config_file +from cycode.cli.apps.ai_guardrails.ides._plugin_utils import build_global_config_files from cycode.cli.apps.ai_guardrails.ides._skill_utils import walk_skill_dirs from cycode.cli.apps.ai_guardrails.ides.base import IDE, DecisionAction, HookDecision from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload @@ -64,6 +64,44 @@ def _load_cursor_mcp_config(config_path: Optional[Path] = None) -> Optional[dict return None +def _cursor_mcp_config_paths(workspace_roots: object) -> list[Path]: + roots = workspace_roots if isinstance(workspace_roots, list) else [] + paths = [Path(root) / _REPO_SUBDIR / _MCP_CONFIG_FILENAME for root in roots if isinstance(root, str) and root] + paths.append(_cursor_mcp_config_path()) + return paths + + +def _server_command_line(server: dict) -> Optional[str]: + command = server.get('command') + if not isinstance(command, str) or not command: + return None + args = server.get('args') + return ' '.join([command, *(str(arg) for arg in args)]) if isinstance(args, list) else command + + +def _resolve_mcp_server_name(raw_payload: dict) -> Optional[str]: + """Cursor sends the server's url or command, but the platform knows it by its mcp.json entry name.""" + url = raw_payload.get('url') + command = raw_payload.get('command') + if not url and not command: + return None + + lowered_command = command.lower() if command else None + for config_path in _cursor_mcp_config_paths(raw_payload.get('workspace_roots')): + servers = (_load_cursor_mcp_config(config_path) or {}).get('mcpServers') + if not isinstance(servers, dict): + continue + for name, server in servers.items(): + if not isinstance(server, dict): + continue + if url and server.get('url') == url: + return name + if command and (command == _server_command_line(server) or lowered_command == name.lower()): + return name + + return command + + class Cursor(IDE): name: ClassVar[str] = 'cursor' display_name: ClassVar[str] = 'Cursor' @@ -95,7 +133,9 @@ def parse_hook_payload(self, raw_payload: dict) -> AIHookPayload: ide_version=raw_payload.get('cursor_version'), prompt=raw_payload.get('prompt', ''), file_path=raw_payload.get('file_path') or raw_payload.get('path'), - mcp_server_name=raw_payload.get('command'), + mcp_server_name=( + _resolve_mcp_server_name(raw_payload) if canonical_event == AiHookEventType.MCP_EXECUTION else None + ), mcp_tool_name=raw_payload.get('tool_name') or raw_payload.get('tool'), mcp_arguments=(raw_payload.get('arguments') or raw_payload.get('tool_input') or raw_payload.get('input')), ) @@ -124,13 +164,11 @@ def build_session_payload(self, raw_payload: dict) -> AIHookPayload: ide_version=raw_payload.get('cursor_version'), ) - def get_session_context(self) -> tuple[Optional[dict], dict]: + def get_session_context(self) -> tuple[list[dict], dict]: config = _load_cursor_mcp_config() if not config: - return None, {} - config_path = _cursor_mcp_config_path() - global_config_file = build_global_config_file(config_path, config.get('mcpServers')) - return global_config_file, {} + return [], {} + return build_global_config_files(_cursor_mcp_config_path(), config.get('mcpServers')), {} def get_skills(self) -> list[dict]: return walk_skill_dirs(_cursor_skills_dir()) diff --git a/cycode/cli/apps/ai_guardrails/scan/guardrail_config.py b/cycode/cli/apps/ai_guardrails/scan/guardrail_config.py index 416262b8..a2e95d61 100644 --- a/cycode/cli/apps/ai_guardrails/scan/guardrail_config.py +++ b/cycode/cli/apps/ai_guardrails/scan/guardrail_config.py @@ -3,7 +3,7 @@ session-start fetches the tenant's resolved guardrail config from the platform and writes it here; scans only read. Per-agent modes and sensitive-path globs are platform-owned - local policy files never carry them. An absent or corrupt cache means built-in defaults (Report -everywhere + the default globs), always synchronous. +everywhere + the default globs, with the unauthorized MCP server guardrail Off), always synchronous. """ import json @@ -23,7 +23,7 @@ GUARDRAILS_CONFIG_FILE_NAME = 'ai-guardrails-config.json' -_DEFAULT_TTL_SECONDS = 900 +DEFAULT_TTL_SECONDS = 900 # Guardrail keys are the CLI's block-reason vocabulary. Anything else in the payload (a future # guardrail this CLI doesn't implement) is ignored - unknown config must never fail closed. @@ -34,14 +34,24 @@ BlockReason.SECRETS_IN_FILE, BlockReason.SENSITIVE_PATH, BlockReason.SECRETS_IN_MCP_ARGS, + BlockReason.UNAUTHORIZED_MCP_SERVER, ) ) +# Off by default (not Report), so they never start enforcing on a tenant that hasn't configured them. +_DEFAULT_OFF_GUARDRAIL_KEYS = frozenset((BlockReason.UNAUTHORIZED_MCP_SERVER.value,)) + def get_config_cache_path() -> Path: return Path.home() / CYCODE_CONFIGURATION_DIRECTORY / GUARDRAILS_CONFIG_FILE_NAME +def default_mode_for(guardrail_key: str) -> str: + if guardrail_key in _DEFAULT_OFF_GUARDRAIL_KEYS: + return GuardrailCellMode.OFF.value + return GuardrailCellMode.REPORT.value + + def _default_sensitive_globs() -> list: return list(DEFAULT_SENSITIVE_PATH_GLOBS) @@ -60,10 +70,20 @@ def __post_init__(self) -> None: if guardrail.get('key') in _KNOWN_GUARDRAIL_KEYS } + @property + def ttl_seconds(self) -> float: + return self.payload.get('ttl_seconds') or DEFAULT_TTL_SECONDS + + def _agents(self, guardrail_key: str) -> dict: + return (self._guardrails.get(guardrail_key) or {}).get('agents') or {} + def mode_for(self, guardrail_key: str, ide_name: Optional[str]) -> str: """The platform keys the cells by our --ide names, so the lookup is direct.""" - agents = (self._guardrails.get(guardrail_key) or {}).get('agents') or {} - return str(agents.get((ide_name or '').lower(), GuardrailCellMode.REPORT.value)).lower() + agents = self._agents(guardrail_key) + return str(agents.get((ide_name or '').lower(), default_mode_for(guardrail_key))).lower() + + def is_off_for_every_agent(self, guardrail_key: str) -> bool: + return all(str(mode).lower() == GuardrailCellMode.OFF for mode in self._agents(guardrail_key).values()) def _modes_for_event(self, event_name: str, ide_name: Optional[str]) -> list: return [ @@ -87,8 +107,7 @@ def sensitive_globs(self) -> list: return globs if isinstance(globs, list) and globs else _default_sensitive_globs() def is_expired(self) -> bool: - ttl = self.payload.get('ttl_seconds') or _DEFAULT_TTL_SECONDS - return time.time() - self.fetched_at > ttl + return time.time() - self.fetched_at > self.ttl_seconds def needs_refresh(self, tenant_id: Optional[str]) -> bool: """Expired, or fetched for another tenant (the user switched tenants since).""" @@ -99,14 +118,14 @@ def apply_platform_config(policy: dict, config: Optional[GuardrailConfig], ide_n """Overlay the platform-owned enforcement config onto the local knobs-only policy. The platform is the only mode source: no cache (cold start) means the built-in defaults - - Report everywhere with the default globs - which equal an unconfigured tenant's platform - config, so behaviour is uniform either way. Each matrix cell lands on its own per-feature - action, so the two FileRead guardrails (content scan vs. sensitive path) keep independent modes. + Report everywhere (unauthorized MCP server Off) with the default globs - which equal an unconfigured + tenant's platform config, so behaviour is uniform either way. Each matrix cell lands on its own + per-feature action, so guardrails sharing an event (e.g. the two FileRead ones) keep independent modes. An all-Off event never reaches here at all: scan_command skips it. """ def cell(guardrail_key: str) -> str: - return config.mode_for(guardrail_key, ide_name) if config is not None else GuardrailCellMode.REPORT.value + return config.mode_for(guardrail_key, ide_name) if config is not None else default_mode_for(guardrail_key) def action(guardrail_key: str) -> str: return PolicyMode.BLOCK.value if cell(guardrail_key) == GuardrailCellMode.BLOCK else PolicyMode.WARN.value @@ -123,7 +142,11 @@ def action(guardrail_key: str) -> str: ) file_read['path_action'] = action(BlockReason.SENSITIVE_PATH) - policy.setdefault('mcp', {})['action'] = action(BlockReason.SECRETS_IN_MCP_ARGS) + mcp = policy.setdefault('mcp', {}) + mcp['scan_args'] = cell(BlockReason.SECRETS_IN_MCP_ARGS) != GuardrailCellMode.OFF + mcp['action'] = action(BlockReason.SECRETS_IN_MCP_ARGS) + mcp['check_server'] = cell(BlockReason.UNAUTHORIZED_MCP_SERVER) != GuardrailCellMode.OFF + mcp['server_action'] = action(BlockReason.UNAUTHORIZED_MCP_SERVER) def save_guardrail_config(payload: dict, tenant_id: Optional[str]) -> None: diff --git a/cycode/cli/apps/ai_guardrails/scan/handlers.py b/cycode/cli/apps/ai_guardrails/scan/handlers.py index c4a37b35..e977291a 100644 --- a/cycode/cli/apps/ai_guardrails/scan/handlers.py +++ b/cycode/cli/apps/ai_guardrails/scan/handlers.py @@ -21,6 +21,7 @@ from cycode.cli.apps.ai_guardrails.consts import GuardrailsMode, PolicyMode from cycode.cli.apps.ai_guardrails.ides.base import HookDecision +from cycode.cli.apps.ai_guardrails.scan.mcp_server_status import is_enforced, load_mcp_server_statuses from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload from cycode.cli.apps.ai_guardrails.scan.policy import get_policy_value from cycode.cli.apps.ai_guardrails.scan.types import ( @@ -217,6 +218,16 @@ class _ArgScanFeature: deny_agent_message: str ask_message: Callable[[str], str] ask_agent_message: str + scan_enabled: bool = True + + +class _PreScanFinding(NamedTuple): + """A finding on the call itself (e.g. an unauthorized MCP server), made before the text scan.""" + + block_reason: BlockReason + mode: GuardrailsMode + deny_message: str + deny_agent_message: str def _handle_arg_scan( @@ -225,8 +236,12 @@ def _handle_arg_scan( policy: dict, feature: _ArgScanFeature, scan_text: str, + pre_scan_finding: Optional[_PreScanFinding] = None, ) -> HookDecision: - """Shared scan + decision flow for MCP_EXECUTION and COMMAND_EXEC events.""" + """Shared scan + decision flow for MCP_EXECUTION and COMMAND_EXEC events. + + A Report-mode pre-scan finding still scans, and a secret found then takes over the response. + """ ai_client = ctx.obj['ai_security_client'] max_bytes = get_policy_value(policy, 'secrets', 'max_bytes', default=200000) @@ -239,6 +254,18 @@ def _handle_arg_scan( error_message = None try: + if pre_scan_finding is not None: + block_reason = pre_scan_finding.block_reason + if pre_scan_finding.mode == GuardrailsMode.BLOCK: + outcome = AIHookOutcome.BLOCKED + return HookDecision.deny( + feature.event_type, pre_scan_finding.deny_message, pre_scan_finding.deny_agent_message + ) + outcome = AIHookOutcome.WARNED + + if not feature.scan_enabled: + return HookDecision.allow(feature.event_type) + scan_outcome = _scan_text_for_secrets( ctx, clipped, @@ -282,8 +309,43 @@ def _handle_arg_scan( ) +def _check_mcp_server_authorization(payload: AIHookPayload, policy: dict) -> Optional[_PreScanFinding]: + """Fails open: no server name or no cached statuses lets the call through.""" + mcp_config = get_policy_value(policy, 'mcp', default={}) + if not get_policy_value(mcp_config, 'check_server', default=False): + return None + + server_name = payload.mcp_server_name + if not server_name: + logger.debug('No MCP server name in the payload; skipping the server authorization check') + return None + + statuses = load_mcp_server_statuses() + if statuses is None: + logger.debug('No cached MCP server statuses; skipping the server authorization check') + return None + + status = statuses.status_of(server_name) + if not is_enforced(status): + return None + + logger.debug('MCP server is not authorized, %s', {'mcp_server_name': server_name, 'status': status}) + return _PreScanFinding( + block_reason=BlockReason.UNAUTHORIZED_MCP_SERVER, + mode=get_effective_mode(mcp_config, action_key='server_action'), + deny_message=( + f"Cycode blocked MCP server '{server_name}': it is not authorized in your organization. " + 'Contact your admin to authorize it.' + ), + deny_agent_message=( + f"The MCP server '{server_name}' is not authorized in this organization. " + 'Do not retry its tools or reach it another way.' + ), + ) + + def handle_before_mcp_execution(ctx: typer.Context, payload: AIHookPayload, policy: dict) -> HookDecision: - """Scan MCP tool arguments for secrets before execution.""" + """Check the MCP server is authorized, then scan the tool arguments for secrets.""" tool = payload.mcp_tool_name or 'unknown' args = payload.mcp_arguments or {} args_text = args if isinstance(args, str) else json.dumps(args) @@ -297,8 +359,10 @@ def handle_before_mcp_execution(ctx: typer.Context, payload: AIHookPayload, poli deny_agent_message='Do not pass secrets to tools. Use secret references (name/id) instead.', ask_message=lambda v: f'Allow MCP tool call "{tool}"? {v}', ask_agent_message='Possible secrets detected in tool arguments; proceed with caution.', + scan_enabled=get_policy_value(policy, 'mcp', 'scan_args', default=True), ), scan_text=args_text, + pre_scan_finding=_check_mcp_server_authorization(payload, policy), ) diff --git a/cycode/cli/apps/ai_guardrails/scan/mcp_server_status.py b/cycode/cli/apps/ai_guardrails/scan/mcp_server_status.py new file mode 100644 index 00000000..f6a19c1a --- /dev/null +++ b/cycode/cli/apps/ai_guardrails/scan/mcp_server_status.py @@ -0,0 +1,103 @@ +"""MCP server authorization status cache: session-start writes it, the pre-MCP-execution hook reads it. + +An absent or corrupt cache means no status is known, so the guardrail fails open. +""" + +import json +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Optional + +from marshmallow import EXCLUDE, Schema, fields, post_load + +from cycode.cli.apps.ai_guardrails.scan.guardrail_config import DEFAULT_TTL_SECONDS +from cycode.cli.consts import CYCODE_CONFIGURATION_DIRECTORY +from cycode.cli.utils.path_utils import atomic_write_text, quarantine_corrupt_file +from cycode.cyclient.models import McpServerAuthorizationStatus, McpServerStatus, McpServerStatusSchema +from cycode.logger import get_logger + +logger = get_logger('AI Guardrails') + +MCP_SERVER_STATUSES_FILE_NAME = 'ai-guardrails-mcp-servers.json' + + +# An alias shared by several servers (e.g. configured differently per device) gets the most restrictive status. +_RESTRICTIVENESS = { + McpServerAuthorizationStatus.AUTHORIZED: 0, + McpServerAuthorizationStatus.UNREVIEWED: 1, + McpServerAuthorizationStatus.UNAUTHORIZED: 2, +} + + +def get_mcp_server_statuses_cache_path() -> Path: + return Path.home() / CYCODE_CONFIGURATION_DIRECTORY / MCP_SERVER_STATUSES_FILE_NAME + + +def is_enforced(status: Optional[McpServerAuthorizationStatus]) -> bool: + """``status`` is None when the platform never saw the server.""" + return status == McpServerAuthorizationStatus.UNAUTHORIZED + + +@dataclass +class McpServerStatuses: + servers: list[McpServerStatus] + fetched_at: float + tenant_id: Optional[str] = None + ttl_seconds: float = DEFAULT_TTL_SECONDS + _by_alias: dict = field(init=False, repr=False) + + def __post_init__(self) -> None: + # The platform matches aliases case-insensitively, so the CLI does too. + self._by_alias = {} + for server in self.servers: + if not server.alias: + continue + key = server.alias.lower() + known = self._by_alias.get(key) + if known is None or _RESTRICTIVENESS[server.status] > _RESTRICTIVENESS[known]: + self._by_alias[key] = server.status + + def status_of(self, alias: str) -> Optional[McpServerAuthorizationStatus]: + return self._by_alias.get(alias.lower()) + + +class McpServerStatusesSchema(Schema): + class Meta: + unknown = EXCLUDE + + servers = fields.List(fields.Nested(McpServerStatusSchema), required=True) + fetched_at = fields.Float(required=True) + tenant_id = fields.String(allow_none=True, load_default=None) + ttl_seconds = fields.Float(allow_none=True, load_default=None) + + @post_load + def build_dto(self, data: dict[str, Any], **_) -> McpServerStatuses: + data['ttl_seconds'] = data['ttl_seconds'] or DEFAULT_TTL_SECONDS + return McpServerStatuses(**data) + + +def save_mcp_server_statuses( + servers: list[McpServerStatus], tenant_id: Optional[str], ttl_seconds: float = DEFAULT_TTL_SECONDS +) -> None: + path = get_mcp_server_statuses_cache_path() + statuses = McpServerStatuses(servers=servers, fetched_at=time.time(), tenant_id=tenant_id, ttl_seconds=ttl_seconds) + try: + path.parent.mkdir(parents=True, exist_ok=True) + atomic_write_text(str(path), json.dumps(McpServerStatusesSchema().dump(statuses))) + except Exception as e: + logger.debug('Failed to save MCP server statuses cache', exc_info=e) + + +def load_mcp_server_statuses() -> Optional[McpServerStatuses]: + path = get_mcp_server_statuses_cache_path() + if not path.exists(): + return None + + try: + with open(path, encoding='UTF-8') as file: + return McpServerStatusesSchema().load(json.load(file)) + except Exception as e: + logger.warning('MCP server statuses cache is corrupt and will be moved aside', exc_info=e) + quarantine_corrupt_file(str(path)) + return None diff --git a/cycode/cli/apps/ai_guardrails/scan/types.py b/cycode/cli/apps/ai_guardrails/scan/types.py index 5d18e07d..3ffb6280 100644 --- a/cycode/cli/apps/ai_guardrails/scan/types.py +++ b/cycode/cli/apps/ai_guardrails/scan/types.py @@ -40,6 +40,7 @@ class BlockReason(StrEnum): SECRETS_IN_FILE = 'secrets_in_file' SECRETS_IN_MCP_ARGS = 'secrets_in_mcp_args' SENSITIVE_PATH = 'sensitive_path' + UNAUTHORIZED_MCP_SERVER = 'unauthorized_mcp_server' SCAN_FAILURE = 'scan_failure' diff --git a/cycode/cli/apps/ai_guardrails/session_start_command.py b/cycode/cli/apps/ai_guardrails/session_start_command.py index 282fab78..a56161b9 100644 --- a/cycode/cli/apps/ai_guardrails/session_start_command.py +++ b/cycode/cli/apps/ai_guardrails/session_start_command.py @@ -15,7 +15,14 @@ collect_all_skills, get_ide, ) -from cycode.cli.apps.ai_guardrails.scan.guardrail_config import load_guardrail_config, save_guardrail_config +from cycode.cli.apps.ai_guardrails.scan.guardrail_config import ( + DEFAULT_TTL_SECONDS, + GuardrailConfig, + load_guardrail_config, + save_guardrail_config, +) +from cycode.cli.apps.ai_guardrails.scan.mcp_server_status import save_mcp_server_statuses +from cycode.cli.apps.ai_guardrails.scan.types import BlockReason from cycode.cli.apps.ai_guardrails.scan.utils import read_stdin_text, safe_json_parse from cycode.cli.apps.auth.auth_common import get_authorization_info from cycode.cli.apps.auth.auth_manager import AuthManager @@ -49,14 +56,21 @@ def _session_context_digest(report: dict) -> str: return hashlib.sha256(canonical.encode('utf-8')).hexdigest() -def _should_skip_report(digest: str, tenant_id: Optional[str]) -> bool: +def _report_ttl_seconds(config: Optional[GuardrailConfig]) -> float: + """The response carries the MCP server statuses, so the guardrail's TTL applies while it is enabled.""" + if config is None or config.is_off_for_every_agent(BlockReason.UNAUTHORIZED_MCP_SERVER): + return _SESSION_CONTEXT_TTL_SECONDS + return config.ttl_seconds + + +def _should_skip_report(digest: str, tenant_id: Optional[str], ttl_seconds: float) -> bool: """Skip when the same payload was already sent for this tenant and the TTL hasn't expired.""" try: cache = json.loads(_session_context_cache_path().read_text(encoding='utf-8')) return ( cache.get('hash') == digest and cache.get('tenant_id') == tenant_id - and time.time() - float(cache.get('sent_at', 0)) < _SESSION_CONTEXT_TTL_SECONDS + and time.time() - float(cache.get('sent_at', 0)) < ttl_seconds ) except Exception: # Missing/corrupt cache reads as a miss - over-sending is harmless @@ -78,6 +92,7 @@ def _report_session_context( ai_client: 'AISecurityManagerClient', user_email: Optional[str], tenant_id: Optional[str], + config: Optional[GuardrailConfig], ) -> None: """Report the device + cross-IDE session context to the AI security manager. Never raises. @@ -86,7 +101,7 @@ def _report_session_context( until the TTL expires. """ try: - config_files_by_ide, enabled_plugins = collect_all_session_contexts() + config_files, enabled_plugins = collect_all_session_contexts() report = { 'hostname': get_hostname(), 'platform_name': get_platform_name(), @@ -94,7 +109,7 @@ def _report_session_context( 'serial_number': get_serial_number(), 'last_login_user': get_last_login_user(), # Sorted by path so the digest is stable regardless of IDE registry order. - 'config_files': sorted(config_files_by_ide.values(), key=lambda f: f['path']), + 'config_files': sorted(config_files, key=lambda f: f['path']), 'enabled_plugins': enabled_plugins, # Already deduplicated and sorted by path, for the same digest-stability reason. # Editing a skill body changes the digest and so re-reports the device's inventory. @@ -103,12 +118,19 @@ def _report_session_context( } digest = _session_context_digest(report) - if _should_skip_report(digest, tenant_id): + if _should_skip_report(digest, tenant_id, _report_ttl_seconds(config)): logger.debug('Session context unchanged; skipping report') return - if ai_client.report_session_context(**report): - _save_report_cache(digest, tenant_id) + response = ai_client.report_session_context(**report) + if response is None: + return + + _save_report_cache(digest, tenant_id) + if response.mcp_servers is not None: + ttl_seconds = config.ttl_seconds if config is not None else DEFAULT_TTL_SECONDS + save_mcp_server_statuses(response.mcp_servers, tenant_id, ttl_seconds) + logger.debug('MCP server statuses cache updated') except Exception as e: logger.debug('Failed to report session context', exc_info=e) @@ -166,14 +188,14 @@ def session_start_command( except Exception as e: logger.debug('Failed to create conversation during session start', exc_info=e) - # Report session context (device + cross-IDE MCP servers and plugins) - _report_session_context(ai_client, session_payload.ide_user_email, auth_info.tenant_id) - # SessionStart precedes the first prompt hook in every IDE, so scans normally find a cache. - _sync_guardrail_config(ai_client, auth_info.tenant_id) + config = _sync_guardrail_config(ai_client, auth_info.tenant_id) + + # Report session context (device + cross-IDE MCP servers and plugins) + _report_session_context(ai_client, session_payload.ide_user_email, auth_info.tenant_id, config) -def _sync_guardrail_config(ai_client: 'AISecurityManagerClient', tenant_id: Optional[str]) -> None: +def _sync_guardrail_config(ai_client: 'AISecurityManagerClient', tenant_id: Optional[str]) -> Optional[GuardrailConfig]: """Refresh the guardrail config cache when it is expired or belongs to another tenant. Every step here swallows its own failures - a broken cache or an unreachable platform @@ -182,9 +204,12 @@ def _sync_guardrail_config(ai_client: 'AISecurityManagerClient', tenant_id: Opti cached = load_guardrail_config() if cached is not None and not cached.needs_refresh(tenant_id): logger.debug('Guardrail config cache is fresh, skipping fetch') - return + return cached resolved = ai_client.get_resolved_guardrails() - if resolved: - save_guardrail_config(resolved, tenant_id) - logger.debug('Guardrail config cache updated') + if not resolved: + return cached + + save_guardrail_config(resolved, tenant_id) + logger.debug('Guardrail config cache updated') + return GuardrailConfig(payload=resolved, fetched_at=time.time(), tenant_id=tenant_id) diff --git a/cycode/cyclient/ai_security_manager_client.py b/cycode/cyclient/ai_security_manager_client.py index c3a9fb95..ef536c91 100644 --- a/cycode/cyclient/ai_security_manager_client.py +++ b/cycode/cyclient/ai_security_manager_client.py @@ -5,6 +5,7 @@ from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError from cycode.cyclient.cycode_client_base import CycodeClientBase from cycode.cyclient.logger import logger +from cycode.cyclient.models import SessionContextResponse, SessionContextResponseSchema if TYPE_CHECKING: from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload @@ -114,8 +115,8 @@ def report_session_context( enabled_plugins: Optional[dict] = None, skill_files: Optional[list[dict]] = None, user_email: Optional[str] = None, - ) -> bool: - """Report session context to the backend. Returns whether the report was accepted.""" + ) -> Optional[SessionContextResponse]: + """Report session context to the backend. Returns None when the report was not accepted.""" body: dict = { 'hostname': hostname, 'platform_name': platform_name, @@ -129,9 +130,14 @@ def report_session_context( } try: - self.client.post(self._build_endpoint_path(self._SESSION_CONTEXT_PATH), body=body) - return True + response = self.client.post(self._build_endpoint_path(self._SESSION_CONTEXT_PATH), body=body) except Exception as e: logger.debug('Failed to report session context', exc_info=e) # Don't fail the session if reporting fails - return False + return None + + try: + return SessionContextResponseSchema().load(response.json()) + except Exception as e: + logger.debug('Failed to parse the session context response', exc_info=e) + return SessionContextResponse() diff --git a/cycode/cyclient/models.py b/cycode/cyclient/models.py index b4bf4f7a..1d20c4ab 100644 --- a/cycode/cyclient/models.py +++ b/cycode/cyclient/models.py @@ -1,7 +1,8 @@ from dataclasses import dataclass +from enum import Enum from typing import Any, Optional -from marshmallow import EXCLUDE, Schema, fields, post_load +from marshmallow import EXCLUDE, Schema, fields, missing, post_load class Detection(Schema): @@ -578,3 +579,60 @@ class Meta: @post_load def build_dto(self, data: dict[str, Any], **_) -> 'ScanConfiguration': return ScanConfiguration(**data) + + +class McpServerAuthorizationStatus(str, Enum): + AUTHORIZED = 'Authorized' + UNREVIEWED = 'Unreviewed' + UNAUTHORIZED = 'Unauthorized' + + @classmethod + def parse(cls, raw_status: object) -> 'McpServerAuthorizationStatus': + """An unknown status reads as Unreviewed: no decision was made on the server.""" + lowered = str(raw_status or '').lower() + return next((status for status in cls if status.value.lower() == lowered), cls.UNREVIEWED) + + +class McpServerAuthorizationStatusField(fields.Field): + def _serialize(self, value: Optional[McpServerAuthorizationStatus], *_, **__) -> Optional[str]: + return value.value if value is not None else None + + def deserialize(self, value: object, *_, **__) -> McpServerAuthorizationStatus: + return McpServerAuthorizationStatus.parse(None if value is missing else value) + + +@dataclass +class McpServerStatus: + alias: Optional[str] = None + normalized_id: Optional[str] = None + status: McpServerAuthorizationStatus = McpServerAuthorizationStatus.UNREVIEWED + + +class McpServerStatusSchema(Schema): + class Meta: + unknown = EXCLUDE + + alias = fields.String(allow_none=True, load_default=None) + normalized_id = fields.String(allow_none=True, load_default=None) + status = McpServerAuthorizationStatusField() + + @post_load + def build_dto(self, data: dict[str, Any], **_) -> McpServerStatus: + return McpServerStatus(**data) + + +@dataclass +class SessionContextResponse: + # None when the response carries no (or malformed) MCP server statuses. + mcp_servers: Optional[list[McpServerStatus]] = None + + +class SessionContextResponseSchema(Schema): + class Meta: + unknown = EXCLUDE + + mcp_servers = fields.List(fields.Nested(McpServerStatusSchema), allow_none=True, load_default=None) + + @post_load + def build_dto(self, data: dict[str, Any], **_) -> SessionContextResponse: + return SessionContextResponse(**data) diff --git a/tests/cli/commands/ai_guardrails/ides/test_claude_code.py b/tests/cli/commands/ai_guardrails/ides/test_claude_code.py index 595452c0..12652160 100644 --- a/tests/cli/commands/ai_guardrails/ides/test_claude_code.py +++ b/tests/cli/commands/ai_guardrails/ides/test_claude_code.py @@ -12,6 +12,7 @@ ClaudeCode, _email_from_config, _read_claude_plugin, + _resolve_mcp_server_name, load_claude_config, resolve_plugins, ) @@ -111,13 +112,17 @@ def test_parse_file_read_payload() -> None: def test_parse_mcp_execution_payload() -> None: args = {'resource_type': 'merge_request', 'parent_id': 'org/repo', 'resource_id': '4'} - unified = ClaudeCode().parse_hook_payload( - { - 'hook_event_name': 'PreToolUse', - 'tool_name': 'mcp__gitlab__discussion_list', - 'tool_input': args, - } - ) + with ( + patch('cycode.cli.apps.ai_guardrails.ides.claude_code.load_claude_config', return_value=None), + patch('cycode.cli.apps.ai_guardrails.ides.claude_code.load_claude_settings', return_value=None), + ): + unified = ClaudeCode().parse_hook_payload( + { + 'hook_event_name': 'PreToolUse', + 'tool_name': 'mcp__gitlab__discussion_list', + 'tool_input': args, + } + ) assert unified.event_name == AiHookEventType.MCP_EXECUTION assert unified.mcp_server_name == 'gitlab' @@ -125,6 +130,59 @@ def test_parse_mcp_execution_payload() -> None: assert unified.mcp_arguments == args +def _create_claude_plugin(fs: FakeFilesystem, plugin: str, servers: list[str]) -> dict: + plugin_dir = Path.home() / '.claude' / 'plugins' / 'cache' / 'dummy-marketplace' / plugin / '1.0.0' + fs.create_file( + plugin_dir / '.mcp.json', + contents=json.dumps({'mcpServers': {server: {'command': 'dummy-command'} for server in servers}}), + ) + return {'enabledPlugins': {f'{plugin}@dummy-marketplace': True}} + + +def test_resolve_mcp_server_name_from_user_and_project_config(fs: FakeFilesystem) -> None: + fs.create_file('/repo/.mcp.json', contents=json.dumps({'mcpServers': {'repo.server': {}}})) + config = {'mcpServers': {'my.server': {}}, 'projects': {'/repo': {'mcpServers': {'Project Server': {}}}}} + + assert _resolve_mcp_server_name('my_server', config, None) == 'my.server' + assert _resolve_mcp_server_name('Project_Server', config, '/repo') == 'Project Server' + assert _resolve_mcp_server_name('repo_server', config, '/repo') == 'repo.server' + assert _resolve_mcp_server_name('Project_Server', config, None) == 'Project_Server' + + +def test_resolve_mcp_server_name_from_plugin(fs: FakeFilesystem) -> None: + settings = _create_claude_plugin(fs, 'cycode-dev-dev', ['sentry']) + config = {'mcpServers': {'dev_sentry': {}}} + + with patch('cycode.cli.apps.ai_guardrails.ides.claude_code.load_claude_settings', return_value=settings): + assert _resolve_mcp_server_name('plugin_cycode-dev-dev_sentry', config, None) == 'sentry' + assert _resolve_mcp_server_name('dev_sentry', config, None) == 'dev_sentry' + assert _resolve_mcp_server_name('plugin_cycode-dev-dev_unknown', config, None) == ( + 'plugin_cycode-dev-dev_unknown' + ) + + +def test_parse_mcp_execution_payload_resolves_plugin_server(fs: FakeFilesystem) -> None: + settings = _create_claude_plugin(fs, 'cycode-dev', ['sentry']) + + with ( + patch('cycode.cli.apps.ai_guardrails.ides.claude_code.load_claude_config', return_value=None), + patch('cycode.cli.apps.ai_guardrails.ides.claude_code.load_claude_settings', return_value=settings), + ): + unified = ClaudeCode().parse_hook_payload( + {'hook_event_name': 'PreToolUse', 'tool_name': 'mcp__plugin_cycode-dev_sentry__search', 'tool_input': {}} + ) + + assert unified.mcp_server_name == 'sentry' + assert unified.mcp_tool_name == 'search' + + +def test_parse_non_mcp_payload_skips_server_resolution() -> None: + with patch('cycode.cli.apps.ai_guardrails.ides.claude_code.load_claude_settings') as mock_settings: + ClaudeCode().parse_hook_payload({'hook_event_name': 'PreToolUse', 'tool_name': 'Read', 'tool_input': {}}) + + mock_settings.assert_not_called() + + def test_parse_empty_payload_defaults() -> None: unified = ClaudeCode().parse_hook_payload({'hook_event_name': 'UserPromptSubmit'}) assert unified.event_name == AiHookEventType.PROMPT @@ -348,8 +406,8 @@ def test_session_context_no_config() -> None: patch('cycode.cli.apps.ai_guardrails.ides.claude_code.load_claude_config', return_value=None), patch('cycode.cli.apps.ai_guardrails.ides.claude_code.load_claude_settings', return_value=None), ): - global_config_file, plugins = ClaudeCode().get_session_context() - assert global_config_file is None + global_config_files, plugins = ClaudeCode().get_session_context() + assert global_config_files == [] assert plugins == {} diff --git a/tests/cli/commands/ai_guardrails/ides/test_codex.py b/tests/cli/commands/ai_guardrails/ides/test_codex.py index 7a05b68e..655c75ee 100644 --- a/tests/cli/commands/ai_guardrails/ides/test_codex.py +++ b/tests/cli/commands/ai_guardrails/ides/test_codex.py @@ -70,6 +70,55 @@ def test_parse_mcp_execution_payload() -> None: assert unified.mcp_arguments == args +_CODEX_HOME = '/dummy/codex-home' + + +def _parse_mcp(tool_name: str) -> tuple: + unified = Codex().parse_hook_payload({'hook_event_name': 'PreToolUse', 'tool_name': tool_name}) + return unified.mcp_server_name, unified.mcp_tool_name + + +@pytest.fixture +def codex_servers(fs: FakeFilesystem, monkeypatch: pytest.MonkeyPatch) -> FakeFilesystem: + monkeypatch.setenv('CODEX_HOME', _CODEX_HOME) + fs.create_file( + f'{_CODEX_HOME}/config.toml', + contents=( + '[mcp_servers.my-server]\ncommand = "a"\n' + '[mcp_servers."docs.example.com"]\ncommand = "b"\n' + '[mcp_servers.a__b]\ncommand = "c"\n' + '[mcp_servers.a]\ncommand = "d"\n' + '[plugins."demo@market"]\nenabled = true\n' + ), + ) + plugin_dir = Path(f'{_CODEX_HOME}/plugins/cache/market/demo/abc123') + fs.create_file(plugin_dir / '.codex-plugin' / 'plugin.json', contents=json.dumps({'mcpServers': '.mcp.json'})) + fs.create_file(plugin_dir / '.mcp.json', contents=json.dumps({'mcpServers': {'plugin-srv': {'command': 'e'}}})) + return fs + + +@pytest.mark.parametrize( + ('tool_name', 'expected'), + [ + ('mcp__my_server__list', ('my-server', 'list')), + ('mcp__docs_example_com__search', ('docs.example.com', 'search')), + ('mcp__plugin_srv__run', ('plugin-srv', 'run')), + ('mcp__my_server_0123456789ab__list', ('my-server', 'list')), + ('mcp__a__b__tool', ('a__b', 'tool')), + ('mcp__a__tool', ('a', 'tool')), + ('mcp__unknown__tool', ('unknown', 'tool')), + ], +) +def test_parse_mcp_resolves_codex_server_name(codex_servers: FakeFilesystem, tool_name: str, expected: tuple) -> None: + assert _parse_mcp(tool_name) == expected + + +def test_parse_non_mcp_event_reads_no_codex_config() -> None: + with patch('cycode.cli.apps.ai_guardrails.ides.codex._load_codex_config') as load_config: + Codex().parse_hook_payload({'hook_event_name': 'UserPromptSubmit', 'prompt': 'hi'}) + load_config.assert_not_called() + + def test_parse_unknown_event_falls_through() -> None: unified = Codex().parse_hook_payload({'hook_event_name': 'Stop'}) assert unified.event_name == 'Stop' @@ -320,17 +369,17 @@ def test_session_context_reads_mcp_servers() -> None: 'cycode.cli.apps.ai_guardrails.ides.codex._load_codex_config', return_value={'mcp_servers': mcp}, ): - global_config_file, plugins = Codex().get_session_context() - assert global_config_file is not None - assert global_config_file['path'].endswith('config.toml') - assert global_config_file['content'] == json.dumps({'mcpServers': mcp}) + global_config_files, plugins = Codex().get_session_context() + assert len(global_config_files) == 1 + assert global_config_files[0]['path'].endswith('config.toml') + assert global_config_files[0]['content'] == json.dumps({'mcpServers': mcp}) assert plugins == {} def test_session_context_no_config() -> None: with patch('cycode.cli.apps.ai_guardrails.ides.codex._load_codex_config', return_value=None): - global_config_file, plugins = Codex().get_session_context() - assert global_config_file is None + global_config_files, plugins = Codex().get_session_context() + assert global_config_files == [] assert plugins == {} diff --git a/tests/cli/commands/ai_guardrails/ides/test_contract.py b/tests/cli/commands/ai_guardrails/ides/test_contract.py index 05596e76..44ba8d73 100644 --- a/tests/cli/commands/ai_guardrails/ides/test_contract.py +++ b/tests/cli/commands/ai_guardrails/ides/test_contract.py @@ -125,12 +125,13 @@ def test_build_session_payload_tags_ide(ide: IDE) -> None: def test_get_session_context_returns_pair(ide: IDE) -> None: - """Session context must be a ``(global_config_file, plugins)`` pair. + """Session context must be a ``(global_config_files, plugins)`` pair. - ``global_config_file`` is ``None`` or a ``{"path", "content"}`` dict; ``plugins`` is a dict. + ``global_config_files`` is a list of ``{"path", "content"}`` dicts; ``plugins`` is a dict. """ - global_config_file, plugins = ide.get_session_context() - assert global_config_file is None or isinstance(global_config_file, dict) + global_config_files, plugins = ide.get_session_context() + assert isinstance(global_config_files, list) + assert all(isinstance(f, dict) for f in global_config_files) assert isinstance(plugins, dict) diff --git a/tests/cli/commands/ai_guardrails/ides/test_copilot.py b/tests/cli/commands/ai_guardrails/ides/test_copilot.py index a0f84fca..be2053bb 100644 --- a/tests/cli/commands/ai_guardrails/ides/test_copilot.py +++ b/tests/cli/commands/ai_guardrails/ides/test_copilot.py @@ -362,19 +362,33 @@ def test_get_session_context_normalizes_servers_key(fs: FakeFilesystem) -> None: contents=json.dumps({'servers': {'gitlab': {'type': 'stdio', 'command': 'dummy-mcp'}}}), ) - global_config_file, plugins = Copilot().get_session_context() + global_config_files, plugins = Copilot().get_session_context() - assert global_config_file is not None - assert global_config_file['path'] == str(config_path) + assert len(global_config_files) == 1 + assert global_config_files[0]['path'] == str(config_path) # VS Code's `servers` key is normalized to the canonical mcpServers shape. - assert json.loads(global_config_file['content']) == { + assert json.loads(global_config_files[0]['content']) == { 'mcpServers': {'gitlab': {'type': 'stdio', 'command': 'dummy-mcp'}} } assert plugins == {} +def test_get_session_context_reports_both_runtimes_configs(fs: FakeFilesystem) -> None: + vscode_path = _vscode_mcp_config_path() + fs.create_file(vscode_path, contents=json.dumps({'servers': {'gitlab': {'command': 'gitlab-mcp'}}})) + agent_path = Path.home() / '.copilot' / 'mcp-config.json' + fs.create_file(agent_path, contents=json.dumps({'mcpServers': {'linear': {'command': 'linear-mcp'}}})) + + global_config_files, _ = Copilot().get_session_context() + + assert global_config_files == [ + {'path': str(vscode_path), 'content': json.dumps({'mcpServers': {'gitlab': {'command': 'gitlab-mcp'}}})}, + {'path': str(agent_path), 'content': json.dumps({'mcpServers': {'linear': {'command': 'linear-mcp'}}})}, + ] + + def test_get_session_context_without_config(fs: FakeFilesystem) -> None: - assert Copilot().get_session_context() == (None, {}) + assert Copilot().get_session_context() == ([], {}) # --- plugins inventory ----------------------------------------------------------- diff --git a/tests/cli/commands/ai_guardrails/ides/test_cursor.py b/tests/cli/commands/ai_guardrails/ides/test_cursor.py index 0547019e..60834020 100644 --- a/tests/cli/commands/ai_guardrails/ides/test_cursor.py +++ b/tests/cli/commands/ai_guardrails/ides/test_cursor.py @@ -49,7 +49,8 @@ def test_parse_file_read_payload() -> None: assert unified.file_path == '/path/to/secret.env' -def test_parse_mcp_execution_payload() -> None: +def test_parse_mcp_execution_payload(fs: FakeFilesystem) -> None: + # No MCP config on disk: the raw command is kept as the server name. args: dict[str, Any] = {'resource_type': 'merge_request', 'parent_id': 'org/repo', 'resource_id': '4'} unified = Cursor().parse_hook_payload( { @@ -66,6 +67,72 @@ def test_parse_mcp_execution_payload() -> None: assert unified.mcp_arguments == args +def _write_mcp_config(fs: FakeFilesystem, path: Path, servers: dict) -> None: + fs.create_file(str(path), contents=json.dumps({'mcpServers': servers})) + + +def _parse_mcp(**fields: Any) -> Any: + return Cursor().parse_hook_payload({'hook_event_name': 'beforeMCPExecution', 'tool_name': 't', **fields}) + + +_USER_MCP_CONFIG = Path.home() / '.cursor' / 'mcp.json' + + +def test_mcp_server_name_resolved_by_url(fs: FakeFilesystem) -> None: + _write_mcp_config( + fs, + _USER_MCP_CONFIG, + {'notion': {'url': 'https://mcp.notion.com/mcp'}, 'other': {'url': 'https://example.com/mcp'}}, + ) + + assert _parse_mcp(url='https://mcp.notion.com/mcp').mcp_server_name == 'notion' + + +def test_mcp_server_name_resolved_by_command_line(fs: FakeFilesystem) -> None: + _write_mcp_config( + fs, + _USER_MCP_CONFIG, + { + 'filesystem': {'command': 'npx', 'args': ['-y', '@modelcontextprotocol/server-filesystem']}, + 'github': {'command': 'npx', 'args': ['-y', '@modelcontextprotocol/server-github']}, + 'local-db': {'command': '/usr/local/bin/db-mcp'}, + }, + ) + + assert _parse_mcp(command='npx -y @modelcontextprotocol/server-github').mcp_server_name == 'github' + assert _parse_mcp(command='/usr/local/bin/db-mcp').mcp_server_name == 'local-db' + + +def test_mcp_server_name_resolved_by_entry_name(fs: FakeFilesystem) -> None: + _write_mcp_config(fs, _USER_MCP_CONFIG, {'gitlab': {'command': '/opt/homebrew/bin/gitlab-mcp'}}) + + assert _parse_mcp(command='GitLab').mcp_server_name == 'gitlab' + + +def test_mcp_server_name_prefers_the_project_config(fs: FakeFilesystem) -> None: + _write_mcp_config(fs, _USER_MCP_CONFIG, {'user-notion': {'url': 'https://mcp.notion.com/mcp'}}) + _write_mcp_config( + fs, Path('/work/repo/.cursor/mcp.json'), {'project-notion': {'url': 'https://mcp.notion.com/mcp'}} + ) + + unified = _parse_mcp(url='https://mcp.notion.com/mcp', workspace_roots=['/work/repo']) + + assert unified.mcp_server_name == 'project-notion' + + +def test_mcp_server_name_without_a_matching_entry_falls_back(fs: FakeFilesystem) -> None: + _write_mcp_config(fs, _USER_MCP_CONFIG, {'notion': {'url': 'https://mcp.notion.com/mcp'}}) + + assert _parse_mcp(command='npx some-unknown-server').mcp_server_name == 'npx some-unknown-server' + assert _parse_mcp(url='https://unknown.example.com/mcp').mcp_server_name is None + + +def test_mcp_server_name_with_a_corrupt_config_falls_back(fs: FakeFilesystem) -> None: + fs.create_file(str(_USER_MCP_CONFIG), contents='not json {') + + assert _parse_mcp(command='GitLab').mcp_server_name == 'GitLab' + + def test_parse_alternative_field_names() -> None: """Cursor's payload has alternative names for some fields.""" fr = Cursor().parse_hook_payload({'hook_event_name': 'beforeReadFile', 'path': '/alt/path.txt'}) @@ -136,24 +203,26 @@ def test_session_payload_carries_cursor_fields() -> None: def test_session_context_loads_mcp_servers() -> None: - """Cursor wraps ~/.cursor/mcp.json into a global_config_file.""" + """Cursor wraps ~/.cursor/mcp.json into a global config file.""" mcp_servers = {'github': {'command': 'npx', 'args': ['-y', '@modelcontextprotocol/server-github']}} with patch('cycode.cli.apps.ai_guardrails.ides.cursor._load_cursor_mcp_config') as load: load.return_value = {'mcpServers': mcp_servers} - global_config_file, plugins = Cursor().get_session_context() + global_config_files, plugins = Cursor().get_session_context() - assert global_config_file == { - 'path': str(Path.home() / '.cursor' / 'mcp.json'), - 'content': json.dumps({'mcpServers': mcp_servers}), - } + assert global_config_files == [ + { + 'path': str(Path.home() / '.cursor' / 'mcp.json'), + 'content': json.dumps({'mcpServers': mcp_servers}), + } + ] assert plugins == {} def test_session_context_no_config_returns_empty() -> None: with patch('cycode.cli.apps.ai_guardrails.ides.cursor._load_cursor_mcp_config', return_value=None): - global_config_file, plugins = Cursor().get_session_context() - assert global_config_file is None + global_config_files, plugins = Cursor().get_session_context() + assert global_config_files == [] assert plugins == {} diff --git a/tests/cli/commands/ai_guardrails/scan/conftest.py b/tests/cli/commands/ai_guardrails/scan/conftest.py index 86c1b2ae..a253fbb5 100644 --- a/tests/cli/commands/ai_guardrails/scan/conftest.py +++ b/tests/cli/commands/ai_guardrails/scan/conftest.py @@ -10,9 +10,10 @@ def resolved_guardrails_payload( sensitive_path: str = 'Report', mcp: str = 'Report', globs: Optional[list] = None, + mcp_server: Optional[str] = None, ) -> dict: """A platform resolved-config payload with the given per-guardrail modes for the cursor agent.""" - return { + payload = { 'ttl_seconds': 900, 'guardrails': [ {'key': 'secrets_in_prompt', 'event_type': 'Prompt', 'agents': {'cursor': prompt, 'claude-code': 'Block'}}, @@ -26,9 +27,19 @@ def resolved_guardrails_payload( {'key': 'secrets_in_mcp_args', 'event_type': 'McpExecution', 'agents': {'cursor': mcp}}, ], } + if mcp_server is not None: + payload['guardrails'].append( + { + 'key': 'unauthorized_mcp_server', + 'policy_type': 'UnauthorizedAiTools', + 'event_type': 'McpExecution', + 'agents': {'cursor': mcp_server}, + } + ) + return payload -def platform_config(fetched_at: Optional[float] = None, **modes: str) -> GuardrailConfig: +def platform_config(fetched_at: Optional[float] = None, **modes: Optional[str]) -> GuardrailConfig: """A cached platform config; keyword args are the per-guardrail modes (see resolved_guardrails_payload).""" return GuardrailConfig( payload=resolved_guardrails_payload(**modes), diff --git a/tests/cli/commands/ai_guardrails/scan/test_guardrail_config.py b/tests/cli/commands/ai_guardrails/scan/test_guardrail_config.py index efde1d26..90065d93 100644 --- a/tests/cli/commands/ai_guardrails/scan/test_guardrail_config.py +++ b/tests/cli/commands/ai_guardrails/scan/test_guardrail_config.py @@ -82,7 +82,7 @@ def test_unknown_guardrail_keys_are_ignored() -> None: payload = _payload() # A future guardrail this CLI doesn't implement must never affect decisions (fail-open). payload['guardrails'].append( - {'key': 'unauthorized_mcp_server', 'event_type': 'McpExecution', 'agents': {'cursor': 'Block'}} + {'key': 'future_guardrail', 'event_type': 'McpExecution', 'agents': {'cursor': 'Block'}} ) config = GuardrailConfig(payload=payload, fetched_at=time.time()) @@ -151,3 +151,55 @@ def test_apply_platform_config_sensitive_path_off_clears_globs() -> None: assert policy['file_read']['deny_globs'] == [] assert policy['file_read']['scan_content'] is True + + +# --- unauthorized MCP server guardrail --- + + +def test_unauthorized_mcp_server_defaults_to_off() -> None: + # Unlike the other guardrails: a missing entry, or a missing agent cell, must not enable it. + missing_entry = _config() + assert missing_entry.mode_for('unauthorized_mcp_server', 'cursor') == 'off' + assert missing_entry.is_off_for_every_agent('unauthorized_mcp_server') is True + + missing_cell = _config(mcp_server='Block') + assert missing_cell.mode_for('unauthorized_mcp_server', 'codex') == 'off' + assert missing_cell.can_event_block('McpExecution', 'codex') is False + + +def test_unauthorized_mcp_server_block_makes_the_event_blockable() -> None: + config = _config(mcp_server='Block') + + assert config.can_event_block('McpExecution', 'cursor') is True + assert config.is_off_for_every_agent('unauthorized_mcp_server') is False + + +def test_mcp_event_is_off_only_when_both_mcp_guardrails_are_off() -> None: + assert _config(mcp='Off').is_event_off('McpExecution', 'cursor') is True + assert _config(mcp='Off', mcp_server='Off').is_event_off('McpExecution', 'cursor') is True + assert _config(mcp='Off', mcp_server='Report').is_event_off('McpExecution', 'cursor') is False + + +def test_apply_platform_config_without_cache_leaves_the_mcp_server_check_off() -> None: + policy: dict = {} + + apply_platform_config(policy, None, 'cursor') + + assert policy['mcp']['check_server'] is False + assert policy['mcp']['scan_args'] is True + + +def test_apply_platform_config_mcp_server_cells() -> None: + policy: dict = {} + apply_platform_config(policy, _config(mcp='Off', mcp_server='Block'), 'cursor') + + assert policy['mcp']['scan_args'] is False + assert policy['mcp']['check_server'] is True + assert policy['mcp']['server_action'] == 'block' + + policy = {} + apply_platform_config(policy, _config(mcp_server='Report'), 'cursor') + + assert policy['mcp']['check_server'] is True + assert policy['mcp']['server_action'] == 'warn' + assert policy['mcp']['action'] == 'warn' diff --git a/tests/cli/commands/ai_guardrails/scan/test_handlers.py b/tests/cli/commands/ai_guardrails/scan/test_handlers.py index 5235f1a3..4c5abe24 100644 --- a/tests/cli/commands/ai_guardrails/scan/test_handlers.py +++ b/tests/cli/commands/ai_guardrails/scan/test_handlers.py @@ -2,8 +2,9 @@ import os import threading +import time from multiprocessing import synchronize -from typing import Any +from typing import Any, Optional from unittest.mock import MagicMock, patch import pytest @@ -22,11 +23,12 @@ handle_before_read_file, handle_before_submit_prompt, ) +from cycode.cli.apps.ai_guardrails.scan.mcp_server_status import McpServerStatuses from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload from cycode.cli.apps.ai_guardrails.scan.types import AiHookEventType, AIHookOutcome, BlockReason from cycode.cli.apps.ai_guardrails.scan.utils import MAX_VIOLATION_DETAIL_LINES, build_violation_summary from cycode.cli.models import Document, DocumentDetections, LocalScanResult -from cycode.cyclient.models import Detection +from cycode.cyclient.models import Detection, McpServerAuthorizationStatus, McpServerStatus @pytest.fixture @@ -582,6 +584,157 @@ def test_handle_before_mcp_execution_with_secrets_warned( assert call_args.args[2] == AIHookOutcome.WARNED +# Tests for the unauthorized MCP server guardrail + + +def _server_check_policy(default_policy: dict[str, Any], server_action: str = 'block') -> dict[str, Any]: + default_policy['mcp'].update(check_server=True, server_action=server_action) + return default_policy + + +def _mcp_payload(server: Optional[str] = 'github') -> AIHookPayload: + return AIHookPayload( + event_name='McpExecution', + conversation_id='conv-1', + ide_provider='claude-code', + mcp_server_name=server, + mcp_tool_name='create_issue', + mcp_arguments={'title': 'hello'}, + ) + + +def _cached_statuses(*rows: tuple[str, str]) -> McpServerStatuses: + servers = [McpServerStatus(alias=a, status=McpServerAuthorizationStatus.parse(s)) for a, s in rows] + return McpServerStatuses(servers=servers, fetched_at=time.time()) + + +def _reported_event(mock_ctx: MagicMock) -> tuple[AIHookOutcome, Optional[BlockReason]]: + call_args = mock_ctx.obj['ai_security_client'].create_event.call_args + return call_args.args[2], call_args.kwargs['block_reason'] + + +@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets') +@patch('cycode.cli.apps.ai_guardrails.scan.handlers.load_mcp_server_statuses') +def test_unauthorized_mcp_server_block_denies_without_scanning( + mock_statuses: MagicMock, mock_scan: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any] +) -> None: + mock_statuses.return_value = _cached_statuses(('github', 'Unauthorized')) + + result = handle_before_mcp_execution(mock_ctx, _mcp_payload(), _server_check_policy(default_policy)) + + assert result.action == DecisionAction.DENY + assert result.user_message == ( + "Cycode blocked MCP server 'github': it is not authorized in your organization. " + 'Contact your admin to authorize it.' + ) + assert "'github'" in result.agent_message + mock_scan.assert_not_called() + assert _reported_event(mock_ctx) == (AIHookOutcome.BLOCKED, BlockReason.UNAUTHORIZED_MCP_SERVER) + + +@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets') +@patch('cycode.cli.apps.ai_guardrails.scan.handlers.load_mcp_server_statuses') +def test_unauthorized_mcp_server_report_allows_warns_and_still_scans( + mock_statuses: MagicMock, mock_scan: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any] +) -> None: + mock_statuses.return_value = _cached_statuses(('github', 'Unauthorized')) + mock_scan.return_value = ScanOutcome(scan_id='scan-1') + + result = handle_before_mcp_execution(mock_ctx, _mcp_payload(), _server_check_policy(default_policy, 'warn')) + + assert result == HookDecision.allow(AiHookEventType.MCP_EXECUTION) + mock_scan.assert_called_once() + assert _reported_event(mock_ctx) == (AIHookOutcome.WARNED, BlockReason.UNAUTHORIZED_MCP_SERVER) + + +@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets') +@patch('cycode.cli.apps.ai_guardrails.scan.handlers.load_mcp_server_statuses') +def test_unauthorized_mcp_server_report_then_secret_block_takes_over( + mock_statuses: MagicMock, mock_scan: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any] +) -> None: + mock_statuses.return_value = _cached_statuses(('github', 'Unauthorized')) + mock_scan.return_value = ScanOutcome('Found 1 secret: token', 'scan-1', GuardrailsMode.BLOCK) + + result = handle_before_mcp_execution(mock_ctx, _mcp_payload(), _server_check_policy(default_policy, 'warn')) + + assert result.action == DecisionAction.DENY + assert 'Found 1 secret: token' in result.user_message + assert _reported_event(mock_ctx) == (AIHookOutcome.BLOCKED, BlockReason.SECRETS_IN_MCP_ARGS) + + +@pytest.mark.parametrize( + ('rows', 'server', 'expected_action'), + [ + # Only explicitly Unauthorized servers are enforced. + ((('github', 'Unreviewed'),), 'github', DecisionAction.ALLOW), + ((('github', 'Authorized'),), 'github', DecisionAction.ALLOW), + ((), 'github', DecisionAction.ALLOW), + # One alias, several servers: the most restrictive status wins. + ((('github', 'Authorized'), ('GitHub', 'Unauthorized')), 'github', DecisionAction.DENY), + # No server name to check: fail open. + ((('github', 'Unauthorized'),), None, DecisionAction.ALLOW), + ], +) +@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets', return_value=ScanOutcome()) +@patch('cycode.cli.apps.ai_guardrails.scan.handlers.load_mcp_server_statuses') +def test_unauthorized_mcp_server_enforcement( + mock_statuses: MagicMock, + mock_scan: MagicMock, + mock_ctx: MagicMock, + default_policy: dict[str, Any], + rows: tuple, + server: Optional[str], + expected_action: DecisionAction, +) -> None: + mock_statuses.return_value = _cached_statuses(*rows) + policy = _server_check_policy(default_policy) + + result = handle_before_mcp_execution(mock_ctx, _mcp_payload(server), policy) + + assert result.action == expected_action + + +@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets', return_value=ScanOutcome()) +@patch('cycode.cli.apps.ai_guardrails.scan.handlers.load_mcp_server_statuses', return_value=None) +def test_unauthorized_mcp_server_without_cached_statuses_fails_open( + mock_statuses: MagicMock, mock_scan: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any] +) -> None: + policy = _server_check_policy(default_policy) + + result = handle_before_mcp_execution(mock_ctx, _mcp_payload(), policy) + + assert result == HookDecision.allow(AiHookEventType.MCP_EXECUTION) + assert _reported_event(mock_ctx) == (AIHookOutcome.ALLOWED, None) + + +@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets', return_value=ScanOutcome()) +@patch('cycode.cli.apps.ai_guardrails.scan.handlers.load_mcp_server_statuses') +def test_unauthorized_mcp_server_off_never_reads_statuses( + mock_statuses: MagicMock, mock_scan: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any] +) -> None: + # The default policy carries no platform overlay: the check is off. + result = handle_before_mcp_execution(mock_ctx, _mcp_payload(), default_policy) + + assert result == HookDecision.allow(AiHookEventType.MCP_EXECUTION) + mock_statuses.assert_not_called() + + +@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets') +@patch('cycode.cli.apps.ai_guardrails.scan.handlers.load_mcp_server_statuses') +def test_args_scan_off_skips_the_scan_when_only_the_server_check_is_on( + mock_statuses: MagicMock, mock_scan: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any] +) -> None: + mock_statuses.return_value = _cached_statuses(('github', 'Authorized')) + policy = _server_check_policy(default_policy) + policy['mcp']['scan_args'] = False + + result = handle_before_mcp_execution(mock_ctx, _mcp_payload(), policy) + + assert result == HookDecision.allow(AiHookEventType.MCP_EXECUTION) + mock_scan.assert_not_called() + assert _reported_event(mock_ctx) == (AIHookOutcome.ALLOWED, None) + + def test_get_effective_mode_reads_the_guardrails_action() -> None: assert get_effective_mode({'action': 'block'}) == GuardrailsMode.BLOCK assert get_effective_mode({'action': 'warn'}) == GuardrailsMode.REPORT diff --git a/tests/cli/commands/ai_guardrails/scan/test_mcp_server_status.py b/tests/cli/commands/ai_guardrails/scan/test_mcp_server_status.py new file mode 100644 index 00000000..b7842384 --- /dev/null +++ b/tests/cli/commands/ai_guardrails/scan/test_mcp_server_status.py @@ -0,0 +1,146 @@ +"""Tests for the MCP server authorization status cache.""" + +import json +import time +from pathlib import Path +from typing import Optional + +import pytest +from pyfakefs.fake_filesystem import FakeFilesystem + +from cycode.cli.apps.ai_guardrails.scan.mcp_server_status import ( + McpServerStatuses, + get_mcp_server_statuses_cache_path, + is_enforced, + load_mcp_server_statuses, + save_mcp_server_statuses, +) +from cycode.cyclient.models import McpServerAuthorizationStatus, McpServerStatus + +_AUTHORIZED = McpServerAuthorizationStatus.AUTHORIZED +_UNREVIEWED = McpServerAuthorizationStatus.UNREVIEWED +_UNAUTHORIZED = McpServerAuthorizationStatus.UNAUTHORIZED + + +@pytest.fixture(autouse=True) +def _fake_home(fs: FakeFilesystem, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv('HOME', '/home/testuser') + fs.create_dir('/home/testuser') + + +def _statuses(*rows: tuple[str, str]) -> McpServerStatuses: + return McpServerStatuses( + servers=[ + McpServerStatus(alias, f'id:{alias}', McpServerAuthorizationStatus.parse(status)) for alias, status in rows + ], + fetched_at=time.time(), + ) + + +def test_cache_file_sits_next_to_the_guardrail_config() -> None: + assert get_mcp_server_statuses_cache_path() == Path.home() / '.cycode' / 'ai-guardrails-mcp-servers.json' + + +@pytest.mark.parametrize( + ('raw', 'expected'), + [ + ('Authorized', _AUTHORIZED), + ('unauthorized', _UNAUTHORIZED), + ('UNREVIEWED', _UNREVIEWED), + ('Pending', _UNREVIEWED), + (None, _UNREVIEWED), + ], +) +def test_parse_status_is_case_insensitive_and_unknown_reads_unreviewed( + raw: Optional[str], expected: McpServerAuthorizationStatus +) -> None: + assert McpServerAuthorizationStatus.parse(raw) == expected + + +@pytest.mark.parametrize( + ('status', 'expected'), + [ + (_UNAUTHORIZED, True), + (_UNREVIEWED, False), + (_AUTHORIZED, False), + (None, False), + ], +) +def test_is_enforced(status: Optional[McpServerAuthorizationStatus], expected: bool) -> None: + assert is_enforced(status) is expected + + +def test_status_of_is_case_insensitive() -> None: + assert _statuses(('GitHub', 'Unauthorized')).status_of('github') == _UNAUTHORIZED + + +def test_status_of_unknown_alias_returns_none() -> None: + assert _statuses(('github', 'Authorized')).status_of('notion') is None + + +def test_the_most_restrictive_status_wins_for_one_alias() -> None: + statuses = _statuses(('github', 'Authorized'), ('github', 'Unauthorized'), ('github', 'Unreviewed')) + assert statuses.status_of('github') == _UNAUTHORIZED + + statuses = _statuses(('notion', 'Authorized'), ('Notion', 'Unreviewed')) + assert statuses.status_of('notion') == _UNREVIEWED + + +def test_rows_without_an_alias_are_ignored() -> None: + statuses = McpServerStatuses( + servers=[McpServerStatus(status=_UNAUTHORIZED), McpServerStatus(alias='', status=_UNAUTHORIZED)], + fetched_at=time.time(), + ) + assert statuses.status_of('') is None + + +def test_save_and_load_round_trip() -> None: + save_mcp_server_statuses([McpServerStatus('github', 'pkg:gh', _UNAUTHORIZED)], 'tenant-a', ttl_seconds=60) + + statuses = load_mcp_server_statuses() + + assert statuses is not None + assert statuses.status_of('github') == _UNAUTHORIZED + assert statuses.servers == [McpServerStatus('github', 'pkg:gh', _UNAUTHORIZED)] + assert statuses.tenant_id == 'tenant-a' + assert statuses.ttl_seconds == 60 + + +def test_saved_cache_keeps_the_file_format() -> None: + save_mcp_server_statuses([McpServerStatus('github', 'pkg:gh', _UNAUTHORIZED)], 'tenant-a', ttl_seconds=60) + + content = json.loads(get_mcp_server_statuses_cache_path().read_text(encoding='UTF-8')) + + assert set(content) == {'servers', 'fetched_at', 'tenant_id', 'ttl_seconds'} + assert content['servers'] == [{'alias': 'github', 'normalized_id': 'pkg:gh', 'status': 'Unauthorized'}] + + +def test_load_ignores_unknown_fields_and_reads_unknown_statuses_as_unreviewed(fs: FakeFilesystem) -> None: + fs.create_file( + str(get_mcp_server_statuses_cache_path()), + contents=json.dumps( + {'fetched_at': time.time(), 'future': 1, 'servers': [{'alias': 'github', 'status': 'Pending', 'x': 1}]} + ), + ) + + statuses = load_mcp_server_statuses() + + assert statuses is not None + assert statuses.status_of('github') == _UNREVIEWED + + +def test_load_missing_cache_returns_none() -> None: + assert load_mcp_server_statuses() is None + + +@pytest.mark.parametrize( + 'contents', + ['{"servers": {}, "fetched_at": 1}', '{"servers": []}', '{"servers": ["garbage"], "fetched_at": 1}', 'not json'], +) +def test_corrupt_cache_is_quarantined(contents: str, fs: FakeFilesystem) -> None: + path = get_mcp_server_statuses_cache_path() + fs.create_file(str(path), contents=contents) + + assert load_mcp_server_statuses() is None + assert not path.exists() + assert Path(f'{path}.corrupt').exists() diff --git a/tests/cli/commands/ai_guardrails/scan/test_scan_command.py b/tests/cli/commands/ai_guardrails/scan/test_scan_command.py index 8acaf1ac..b756d534 100644 --- a/tests/cli/commands/ai_guardrails/scan/test_scan_command.py +++ b/tests/cli/commands/ai_guardrails/scan/test_scan_command.py @@ -503,6 +503,55 @@ def test_detach_decision_is_per_event( scan_command(mock_ctx, ide='cursor') mock_respawn.assert_called_once() + @pytest.mark.parametrize( + ('mcp_server', 'expect_detach'), + [(None, True), ('off', True), ('report', True), ('block', False)], + ) + def test_unauthorized_mcp_server_block_keeps_the_mcp_event_synchronous( + self, + mock_ctx: MagicMock, + mocker: MockerFixture, + mock_respawn: MagicMock, + not_detached: None, + mcp_server: Optional[str], + expect_detach: bool, + ) -> None: + config = platform_config(mcp='report', mcp_server=mcp_server) + mocker.patch('cycode.cli.apps.ai_guardrails.scan.scan_command._initialize_clients') + handler = MagicMock(return_value=HookDecision.allow(AiHookEventType.MCP_EXECUTION)) + mocker.patch('cycode.cli.apps.ai_guardrails.scan.scan_command.get_handler_for_event', return_value=handler) + mocker.patch('cycode.cli.apps.ai_guardrails.scan.scan_command.load_guardrail_config', return_value=config) + mocker.patch('cycode.cli.apps.ai_guardrails.scan.scan_command.load_policy', return_value={'fail_open': True}) + + mocker.patch('sys.stdin', StringIO(json.dumps({'hook_event_name': 'beforeMCPExecution', 'tool_name': 't'}))) + scan_command(mock_ctx, ide='cursor') + + assert mock_respawn.called is expect_detach + assert handler.called is not expect_detach + + def test_mcp_event_runs_when_only_the_server_check_is_on( + self, + mock_ctx: MagicMock, + mocker: MockerFixture, + mock_respawn: MagicMock, + not_detached: None, + ) -> None: + config = platform_config(mcp='off', mcp_server='block') + mocker.patch('cycode.cli.apps.ai_guardrails.scan.scan_command._initialize_clients') + handler = MagicMock(return_value=HookDecision.allow(AiHookEventType.MCP_EXECUTION)) + mocker.patch('cycode.cli.apps.ai_guardrails.scan.scan_command.get_handler_for_event', return_value=handler) + mocker.patch('cycode.cli.apps.ai_guardrails.scan.scan_command.load_guardrail_config', return_value=config) + mocker.patch('cycode.cli.apps.ai_guardrails.scan.scan_command.load_policy', return_value={'fail_open': True}) + + mocker.patch('sys.stdin', StringIO(json.dumps({'hook_event_name': 'beforeMCPExecution', 'tool_name': 't'}))) + scan_command(mock_ctx, ide='cursor') + + handler.assert_called_once() + policy = handler.call_args.args[2] + assert policy['mcp']['scan_args'] is False + assert policy['mcp']['check_server'] is True + assert policy['mcp']['server_action'] == 'block' + def test_failed_respawn_falls_back_to_synchronous_scan( self, mock_ctx: MagicMock, diff --git a/tests/cli/commands/ai_guardrails/test_session_start_command.py b/tests/cli/commands/ai_guardrails/test_session_start_command.py index 2ace737a..bd63dfef 100644 --- a/tests/cli/commands/ai_guardrails/test_session_start_command.py +++ b/tests/cli/commands/ai_guardrails/test_session_start_command.py @@ -4,6 +4,7 @@ import time from io import StringIO from pathlib import Path +from typing import Optional from unittest.mock import ANY, MagicMock, patch import pytest @@ -17,6 +18,7 @@ from cycode.cli.apps.ai_guardrails.ides import cursor as _cursor_mod from cycode.cli.apps.ai_guardrails.scan.guardrail_config import GuardrailConfig from cycode.cli.apps.ai_guardrails.session_start_command import session_start_command +from cycode.cyclient.models import McpServerAuthorizationStatus, McpServerStatus, SessionContextResponse @pytest.fixture @@ -36,6 +38,14 @@ def mock_save_guardrail_config(monkeypatch: pytest.MonkeyPatch) -> MagicMock: return save_mock +@pytest.fixture(autouse=True) +def mock_save_mcp_server_statuses(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """Keep tests hermetic: never write the real MCP server statuses cache.""" + save_mock = MagicMock() + monkeypatch.setattr(_session_start_mod, 'save_mcp_server_statuses', save_mock) + return save_mock + + @pytest.fixture(autouse=True) def _isolated_session_context_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """Keep the dedup cache away from the real ~/.cycode in every test.""" @@ -266,7 +276,7 @@ def test_reports_cross_ide_session_context( cursor_file = {'path': '/home/u/.cursor/mcp.json', 'content': '{"mcpServers": {}}'} claude_file = {'path': '/home/u/.claude.json', 'content': '{"mcpServers": {}}'} plugins = {'dummy-plugin@dummy-marketplace': {'enabled': True}} - mock_collect.return_value = ({'cursor': cursor_file, 'claude-code': claude_file}, plugins) + mock_collect.return_value = ([cursor_file, claude_file], plugins) payload = {'session_id': 'session-123'} @@ -302,7 +312,7 @@ def test_no_mcp_anywhere_still_reports_device( mock_get_auth.return_value = MagicMock(tenant_id='tenant-1') mock_ai_client = MagicMock() mock_get_client.return_value = mock_ai_client - mock_collect.return_value = ({}, {}) + mock_collect.return_value = ([], {}) payload = {'session_id': 'session-123'} @@ -484,7 +494,7 @@ def test_unchanged_context_skips_second_report( mock_get_auth.return_value = MagicMock(tenant_id='tenant-1') mock_ai_client = MagicMock() mock_get_client.return_value = mock_ai_client - mock_collect.return_value = ({'cursor': {'path': '/p', 'content': 'c'}}, {}) + mock_collect.return_value = ([{'path': '/p', 'content': 'c'}], {}) _run_session_start(mock_ctx, {'session_id': 'session-1'}) _run_session_start(mock_ctx, {'session_id': 'session-2'}) @@ -506,10 +516,10 @@ def test_changed_context_resends( mock_ai_client = MagicMock() mock_get_client.return_value = mock_ai_client - mock_collect.return_value = ({'cursor': {'path': '/p', 'content': 'c1'}}, {}) + mock_collect.return_value = ([{'path': '/p', 'content': 'c1'}], {}) _run_session_start(mock_ctx, {'session_id': 'session-1'}) - mock_collect.return_value = ({'cursor': {'path': '/p', 'content': 'c2'}}, {}) + mock_collect.return_value = ([{'path': '/p', 'content': 'c2'}], {}) _run_session_start(mock_ctx, {'session_id': 'session-2'}) assert mock_ai_client.report_session_context.call_count == 2 @@ -527,7 +537,7 @@ def test_tenant_change_resends( """Re-authenticating against a different tenant must re-send the same inventory.""" mock_ai_client = MagicMock() mock_get_client.return_value = mock_ai_client - mock_collect.return_value = ({'cursor': {'path': '/p', 'content': 'c'}}, {}) + mock_collect.return_value = ([{'path': '/p', 'content': 'c'}], {}) mock_get_auth.return_value = MagicMock(tenant_id='tenant-1') _run_session_start(mock_ctx, {'session_id': 'session-1'}) @@ -550,9 +560,9 @@ def test_failed_report_is_not_cached( """A failed send must not populate the cache - the next session retries.""" mock_get_auth.return_value = MagicMock(tenant_id='tenant-1') mock_ai_client = MagicMock() - mock_ai_client.report_session_context.return_value = False + mock_ai_client.report_session_context.return_value = None mock_get_client.return_value = mock_ai_client - mock_collect.return_value = ({'cursor': {'path': '/p', 'content': 'c'}}, {}) + mock_collect.return_value = ([{'path': '/p', 'content': 'c'}], {}) _run_session_start(mock_ctx, {'session_id': 'session-1'}) _run_session_start(mock_ctx, {'session_id': 'session-2'}) @@ -573,7 +583,7 @@ def test_expired_ttl_resends( mock_get_auth.return_value = MagicMock(tenant_id='tenant-1') mock_ai_client = MagicMock() mock_get_client.return_value = mock_ai_client - mock_collect.return_value = ({'cursor': {'path': '/p', 'content': 'c'}}, {}) + mock_collect.return_value = ([{'path': '/p', 'content': 'c'}], {}) _run_session_start(mock_ctx, {'session_id': 'session-1'}) @@ -597,10 +607,10 @@ def test_collect_all_session_contexts_merges_plugins_first_wins() -> None: codex_plugin = {'enabled': True, 'version': '2.0.0'} with ( - patch.object(IDES['cursor'], 'get_session_context', return_value=(None, {})), - patch.object(IDES['claude-code'], 'get_session_context', return_value=(None, {'plug@m': claude_plugin})), - patch.object(IDES['codex'], 'get_session_context', return_value=(None, {'plug@m': codex_plugin})), - patch.object(IDES['copilot'], 'get_session_context', return_value=(None, {})), + patch.object(IDES['cursor'], 'get_session_context', return_value=([], {})), + patch.object(IDES['claude-code'], 'get_session_context', return_value=([], {'plug@m': claude_plugin})), + patch.object(IDES['codex'], 'get_session_context', return_value=([], {'plug@m': codex_plugin})), + patch.object(IDES['copilot'], 'get_session_context', return_value=([], {})), ): _, plugins = collect_all_session_contexts() @@ -702,6 +712,145 @@ def test_fresh_guardrail_config_cache_skips_fetch( mock_save_guardrail_config.assert_called_once_with(mock_ai_client.get_resolved_guardrails.return_value, 'tenant-b') +# MCP server statuses (unauthorized MCP server guardrail) + + +def _guardrail_config_with_mcp_server(agents: dict) -> GuardrailConfig: + payload = { + 'ttl_seconds': 600, + 'guardrails': [ + { + 'key': 'unauthorized_mcp_server', + 'event_type': 'McpExecution', + 'agents': agents, + } + ], + } + return GuardrailConfig(payload=payload, fetched_at=time.time(), tenant_id='tenant-a') + + +_SERVERS = [McpServerStatus('github', 'pkg:gh', McpServerAuthorizationStatus.UNAUTHORIZED)] + + +def _run_with_config( + mock_ctx: MagicMock, + monkeypatch: pytest.MonkeyPatch, + ai_client: MagicMock, + config: Optional[GuardrailConfig], + tenant_id: str = 'tenant-a', +) -> None: + monkeypatch.setattr(_session_start_mod, 'load_guardrail_config', MagicMock(return_value=config)) + with ( + patch.object(_session_start_mod, 'get_authorization_info', return_value=MagicMock(tenant_id=tenant_id)), + patch.object(_session_start_mod, 'get_ai_security_manager_client', return_value=ai_client), + patch.object( + _session_start_mod, + 'collect_all_session_contexts', + return_value=([{'path': '/p', 'content': 'c'}], {}), + ), + patch('sys.stdin', new=StringIO(json.dumps({'conversation_id': 'conv-1'}))), + ): + session_start_command(mock_ctx, ide='cursor') + + +def _age_report_cache(seconds: float) -> None: + cache_path = _session_start_mod._session_context_cache_path() + cache = json.loads(cache_path.read_text(encoding='utf-8')) + cache['sent_at'] -= seconds + cache_path.write_text(json.dumps(cache), encoding='utf-8') + + +def _statuses_client(mcp_servers: Optional[list] = _SERVERS) -> MagicMock: + ai_client = MagicMock() + ai_client.report_session_context.return_value = SessionContextResponse(mcp_servers=mcp_servers) + return ai_client + + +def test_session_context_response_saves_the_mcp_server_statuses( + mock_ctx: MagicMock, monkeypatch: pytest.MonkeyPatch, mock_save_mcp_server_statuses: MagicMock +) -> None: + config = _guardrail_config_with_mcp_server({'cursor': 'Block'}) + + _run_with_config(mock_ctx, monkeypatch, _statuses_client(), config) + + # Cached with the guardrail config's TTL. + mock_save_mcp_server_statuses.assert_called_once_with(_SERVERS, 'tenant-a', 600) + + +@pytest.mark.parametrize('response', [None, SessionContextResponse(mcp_servers=None)]) +def test_failed_or_malformed_session_context_response_keeps_the_cache( + response: Optional[SessionContextResponse], + mock_ctx: MagicMock, + monkeypatch: pytest.MonkeyPatch, + mock_save_mcp_server_statuses: MagicMock, +) -> None: + ai_client = MagicMock() + ai_client.report_session_context.return_value = response + + _run_with_config(mock_ctx, monkeypatch, ai_client, _guardrail_config_with_mcp_server({'cursor': 'Block'})) + + ai_client.report_session_context.assert_called_once() + mock_save_mcp_server_statuses.assert_not_called() + + +@pytest.mark.parametrize( + 'config', + [ + None, + GuardrailConfig(payload={'guardrails': []}, fetched_at=time.time(), tenant_id='tenant-a'), + _guardrail_config_with_mcp_server({'cursor': 'Off', 'claude-code': 'Off'}), + ], +) +def test_unchanged_context_is_skipped_for_the_report_ttl_while_the_guardrail_is_off( + config: Optional[GuardrailConfig], mock_ctx: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + ai_client = _statuses_client() + + _run_with_config(mock_ctx, monkeypatch, ai_client, config) + _age_report_cache(601) + _run_with_config(mock_ctx, monkeypatch, ai_client, config) + + ai_client.report_session_context.assert_called_once() + + +def test_unchanged_context_is_skipped_within_the_guardrail_ttl( + mock_ctx: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + config = _guardrail_config_with_mcp_server({'cursor': 'Off', 'claude-code': 'Report'}) + ai_client = _statuses_client() + + _run_with_config(mock_ctx, monkeypatch, ai_client, config) + _age_report_cache(599) + _run_with_config(mock_ctx, monkeypatch, ai_client, config) + + ai_client.report_session_context.assert_called_once() + + +def test_unchanged_context_is_resent_after_the_guardrail_ttl_to_refresh_statuses( + mock_ctx: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + config = _guardrail_config_with_mcp_server({'cursor': 'Block'}) + ai_client = _statuses_client() + + _run_with_config(mock_ctx, monkeypatch, ai_client, config) + _age_report_cache(601) + _run_with_config(mock_ctx, monkeypatch, ai_client, config) + + assert ai_client.report_session_context.call_count == 2 + + +def test_unchanged_context_is_resent_for_another_tenant_while_the_guardrail_is_enabled( + mock_ctx: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + config = _guardrail_config_with_mcp_server({'cursor': 'Block'}) + ai_client = _statuses_client() + + _run_with_config(mock_ctx, monkeypatch, ai_client, config) + _run_with_config(mock_ctx, monkeypatch, ai_client, config, tenant_id='tenant-b') + + assert ai_client.report_session_context.call_count == 2 + + # Skills reporting @@ -721,7 +870,7 @@ def test_reports_skill_files( mock_get_auth.return_value = MagicMock(tenant_id='tenant-1') mock_ai_client = MagicMock() mock_get_client.return_value = mock_ai_client - mock_collect.return_value = ({}, {}) + mock_collect.return_value = ([], {}) content = '---\nname: dummy-skill\n---\nBody.\n' skill_file = _write_claude_skill(isolated_home, 'dummy-skill', content) skills = [{'path': str(skill_file), 'content': content}] @@ -760,7 +909,7 @@ def test_editing_a_skill_re_reports( mock_get_auth.return_value = MagicMock(tenant_id='tenant-1') mock_ai_client = MagicMock() mock_get_client.return_value = mock_ai_client - mock_collect.return_value = ({}, {}) + mock_collect.return_value = ([], {}) payload = json.dumps({'session_id': 'session-123'}) _write_claude_skill(isolated_home, 'dummy-skill', 'first') diff --git a/tests/cyclient/test_ai_security_manager_client.py b/tests/cyclient/test_ai_security_manager_client.py index 4d37e020..8e3e1da2 100644 --- a/tests/cyclient/test_ai_security_manager_client.py +++ b/tests/cyclient/test_ai_security_manager_client.py @@ -1,8 +1,11 @@ from unittest.mock import MagicMock +import pytest + from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload from cycode.cli.apps.ai_guardrails.scan.types import AiHookEventType, AIHookOutcome from cycode.cyclient.ai_security_manager_client import AISecurityManagerClient +from cycode.cyclient.models import McpServerAuthorizationStatus, McpServerStatus def _build_client() -> tuple[AISecurityManagerClient, MagicMock]: @@ -74,3 +77,46 @@ def test_create_event_without_a_conversation_posts_nothing() -> None: client.create_event(AIHookPayload(event_name='Prompt'), AiHookEventType.PROMPT, AIHookOutcome.ALLOWED) http_client.post.assert_not_called() + + +def test_report_session_context_returns_the_mcp_server_statuses() -> None: + client, http_client = _build_client() + http_client.post.return_value.json.return_value = { + 'mcp_servers': [ + {'alias': 'github', 'normalized_id': 'pkg:gh', 'status': 'Unauthorized', 'unknown_field': 1}, + {'alias': 'notion', 'normalized_id': 'pkg:notion', 'status': 'SomethingNew'}, + ], + 'unknown_field': 1, + } + + response = client.report_session_context(hostname='host') + + assert response is not None + assert response.mcp_servers == [ + McpServerStatus('github', 'pkg:gh', McpServerAuthorizationStatus.UNAUTHORIZED), + McpServerStatus('notion', 'pkg:notion', McpServerAuthorizationStatus.UNREVIEWED), + ] + assert http_client.post.call_args.args[0] == 'v4/ai-security/interactions/session-context' + + +def test_report_session_context_failure_returns_none() -> None: + client, http_client = _build_client() + http_client.post.side_effect = RuntimeError('boom') + + assert client.report_session_context(hostname='host') is None + + +@pytest.mark.parametrize( + 'body', [{}, [], {'mcp_servers': {'github': 'Unauthorized'}}, {'mcp_servers': ['github']}, ValueError('not json')] +) +def test_report_session_context_malformed_response_has_no_statuses(body: object) -> None: + client, http_client = _build_client() + if isinstance(body, Exception): + http_client.post.return_value.json.side_effect = body + else: + http_client.post.return_value.json.return_value = body + + response = client.report_session_context(hostname='host') + + assert response is not None + assert response.mcp_servers is None