From 420dc9077b204c1ad75bde25f3f23ff8c11db770 Mon Sep 17 00:00:00 2001 From: Aashish Tamsya Date: Thu, 16 Jul 2026 05:51:52 +0900 Subject: [PATCH] feat(grok-build): add Grok Build wrap command and MCP integration (#1629) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Adds first-class Grok Build support to Headroom so Grok CLI sessions can route through the local proxy for context compression and savings tracking. This PR introduces `headroom wrap grok-build` / `headroom unwrap grok-build`, a `grok_build` provider slice, Grok MCP registrar support, and install/telemetry wiring so Grok traffic is attributed correctly in the proxy and dashboard. Review follow-up (`9368c413`): when users already own `[model.grok-build]` in `~/.grok/config.toml`, wrap rewrites `base_url` in that table in place instead of appending a duplicate header (invalid TOML). ## Type of Change - [x] New feature (non-breaking change that adds functionality) ## Changes Made - Added `headroom/providers/grok_build/` with runtime helpers, reversible `~/.grok/config.toml` injection, and install env builders. - Added `headroom wrap grok-build` and `headroom unwrap grok-build` CLI commands. - Added `GrokRegistrar` for Headroom MCP registration in Grok config. - Wired `grok_build` into install planner/registry, agent savings, telemetry, and proxy client detection (`grok/` user agent). - **Review fix:** rewrite `base_url` inside an existing user-owned `[model.grok-build]` table in place (`# was: …` metadata). - Added regression tests + docs (`grok-build.mdx`, `proxy.mdx`) and CHANGELOG entry. ## Testing - [x] Unit tests pass (`pytest`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ pytest -q tests/test_provider_grok_build.py tests/test_mcp_registry/test_grok_registrar.py ============================== 12 passed in 1.13s ============================== ``` See **Screenshots** below for terminal captures (pytest, review-fix in-place rewrite, proxy `/readyz`, unwrap). ## Real Behavior Proof - Environment: macOS, Python 3.11.12 venv, feat/grok-build @ `9368c413`, isolated `GROK_HOME` temp dirs, proxy port 8799 - Exact command / steps: see screenshot evidence (wrap/unwrap, in-place table rewrite, `/readyz`) - Observed result: see screenshots — 12 tests pass; single `[model.grok-build]` table after wrap on pre-existing config; proxy healthy; unwrap restores backup - Not tested: Live interactive Grok chat with xAI auth through the proxy ## 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 - [ ] 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 - [x] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) Terminal captures from local verification (`9368c413`). Assets hosted on fork prerelease only — **not** in the source tree. **1. Pytest — 12 passed (incl. review-fix regression)** ![pytest 12 passed](https://github.com/aashishtamsya/headroom/releases/download/pr-1629-evidence/01-pytest.png) **2. Review fix — in-place `[model.grok-build]` rewrite (single table, `# was:` metadata)** ![review fix in-place rewrite](https://github.com/aashishtamsya/headroom/releases/download/pr-1629-evidence/02-review-fix-in-place.png) **3. Proxy health — `/readyz` healthy on port 8799** ![proxy readyz healthy](https://github.com/aashishtamsya/headroom/releases/download/pr-1629-evidence/03-proxy-health.png) **4. Unwrap — restores pre-wrap backup** ![unwrap restores backup](https://github.com/aashishtamsya/headroom/releases/download/pr-1629-evidence/04-unwrap.png) ## Additional Notes Screenshot assets: https://github.com/aashishtamsya/headroom/releases/tag/pr-1629-evidence (temporary prerelease; safe to delete after merge). --------- Co-authored-by: JerrettDavis --- .gitignore | 3 + CHANGELOG.md | 1 + docs/content/docs/grok-build.mdx | 62 +++++++ docs/content/docs/meta.json | 1 + docs/content/docs/proxy.mdx | 7 + headroom/cli/install.py | 4 +- headroom/cli/wrap.py | 148 ++++++++++++++++- headroom/install/models.py | 1 + headroom/install/planner.py | 4 + headroom/providers/grok_build/__init__.py | 5 + headroom/providers/grok_build/config.py | 153 ++++++++++++++++++ headroom/providers/grok_build/install.py | 14 ++ headroom/providers/grok_build/runtime.py | 57 +++++++ headroom/providers/install_registry.py | 2 + headroom/proxy/auth_policy.py | 2 + headroom/proxy/server.py | 1 + headroom/telemetry/context.py | 12 +- pyproject.toml | 2 +- tests/test_auth_policy.py | 7 + tests/test_cli/test_wrap_bridge.py | 21 +++ .../test_mcp_registry/test_grok_registrar.py | 43 +++++ tests/test_provider_grok_build.py | 130 +++++++++++++++ 22 files changed, 675 insertions(+), 5 deletions(-) create mode 100644 docs/content/docs/grok-build.mdx create mode 100644 headroom/providers/grok_build/__init__.py create mode 100644 headroom/providers/grok_build/config.py create mode 100644 headroom/providers/grok_build/install.py create mode 100644 headroom/providers/grok_build/runtime.py create mode 100644 tests/test_mcp_registry/test_grok_registrar.py create mode 100644 tests/test_provider_grok_build.py diff --git a/.gitignore b/.gitignore index ebba5f120..4f7fcc7db 100644 --- a/.gitignore +++ b/.gitignore @@ -234,6 +234,9 @@ pyrightconfig.json \#*\# .\#* +# Local git worktrees (isolated feature branches) +.worktrees/ + # Local development configuration CLAUDE.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 2028457f0..7d96a5d4b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -109,6 +109,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Features +* **grok-build:** add first-class Grok Build support — `headroom wrap grok-build` / `headroom unwrap grok-build`, reversible `~/.grok/config.toml` injection (in-place `base_url` rewrite when `[model.grok-build]` already exists), `GrokRegistrar` MCP install, and install/telemetry wiring ([#1629](https://github.com/headroomlabs-ai/headroom/pull/1629)). * **wrap:** add `headroom wrap omp` / `headroom unwrap omp` for Oh My Pi — points omp's built-in `anthropic` provider at the local proxy via a marker-fenced `providers.anthropic.baseUrl` override in `~/.omp/agent/models.yml`, snapshotting the pre-wrap file byte-for-byte and restoring it on unwrap. omp resolves its Anthropic chat endpoint from models.yml (`ANTHROPIC_BASE_URL` only feeds its web-search helper), and a same-ID override keeps omp's bundled model catalog and stored credentials ([#1149](https://github.com/headroomlabs-ai/headroom/issues/1149)) * **compress:** expose `frozen_message_count` in library-mode `compress()` via a new `CompressConfig` field (default `0`, unchanged behavior). `read_lifecycle.apply()` already skips stale-Read replacements inside a frozen message prefix, but only the proxy handlers could pass it — `ContentRouter` reads it from transform kwargs and the public API never forwarded it. Library-mode callers that manage their own conversation loop can now stop transforms from rewriting messages already anchored in the provider's prompt cache, which would otherwise convert 0.1x cached prefix reads into full-price cache writes ([#2178](https://github.com/headroomlabs-ai/headroom/pull/2178)). diff --git a/docs/content/docs/grok-build.mdx b/docs/content/docs/grok-build.mdx new file mode 100644 index 000000000..e11b658c0 --- /dev/null +++ b/docs/content/docs/grok-build.mdx @@ -0,0 +1,62 @@ +--- +title: Grok Build Integration +description: Route Grok Build traffic through Headroom for token compression and per-project savings attribution. +--- + +Use `headroom wrap grok-build` to route Grok Build LLM traffic through the local Headroom proxy. The wrapper starts or reuses the proxy, injects a reversible `[model.grok-build]` override into `~/.grok/config.toml` (or `$GROK_HOME/config.toml`), optionally sets up RTK or `lean-ctx`, and prints next steps for launching `grok`. + +## Quick Start + +```bash +headroom wrap grok-build +``` + +In another terminal, from the same project directory: + +```bash +grok +``` + +When you are done: + +```bash +headroom unwrap grok-build +``` + +## What `wrap grok-build` Does + +| Step | What happens | +|---|---| +| Proxy | Starts the Headroom proxy unless `--no-proxy` is set | +| Model config | Writes or updates `[model.grok-build] base_url` in Grok's `config.toml`, pointing at `http://127.0.0.1:/v1` (with optional `/p/` prefix for savings attribution) | +| Existing config | If you already have a `[model.grok-build]` table, Headroom rewrites `base_url` in place instead of appending a duplicate table (invalid TOML) | +| Context tool | Injects RTK or `lean-ctx` guidance into project `AGENTS.md` unless `--no-context-tool` is set | +| MCP install | `headroom mcp install` can register Headroom MCP via `GrokRegistrar` | +| Backup | Snapshots `config.toml` to `config.toml.headroom-backup` before the first injection | + +## Options + +```bash +headroom wrap grok-build \ + --port 8787 \ # Proxy port (default: 8787) + --no-context-tool \ # Skip RTK / lean-ctx setup + --no-proxy \ # Use an existing proxy instead of starting one + --learn \ # Enable live traffic learning + --memory # Enable persistent memory +``` + +## Environment Variables + +| Variable | Description | +|---|---| +| `GROK_HOME` | Override Grok config directory (default: `~/.grok`) | +| `HEADROOM_CONTEXT_TOOL` | Set to `lean-ctx` to use lean-ctx instead of RTK | +| `XAI_API_KEY` | Grok API key (also accepts `GROK_CODE_XAI_API_KEY`) | + +## Persistent Install + +`grok_build` is an install target for `headroom install apply --providers manual --target grok_build`. The install manifest records proxy env values for Grok Build alongside other wrapped agents. + +## Unwrap + +`headroom unwrap grok-build` restores the pre-wrap `config.toml` from backup when available, or strips Headroom marker blocks and in-place `base_url` rewrites when no backup exists. \ No newline at end of file diff --git a/docs/content/docs/meta.json b/docs/content/docs/meta.json index 2089afcff..430384417 100644 --- a/docs/content/docs/meta.json +++ b/docs/content/docs/meta.json @@ -37,6 +37,7 @@ "claude-code-vertex", "claude-code-azure-foundry", "opencode", + "grok-build", "mcp", "---Configuration---", "configuration", diff --git a/docs/content/docs/proxy.mdx b/docs/content/docs/proxy.mdx index c19861be5..f70e674c0 100644 --- a/docs/content/docs/proxy.mdx +++ b/docs/content/docs/proxy.mdx @@ -320,12 +320,19 @@ headroom wrap aider # Cursor (starts the proxy and prints settings to paste into Cursor) headroom wrap cursor + +# Grok Build (updates ~/.grok/config.toml and starts the proxy) +headroom wrap grok-build ``` Cursor reads model endpoints from its settings UI, so `headroom wrap cursor` does not rewrite Cursor configuration or launch the app. After it starts the proxy, copy the printed base URL into Cursor's model settings. +Grok Build reads model endpoints from `~/.grok/config.toml`. `headroom wrap grok-build` +injects or updates `[model.grok-build] base_url` to point at the local proxy, then +run `grok` from the same project directory. See [Grok Build Integration](/docs/grok-build). + For environment-driven clients, you can also set the base URL manually: ```bash diff --git a/headroom/cli/install.py b/headroom/cli/install.py index f2b2c7189..0c0b6bae5 100644 --- a/headroom/cli/install.py +++ b/headroom/cli/install.py @@ -389,7 +389,9 @@ def _echo_installed(manifest: DeploymentManifest, *, prefix: str = "Installed pe "--target", "targets", multiple=True, - type=click.Choice(["claude", "copilot", "codex", "aider", "cursor", "openclaw", "opencode"]), + type=click.Choice( + ["claude", "copilot", "codex", "aider", "cursor", "grok_build", "openclaw", "opencode"] + ), help="Tool target to configure when --providers manual is used.", ) @click.option("--profile", default="default", show_default=True, help="Deployment profile name.") diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 09f4d7603..3a53bcfd1 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -9,6 +9,7 @@ Usage: 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 grok-build # Start proxy + configure Grok Build headroom wrap openclaw # Install + configure OpenClaw plugin headroom wrap claude --no-context-tool # Without CLI context-tool setup headroom wrap claude --port 9999 # Custom proxy port @@ -114,6 +115,11 @@ from headroom.providers.copilot import ( ) 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.grok_build import render_setup_lines as _render_grok_build_setup_lines +from headroom.providers.grok_build.config import ( + inject_grok_provider_config, + restore_grok_provider_config, +) from headroom.providers.mistral_vibe import build_launch_env as _build_mistral_vibe_launch_env from headroom.providers.omp import build_launch_env as _build_omp_launch_env from headroom.providers.omp import inject_models_override as _inject_omp_models_override @@ -189,7 +195,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", "grok", "opencode"} +_AGENT_SAVINGS_TARGET_AGENTS = {"claude", "codex", "cursor", "grok", "grok_build", "opencode"} _WRAP_PROXY_TIMEOUT_ENV = "HEADROOM_WRAP_PROXY_TIMEOUT" _WRAP_PROXY_TIMEOUT_DEFAULT_SECONDS = 45 _WRAP_PROXY_TIMEOUT_ML_DEFAULT_SECONDS = 90 @@ -204,7 +210,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", "grok"} +_AGENT_SAVINGS_WRAP_AGENTS = {"claude", "codex", "cursor", "grok", "grok_build"} # 1M context window for `wrap claude` (#1158). Claude Code only sends the # `context-1m` beta header — unlocking the 1M window for entitled subscription @@ -4205,6 +4211,7 @@ def wrap() -> None: headroom wrap vibe # Mistral Vibe headroom wrap grok # Grok CLI (xAI) headroom wrap cursor # Cursor (prints config instructions) + headroom wrap grok-build # Grok Build (updates ~/.grok/config.toml) headroom wrap cline # Cline (VS Code; prints config instructions) headroom wrap continue # Continue (VS Code/JetBrains; injects systemMessage) headroom wrap goose # Goose (Block) CLI @@ -5951,6 +5958,102 @@ def cursor( ) +# ============================================================================= +# Grok Build +# ============================================================================= + + +@wrap.command("grok-build", 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-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("--verbose", "-v", is_flag=True, help="Verbose output") +@click.option("--prepare-only", is_flag=True, hidden=True) +def grok_build( + port: int, + no_rtk: bool, + no_proxy: bool, + learn: bool, + memory: bool, + verbose: bool, + prepare_only: bool, +) -> None: + """Start Headroom proxy for use with Grok Build. + + \b + Grok Build reads model endpoints from ``~/.grok/config.toml``. This + command starts the proxy, optionally sets up the selected CLI context + tool, injects a Headroom-managed ``[model.grok-build]`` override, and + prints next steps. + + \b + Example: + headroom wrap grok-build + headroom wrap grok-build --no-context-tool + headroom wrap grok-build --port 9999 + """ + project = _project_name_from_cwd() + 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_build", + agent_display="Grok Build", + marker_path=agents_md, + on_rtk_ready=lambda _rtk: ( + _inject_rtk_instructions(cast(Path, agents_md), verbose=verbose) + if agents_md is not None + else None + ), + verbose=verbose, + ) + + if prepare_only: + try: + config_file = inject_grok_provider_config(port, project=project) + click.echo(f" Grok config: injected Headroom proxy override into {config_file}") + except Exception as e: + click.echo(f" Warning: could not update Grok config: {e}") + return + + def _print_grok_build_setup(actual_port: int) -> None: + try: + config_file = inject_grok_provider_config(actual_port, project=project) + click.echo(f" Grok config: injected Headroom proxy override into {config_file}") + click.echo() + except Exception as e: + click.echo(f" Warning: could not update Grok config: {e}") + click.echo() + for line in _render_grok_build_setup_lines(actual_port, project=project): + click.echo(line) + if not no_rtk: + click.echo() + if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX: + click.echo(" lean-ctx configured for Grok Build") + else: + click.echo(" rtk instructions injected into AGENTS.md") + click.echo(" Grok Build will use token-optimized commands automatically.") + + _run_proxy_only_watcher( + agent_label="grok-build", + port=port, + no_proxy=no_proxy, + learn=learn, + memory=memory, + agent_type="grok_build", + print_setup_lines=_print_grok_build_setup, + ) + + # ============================================================================= # Cline (VS Code extension) # ============================================================================= @@ -7171,6 +7274,47 @@ def unwrap_openclaw( # ============================================================================= +@unwrap.command("grok-build") +@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_build(port: int, no_stop_proxy: bool) -> None: + """Undo ``headroom wrap grok-build`` edits to the active Grok config file.""" + click.echo() + click.echo(" ╔═══════════════════════════════════════════════╗") + click.echo(" ║ HEADROOM UNWRAP: GROK BUILD ║") + click.echo(" ╚═══════════════════════════════════════════════╝") + click.echo() + + try: + status, config_file = restore_grok_provider_config() + except Exception as e: # pragma: no cover - filesystem-level errors + raise click.ClickException(f"could not unwrap Grok Build config: {e}") from e + + if status == "restored": + click.echo(f" Restored prior {config_file} from pre-wrap backup.") + elif status == "cleaned": + click.echo(f" Removed Headroom block from {config_file}; other content preserved.") + elif status == "removed": + click.echo(f" Removed {config_file} (contained only Headroom-written config).") + else: + if not os.environ.get("GROK_HOME"): + click.echo( + " Warning: found no Headroom wrap markers in the default Grok config. " + "If you wrapped Grok Build with GROK_HOME, rerun unwrap with the same " + "environment variable, e.g. GROK_HOME=/path/to/grok-home " + "headroom unwrap grok-build." + ) + click.echo(f" Nothing to undo: {config_file} has no Headroom wrap markers.") + + click.echo() + click.echo("✓ Grok Build is no longer routed through the Headroom proxy.") + if not no_stop_proxy and status != "noop": + _echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port) + click.echo() + + @unwrap.command("codex") @click.option( "--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)" diff --git a/headroom/install/models.py b/headroom/install/models.py index 3b13ad950..1f08e09c1 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_BUILD = "grok_build" GROK = "grok" OPENCLAW = "openclaw" OPENCODE = "opencode" diff --git a/headroom/install/planner.py b/headroom/install/planner.py index aba198eac..79c11ed31 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_BUILD, ToolTarget.GROK, ToolTarget.OPENCLAW, ToolTarget.OPENCODE, @@ -55,6 +56,9 @@ def detect_targets() -> list[str]: continue if target == ToolTarget.CURSOR and shutil.which("cursor"): detected.append(target.value) + continue + if target == ToolTarget.GROK_BUILD and shutil.which("grok"): + detected.append(target.value) return detected diff --git a/headroom/providers/grok_build/__init__.py b/headroom/providers/grok_build/__init__.py new file mode 100644 index 000000000..2099af6dd --- /dev/null +++ b/headroom/providers/grok_build/__init__.py @@ -0,0 +1,5 @@ +"""Grok Build provider helpers.""" + +from .runtime import build_proxy_targets, proxy_base_url, render_setup_lines + +__all__ = ["build_proxy_targets", "proxy_base_url", "render_setup_lines"] diff --git a/headroom/providers/grok_build/config.py b/headroom/providers/grok_build/config.py new file mode 100644 index 000000000..ecd570321 --- /dev/null +++ b/headroom/providers/grok_build/config.py @@ -0,0 +1,153 @@ +"""Grok Build config.toml helpers for wrap and persistent install.""" + +from __future__ import annotations + +import os +import re +import shutil +from pathlib import Path + +from headroom import fsutil + +from .runtime import build_proxy_targets + +_MARKER_START = "# --- headroom:grok-build:start ---" +_MARKER_END = "# --- headroom:grok-build:end ---" +_BLOCK_RE = re.compile( + re.escape(_MARKER_START) + r".*?" + re.escape(_MARKER_END) + r"\n?", + re.DOTALL, +) +_GROK_BUILD_TABLE_RE = re.compile(r"(?m)^\[model\.grok-build\]\s*$") +_NEXT_TABLE_RE = re.compile(r"(?m)^\[") +_BASE_URL_LINE_RE = re.compile( + r'(?m)^(?P[ \t]*)base_url[ \t]*=[ \t]*"(?P[^"\n]*)".*$' +) + + +def grok_home_dir() -> Path: + """Return the Grok home/config directory.""" + env_path = os.environ.get("GROK_HOME", "").strip() + if env_path: + return Path(env_path).expanduser() + return Path.home() / ".grok" + + +def grok_config_paths() -> tuple[Path, Path]: + """Return ``(config_file, backup_file)`` for Grok Build.""" + config_file = grok_home_dir() / "config.toml" + backup_file = config_file.with_suffix(".toml.headroom-backup") + return config_file, backup_file + + +def snapshot_grok_config_if_unwrapped(config_file: Path, backup_file: Path) -> None: + """Snapshot ``config.toml`` before the first Headroom injection.""" + if backup_file.exists(): + return + if not config_file.exists(): + return + try: + content = fsutil.read_text(config_file) + except OSError: + return + if _MARKER_START in content: + return + backup_file.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(config_file, backup_file) + + +def strip_grok_headroom_blocks(content: str) -> str: + """Remove Headroom-managed Grok config blocks.""" + content = _BLOCK_RE.sub("", content) + content = re.sub(r"\n{3,}", "\n\n", content) + return content.strip() + + +def has_user_grok_build_model_table(content: str) -> bool: + """Return True when ``content`` already declares ``[model.grok-build]``.""" + return _GROK_BUILD_TABLE_RE.search(content) is not None + + +def redirect_existing_grok_build_base_url(content: str, base_url: str) -> tuple[str, bool]: + """Rewrite ``base_url`` inside an existing ``[model.grok-build]`` table. + + TOML rejects duplicate table headers, so when the user already owns + ``[model.grok-build]`` we update that table in place instead of appending + a second one. The previous ``base_url`` value is preserved in a trailing + ``# was: …`` comment for visibility; the pre-wrap snapshot still enables + byte-for-byte restore on ``headroom unwrap grok-build``. + """ + match = _GROK_BUILD_TABLE_RE.search(content) + if match is None: + return content, False + + section_start = match.end() + next_table = _NEXT_TABLE_RE.search(content, section_start) + section_end = next_table.start() if next_table else len(content) + section = content[section_start:section_end] + + if _BASE_URL_LINE_RE.search(section): + + def _replace(match_obj: re.Match[str]) -> str: + original_value = match_obj.group("value") + if original_value == base_url: + return match_obj.group(0) + indent = match_obj.group("indent") + return f'{indent}base_url = "{base_url}" # was: {original_value}' + + section = _BASE_URL_LINE_RE.sub(_replace, section, count=1) + else: + section = f'\nbase_url = "{base_url}"' + section + + updated = content[:section_start] + section + content[section_end:] + return updated, updated != content + + +def render_headroom_block(port: int, project: str | None = None) -> str: + """Render the Headroom-managed ``[model.grok-build]`` override block.""" + target = build_proxy_targets(port, project) + return f'{_MARKER_START}\n[model.grok-build]\nbase_url = "{target.base_url}"\n{_MARKER_END}\n' + + +def inject_grok_provider_config(port: int, project: str | None = None) -> Path: + """Inject or refresh the Headroom proxy override into Grok config.""" + config_file, backup_file = grok_config_paths() + config_file.parent.mkdir(parents=True, exist_ok=True) + snapshot_grok_config_if_unwrapped(config_file, backup_file) + + if config_file.exists(): + content = strip_grok_headroom_blocks(fsutil.read_text(config_file)) + else: + content = "" + + target = build_proxy_targets(port, project) + if has_user_grok_build_model_table(content): + content, _ = redirect_existing_grok_build_base_url(content, target.base_url) + else: + block = render_headroom_block(port, project) + if content: + content = content.rstrip() + "\n\n" + block + else: + content = block + + fsutil.write_text(config_file, content) + return config_file + + +def restore_grok_provider_config() -> tuple[str, Path]: + """Undo ``inject_grok_provider_config`` for the active Grok config file.""" + config_file, backup_file = grok_config_paths() + if backup_file.exists(): + shutil.copy2(backup_file, config_file) + backup_file.unlink() + return "restored", config_file + + if not config_file.exists(): + return "noop", config_file + + content = strip_grok_headroom_blocks(fsutil.read_text(config_file)) + if content: + fsutil.write_text(config_file, content) + return "cleaned", config_file + + config_file.unlink(missing_ok=True) + return "removed", config_file diff --git a/headroom/providers/grok_build/install.py b/headroom/providers/grok_build/install.py new file mode 100644 index 000000000..b78f1a670 --- /dev/null +++ b/headroom/providers/grok_build/install.py @@ -0,0 +1,14 @@ +"""Grok Build install-time helpers.""" + +from __future__ import annotations + +from .runtime import build_proxy_targets + + +def build_install_env(*, port: int, backend: str) -> dict[str, str]: + """Build the persistent install environment for Grok Build.""" + del backend + target = build_proxy_targets(port) + return { + "GROK_MODEL_GROK_BUILD_BASE_URL": target.base_url, + } diff --git a/headroom/providers/grok_build/runtime.py b/headroom/providers/grok_build/runtime.py new file mode 100644 index 000000000..f8bc7f123 --- /dev/null +++ b/headroom/providers/grok_build/runtime.py @@ -0,0 +1,57 @@ +"""Runtime helpers for Grok Build integrations.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from headroom.proxy.project_context import with_project_prefix + + +def proxy_base_url(port: int) -> str: + """Return the local proxy base URL for OpenAI-compatible Grok traffic.""" + return f"http://127.0.0.1:{port}/v1" + + +@dataclass(frozen=True) +class GrokBuildProxyTarget: + """Resolved local proxy target shown in Grok Build setup instructions.""" + + base_url: str + + +def build_proxy_targets(port: int, project: str | None = None) -> GrokBuildProxyTarget: + """Build the local proxy URL shown to Grok Build users. + + ``project`` (the wrap launch directory) is encoded as a ``/p/`` + base-URL prefix because Grok cannot send custom headers; the proxy + strips it and attributes savings per project. + """ + return GrokBuildProxyTarget( + base_url=with_project_prefix(proxy_base_url(port), project), + ) + + +def render_setup_lines(port: int, project: str | None = None) -> list[str]: + """Render Grok Build setup instructions for the local proxy.""" + target = build_proxy_targets(port, project) + lines = [ + " Headroom proxy is running. Configure Grok Build:", + "", + " ~/.grok/config.toml has been updated with:", + " [model.grok-build]", + f' base_url = "{target.base_url}"', + "", + " Start Grok Build in this project directory:", + " grok", + "", + " Or switch models in an existing session:", + " /model grok-build", + ] + if project: + lines += [ + "", + f" Dashboard savings will be attributed to project '{project}'", + " (the directory this command was run from). Re-run from another", + " project directory to get that project's URL.", + ] + return lines diff --git a/headroom/providers/install_registry.py b/headroom/providers/install_registry.py index 130126c59..a3e053aa1 100644 --- a/headroom/providers/install_registry.py +++ b/headroom/providers/install_registry.py @@ -30,6 +30,7 @@ from headroom.providers.cortex_code.install import ( ) 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.grok_build.install import build_install_env as _build_grok_build_install_env from headroom.providers.openclaw.install import ( apply_provider_scope as _apply_openclaw_provider_scope, ) @@ -55,6 +56,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": _build_grok_build_install_env, "grok": _build_grok_install_env, "opencode": _build_opencode_install_env, } diff --git a/headroom/proxy/auth_policy.py b/headroom/proxy/auth_policy.py index 0c3cc9d4a..73f6cf91d 100644 --- a/headroom/proxy/auth_policy.py +++ b/headroom/proxy/auth_policy.py @@ -19,6 +19,7 @@ SUBSCRIPTION_UA_PREFIXES: tuple[str, ...] = ( "claude-code/", "codex-cli/", "cursor/", + "grok/", "claude-vscode/", "github-copilot/", "anthropic-cli/", @@ -33,6 +34,7 @@ CLIENT_UA_MAP: tuple[tuple[str, str], ...] = ( ("anthropic-cli/", "anthropic-cli"), ("codex-cli/", "codex"), ("cursor/", "cursor"), + ("grok/", "grok_build"), ("zed/", "zed"), ("aider/", "aider"), ("droid/", "droid"), diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 6c5511d46..9673476d9 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -218,6 +218,7 @@ _AGENT_LABELS: dict[str, str] = { "codex": "Codex", "codex-cli": "Codex", "cursor": "Cursor", + "grok_build": "Grok Build", "copilot": "GitHub Copilot", "github-copilot": "GitHub Copilot", "aider": "Aider", diff --git a/headroom/telemetry/context.py b/headroom/telemetry/context.py index e68ec5e75..e845b1ed7 100644 --- a/headroom/telemetry/context.py +++ b/headroom/telemetry/context.py @@ -22,7 +22,17 @@ logger = logging.getLogger(__name__) _KNOWN_WRAP_AGENTS = frozenset( - {"claude", "copilot", "codex", "aider", "cursor", "omp", "openclaw", "opencode"} + { + "claude", + "copilot", + "codex", + "aider", + "cursor", + "grok_build", + "omp", + "openclaw", + "opencode", + } ) # Stack slugs must start with a letter and contain only [a-z0-9_], max 64 chars. diff --git a/pyproject.toml b/pyproject.toml index 4a4915a76..a46be9b74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "headroom-ai" version = "0.32.0" description = "The Context Optimization Layer for LLM Applications - Cut costs by 50-90%" readme = "README.md" -license = "Apache-2.0" +license = "Apache-2.0" requires-python = ">=3.10" authors = [ { name = "Headroom Contributors" } diff --git a/tests/test_auth_policy.py b/tests/test_auth_policy.py index 02ae27a63..e15c3552b 100644 --- a/tests/test_auth_policy.py +++ b/tests/test_auth_policy.py @@ -43,6 +43,13 @@ def test_client_explicit_override_wins_over_user_agent() -> None: assert classify_client_signals(signals) == "aider" +def test_grok_build_user_agent_is_subscription_client() -> None: + signals = AuthSignals(user_agent="grok/1.2.3") + + assert classify_auth_signals(signals) is AuthMode.SUBSCRIPTION + assert classify_client_signals(signals) == "grok_build" + + def test_codex_stamp_only_for_unidentified_responses_callers() -> None: assert should_stamp_codex_client_signals("/v1/responses", AuthSignals()) is True assert ( diff --git a/tests/test_cli/test_wrap_bridge.py b/tests/test_cli/test_wrap_bridge.py index c160ae3e0..a3e21c91b 100644 --- a/tests/test_cli/test_wrap_bridge.py +++ b/tests/test_cli/test_wrap_bridge.py @@ -94,6 +94,27 @@ def test_wrap_codex_prepare_only_updates_config(monkeypatch, tmp_path: Path) -> assert 'base_url = "http://127.0.0.1:8787/v1"' in content +def test_wrap_grok_build_uses_actual_proxy_port(monkeypatch, tmp_path: Path) -> None: + _set_test_home(monkeypatch, tmp_path) + runner = CliRunner() + + def fake_watcher(**kwargs) -> None: + kwargs["print_setup_lines"](9999) + + monkeypatch.setattr("headroom.cli.wrap._run_proxy_only_watcher", fake_watcher) + + result = runner.invoke(main, ["wrap", "grok-build", "--no-context-tool", "--port", "8787"]) + + assert result.exit_code == 0, result.output + config_file = tmp_path / ".grok" / "config.toml" + assert config_file.exists() + content = config_file.read_text(encoding="utf-8") + assert 'base_url = "http://127.0.0.1:9999/' in content + assert "http://127.0.0.1:8787/" not in content + assert "http://127.0.0.1:9999/" in result.output + assert "http://127.0.0.1:8787/" not in result.output + + def test_wrap_codex_prepare_only_uses_lean_ctx_when_configured(monkeypatch, tmp_path: Path) -> None: _set_test_home(monkeypatch, tmp_path) monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx") diff --git a/tests/test_mcp_registry/test_grok_registrar.py b/tests/test_mcp_registry/test_grok_registrar.py new file mode 100644 index 000000000..c09058917 --- /dev/null +++ b/tests/test_mcp_registry/test_grok_registrar.py @@ -0,0 +1,43 @@ +"""Tests for the Grok Build MCP registrar.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from headroom.mcp_registry.base import RegisterStatus, ServerSpec +from headroom.mcp_registry.grok import GrokRegistrar + + +def _make_registrar(tmp_path: Path) -> GrokRegistrar: + return GrokRegistrar(home_dir=tmp_path) + + +def _spec() -> ServerSpec: + return ServerSpec( + name="headroom", + command="/usr/bin/python", + args=("-m", "headroom.cli", "mcp", "serve"), + ) + + +def test_detect_true_when_grok_dir_exists(tmp_path: Path) -> None: + (tmp_path / ".grok").mkdir() + assert _make_registrar(tmp_path).detect() is True + + +def test_detect_false_when_grok_dir_missing(tmp_path: Path) -> None: + assert _make_registrar(tmp_path).detect() is False + + +def test_register_uses_grok_home_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + grok_home = tmp_path / "custom-grok-home" + monkeypatch.setenv("GROK_HOME", str(grok_home)) + + result = GrokRegistrar().register_server(_spec()) + + assert result.status == RegisterStatus.REGISTERED + config = grok_home / "config.toml" + assert config.exists() + assert "[mcp_servers.headroom]" in config.read_text() diff --git a/tests/test_provider_grok_build.py b/tests/test_provider_grok_build.py new file mode 100644 index 000000000..43ae61e7e --- /dev/null +++ b/tests/test_provider_grok_build.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +from headroom.providers.grok_build import build_proxy_targets, render_setup_lines +from headroom.providers.grok_build.config import ( + inject_grok_provider_config, + redirect_existing_grok_build_base_url, + render_headroom_block, + restore_grok_provider_config, + strip_grok_headroom_blocks, +) +from headroom.providers.grok_build.install import build_install_env + +if sys.version_info >= (3, 11): + import tomllib +else: # pragma: no cover + import tomli as tomllib # type: ignore[no-redef] + + +def _assert_valid_toml(content: str) -> None: + payload = content.encode("utf-8") + try: + tomllib.loads(payload) + except TypeError: + # Some environments expose a str-accepting TOML parser shim. + tomllib.loads(content) # type: ignore[arg-type] + + +def _count_grok_build_tables(content: str) -> int: + return content.count("[model.grok-build]") + + +def test_grok_build_proxy_targets_use_local_headroom_proxy() -> None: + target = build_proxy_targets(9999) + + assert target.base_url == "http://127.0.0.1:9999/v1" + + +def test_grok_build_setup_lines_include_proxy_url() -> None: + lines = render_setup_lines(8787) + joined = "\n".join(lines) + + assert "http://127.0.0.1:8787/v1" in joined + assert "[model.grok-build]" in joined + + +def test_grok_build_build_install_env_returns_proxy_url() -> None: + env = build_install_env(port=7654, backend="ignored") + + assert env == {"GROK_MODEL_GROK_BUILD_BASE_URL": "http://127.0.0.1:7654/v1"} + + +def test_grok_build_proxy_targets_apply_project_path_prefix() -> None: + target = build_proxy_targets(9999, project="frontend") + + assert target.base_url == "http://127.0.0.1:9999/p/frontend/v1" + + +def test_grok_build_setup_lines_mention_project_attribution() -> None: + lines = render_setup_lines(8787, project="frontend") + joined = "\n".join(lines) + + assert "http://127.0.0.1:8787/p/frontend/v1" in joined + assert "attributed to project 'frontend'" in joined + + +def test_grok_build_config_inject_and_restore_round_trip(tmp_path: Path, monkeypatch) -> None: + grok_home = tmp_path / ".grok" + grok_home.mkdir() + monkeypatch.setenv("GROK_HOME", str(grok_home)) + + config_file = inject_grok_provider_config(8787, project="demo") + content = config_file.read_text(encoding="utf-8") + + assert render_headroom_block(8787, project="demo").strip() in content + assert 'base_url = "http://127.0.0.1:8787/p/demo/v1"' in content + + status, _ = restore_grok_provider_config() + assert status == "removed" + assert not config_file.exists() + + +def test_grok_build_config_strip_preserves_user_content() -> None: + original = f'[models]\ndefault = "grok-build"\n\n{render_headroom_block(8787)}' + cleaned = strip_grok_headroom_blocks(original) + + assert "[models]" in cleaned + assert "headroom:grok-build" not in cleaned + + +def test_grok_build_inject_updates_existing_user_table_without_duplicate( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + grok_home = tmp_path / ".grok" + grok_home.mkdir() + monkeypatch.setenv("GROK_HOME", str(grok_home)) + + original = ( + "[models]\n" + 'default = "grok-build"\n\n' + "[model.grok-build]\n" + 'model = "grok-build"\n' + 'base_url = "https://api.x.ai/v1"\n' + "temperature = 0.5\n" + ) + config_file = grok_home / "config.toml" + config_file.write_text(original, encoding="utf-8") + + inject_grok_provider_config(8787, project="demo") + content = config_file.read_text(encoding="utf-8") + + assert _count_grok_build_tables(content) == 1 + assert 'base_url = "http://127.0.0.1:8787/p/demo/v1" # was: https://api.x.ai/v1' in content + assert "temperature = 0.5" in content + assert "headroom:grok-build" not in content + _assert_valid_toml(content) + + +def test_grok_build_redirect_existing_base_url_is_idempotent() -> None: + original = '[model.grok-build]\nbase_url = "http://127.0.0.1:8787/v1"\ntemperature = 0.2\n' + + updated, changed = redirect_existing_grok_build_base_url(original, "http://127.0.0.1:8787/v1") + + assert changed is False + assert updated == original + _assert_valid_toml(updated)