From cb388f6af22a1f73459a72951e7c0c85e7ef11c9 Mon Sep 17 00:00:00 2001 From: roman-t3a Date: Wed, 15 Jul 2026 22:51:38 +0400 Subject: [PATCH] feat(wrap): add first-class Grok CLI support (#1823) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Adds first-class Grok CLI integration so Headroom can wrap, compress, and learn from Grok sessions the same way it does for Claude Code and Codex. Grok routes inference through `GROK_CLI_CHAT_PROXY_BASE_URL`; Headroom sets that to the local proxy so chat traffic is compressed before forwarding to xAI. MCP retrieval and session learning follow the same patterns as Codex. Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [x] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Add `headroom/providers/grok/` provider slice (`runtime.py`, `install.py`) with `GROK_CLI_CHAT_PROXY_BASE_URL` routing and project attribution prefix - Add `headroom wrap grok` and `headroom unwrap grok` CLI commands (MCP registration, RTK/context-tool setup, proxy launch) - Add `GrokRegistrar` for marker-delimited `[mcp_servers.headroom]` injection in `~/.grok/config.toml` - Add `headroom learn --agent grok` plugin parsing `~/.grok/sessions/*/updates.jsonl` and `GrokWriter` targeting `GROK.md` - Register Grok in install planner, install registry, MCP install list, and agent savings tracking - Update README agent compatibility matrix and unwrap support list - Add unit tests for provider, wrap CLI, MCP registrar, and learn plugin ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ uv run pytest tests/test_provider_grok.py tests/test_cli/test_wrap_grok.py tests/test_mcp_registry_grok.py tests/test_learn_grok_plugin.py -q ============================== 10 passed in 0.15s ============================== $ uv run ruff check headroom/providers/grok headroom/mcp_registry/grok.py headroom/learn/plugins/grok.py headroom/cli/wrap.py headroom/learn/writer.py All checks passed! $ uv run mypy headroom/providers/grok headroom/mcp_registry/grok.py headroom/learn/plugins/grok.py Success: no issues found in 5 source files ``` ## Real Behavior Proof - Environment: macOS (darwin), Python 3.13.14 via `uv`, repo at `/Users/s/dev/headroom`, Grok CLI at `/Users/s/.grok/bin/grok` - Exact command / steps: `cd /Users/s/dev/headroom && uv run pytest tests/test_provider_grok.py tests/test_cli/test_wrap_grok.py tests/test_mcp_registry_grok.py tests/test_learn_grok_plugin.py -q`; `uv run ruff check headroom/providers/grok headroom/mcp_registry/grok.py headroom/learn/plugins/grok.py`; `uv run python -c "from headroom.providers.grok import build_launch_env; env, display = build_launch_env(8787, environ={}); print(display[0])"`; `uv run python -c "from pathlib import Path; import tempfile; from headroom.mcp_registry.grok import GrokRegistrar; from headroom.mcp_registry.install import build_headroom_spec; td=tempfile.mkdtemp(); reg=GrokRegistrar(home_dir=Path(td)); print(reg.register_server(build_headroom_spec('http://127.0.0.1:8787'), force=True).status.value)"` - Observed result: pytest reported `10 passed`; ruff reported `All checks passed!`; `build_launch_env` printed `GROK_CLI_CHAT_PROXY_BASE_URL=http://127.0.0.1:8787/v1`; MCP registrar returned `registered` and wrote the Headroom marker block to `config.toml` - Not tested: live `headroom wrap grok` session with authenticated Grok API traffic through a running Headroom proxy (requires maintainer environment with active `grok login`) ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A — CLI/integration change only. ## Additional Notes - Follows the provider-slice pattern from `b17c6d81` / `93a1f211` (Codex/Cursor/Aider extraction). - Routing uses the session env var only (not `config.toml` endpoint override) so `grok login` session auth continues to work. - Manual E2E wrap/unwrap with real Grok sessions is left for maintainer verification. --------- Co-authored-by: JerrettDavis Co-authored-by: Tejas Chopra --- README.md | 11 +- headroom/cli/learn.py | 2 +- headroom/cli/wrap.py | 212 ++++++++++++++++++++++- headroom/install/models.py | 1 + headroom/install/planner.py | 1 + headroom/learn/__init__.py | 3 +- headroom/learn/plugins/grok.py | 202 ++++++++++++++++++++++ headroom/learn/writer.py | 30 ++++ headroom/mcp_registry/__init__.py | 2 + headroom/mcp_registry/grok.py | 229 +++++++++++++++++++++++++ headroom/mcp_registry/install.py | 3 +- headroom/providers/grok/__init__.py | 5 + headroom/providers/grok/install.py | 11 ++ headroom/providers/grok/runtime.py | 37 ++++ headroom/providers/install_registry.py | 2 + tests/test_cli/test_wrap_grok.py | 64 +++++++ tests/test_learn_grok_plugin.py | 59 +++++++ tests/test_mcp_registry_grok.py | 31 ++++ tests/test_provider_grok.py | 27 +++ 19 files changed, 922 insertions(+), 10 deletions(-) create mode 100644 headroom/learn/plugins/grok.py create mode 100644 headroom/mcp_registry/grok.py create mode 100644 headroom/providers/grok/__init__.py create mode 100644 headroom/providers/grok/install.py create mode 100644 headroom/providers/grok/runtime.py create mode 100644 tests/test_cli/test_wrap_grok.py create mode 100644 tests/test_learn_grok_plugin.py create mode 100644 tests/test_mcp_registry_grok.py create mode 100644 tests/test_provider_grok.py diff --git a/README.md b/README.md index 73cf7fa8e..4d8a52605 100644 --- a/README.md +++ b/README.md @@ -49,10 +49,10 @@ Headroom compresses everything your AI agent reads — tool outputs, logs, RAG c - **Library** — `compress(messages)` in Python or TypeScript, inline in any app - **Proxy** — `headroom proxy --port 8787`, zero code changes, any language -- **Agent wrap** — `headroom wrap claude|codex|copilot|cursor|aider|opencode|cline|continue|goose|openhands|openclaw|vibe|zcode` in one command; undo with `headroom unwrap ` +- **Agent wrap** — `headroom wrap claude|codex|grok|copilot|cursor|aider|opencode|cline|continue|goose|openhands|openclaw|vibe|zcode` in one command; undo with `headroom unwrap ` - **MCP server** — `headroom_compress`, `headroom_retrieve`, `headroom_stats` for any MCP client -- **Cross-agent memory** — shared store across Claude, Codex, Gemini, auto-dedup -- **`headroom learn`** — mines failed sessions, writes corrections to `CLAUDE.local.md` (default, gitignored) or `CLAUDE.md` / `AGENTS.md` / `GEMINI.md` +- **Cross-agent memory** — shared store across Claude, Codex, Gemini, Grok, auto-dedup +- **`headroom learn`** — mines failed sessions, writes corrections to `CLAUDE.local.md` (default, gitignored) or `CLAUDE.md` / `AGENTS.md` / `GEMINI.md` / `GROK.md` - **Output token reduction** — trims what the model *writes back* (not just what you send): drops ceremony/restated code and skips deep "thinking" on routine steps. See [Output token reduction](#output-token-reduction-cut-what-the-model-writes-back). - **Reversible (CCR)** — originals are cached for retrieval on demand @@ -227,6 +227,7 @@ shows an **Output Tokens Saved** card next to input compression, labelled |--------------|:---------------:|----------------------------------| | Claude Code | ✅ | `--memory` · `--code-graph` · `--1m` · `--tool-search` | | Codex | ✅ | shares memory with Claude | +| Grok CLI | ✅ | routes via `GROK_CLI_CHAT_PROXY_BASE_URL` | | Cursor | Manual setup | starts proxy and prints base URLs for Cursor settings | | Aider | ✅ | starts proxy + launches | | Copilot CLI | ✅ | starts proxy + launches | @@ -241,7 +242,7 @@ shows an **Output Tokens Saved** card next to input compression, labelled | ZCode | ✅ | starts proxy and prints base URLs for ZCode settings | Any OpenAI-compatible client works via `headroom proxy`. MCP-native: `headroom mcp install`. -Undo durable wrapping with `headroom unwrap ` (supports: `claude`, `copilot`, `codex`, `opencode`, `openclaw`, `zcode`). +Undo durable wrapping with `headroom unwrap ` (supports: `claude`, `copilot`, `codex`, `grok`, `opencode`, `openclaw`, `zcode`). Registry authors can use the canonical [`server.json`](server.json) in the repo root instead of reconstructing the `headroom mcp serve` contract from prose. ### GitHub Copilot CLI subscription mode @@ -339,7 +340,7 @@ Headroom exposes one stable request lifecycle across `compress()`, the SDK, and Provider and tool-specific behavior lives under `headroom/providers/` so core orchestration stays focused on lifecycle, sequencing, and policy. -- **CLI/tool slices**: `headroom/providers/claude`, `copilot`, `codex`, `openclaw` +- **CLI/tool slices**: `headroom/providers/claude`, `copilot`, `codex`, `grok`, `openclaw` - **Provider runtime slices**: `headroom/providers/claude`, `gemini`, plus shared backend/runtime dispatch in `headroom/providers/registry.py` - **Core files stay orchestration-first**: `wrap.py`, `client.py`, `cli/proxy.py`, and `proxy/server.py` delegate provider-specific env shaping, API target normalization, backend selection, and transport dispatch. diff --git a/headroom/cli/learn.py b/headroom/cli/learn.py index 04cbda325..bcbf51466 100644 --- a/headroom/cli/learn.py +++ b/headroom/cli/learn.py @@ -52,7 +52,7 @@ class _AgentChoice(click.ParamType): _AGENT_HELP = """Which coding agent to analyze. Auto-detects by default. \b -Built-in: claude, codex, gemini. +Built-in: claude, codex, gemini, grok. External plugins register via 'headroom.learn_plugin' entry point. Use 'auto' (default) to scan all detected agents.""" diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 5ebcdf78a..a2af2ecee 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -7,6 +7,7 @@ Usage: headroom wrap aider # Start proxy + aider headroom wrap openclaude # Start proxy + OpenClaude headroom wrap vibe # Start proxy + Mistral Vibe + headroom wrap grok # Start proxy + Grok CLI headroom wrap cursor # Start proxy + print Cursor config instructions headroom wrap openclaw # Install + configure OpenClaw plugin headroom wrap claude --no-context-tool # Without CLI context-tool setup @@ -108,6 +109,7 @@ from headroom.providers.copilot import ( validate_configuration as _validate_copilot_configuration, ) from headroom.providers.cursor import render_setup_lines as _render_cursor_setup_lines +from headroom.providers.grok import build_launch_env as _build_grok_launch_env from headroom.providers.mistral_vibe import build_launch_env as _build_mistral_vibe_launch_env from headroom.providers.openclaw import ( OPENCLAW_NPM_PACKAGE, @@ -179,7 +181,7 @@ _CONTEXT_TOOL_ENV = "HEADROOM_CONTEXT_TOOL" _CONTEXT_TOOL_RTK = "rtk" _CONTEXT_TOOL_LEAN_CTX = "lean-ctx" _VALID_CONTEXT_TOOLS = {_CONTEXT_TOOL_RTK, _CONTEXT_TOOL_LEAN_CTX} -_AGENT_SAVINGS_TARGET_AGENTS = {"claude", "codex", "cursor", "opencode"} +_AGENT_SAVINGS_TARGET_AGENTS = {"claude", "codex", "cursor", "grok", "opencode"} _WRAP_PROXY_TIMEOUT_ENV = "HEADROOM_WRAP_PROXY_TIMEOUT" _WRAP_PROXY_TIMEOUT_DEFAULT_SECONDS = 45 _WRAP_PROXY_TIMEOUT_ML_DEFAULT_SECONDS = 90 @@ -194,7 +196,7 @@ _WRAP_PROXY_TIMEOUT_ML_MODULES = ("torch", "sentence_transformers", "spacy") # `init` and `install` via the Claude provider package to prevent drift. _TOOL_SEARCH_ENV = TOOL_SEARCH_ENV _TOOL_SEARCH_DEFAULT = TOOL_SEARCH_DEFAULT -_AGENT_SAVINGS_WRAP_AGENTS = {"claude", "codex", "cursor"} +_AGENT_SAVINGS_WRAP_AGENTS = {"claude", "codex", "cursor", "grok"} # 1M context window for `wrap claude` (#1158). Claude Code only sends the # `context-1m` beta header — unlocking the 1M window for entitled subscription @@ -3855,6 +3857,7 @@ def wrap() -> None: headroom wrap aider # Aider headroom wrap openclaude # OpenClaude headroom wrap vibe # Mistral Vibe + headroom wrap grok # Grok CLI (xAI) headroom wrap cursor # Cursor (prints config instructions) headroom wrap cline # Cline (VS Code; prints config instructions) headroom wrap continue # Continue (VS Code/JetBrains; injects systemMessage) @@ -5339,6 +5342,149 @@ def vibe( ) +# ============================================================================= +# Grok CLI +# ============================================================================= + + +@wrap.command(context_settings={"ignore_unknown_options": True}) +@click.option( + "--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)" +) +@click.option( + "--no-context-tool", + "--no-rtk", + "no_rtk", + is_flag=True, + help="Skip CLI context-tool setup", +) +@click.option("--no-mcp", is_flag=True, help="Skip headroom MCP server registration") +@click.option( + "--no-tokensave", + is_flag=True, + help="Skip the tokensave code-graph MCP server (primary coding-task compressor)", +) +@click.option( + "--serena", + is_flag=True, + help="Force the Serena MCP backup compressor on (registered automatically when " + "tokensave is unavailable)", +) +@click.option("--no-serena", is_flag=True, help="Never register the Serena backup compressor") +@click.option( + "--code-graph", + is_flag=True, + help="Force a tokensave code-graph index now (tokensave is the default compressor)", +) +@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)") +@click.option("--learn", is_flag=True, help="Enable live traffic learning") +@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory") +@click.option( + "--backend", + default=None, + help="API backend for the proxy: 'anthropic' (default), 'litellm-xai', etc.", +) +@click.option("--anyllm-provider", default=None, help="Provider for any-llm backend") +@click.option("--region", default=None, help="Cloud region for Vertex/Bedrock backends") +@click.option("--verbose", "-v", is_flag=True, help="Verbose output") +@click.option("--prepare-only", is_flag=True, hidden=True) +@click.argument("grok_args", nargs=-1, type=click.UNPROCESSED) +def grok( + port: int, + no_rtk: bool, + no_mcp: bool, + no_tokensave: bool, + serena: bool, + no_serena: bool, + code_graph: bool, + no_proxy: bool, + learn: bool, + memory: bool, + backend: str | None, + anyllm_provider: str | None, + region: str | None, + verbose: bool, + prepare_only: bool, + grok_args: tuple, +) -> None: + """Launch Grok CLI through Headroom proxy. + + \b + Sets ``GROK_CLI_CHAT_PROXY_BASE_URL`` so Grok routes inference traffic + through Headroom. Registers the headroom MCP server in ``~/.grok/config.toml`` + so Grok can call ``headroom_retrieve`` on compression markers. + + \b + Examples: + headroom wrap grok # Start proxy + context tool + grok + headroom wrap grok -- -p "fix the bug" # Pass prompt to grok + headroom wrap grok --no-context-tool # Skip CLI context-tool setup + headroom wrap grok --no-mcp # Skip MCP retrieve tool registration + headroom wrap grok --port 9999 # Custom proxy port + """ + agents_md: Path | None = Path.cwd() / "AGENTS.md" if not no_rtk else None + if not no_rtk: + _setup_context_tool_for_agent( + agent="grok", + agent_display="Grok", + marker_path=agents_md, + on_rtk_ready=lambda _rtk: _inject_rtk_instructions( + cast(Path, agents_md), verbose=verbose + ), + verbose=verbose, + ) + + if not no_mcp: + from headroom.mcp_registry import GrokRegistrar + + _setup_headroom_mcp(GrokRegistrar(), port, verbose=verbose, force=True) + elif verbose: + click.echo(" Skipping MCP retrieve tool (--no-mcp)") + + from headroom.mcp_registry import GrokRegistrar + + _setup_coding_compressor( + GrokRegistrar(), + serena_context="grok", + serena=serena, + no_serena=no_serena, + no_tokensave=no_tokensave, + verbose=verbose, + force=True, + ) + + if prepare_only: + return + + grok_bin = shutil.which("grok") + if not grok_bin: + click.echo("Error: 'grok' not found in PATH.") + click.echo("Install Grok CLI: https://docs.x.ai/docs/grok-cli") + raise SystemExit(1) + + env, env_vars_display = _build_grok_launch_env( + port, os.environ, project=_project_name_from_cwd() + ) + + _launch_tool( + binary=grok_bin, + args=grok_args, + env=env, + port=port, + no_proxy=no_proxy, + tool_label="GROK", + env_vars_display=env_vars_display, + learn=learn, + memory=memory, + agent_type="grok", + code_graph=code_graph, + backend=backend, + anyllm_provider=anyllm_provider, + region=region, + openai_api_url="https://api.x.ai", + ) + + # ============================================================================= # Cursor # ============================================================================= @@ -6758,6 +6904,68 @@ def unwrap_codex(port: int, no_stop_proxy: bool) -> None: click.echo() +# ============================================================================= +# Grok CLI (unwrap) +# ============================================================================= + + +@unwrap.command("grok") +@click.option( + "--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)" +) +@click.option("--no-stop-proxy", is_flag=True, help="Do not stop the local Headroom proxy") +def unwrap_grok(port: int, no_stop_proxy: bool) -> None: + """Undo durable ``headroom wrap grok`` MCP and guidance edits. + + Grok API routing is session-scoped via ``GROK_CLI_CHAT_PROXY_BASE_URL`` and + does not require config restoration. This command removes Headroom MCP + servers from ``~/.grok/config.toml`` and strips injected RTK guidance from + the project ``AGENTS.md``. + """ + click.echo() + click.echo(" ╔═══════════════════════════════════════════════╗") + click.echo(" ║ HEADROOM UNWRAP: GROK ║") + click.echo(" ╚═══════════════════════════════════════════════╝") + click.echo() + + from headroom.mcp_registry import GrokRegistrar + + grok_registrar = GrokRegistrar() + removed_any = False + if grok_registrar.detect(): + tokensave_status = _remove_headroom_installed_tokensave_mcp(grok_registrar) + if tokensave_status == "removed": + click.echo(" Removed Headroom-installed tokensave MCP server from Grok.") + removed_any = True + elif tokensave_status == "failed": + click.echo(" tokensave MCP server matched Headroom ledger but could not be removed.") + + serena_status = _remove_headroom_installed_serena_mcp(grok_registrar) + if serena_status == "removed": + click.echo(" Removed Headroom-installed Serena MCP server from Grok.") + removed_any = True + elif serena_status == "failed": + click.echo(" Serena MCP server matched Headroom ledger but could not be removed.") + + if grok_registrar.unregister_server("headroom"): + click.echo(" Removed Headroom MCP server from Grok config.") + removed_any = True + + if _remove_rtk_instructions(Path.cwd() / "AGENTS.md"): + click.echo(" Removed Headroom rtk instructions from project AGENTS.md.") + removed_any = True + + if not removed_any: + click.echo(" Nothing to undo: no Headroom MCP markers or rtk guidance found.") + + click.echo() + click.echo("✓ Grok is no longer configured for Headroom MCP retrieval.") + click.echo(" Start Grok without `headroom wrap grok` so API traffic skips the proxy.") + if not no_stop_proxy and removed_any: + _echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port) + click.echo() + + # ============================================================================= # ZCode (unwrap) # ============================================================================= diff --git a/headroom/install/models.py b/headroom/install/models.py index dc5f0bd85..3b13ad950 100644 --- a/headroom/install/models.py +++ b/headroom/install/models.py @@ -55,6 +55,7 @@ class ToolTarget(str, Enum): CODEX = "codex" AIDER = "aider" CURSOR = "cursor" + GROK = "grok" OPENCLAW = "openclaw" OPENCODE = "opencode" diff --git a/headroom/install/planner.py b/headroom/install/planner.py index c4ac14179..aba198eac 100644 --- a/headroom/install/planner.py +++ b/headroom/install/planner.py @@ -26,6 +26,7 @@ SUPPORTED_TARGETS = [ ToolTarget.CODEX, ToolTarget.AIDER, ToolTarget.CURSOR, + ToolTarget.GROK, ToolTarget.OPENCLAW, ToolTarget.OPENCODE, ] diff --git a/headroom/learn/__init__.py b/headroom/learn/__init__.py index e13152107..36e1c4ca9 100644 --- a/headroom/learn/__init__.py +++ b/headroom/learn/__init__.py @@ -7,7 +7,8 @@ prevents future token waste. Plugin architecture: plugins/claude.py ─┐ plugins/codex.py ─┤→ Analyzer (LLM) → Writer (adapter) - plugins/gemini.py ─┘ + plugins/gemini.py ─┤ + plugins/grok.py ─┘ Built-in plugins are auto-discovered from headroom.learn.plugins.*. External plugins register via the ``headroom.learn_plugin`` entry point. diff --git a/headroom/learn/plugins/grok.py b/headroom/learn/plugins/grok.py new file mode 100644 index 000000000..c274dcf29 --- /dev/null +++ b/headroom/learn/plugins/grok.py @@ -0,0 +1,202 @@ +"""Grok CLI plugin for headroom learn. + +Reads session logs from ~/.grok/sessions///updates.jsonl. +""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path +from urllib.parse import unquote + +from .._shared import classify_error, is_error_content, normalize_tool_name +from ..base import ConversationScanner, LearnPlugin +from ..models import ErrorCategory, ProjectInfo, SessionData, ToolCall +from ..writer import ContextWriter, GrokWriter + +logger = logging.getLogger(__name__) + + +class GrokPlugin(LearnPlugin, ConversationScanner): + """Reads Grok CLI session logs from ~/.grok/sessions/.""" + + def __init__(self, grok_dir: Path | None = None): + self.grok_dir = grok_dir or Path.home() / ".grok" + self.sessions_dir = self.grok_dir / "sessions" + + @property + def name(self) -> str: + return "grok" + + @property + def display_name(self) -> str: + return "Grok CLI" + + @property + def description(self) -> str: + return "Grok CLI (~/.grok/sessions/)" + + def detect(self) -> bool: + if not self.sessions_dir.exists(): + return False + return any(self.sessions_dir.rglob("updates.jsonl")) + + def create_writer(self) -> ContextWriter: + return GrokWriter() + + def discover_projects(self) -> list[ProjectInfo]: + if not self.sessions_dir.exists(): + return [] + + projects: list[ProjectInfo] = [] + for workspace_dir in sorted(self.sessions_dir.iterdir()): + if not workspace_dir.is_dir(): + continue + session_files = list(workspace_dir.glob("*/updates.jsonl")) + if not session_files: + continue + + decoded = unquote(workspace_dir.name) + project_path = Path(decoded) if decoded.startswith("/") else Path.cwd() + agents_md = project_path / "AGENTS.md" + grok_md = project_path / "GROK.md" + + projects.append( + ProjectInfo( + name=workspace_dir.name, + project_path=project_path, + data_path=workspace_dir, + context_file=grok_md + if grok_md.exists() + else agents_md + if agents_md.exists() + else None, + ) + ) + return projects + + def scan_project( + self, project: ProjectInfo, max_workers: int = 1, include_subagents: bool = True + ) -> list[SessionData]: + del include_subagents + session_files = sorted(project.data_path.glob("*/updates.jsonl")) + if not session_files: + return [] + + if max_workers <= 1 or len(session_files) <= 1: + return [s for f in session_files if (s := self._scan_session(f)) and s.tool_calls] + + from concurrent.futures import ThreadPoolExecutor, as_completed + + sessions: list[SessionData] = [] + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = {executor.submit(self._scan_session, f): f for f in session_files} + for future in as_completed(futures): + session = future.result() + if session and session.tool_calls: + sessions.append(session) + return sessions + + def _scan_session(self, jsonl_path: Path) -> SessionData | None: + session_id = jsonl_path.parent.name + pending_calls: dict[str, tuple[str, dict]] = {} + tool_calls: list[ToolCall] = [] + msg_index = 0 + + try: + with open(jsonl_path, encoding="utf-8", errors="replace") as f: + for line in f: + try: + entry = json.loads(line) + except json.JSONDecodeError: + continue + + params = entry.get("params", {}) + if not isinstance(params, dict): + continue + update = params.get("update", {}) + if not isinstance(update, dict): + continue + + session_update = update.get("sessionUpdate", "") + if session_update == "tool_call": + msg_index += 1 + call_id = str(update.get("toolCallId", "")) + title = str(update.get("title", "tool")) + raw_input = update.get("rawInput", {}) + if not isinstance(raw_input, dict): + raw_input = {"raw": raw_input} + name = normalize_tool_name(title) + if call_id: + pending_calls[call_id] = (name, raw_input) + continue + + if session_update != "tool_call_update": + continue + + status = update.get("status") + if status not in ("completed", "failed"): + continue + + call_id = str(update.get("toolCallId", "")) + if call_id not in pending_calls: + continue + + msg_index += 1 + name, inp = pending_calls[call_id] + result_content = _extract_tool_output(update) + is_err = status == "failed" or is_error_content(result_content) + error_cat = classify_error(result_content) if is_err else ErrorCategory.UNKNOWN + + tool_calls.append( + ToolCall( + name=name, + tool_call_id=call_id, + input_data=inp, + output=result_content, + is_error=is_err, + error_category=error_cat, + msg_index=msg_index, + output_bytes=len(result_content.encode("utf-8")), + ) + ) + except OSError as exc: + logger.debug("Failed to read Grok session %s: %s", jsonl_path, exc) + return None + + if not tool_calls: + return None + return SessionData(session_id=session_id, tool_calls=tool_calls) + + +def _extract_tool_output(update: dict) -> str: + raw_output = update.get("rawOutput") + if isinstance(raw_output, dict): + for key in ("output_for_prompt", "output", "FileContent"): + value = raw_output.get(key) + if isinstance(value, str) and value.strip(): + return value + if isinstance(value, dict): + content = value.get("content") + if isinstance(content, str) and content.strip(): + return content + + content_blocks = update.get("content") + if isinstance(content_blocks, list): + parts: list[str] = [] + for block in content_blocks: + if not isinstance(block, dict): + continue + inner = block.get("content") + if isinstance(inner, dict): + text = inner.get("text") + if isinstance(text, str): + parts.append(text) + if parts: + return "\n".join(parts) + + return "" + + +plugin = GrokPlugin() diff --git a/headroom/learn/writer.py b/headroom/learn/writer.py index 3a5e8a03b..6fef4362c 100644 --- a/headroom/learn/writer.py +++ b/headroom/learn/writer.py @@ -408,3 +408,33 @@ class GeminiWriter(ContextWriter): gemini_md.write_text(full_content, encoding="utf-8") return result + + +# ============================================================================= +# Grok Writer (Grok CLI) +# ============================================================================= + + +class GrokWriter(ContextWriter): + """Writes learned patterns to GROK.md for Grok CLI.""" + + def write( + self, + recommendations: list[Recommendation], + project: ProjectInfo, + dry_run: bool = True, + ) -> WriteResult: + result = WriteResult() + result.dry_run = dry_run + + if not recommendations: + return result + + grok_md = project.context_file or (project.project_path / "GROK.md") + full_content = _merge_into_file(grok_md, recommendations) + result.add(grok_md, full_content) + if not dry_run: + grok_md.parent.mkdir(parents=True, exist_ok=True) + grok_md.write_text(full_content, encoding="utf-8") + + return result diff --git a/headroom/mcp_registry/__init__.py b/headroom/mcp_registry/__init__.py index e812359f2..c9aec0635 100644 --- a/headroom/mcp_registry/__init__.py +++ b/headroom/mcp_registry/__init__.py @@ -17,6 +17,7 @@ from .base import MCPRegistrar, RegisterResult, RegisterStatus, ServerSpec from .claude import ClaudeRegistrar from .codex import CodexRegistrar from .display import any_succeeded, format_result, format_results +from .grok import GrokRegistrar from .install import ( DEFAULT_PROXY_URL, build_headroom_spec, @@ -32,6 +33,7 @@ __all__ = [ "DEFAULT_PROXY_URL", "ClaudeRegistrar", "CodexRegistrar", + "GrokRegistrar", "MCPRegistrar", "OpencodeRegistrar", "RegisterResult", diff --git a/headroom/mcp_registry/grok.py b/headroom/mcp_registry/grok.py new file mode 100644 index 000000000..37ee764c8 --- /dev/null +++ b/headroom/mcp_registry/grok.py @@ -0,0 +1,229 @@ +"""Grok CLI MCP registrar. + +Grok stores MCP server config in ``$GROK_HOME/config.toml`` (default +``~/.grok/config.toml``) as ``[mcp_servers.]`` tables. There is no +general-purpose CLI for adding entries, so we edit the file in place using +marker-delimited blocks so we can idempotently inject, replace, and remove +our entry without disturbing anything else the user has configured. +""" + +from __future__ import annotations + +import logging +import os +import sys +from pathlib import Path +from typing import Any + +from headroom import fsutil + +from .base import MCPRegistrar, RegisterResult, RegisterStatus, ServerSpec + +if sys.version_info >= (3, 11): + import tomllib +else: # pragma: no cover — exercised only on 3.10 + import tomli as tomllib # type: ignore[no-redef] + +logger = logging.getLogger(__name__) + +_MARKER_START = "# --- Headroom MCP server ---" +_MARKER_END = "# --- end Headroom MCP server ---" + + +def _marker_start(server_name: str) -> str: + if server_name == "headroom": + return _MARKER_START + return f"# --- Headroom MCP server: {server_name} ---" + + +def _marker_end(server_name: str) -> str: + if server_name == "headroom": + return _MARKER_END + return f"# --- end Headroom MCP server: {server_name} ---" + + +class GrokRegistrar(MCPRegistrar): + """Register MCP servers with the Grok CLI.""" + + name = "grok" + display_name = "Grok CLI" + + def __init__(self, *, home_dir: Path | None = None) -> None: + if home_dir is not None: + self._grok_dir = home_dir / ".grok" + elif os.environ.get("GROK_HOME"): + self._grok_dir = Path(os.environ["GROK_HOME"]).expanduser() + else: + self._grok_dir = Path.home() / ".grok" + self._config_file = self._grok_dir / "config.toml" + + def detect(self) -> bool: + return self._grok_dir.is_dir() + + def get_server(self, server_name: str) -> ServerSpec | None: + data = self._load_toml() + servers = data.get("mcp_servers", {}) + if not isinstance(servers, dict): + return None + entry = servers.get(server_name) + if not isinstance(entry, dict): + return None + return _entry_to_spec(server_name, entry) + + def register_server(self, spec: ServerSpec, *, force: bool = False) -> RegisterResult: + existing = self.get_server(spec.name) + + if existing is not None and _specs_equivalent(existing, spec): + return RegisterResult(RegisterStatus.ALREADY, "matches current configuration") + + if existing is not None and not force: + content = self._read_text() + if _marker_start(spec.name) not in content: + return RegisterResult( + RegisterStatus.MISMATCH, + "user-managed [mcp_servers." + f"{spec.name}] entry outside Headroom markers; " + f"{_diff_specs(existing, spec)}", + ) + return RegisterResult(RegisterStatus.MISMATCH, _diff_specs(existing, spec)) + + if existing is not None and force: + content = self._read_text() + if _marker_start(spec.name) not in content: + return RegisterResult( + RegisterStatus.MISMATCH, + "user-managed [mcp_servers." + f"{spec.name}] entry outside Headroom markers; " + f"{_diff_specs(existing, spec)}", + ) + self.unregister_server(spec.name) + + return self._write_block(spec) + + def unregister_server(self, server_name: str) -> bool: + if not self._config_file.exists(): + return False + content = self._read_text() + marker_start = _marker_start(server_name) + marker_end = _marker_end(server_name) + if marker_start not in content or marker_end not in content: + return False + try: + start = content.index(marker_start) + end = content.index(marker_end) + len(marker_end) + except ValueError: + return False + before = content[:start].rstrip("\n") + after = content[end:].lstrip("\n") + if before and after: + new_content = before + "\n\n" + after + else: + new_content = (before or after).rstrip("\n") + ("\n" if (before or after) else "") + try: + fsutil.write_text(self._config_file, new_content) + except OSError: + return False + return True + + def _load_toml(self) -> dict[str, Any]: + if not self._config_file.exists(): + return {} + try: + data = tomllib.loads(fsutil.read_text(self._config_file)) + except (tomllib.TOMLDecodeError, OSError): + return {} + return data if isinstance(data, dict) else {} + + def _read_text(self) -> str: + return fsutil.read_text(self._config_file, default="") + + def _write_block(self, spec: ServerSpec) -> RegisterResult: + block = _render_block(spec) + try: + self._grok_dir.mkdir(parents=True, exist_ok=True) + content = self._read_text() + marker_start = _marker_start(spec.name) + marker_end = _marker_end(spec.name) + if marker_start in content and marker_end in content: + start = content.index(marker_start) + end = content.index(marker_end) + len(marker_end) + content = ( + content[:start].rstrip("\n") + + ("\n\n" if content[:start].rstrip("\n") else "") + + block + + "\n" + + content[end:].lstrip("\n") + ) + elif content.strip(): + content = content.rstrip("\n") + "\n\n" + block + "\n" + else: + content = block + "\n" + fsutil.write_text(self._config_file, content) + except OSError as exc: + return RegisterResult( + RegisterStatus.FAILED, f"could not write {self._config_file}: {exc}" + ) + return RegisterResult(RegisterStatus.REGISTERED, f"wrote to {self._config_file}") + + +def _render_block(spec: ServerSpec) -> str: + lines: list[str] = [ + _marker_start(spec.name), + f"[mcp_servers.{spec.name}]", + f"command = {_toml_str(spec.command)}", + ] + if spec.args: + items = ", ".join(_toml_str(a) for a in spec.args) + lines.append(f"args = [{items}]") + if spec.env: + lines.append("") + lines.append(f"[mcp_servers.{spec.name}.env]") + for k, v in spec.env.items(): + lines.append(f"{k} = {_toml_str(v)}") + lines.append(_marker_end(spec.name)) + return "\n".join(lines) + + +def _toml_str(s: str) -> str: + escaped = s.replace("\\", "\\\\").replace('"', '\\"') + return f'"{escaped}"' + + +def _entry_to_spec(name: str, entry: dict[str, Any]) -> ServerSpec: + args_value = entry.get("args", []) + if isinstance(args_value, list): + args = tuple(str(x) for x in args_value) + else: + args = () + env_value = entry.get("env", {}) + env: dict[str, str] = {} + if isinstance(env_value, dict): + env = {str(k): str(v) for k, v in env_value.items()} + return ServerSpec( + name=name, + command=str(entry.get("command", "")), + args=args, + env=env, + ) + + +def _specs_equivalent(a: ServerSpec, b: ServerSpec) -> bool: + return ( + a.name == b.name + and a.command == b.command + and tuple(a.args) == tuple(b.args) + and dict(a.env) == dict(b.env) + ) + + +def _diff_specs(existing: ServerSpec, requested: ServerSpec) -> str: + parts: list[str] = [] + if existing.command != requested.command: + parts.append(f"command {existing.command!r} -> {requested.command!r}") + if tuple(existing.args) != tuple(requested.args): + parts.append(f"args {list(existing.args)} -> {list(requested.args)}") + if dict(existing.env) != dict(requested.env): + parts.append(f"env {dict(existing.env)} -> {dict(requested.env)}") + if not parts: + return "spec differs in unidentified field(s)" + return "; ".join(parts) diff --git a/headroom/mcp_registry/install.py b/headroom/mcp_registry/install.py index 49834e1b8..927d90268 100644 --- a/headroom/mcp_registry/install.py +++ b/headroom/mcp_registry/install.py @@ -9,6 +9,7 @@ from headroom.install.runtime import resolve_headroom_command from .base import MCPRegistrar, RegisterResult, RegisterStatus, ServerSpec from .claude import ClaudeRegistrar from .codex import CodexRegistrar +from .grok import GrokRegistrar from .opencode import OpencodeRegistrar #: Default proxy URL used when none is given. @@ -20,7 +21,7 @@ def get_all_registrars() -> list[MCPRegistrar]: The list grows as we add adapters for Cursor, Continue, Cline, etc. """ - return [ClaudeRegistrar(), CodexRegistrar(), OpencodeRegistrar()] + return [ClaudeRegistrar(), CodexRegistrar(), GrokRegistrar(), OpencodeRegistrar()] def build_headroom_spec(proxy_url: str = DEFAULT_PROXY_URL) -> ServerSpec: diff --git a/headroom/providers/grok/__init__.py b/headroom/providers/grok/__init__.py new file mode 100644 index 000000000..9a3abaf3c --- /dev/null +++ b/headroom/providers/grok/__init__.py @@ -0,0 +1,5 @@ +"""Grok CLI provider helpers.""" + +from .runtime import DEFAULT_API_URL, PROXY_ENV_KEY, build_launch_env, proxy_base_url + +__all__ = ["DEFAULT_API_URL", "PROXY_ENV_KEY", "build_launch_env", "proxy_base_url"] diff --git a/headroom/providers/grok/install.py b/headroom/providers/grok/install.py new file mode 100644 index 000000000..1e9239613 --- /dev/null +++ b/headroom/providers/grok/install.py @@ -0,0 +1,11 @@ +"""Grok install-time helpers.""" + +from __future__ import annotations + +from .runtime import PROXY_ENV_KEY, proxy_base_url + + +def build_install_env(*, port: int, backend: str) -> dict[str, str]: + """Build the persistent install environment for Grok CLI.""" + del backend + return {PROXY_ENV_KEY: proxy_base_url(port)} diff --git a/headroom/providers/grok/runtime.py b/headroom/providers/grok/runtime.py new file mode 100644 index 000000000..eb9480312 --- /dev/null +++ b/headroom/providers/grok/runtime.py @@ -0,0 +1,37 @@ +"""Runtime helpers for Grok CLI integrations.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping + +from headroom.proxy.project_context import with_project_prefix + +DEFAULT_API_URL = "https://api.x.ai" +PROXY_ENV_KEY = "GROK_CLI_CHAT_PROXY_BASE_URL" + + +def proxy_base_url(port: int) -> str: + """Return the local proxy base URL used by Grok CLI integrations.""" + return f"http://127.0.0.1:{port}/v1" + + +def build_launch_env( + port: int, + environ: Mapping[str, str] | None = None, + project: str | None = None, +) -> tuple[dict[str, str], list[str]]: + """Build environment variables for Grok CLI through the local proxy. + + Grok routes inference traffic through ``GROK_CLI_CHAT_PROXY_BASE_URL`` + when set (see Grok CLI auth/custom-models docs). The proxy forwards + OpenAI-compatible chat requests upstream to xAI. + + ``project`` (the wrap launch directory) is encoded as a ``/p/`` + base-URL prefix because Grok cannot send custom attribution headers; + the proxy strips it and attributes savings per project. + """ + env = dict(environ or os.environ) + base_url = with_project_prefix(proxy_base_url(port), project) + env[PROXY_ENV_KEY] = base_url + return env, [f"{PROXY_ENV_KEY}={base_url}"] diff --git a/headroom/providers/install_registry.py b/headroom/providers/install_registry.py index ba7a191f3..130126c59 100644 --- a/headroom/providers/install_registry.py +++ b/headroom/providers/install_registry.py @@ -29,6 +29,7 @@ from headroom.providers.cortex_code.install import ( build_install_env as _build_cortex_code_install_env, ) from headroom.providers.cursor.install import build_install_env as _build_cursor_install_env +from headroom.providers.grok.install import build_install_env as _build_grok_install_env from headroom.providers.openclaw.install import ( apply_provider_scope as _apply_openclaw_provider_scope, ) @@ -54,6 +55,7 @@ _ENV_BUILDERS: dict[str, _InstallEnvBuilder] = { "aider": _build_aider_install_env, "cortex-code": _build_cortex_code_install_env, "cursor": _build_cursor_install_env, + "grok": _build_grok_install_env, "opencode": _build_opencode_install_env, } diff --git a/tests/test_cli/test_wrap_grok.py b/tests/test_cli/test_wrap_grok.py new file mode 100644 index 000000000..1cd278809 --- /dev/null +++ b/tests/test_cli/test_wrap_grok.py @@ -0,0 +1,64 @@ +"""Tests for `headroom wrap grok` command.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import patch +from urllib.parse import quote + +import pytest +from click.testing import CliRunner + +from headroom.cli.main import main +from headroom.providers.grok import PROXY_ENV_KEY + + +def _expected_project_prefix() -> str: + return f"/p/{quote(Path.cwd().name, safe='')}" + + +@pytest.fixture +def runner() -> CliRunner: + return CliRunner() + + +def test_wrap_grok_sets_proxy_env( + runner: CliRunner, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + captured: dict[str, object] = {} + + def fake_launch_tool(**kwargs): # noqa: ANN003 + captured.update(kwargs) + + with patch("headroom.cli.wrap.shutil.which", return_value="grok"): + with patch("headroom.cli.wrap._setup_context_tool_for_agent"): + with patch("headroom.cli.wrap._setup_headroom_mcp"): + with patch("headroom.cli.wrap._setup_coding_compressor"): + with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool): + result = runner.invoke( + main, ["wrap", "grok", "--no-rtk", "--no-mcp", "--", "-p", "hello"] + ) + + assert result.exit_code == 0, result.output + env = captured["env"] + assert isinstance(env, dict) + assert env[PROXY_ENV_KEY] == f"http://127.0.0.1:8787{_expected_project_prefix()}/v1" + assert captured["tool_label"] == "GROK" + assert captured["agent_type"] == "grok" + assert captured["args"] == ("-p", "hello") + + +def test_wrap_grok_missing_binary_exits( + runner: CliRunner, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + + with patch("headroom.cli.wrap.shutil.which", return_value=None): + with patch("headroom.cli.wrap._setup_context_tool_for_agent"): + with patch("headroom.cli.wrap._setup_headroom_mcp"): + with patch("headroom.cli.wrap._setup_coding_compressor"): + result = runner.invoke(main, ["wrap", "grok", "--no-rtk", "--no-mcp"]) + + assert result.exit_code == 1 + assert "grok" in result.output.lower() diff --git a/tests/test_learn_grok_plugin.py b/tests/test_learn_grok_plugin.py new file mode 100644 index 000000000..3d774b5b7 --- /dev/null +++ b/tests/test_learn_grok_plugin.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import json +from pathlib import Path + +from headroom.learn.plugins.grok import GrokPlugin + + +def test_grok_plugin_detects_updates_jsonl(tmp_path: Path) -> None: + grok_dir = tmp_path / ".grok" + session_dir = grok_dir / "sessions" / "%2Ftmp%2Fproject" / "session-1" + session_dir.mkdir(parents=True) + (session_dir / "updates.jsonl").write_text("{}\n", encoding="utf-8") + + plugin = GrokPlugin(grok_dir=grok_dir) + + assert plugin.detect() is True + + +def test_grok_plugin_scans_tool_calls(tmp_path: Path) -> None: + grok_dir = tmp_path / ".grok" + session_dir = grok_dir / "sessions" / "%2Ftmp%2Fproject" / "session-1" + session_dir.mkdir(parents=True) + + lines = [ + { + "params": { + "update": { + "sessionUpdate": "tool_call", + "toolCallId": "call-1", + "title": "Shell", + "rawInput": {"command": "false"}, + } + } + }, + { + "params": { + "update": { + "sessionUpdate": "tool_call_update", + "toolCallId": "call-1", + "status": "failed", + "rawOutput": {"output_for_prompt": "Exit code: 1"}, + } + } + }, + ] + (session_dir / "updates.jsonl").write_text( + "\n".join(json.dumps(line) for line in lines) + "\n", + encoding="utf-8", + ) + + plugin = GrokPlugin(grok_dir=grok_dir) + projects = plugin.discover_projects() + assert len(projects) == 1 + + sessions = plugin.scan_project(projects[0]) + assert len(sessions) == 1 + assert len(sessions[0].tool_calls) == 1 + assert sessions[0].tool_calls[0].is_error is True diff --git a/tests/test_mcp_registry_grok.py b/tests/test_mcp_registry_grok.py new file mode 100644 index 000000000..2544d8b5c --- /dev/null +++ b/tests/test_mcp_registry_grok.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from pathlib import Path + +from headroom.mcp_registry.grok import GrokRegistrar +from headroom.mcp_registry.install import build_headroom_spec + + +def test_grok_registrar_writes_marker_block(tmp_path: Path) -> None: + registrar = GrokRegistrar(home_dir=tmp_path) + spec = build_headroom_spec("http://127.0.0.1:9999") + + result = registrar.register_server(spec, force=True) + + assert result.status.value == "registered" + content = registrar._config_file.read_text(encoding="utf-8") + assert "[mcp_servers.headroom]" in content + assert "# --- Headroom MCP server ---" in content + + +def test_grok_registrar_unregister_removes_marker_block(tmp_path: Path) -> None: + registrar = GrokRegistrar(home_dir=tmp_path) + spec = build_headroom_spec() + registrar.register_server(spec, force=True) + + removed = registrar.unregister_server("headroom") + + assert removed is True + assert ( + not registrar._config_file.exists() or "Headroom MCP server" not in registrar._read_text() + ) diff --git a/tests/test_provider_grok.py b/tests/test_provider_grok.py new file mode 100644 index 000000000..08793332d --- /dev/null +++ b/tests/test_provider_grok.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from headroom.providers.grok import PROXY_ENV_KEY, build_launch_env, proxy_base_url +from headroom.providers.grok.install import build_install_env + + +def test_grok_proxy_base_url_uses_local_headroom_proxy() -> None: + assert proxy_base_url(8787) == "http://127.0.0.1:8787/v1" + + +def test_grok_build_launch_env_sets_chat_proxy_base_url() -> None: + env, display = build_launch_env(9999, environ={}) + + assert env[PROXY_ENV_KEY] == "http://127.0.0.1:9999/v1" + assert display == [f"{PROXY_ENV_KEY}=http://127.0.0.1:9999/v1"] + + +def test_grok_build_launch_env_applies_project_prefix() -> None: + env, _display = build_launch_env(8787, environ={}, project="frontend") + + assert env[PROXY_ENV_KEY] == "http://127.0.0.1:8787/p/frontend/v1" + + +def test_grok_build_install_env_returns_proxy_url() -> None: + assert build_install_env(port=7654, backend="ignored") == { + PROXY_ENV_KEY: "http://127.0.0.1:7654/v1", + }