From 0932b8bef4db9109665382b6d7c079a368f08d52 Mon Sep 17 00:00:00 2001 From: Brian Toye Date: Tue, 16 Jun 2026 20:59:51 +0100 Subject: [PATCH] feat: Add support for Mistral Vibe CLI (#935) ## Description Add `headroom wrap vibe` / `headroom unwrap vibe` support for Mistral Vibe CLI so Vibe can launch through Headroom's proxy, compression, and observability path. ## Type of Change - [x] New feature (non-breaking change that adds functionality) ## Changes Made - Added `headroom.providers.mistral_vibe` provider runtime helpers. - Added `headroom wrap vibe` command support and matching unwrap handling. - Configured `VIBE_PROVIDERS` so Vibe routes through the Headroom proxy. - Added tests covering launch, custom ports, no-proxy behavior, code-graph/learn-memory flags, verbose mode, invalid-command handling, and provider JSON structure. - Updated `CHANGELOG.md`. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text pytest -v tests/test_cli/test_wrap_vibe.py # 10 passed ``` ## Real Behavior Proof - Environment: Linux, Python 3.13.13, local checkout from the PR branch. - Exact command / steps: Ran the Vibe wrapper tests and manually launched Mistral Vibe through `headroom wrap vibe` with `VIBE_PROVIDERS` pointing at the Headroom proxy. - Observed result: Vibe launched through Headroom's proxy configuration, and the wrapper tests passed. - Not tested: RTK hook support for Vibe. Persistent installs may eventually hold an expired Vibe auth token because Vibe reads its auth token from the environment at startup; opening another port or removing the persistent install is the current workaround. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review --------- Co-authored-by: Vibe Nuage Agent Co-authored-by: JerrettDavis --- CHANGELOG.md | 1 + headroom/cli/wrap.py | 80 ++++++ headroom/providers/mistral_vibe/__init__.py | 5 + headroom/providers/mistral_vibe/runtime.py | 54 ++++ tests/test_cli/test_wrap_vibe.py | 281 ++++++++++++++++++++ 5 files changed, 421 insertions(+) create mode 100644 headroom/providers/mistral_vibe/__init__.py create mode 100644 headroom/providers/mistral_vibe/runtime.py create mode 100644 tests/test_cli/test_wrap_vibe.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 08fcd9382..627f2d9fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Features +* **vibe:** add Mistral Vibe CLI support with `headroom wrap vibe`. * **proxy:** per-project savings breakdown on the dashboard for all wrapped agents — Claude Code, Codex, aider, Copilot, and Cursor ([#802](https://github.com/chopratejas/headroom/issues/802)). `headroom wrap claude`/`codex` tag requests with an `X-Headroom-Project` header (launch-directory name); `wrap aider`/`copilot`/`cursor` — whose clients cannot send custom headers — use a `/p/` base-URL prefix the proxy strips. Savings are aggregated per project (persisted, schema v3 with transparent v2 migration), exposed as `savings.per_project` in `/stats` and `projects` in `/stats-history`, and shown in a Per-Project Savings dashboard table. ### Features diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 4860ff4b0..49d9c843d 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -5,6 +5,7 @@ Usage: headroom wrap copilot -- --model ... # Start proxy + launch GitHub Copilot CLI headroom wrap codex # Start proxy + OpenAI Codex CLI headroom wrap aider # Start proxy + aider + headroom wrap vibe # Start proxy + Mistral Vibe 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 @@ -85,6 +86,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.mistral_vibe import build_launch_env as _build_mistral_vibe_launch_env from headroom.providers.openclaw import ( build_plugin_entry as _build_openclaw_plugin_entry_impl, ) @@ -2711,6 +2713,7 @@ def wrap() -> None: headroom wrap codex # OpenAI Codex CLI headroom wrap copilot -- --model claude-sonnet-4-20250514 headroom wrap aider # Aider + headroom wrap vibe # Mistral Vibe 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) @@ -3655,6 +3658,83 @@ def aider( ) +# ============================================================================= +# Mistral Vibe +# ============================================================================= + + +@wrap.command(context_settings={"ignore_unknown_options": True}) +@click.option("--port", "-p", default=8787, type=int, help="Proxy port (default: 8787)") +@click.option( + "--no-context-tool", + "--no-rtk", + "no_rtk", + is_flag=True, + help="Skip CLI context-tool setup (no effect for vibe)", +) +@click.option( + "--code-graph", + is_flag=True, + help="Enable code graph indexing via codebase-memory-mcp (optional)", +) +@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) +@click.argument("vibe_args", nargs=-1, type=click.UNPROCESSED) +def vibe( + port: int, + no_rtk: bool, + code_graph: bool, + no_proxy: bool, + learn: bool, + memory: bool, + verbose: bool, + prepare_only: bool, + vibe_args: tuple, +) -> None: + """Launch Mistral Vibe through Headroom proxy. + + \b + Sets VIBE_PROVIDERS to route all Mistral API calls through Headroom. + + \b + Examples: + headroom wrap vibe # Start proxy + vibe + headroom wrap vibe -- "fix the bug" # Pass prompt to vibe + headroom wrap vibe --port 9999 # Custom proxy port + headroom wrap vibe --no-context-tool # Skip CLI context-tool setup + """ + if prepare_only: + return + + vibe_bin = shutil.which("vibe") + if not vibe_bin: + click.echo("Error: 'vibe' not found in PATH.") + click.echo("Install Mistral Vibe: https://github.com/mistralai/mistral-vibe") + raise SystemExit(1) + + env, env_vars_display = _build_mistral_vibe_launch_env( + port, os.environ, project=_project_name_from_cwd() + ) + + _launch_tool( + binary=vibe_bin, + args=vibe_args, + env=env, + port=port, + no_proxy=no_proxy, + tool_label="VIBE", + env_vars_display=env_vars_display, + learn=learn, + memory=memory, + agent_type="vibe", + code_graph=code_graph, + openai_api_url="https://api.mistral.ai", + ) + + # ============================================================================= # Cursor # ============================================================================= diff --git a/headroom/providers/mistral_vibe/__init__.py b/headroom/providers/mistral_vibe/__init__.py new file mode 100644 index 000000000..c050ca015 --- /dev/null +++ b/headroom/providers/mistral_vibe/__init__.py @@ -0,0 +1,5 @@ +"""Mistral Vibe-specific provider helpers.""" + +from .runtime import build_launch_env + +__all__ = ["build_launch_env"] diff --git a/headroom/providers/mistral_vibe/runtime.py b/headroom/providers/mistral_vibe/runtime.py new file mode 100644 index 000000000..016a3cd60 --- /dev/null +++ b/headroom/providers/mistral_vibe/runtime.py @@ -0,0 +1,54 @@ +"""Runtime helpers for Mistral Vibe integrations.""" + +from __future__ import annotations + +import json +import os +from collections.abc import Mapping + +from headroom.providers.codex import proxy_base_url as codex_proxy_base_url +from headroom.proxy.project_context import with_project_prefix + + +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 Mistral Vibe through the local proxy. + + Mistral Vibe uses a provider configuration system with `api_base` field. + It supports overriding providers via the `VIBE_PROVIDERS` environment variable + as a JSON array. When routing through Headroom, we set the mistral provider's + `api_base` to the local proxy URL. The proxy will then forward requests to + the actual Mistral API. + + ``project`` (the wrap launch directory) is encoded as a ``/p/`` + base-URL prefix because Vibe cannot send custom headers; the proxy + strips it and attributes savings per project. + """ + env = dict(environ or os.environ) + # NOTE: With a persistent Headroom deployment (`headroom install`), the proxy + # process captures its environment at startup. Vibe reads `MISTRAL_API_KEY` + # from its own process environment (via `api_key_env_var` below), so if the + # token changes you may need to restart the Vibe process (and, for persistent + # installs, the proxy) to pick up the new value. + base_url = with_project_prefix(codex_proxy_base_url(port), project) + + # Build the providers JSON with mistral provider pointing to Headroom proxy + # We need to override the default mistral provider's api_base + providers = [ + { + "name": "mistral", + "api_base": base_url, + "api_key_env_var": "MISTRAL_API_KEY", + "browser_auth_base_url": "https://console.mistral.ai", + "browser_auth_api_base_url": "https://console.mistral.ai/api", + "backend": "mistral", + } + ] + + providers_json = json.dumps(providers) + env["VIBE_PROVIDERS"] = providers_json + + return env, [f"VIBE_PROVIDERS={providers_json}"] diff --git a/tests/test_cli/test_wrap_vibe.py b/tests/test_cli/test_wrap_vibe.py new file mode 100644 index 000000000..2d3c995a7 --- /dev/null +++ b/tests/test_cli/test_wrap_vibe.py @@ -0,0 +1,281 @@ +"""Tests for `headroom wrap vibe` command.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from headroom.cli import wrap as wrap_mod +from headroom.cli.main import main + + +@pytest.fixture +def runner() -> CliRunner: + return CliRunner() + + +def test_wrap_vibe_launch( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Vibe launches with correct configuration.""" + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + + captured: dict[str, Any] = {} + + def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003 + captured.update(kwargs) + + with patch.object(wrap_mod.shutil, "which", return_value="vibe"): + with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool): + with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None): + result = runner.invoke( + main, ["wrap", "vibe", "--port", "9000", "--", "--prompt", "test"] + ) + + assert result.exit_code == 0, result.output + env = captured["env"] + assert isinstance(env, dict) + assert captured["tool_label"] == "VIBE" + assert captured["agent_type"] == "vibe" + assert captured["args"] == ("--prompt", "test") + + +def test_wrap_vibe_with_project_name( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Project name is encoded in the URL when running from a project directory.""" + project_dir = tmp_path / "my-project" + project_dir.mkdir() + monkeypatch.chdir(project_dir) + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + + captured: dict[str, Any] = {} + + def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003 + captured.update(kwargs) + + with patch.object(wrap_mod.shutil, "which", return_value="vibe"): + with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool): + result = runner.invoke(main, ["wrap", "vibe", "--port", "7000"]) + + assert result.exit_code == 0, result.output + env = captured["env"] + assert isinstance(env, dict) + + providers: list[dict[str, Any]] = json.loads(env["VIBE_PROVIDERS"]) + assert providers[0]["api_base"] == "http://127.0.0.1:7000/p/my-project/v1" + + +def test_wrap_vibe_not_found( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Error message when vibe binary is not found.""" + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + + with patch.object(wrap_mod.shutil, "which", return_value=None): + result = runner.invoke(main, ["wrap", "vibe"]) + + assert result.exit_code == 1 + assert "Error: 'vibe' not found in PATH" in result.output + assert "Install Mistral Vibe: https://github.com/mistralai/mistral-vibe" in result.output + + +def test_wrap_vibe_custom_port( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Custom --port is passed to _launch_tool and appears in VIBE_PROVIDERS.""" + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + + captured: dict[str, Any] = {} + + def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003 + captured.update(kwargs) + + with patch.object(wrap_mod.shutil, "which", return_value="vibe"): + with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool): + with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None): + result = runner.invoke(main, ["wrap", "vibe", "--port", "9999"]) + + assert result.exit_code == 0, result.output + assert captured["port"] == 9999 + env = captured["env"] + providers: list[dict[str, Any]] = json.loads(env["VIBE_PROVIDERS"]) + assert providers[0]["api_base"] == "http://127.0.0.1:9999/v1" + + +def test_wrap_vibe_no_proxy( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """--no-proxy flag prevents proxy startup.""" + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + + captured: dict[str, Any] = {} + + def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003 + captured.update(kwargs) + + with patch.object(wrap_mod.shutil, "which", return_value="vibe"): + with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool): + with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None): + result = runner.invoke(main, ["wrap", "vibe", "--no-proxy"]) + + assert result.exit_code == 0, result.output + assert captured["no_proxy"] is True + + +def test_wrap_vibe_code_graph( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """--code-graph flag is passed to _launch_tool.""" + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + + captured: dict[str, Any] = {} + + def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003 + captured.update(kwargs) + + with patch.object(wrap_mod.shutil, "which", return_value="vibe"): + with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool): + with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None): + result = runner.invoke(main, ["wrap", "vibe", "--code-graph"]) + + assert result.exit_code == 0, result.output + assert captured["code_graph"] is True + + +def test_wrap_vibe_learn_memory( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """--learn and --memory flags are passed to _launch_tool.""" + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + + captured: dict[str, Any] = {} + + def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003 + captured.update(kwargs) + + with patch.object(wrap_mod.shutil, "which", return_value="vibe"): + with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool): + with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None): + result = runner.invoke(main, ["wrap", "vibe", "--learn", "--memory"]) + + assert result.exit_code == 0, result.output + assert captured["learn"] is True + assert captured["memory"] is True + + +def test_wrap_vibe_verbose( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """--verbose flag is accepted by vibe command.""" + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + + captured: dict[str, Any] = {} + + def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003 + captured.update(kwargs) + + with patch.object(wrap_mod.shutil, "which", return_value="vibe"): + with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool): + with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None): + result = runner.invoke(main, ["wrap", "vibe", "--verbose"]) + + assert result.exit_code == 0, result.output + + +def test_wrap_vibe_providers_json_structure( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """VIBE_PROVIDERS env var has correct JSON structure.""" + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + + captured: dict[str, Any] = {} + + def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003 + captured.update(kwargs) + + with patch.object(wrap_mod.shutil, "which", return_value="vibe"): + with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool): + with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None): + result = runner.invoke(main, ["wrap", "vibe"]) + + assert result.exit_code == 0, result.output + env = captured["env"] + assert "VIBE_PROVIDERS" in env + + providers: list[dict[str, Any]] = json.loads(env["VIBE_PROVIDERS"]) + assert isinstance(providers, list) + assert len(providers) == 1 + assert providers[0]["name"] == "mistral" + assert providers[0]["api_key_env_var"] == "MISTRAL_API_KEY" + assert providers[0]["backend"] == "mistral" + assert "api_base" in providers[0] + assert providers[0]["browser_auth_base_url"] == "https://console.mistral.ai" + assert providers[0]["browser_auth_api_base_url"] == "https://console.mistral.ai/api" + + +def test_wrap_vibe_no_context_tool( + runner: CliRunner, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """--no-context-tool and --no-rtk flags are accepted and not passed to vibe.""" + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + + captured: dict[str, Any] = {} + + def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003 + captured.update(kwargs) + + with patch.object(wrap_mod.shutil, "which", return_value="vibe"): + with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool): + with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None): + # Test --no-context-tool + result = runner.invoke(main, ["wrap", "vibe", "--no-context-tool", "--", "test"]) + + assert result.exit_code == 0, result.output + assert captured["args"] == ("test",) + assert "--no-context-tool" not in captured["args"] + + captured.clear() + with patch.object(wrap_mod.shutil, "which", return_value="vibe"): + with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool): + with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None): + # Test --no-rtk + result = runner.invoke(main, ["wrap", "vibe", "--no-rtk", "--", "test"]) + + assert result.exit_code == 0, result.output + assert captured["args"] == ("test",) + assert "--no-rtk" not in captured["args"]