mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(grok-build): add Grok Build wrap command and MCP integration (#1629)
## 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)**

**2. Review fix — in-place `[model.grok-build]` rewrite (single table,
`# was:` metadata)**

**3. Proxy health — `/readyz` healthy on port 8799**

**4. Unwrap — restores pre-wrap backup**

## 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 <mxjerrett@gmail.com>
This commit is contained in:
parent
ecf086f1f8
commit
420dc9077b
22 changed files with 675 additions and 5 deletions
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -234,6 +234,9 @@ pyrightconfig.json
|
|||
\#*\#
|
||||
.\#*
|
||||
|
||||
# Local git worktrees (isolated feature branches)
|
||||
.worktrees/
|
||||
|
||||
# Local development configuration
|
||||
CLAUDE.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)).
|
||||
|
||||
|
|
|
|||
62
docs/content/docs/grok-build.mdx
Normal file
62
docs/content/docs/grok-build.mdx
Normal file
|
|
@ -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:<port>/v1` (with optional `/p/<project>` 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.
|
||||
|
|
@ -37,6 +37,7 @@
|
|||
"claude-code-vertex",
|
||||
"claude-code-azure-foundry",
|
||||
"opencode",
|
||||
"grok-build",
|
||||
"mcp",
|
||||
"---Configuration---",
|
||||
"configuration",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.")
|
||||
|
|
|
|||
|
|
@ -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)"
|
||||
|
|
|
|||
|
|
@ -55,6 +55,7 @@ class ToolTarget(str, Enum):
|
|||
CODEX = "codex"
|
||||
AIDER = "aider"
|
||||
CURSOR = "cursor"
|
||||
GROK_BUILD = "grok_build"
|
||||
GROK = "grok"
|
||||
OPENCLAW = "openclaw"
|
||||
OPENCODE = "opencode"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
5
headroom/providers/grok_build/__init__.py
Normal file
5
headroom/providers/grok_build/__init__.py
Normal file
|
|
@ -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"]
|
||||
153
headroom/providers/grok_build/config.py
Normal file
153
headroom/providers/grok_build/config.py
Normal file
|
|
@ -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<indent>[ \t]*)base_url[ \t]*=[ \t]*"(?P<value>[^"\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
|
||||
14
headroom/providers/grok_build/install.py
Normal file
14
headroom/providers/grok_build/install.py
Normal file
|
|
@ -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,
|
||||
}
|
||||
57
headroom/providers/grok_build/runtime.py
Normal file
57
headroom/providers/grok_build/runtime.py
Normal file
|
|
@ -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/<name>``
|
||||
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
|
||||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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" }
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
43
tests/test_mcp_registry/test_grok_registrar.py
Normal file
43
tests/test_mcp_registry/test_grok_registrar.py
Normal file
|
|
@ -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()
|
||||
130
tests/test_provider_grok_build.py
Normal file
130
tests/test_provider_grok_build.py
Normal file
|
|
@ -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)
|
||||
Loading…
Add table
Add a link
Reference in a new issue