mtrx-cli 0.1.21 → 0.1.22

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mtrx-cli",
-  "version": "0.1.21",
+  "version": "0.1.22",
   "description": "MATRX CLI for routing Codex, Claude, and Cursor through Matrx",
   "homepage": "https://mtrx.so",
   "repository": {

package/src/matrx/__init__.py

@@ -1 +1 @@
-__version__ = "0.1.21"
+__version__ = "0.1.22"

package/src/matrx/cli/launcher.py

@@ -579,33 +579,21 @@ def _build_gemini_env(
     orchestration: dict | None = None,
 ) -> tuple[dict[str, str], str]:
     matrx = state["auth"]["matrx"]
-    # Assuming we might store Gemini-specific keys in future, or use OpenAI key fallback
-    # For now, we don't have a specific 'gemini' auth section in state.py, but we can assume
-    # if direct route, we use env var.
+    proxy_root = ensure_root_url(matrx.get("base_url"))
     proxy_base = ensure_v1_url(matrx.get("base_url"))
     mx_key, matrx_auth_source = _resolve_matrx_route_key(state, env)
-    
-    # Check for direct key in env or potentially saved elsewhere
-    direct_key = (env.get("GOOGLE_API_KEY") or "").strip()
+    direct_key = (env.get("GEMINI_API_KEY") or env.get("GOOGLE_API_KEY") or "").strip()
 
     if route == "matrx":
         if not mx_key:
             raise ValueError("No Matrx key available. Run: mtrx login matrx --key mx_... or set MTRX_KEY")
-
-        # Clear existing Gemini config to force proxy usage
         env.pop("MTRX_KEY", None)
-
         group_id, project_id = _resolve_matrx_context_overrides(state, env)
         session_id = str(uuid.uuid4())
         runtime_agent_id = (
             (orchestration or {}).get("agent_id")
             or _runtime_agent_basename("gemini")[0]
         )
-
-        # Build a project-aware proxy base URL so the proxy can identify the project
-        # even though Gemini CLI doesn't support custom request headers via env vars.
-        # We append matrx context as query params which the proxy reads.
-        proxy_base_with_ctx = proxy_base
         ctx_params: list[str] = []
         if project_id:
             ctx_params.append(f"mtrx_project={project_id}")
@@ -613,31 +601,65 @@ def _build_gemini_env(
             ctx_params.append(f"mtrx_session={session_id}")
         if runtime_agent_id:
             ctx_params.append(f"mtrx_agent={runtime_agent_id}")
-        if ctx_params:
-            proxy_base_with_ctx = f"{proxy_base}?{'&'.join(ctx_params)}"
+        git_branch, git_commit = _capture_git_context(_workspace_cwd(env))
+        if git_branch:
+            ctx_params.append(f"mtrx_branch={git_branch}")
+        if git_commit:
+            ctx_params.append(f"mtrx_commit={git_commit}")
+
+        query_suffix = f"?{'&'.join(ctx_params)}" if ctx_params else ""
+        env_snap = _capture_env_snapshot()
+        env_b64 = base64.b64encode(json.dumps(env_snap).encode()).decode() if env_snap else ""
+        custom_headers = [
+            f"x-matrx-key: {mx_key}",
+            f"x-matrx-agent-id: {runtime_agent_id}",
+            "x-matrx-provider: gemini_code",
+            f"x-matrx-session-id: {session_id}",
+        ]
+        if group_id:
+            custom_headers.append(f"x-matrx-group: {group_id}")
+        if project_id:
+            custom_headers.append(f"x-matrx-project-id: {project_id}")
+        if git_branch:
+            custom_headers.append(f"x-matrx-branch: {git_branch}")
+        if git_commit:
+            custom_headers.append(f"x-matrx-commit: {git_commit}")
+        if env_b64:
+            custom_headers.append(f"x-matrx-env: {env_b64}")
 
-        # Set Proxy Config
-        env["GOOGLE_GEMINI_BASE_URL"] = proxy_base_with_ctx
-        env["GEMINI_API_ENDPOINT"] = proxy_base_with_ctx
-        env["GOOGLE_API_KEY"] = mx_key
+        env["GOOGLE_GEMINI_BASE_URL"] = f"{proxy_base}/v1beta{query_suffix}"
+        env["GOOGLE_VERTEX_BASE_URL"] = f"{proxy_base}/v1beta{query_suffix}"
+        env["GEMINI_API_ENDPOINT"] = env["GOOGLE_GEMINI_BASE_URL"]
+        env["CODE_ASSIST_ENDPOINT"] = proxy_base
+        env["GEMINI_CLI_CUSTOM_HEADERS"] = ", ".join(custom_headers)
+        env["GEMINI_API_KEY_AUTH_MECHANISM"] = "bearer"
 
         return env, matrx_auth_source
 
     # Direct route: clear any matrx-managed env vars
     env.pop("MTRX_KEY", None)
-    
-    # Clear proxy overrides
-    _clear_if_matches(env, "GOOGLE_GEMINI_BASE_URL", proxy_base)
-    _clear_if_matches(env, "GEMINI_API_ENDPOINT", proxy_base)
-    
-    # Clear key if it was the Matrx key
-    current_key = (env.get("GOOGLE_API_KEY") or "").strip()
-    if current_key == mx_key or current_key.startswith("mx_"):
-        env.pop("GOOGLE_API_KEY", None)
-
-    if env.get("GOOGLE_API_KEY"):
+    for key in (
+        "GOOGLE_GEMINI_BASE_URL",
+        "GOOGLE_VERTEX_BASE_URL",
+        "GEMINI_API_ENDPOINT",
+        "CODE_ASSIST_ENDPOINT",
+    ):
+        value = (env.get(key) or "").strip()
+        if "matrx" in value.lower() or "mtrx.so" in value.lower():
+            env.pop(key, None)
+
+    custom_headers = (env.get("GEMINI_CLI_CUSTOM_HEADERS") or "").strip().lower()
+    if "x-matrx-" in custom_headers:
+        env.pop("GEMINI_CLI_CUSTOM_HEADERS", None)
+    if (env.get("GEMINI_API_KEY_AUTH_MECHANISM") or "").strip().lower() == "bearer":
+        env.pop("GEMINI_API_KEY_AUTH_MECHANISM", None)
+
+    if env.get("GEMINI_API_KEY") or env.get("GOOGLE_API_KEY"):
         return env, "existing_google_env"
-        
+
+    if direct_key:
+        return env, "existing_gemini_auth"
+
     return env, "missing_auth"
 
 
@@ -906,26 +928,40 @@ def _validate_gemini_launch_plan(plan: LaunchPlan, state: dict) -> None:
         return
 
     expected_base_url = ensure_v1_url(state.get("auth", {}).get("matrx", {}).get("base_url"))
+    expected_gemini_base = f"{expected_base_url}/v1beta"
 
     base_url = (plan.env.get("GOOGLE_GEMINI_BASE_URL") or "").strip()
     if not base_url:
         base_url = (plan.env.get("GEMINI_API_ENDPOINT") or "").strip()
 
     if not base_url:
-        raise ValueError("Gemini Matrx route is missing GOOGLE_GEMINI_BASE_URL or GEMINI_API_ENDPOINT")
+        raise ValueError("Gemini Matrx route is missing GOOGLE_GEMINI_BASE_URL")
 
-    # Strip query params before comparing (project_id context may be appended as ?mtrx_project=...)
     base_url_no_qs = base_url.split("?")[0].rstrip("/")
-    expected_no_qs = expected_base_url.rstrip("/")
+    expected_no_qs = expected_gemini_base.rstrip("/")
     if base_url_no_qs != expected_no_qs:
         raise ValueError(
-            "Gemini Matrx route must use the Matrx /v1 base URL. "
+            "Gemini Matrx route must use the Matrx Gemini-native base URL. "
             f"Got: {base_url}"
         )
 
-    mx_key = (plan.env.get("GOOGLE_API_KEY") or "").strip()
-    if not mx_key.startswith("mx_"):
-        raise ValueError("Gemini Matrx route is missing a valid GOOGLE_API_KEY (should be mx_...)")
+    vertex_base = (plan.env.get("GOOGLE_VERTEX_BASE_URL") or "").strip()
+    if vertex_base and vertex_base.split("?")[0].rstrip("/") != expected_no_qs:
+        raise ValueError("Gemini Matrx route is missing a Matrx GOOGLE_VERTEX_BASE_URL")
+
+    code_assist_endpoint = (plan.env.get("CODE_ASSIST_ENDPOINT") or "").strip().rstrip("/")
+    if code_assist_endpoint != expected_base_url.rstrip("/"):
+        raise ValueError("Gemini Matrx route is missing a Matrx CODE_ASSIST_ENDPOINT")
+
+    custom_headers = (plan.env.get("GEMINI_CLI_CUSTOM_HEADERS") or "").strip().lower()
+    if "x-matrx-key:" not in custom_headers:
+        raise ValueError("Gemini Matrx route is missing GEMINI_CLI_CUSTOM_HEADERS with X-Matrx-Key")
+    if "x-matrx-provider: gemini_code" not in custom_headers:
+        raise ValueError("Gemini Matrx route is missing GEMINI_CLI_CUSTOM_HEADERS with X-Matrx-Provider=gemini_code")
+    if "x-matrx-session-id:" not in custom_headers:
+        raise ValueError("Gemini Matrx route is missing GEMINI_CLI_CUSTOM_HEADERS with X-Matrx-Session-Id")
+    if "x-matrx-agent-id:" not in custom_headers:
+        raise ValueError("Gemini Matrx route is missing GEMINI_CLI_CUSTOM_HEADERS with X-Matrx-Agent-Id")
 
 
 def _validate_codex_launch_plan(plan: LaunchPlan, state: dict) -> None:
@@ -991,10 +1027,13 @@ def describe_launch_plan(plan: LaunchPlan, state: dict) -> list[str]:
     if plan.tool == "gemini":
         base_url = ensure_v1_url(state.get("auth", {}).get("matrx", {}).get("base_url"))
         _, project_id = _resolve_matrx_context_overrides(state, dict(plan.env))
+        custom_headers = (plan.env.get("GEMINI_CLI_CUSTOM_HEADERS") or "").strip()
         lines = [
             "Launching gemini via Matrx",
-            f"  base_url: {base_url}",
+            f"  gemini_base_url: {plan.env.get('GOOGLE_GEMINI_BASE_URL') or ''}",
+            f"  code_assist_endpoint: {plan.env.get('CODE_ASSIST_ENDPOINT') or ''}",
             f"  auth_source: {plan.auth_source}",
+            f"  custom_headers_present: {bool(custom_headers)}",
             "  runtime_route: env injection",
             "  persistent_route: disabled",
         ]