mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(wrap): add first-class Grok CLI support (#1823)
## 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 <mxjerrett@gmail.com>
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
This commit is contained in:
roman-t3a
GitHub
parent
7bfb1d7f38
commit
cb388f6af2
19 changed files with 922 additions and 10 deletions
11
README.md
11
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 <tool>`
|
||||
- **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 <tool>`
|
||||
- **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 <tool>` (supports: `claude`, `copilot`, `codex`, `opencode`, `openclaw`, `zcode`).
|
||||
Undo durable wrapping with `headroom unwrap <tool>` (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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
# =============================================================================
|
||||
|
|
|
|||
|
|
@ -55,6 +55,7 @@ class ToolTarget(str, Enum):
|
|||
CODEX = "codex"
|
||||
AIDER = "aider"
|
||||
CURSOR = "cursor"
|
||||
GROK = "grok"
|
||||
OPENCLAW = "openclaw"
|
||||
OPENCODE = "opencode"
|
||||
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ SUPPORTED_TARGETS = [
|
|||
ToolTarget.CODEX,
|
||||
ToolTarget.AIDER,
|
||||
ToolTarget.CURSOR,
|
||||
ToolTarget.GROK,
|
||||
ToolTarget.OPENCLAW,
|
||||
ToolTarget.OPENCODE,
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
202
headroom/learn/plugins/grok.py
Normal file
202
headroom/learn/plugins/grok.py
Normal file
|
|
@ -0,0 +1,202 @@
|
|||
"""Grok CLI plugin for headroom learn.
|
||||
|
||||
Reads session logs from ~/.grok/sessions/<workspace>/<session-id>/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()
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
229
headroom/mcp_registry/grok.py
Normal file
229
headroom/mcp_registry/grok.py
Normal file
|
|
@ -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.<name>]`` 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)
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
5
headroom/providers/grok/__init__.py
Normal file
5
headroom/providers/grok/__init__.py
Normal file
|
|
@ -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"]
|
||||
11
headroom/providers/grok/install.py
Normal file
11
headroom/providers/grok/install.py
Normal file
|
|
@ -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)}
|
||||
37
headroom/providers/grok/runtime.py
Normal file
37
headroom/providers/grok/runtime.py
Normal file
|
|
@ -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/<name>``
|
||||
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}"]
|
||||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
64
tests/test_cli/test_wrap_grok.py
Normal file
64
tests/test_cli/test_wrap_grok.py
Normal file
|
|
@ -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()
|
||||
59
tests/test_learn_grok_plugin.py
Normal file
59
tests/test_learn_grok_plugin.py
Normal file
|
|
@ -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
|
||||
31
tests/test_mcp_registry_grok.py
Normal file
31
tests/test_mcp_registry_grok.py
Normal file
|
|
@ -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()
|
||||
)
|
||||
27
tests/test_provider_grok.py
Normal file
27
tests/test_provider_grok.py
Normal file
|
|
@ -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",
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue