|
2 | 2 |
|
3 | 3 | import json |
4 | 4 | import os |
| 5 | +import re |
5 | 6 | import sys |
| 7 | +from collections.abc import Iterator |
6 | 8 | from pathlib import Path |
7 | 9 | from typing import ClassVar, Optional |
8 | 10 |
|
|
15 | 17 |
|
16 | 18 | from cycode.cli.apps.ai_guardrails.consts import CYCODE_SCAN_PROMPT_COMMAND, CYCODE_SESSION_START_COMMAND |
17 | 19 | from cycode.cli.apps.ai_guardrails.ides._plugin_utils import ( |
18 | | - build_global_config_file, |
| 20 | + build_global_config_files, |
19 | 21 | load_plugin_json, |
20 | 22 | resolve_cached_plugin_dir, |
21 | 23 | walk_enabled_plugins, |
@@ -158,6 +160,42 @@ def _resolve_codex_plugins(config: dict) -> dict: |
158 | 160 | ) |
159 | 161 |
|
160 | 162 |
|
| 163 | +def _sanitize_mcp_name(name: str) -> str: |
| 164 | + """Mirrors how Codex sanitizes a server name into a tool namespace.""" |
| 165 | + return re.sub(r'[^A-Za-z0-9_]', '_', name) or '_' |
| 166 | + |
| 167 | + |
| 168 | +def _local_mcp_server_names() -> Iterator[str]: |
| 169 | + """Yield the config key of each locally configured Codex MCP server, plugin servers included.""" |
| 170 | + config = _load_codex_config() |
| 171 | + if not config: |
| 172 | + return |
| 173 | + servers = config.get('mcp_servers') |
| 174 | + if isinstance(servers, dict): |
| 175 | + yield from servers |
| 176 | + for entry in _resolve_codex_plugins(config).values(): |
| 177 | + yield from entry.get('mcp_server_names') or [] |
| 178 | + |
| 179 | + |
| 180 | +def _split_mcp_tool_name(tool_name: str) -> tuple[Optional[str], Optional[str]]: |
| 181 | + """Split ``mcp__<namespace>__<tool>`` into ``(server config key, tool)``; unknown namespaces split on ``__``.""" |
| 182 | + rest = tool_name[len('mcp__') :] |
| 183 | + names = list(_local_mcp_server_names()) |
| 184 | + # Codex appends `_<12 hex>` to a namespace on collision or length overflow. |
| 185 | + for suffix in ('__', r'_+[0-9a-f]{12}__'): |
| 186 | + best: Optional[tuple[int, str]] = None |
| 187 | + for name in names: |
| 188 | + namespace = _sanitize_mcp_name(name).rstrip('_') |
| 189 | + match = namespace and re.match(re.escape(namespace) + suffix, rest) |
| 190 | + if match and (best is None or match.end() > best[0]): |
| 191 | + best = (match.end(), name) |
| 192 | + if best: |
| 193 | + return best[1], rest[best[0] :] |
| 194 | + |
| 195 | + parts = tool_name.split('__') |
| 196 | + return (parts[1] if len(parts) >= 2 else None), (parts[2] if len(parts) >= 3 else None) |
| 197 | + |
| 198 | + |
161 | 199 | def _enable_codex_hooks_feature(scope: str, repo_path: Optional[Path] = None) -> tuple[bool, str]: |
162 | 200 | """Set ``[features] hooks = true`` in Codex's ``config.toml``. |
163 | 201 |
|
@@ -247,12 +285,8 @@ def parse_hook_payload(self, raw_payload: dict) -> AIHookPayload: |
247 | 285 | mcp_server_name = None |
248 | 286 | mcp_tool_name = None |
249 | 287 | mcp_arguments = None |
250 | | - if tool_name.startswith('mcp__'): |
251 | | - parts = tool_name.split('__') |
252 | | - if len(parts) >= 2: |
253 | | - mcp_server_name = parts[1] |
254 | | - if len(parts) >= 3: |
255 | | - mcp_tool_name = parts[2] |
| 288 | + if canonical_event == AiHookEventType.MCP_EXECUTION: |
| 289 | + mcp_server_name, mcp_tool_name = _split_mcp_tool_name(tool_name) |
256 | 290 | mcp_arguments = tool_input |
257 | 291 |
|
258 | 292 | return AIHookPayload( |
@@ -302,17 +336,17 @@ def build_session_payload(self, raw_payload: dict) -> AIHookPayload: |
302 | 336 | def get_user_email(self) -> Optional[str]: |
303 | 337 | return _email_from_auth() |
304 | 338 |
|
305 | | - def get_session_context(self) -> tuple[Optional[dict], dict]: |
| 339 | + def get_session_context(self) -> tuple[list[dict], dict]: |
306 | 340 | config = _load_codex_config() |
307 | 341 | if not config: |
308 | | - return None, {} |
| 342 | + return [], {} |
309 | 343 | # Codex stores MCP servers under `[mcp_servers.<name>]`; the global config |
310 | 344 | # file becomes its own session-context file. Plugins (via |
311 | 345 | # `[plugins."<plugin>@<marketplace>"]`) carry their own config files. |
312 | 346 | config_path = _codex_config_toml_path('user') |
313 | | - global_config_file = build_global_config_file(config_path, config.get('mcp_servers')) |
| 347 | + global_config_files = build_global_config_files(config_path, config.get('mcp_servers')) |
314 | 348 | enriched_plugins = _resolve_codex_plugins(config) |
315 | | - return global_config_file, enriched_plugins |
| 349 | + return global_config_files, enriched_plugins |
316 | 350 |
|
317 | 351 | def get_skills(self) -> list[dict]: |
318 | 352 | return walk_skill_dirs(_codex_skills_dir()) |
0 commit comments