mtrx-cli 0.1.25 → 0.1.27

package/src/matrx/cli/cursor_reroute.py

@@ -17,20 +17,80 @@ import asyncio
 import json
 import logging
 import re
-from typing import Any
+from typing import Any, AsyncGenerator
 
 import httpx
 
 logger = logging.getLogger(__name__)
 
-# Cursor AI RPC paths (Connect protocol). RunSSE = main chat, StreamCpp = code completion.
+# Cursor AI RPC paths (Connect protocol).
+# Cursor uses aiserver.v1.AiServerService for all AI endpoints.
 _AI_PATH_PATTERNS = (
+    r"AiServerService",
+    # r"AiService" intentionally omitted — too broad, matches non-inference endpoints like
+    # KnowledgeBaseList, UpdateVscodeProfile, GetDefaultModel. Actual inference methods
+    # on AiService are all covered by their specific method-level patterns below.
+    r"ChatService",
+    r"StreamUnifiedChat",
+    r"StreamDiff",
+    r"GetCompletion",
     r"RunSSE",
     r"StreamCpp",
     r"BidiAppend",
     r"AgentService",
-    r"AiService",
+    r"/v1/messages",           # Anthropic Messages API
+    r"/v1/chat/completions",   # OpenAI Chat Completions
 )
+_REROUTABLE_AI_PATH_PATTERNS = (
+    r"AiService/.+(Stream|Run|Chat|Completion)",
+    r"StreamUnifiedChat",
+    r"RunSSE",
+    r"AgentService",
+    r"ChatService",
+    r"CppService/.+(Complete|Stream)",
+    r"/v1/messages",
+    r"/v1/chat/completions",
+)
+_AI_SERVICE_CANDIDATE_PATTERNS = (
+    r"AiServerService",
+    # r"AiService" intentionally omitted — see note in _AI_PATH_PATTERNS above.
+    r"AgentService",
+    r"ChatService",
+    r"CppService",
+    r"composer",
+)
+
+
+def is_ai_candidate_request(method: str, path: str, headers: dict[str, str] | None = None) -> bool:
+    """Return True when the request looks like Cursor/model traffic worth inspecting."""
+    if method.upper() != "POST" or not path:
+        return False
+    if is_ai_path(path):
+        return True
+    content_type = ((headers or {}).get("content-type") or "").lower()
+    return any(re.search(p, path, re.IGNORECASE) for p in _AI_SERVICE_CANDIDATE_PATTERNS) and (
+        "connect+proto" in content_type
+        or "application/proto" in content_type
+        or "grpc" in content_type
+        or "application/json" in content_type
+    )
+
+
+def is_reroutable_ai_path(path: str) -> bool:
+    """Return True if the path is a supported AI endpoint for active reroute."""
+    if not path:
+        return False
+    return any(re.search(p, path, re.IGNORECASE) for p in _REROUTABLE_AI_PATH_PATTERNS)
+
+
+def classify_ai_request(method: str, path: str, headers: dict[str, str] | None = None) -> dict[str, bool]:
+    """Classify Cursor requests for logging and reroute decisions."""
+    candidate = is_ai_candidate_request(method, path, headers)
+    reroutable = candidate and is_reroutable_ai_path(path)
+    return {
+        "candidate": candidate,
+        "reroutable": reroutable,
+    }
 
 
 def is_ai_path(path: str) -> bool:
@@ -40,6 +100,122 @@ def is_ai_path(path: str) -> bool:
     return any(re.search(p, path, re.IGNORECASE) for p in _AI_PATH_PATTERNS)
 
 
+def _detect_provider_from_model(model: str) -> str:
+    normalized = (model or "").strip().lower()
+    if normalized.startswith("claude-"):
+        return "anthropic"
+    if normalized.startswith("gemini"):
+        return "google"
+    return "openai"
+
+
+def _extract_request_prompt(*, req_proto: Any, extracted: dict[str, Any]) -> str:
+    parts: list[str] = []
+
+    summary = (extracted.get("conversation_summary") or "").strip()
+    if summary:
+        parts.append(summary)
+
+    debug_info = (getattr(getattr(req_proto, "cmd_k_debug_info", None), "debug_info", "") or "").strip()
+    if debug_info:
+        parts.append(debug_info)
+
+    rules = [
+        (getattr(rule, "rule_definition", "") or "").strip()
+        for rule in list(getattr(req_proto, "rules", []) or [])[:3]
+        if (getattr(rule, "rule_definition", "") or "").strip()
+    ]
+    if rules:
+        parts.append("Project rules:\n" + "\n".join(f"- {rule}" for rule in rules))
+
+    doc_ids = [
+        value.strip()
+        for value in list(getattr(getattr(req_proto, "legacy_context", None), "documentation_identifiers", []) or [])[:5]
+        if isinstance(value, str) and value.strip()
+    ]
+    if doc_ids:
+        parts.append("Relevant docs:\n" + "\n".join(f"- {value}" for value in doc_ids))
+
+    files = [
+        entry.get("path", "").strip()
+        for entry in extracted.get("files", [])[:8]
+        if entry.get("path")
+    ]
+    if files:
+        parts.append("Active files:\n" + "\n".join(f"- {path}" for path in files))
+
+    edits = [
+        entry.get("path", "").strip()
+        for entry in extracted.get("edits", [])[:5]
+        if entry.get("path")
+    ]
+    if edits:
+        parts.append("Recent edits:\n" + "\n".join(f"- {path}" for path in edits))
+
+    return "\n\n".join(part for part in parts if part).strip()
+
+
+def _build_matrx_upstream_request(
+    *,
+    req_proto: Any,
+    extracted: dict[str, Any],
+) -> tuple[str, dict[str, str], dict[str, Any]] | None:
+    model = (extracted.get("model") or "").strip()
+    prompt = _extract_request_prompt(req_proto=req_proto, extracted=extracted)
+    if not model or not prompt:
+        return None
+
+    provider = _detect_provider_from_model(model)
+    if provider == "anthropic":
+        return (
+            "/v1/messages",
+            {"x-api-key": ""},
+            {
+                "model": model,
+                "max_tokens": 2048,
+                "messages": [{"role": "user", "content": prompt}],
+            },
+        )
+
+    return (
+        "/v1/chat/completions",
+        {"authorization": ""},
+        {
+            "model": model,
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": 0.2,
+        },
+    )
+
+
+def _build_cursor_response_bytes(*, text: str, usage: dict[str, Any] | None = None) -> bytes | None:
+    try:
+        from matrx.cli.cursor_connect import build_connect_frame
+        from matrx.cli.cursor_proto import _PROTOS_AVAILABLE, server_chat_pb2  # type: ignore[import]
+    except Exception:
+        return None
+
+    if not _PROTOS_AVAILABLE:
+        return None
+
+    frames: list[bytes] = []
+    text = (text or "").strip()
+    if text:
+        content_resp = server_chat_pb2.StreamUnifiedChatWithToolsResponse()
+        content_resp.content.text = text
+        frames.append(build_connect_frame(0x00, content_resp.SerializeToString()))
+
+    if usage:
+        usage_resp = server_chat_pb2.StreamUnifiedChatWithToolsResponse()
+        usage_resp.usage.input_tokens = int(usage.get("input_tokens", 0) or 0)
+        usage_resp.usage.output_tokens = int(usage.get("output_tokens", 0) or 0)
+        if usage_resp.usage.input_tokens or usage_resp.usage.output_tokens:
+            frames.append(build_connect_frame(0x00, usage_resp.SerializeToString()))
+
+    frames.append(build_connect_frame(0x02, b"{}"))
+    return b"".join(frames)
+
+
 def _cursor_model_to_openai(cursor_model: str) -> str:
     """Map Cursor model names to OpenAI-style names MTRX expects."""
     # Cursor uses names like "claude-sonnet-4" or "gpt-4o" - usually compatible
@@ -54,7 +230,12 @@ def _build_search_query(extracted: dict[str, Any]) -> str:
     query_parts = [f.get("path", "").strip() for f in files[:3] if f.get("path")]
     if query_parts:
         return " ".join(query_parts)
-    return (extracted.get("conversation_summary") or "").strip()
+    summary = extracted.get("conversation_summary") or ""
+    if isinstance(summary, str):
+        return summary.strip()
+    if isinstance(summary, dict):
+        return str(summary.get("summary") or "").strip()
+    return str(summary).strip()
 
 
 def _prepend_context_items(context_items: Any, injected_items: list[Any]) -> None:
@@ -110,6 +291,85 @@ def _inject_memory_context_items(
     return len(injected_items)
 
 
+# Pre-built Connect end-of-stream frame: flags=0x02, payload=b"{}"
+# Frame format: [flags:1][length:4 BE][payload]  →  \x02 \x00\x00\x00\x02 {}
+_EOS_FRAME = b"\x02\x00\x00\x00\x02{}"
+
+
+async def _stream_rerouted_frames(
+    url: str,
+    payload: dict[str, Any],
+    headers: dict[str, str],
+    provider: str,
+) -> AsyncGenerator[bytes, None]:
+    """POST to MTRX and yield Connect-framed protobuf text deltas as SSE events arrive.
+
+    Parses each ``data: `` line from the SSE stream, extracts the text delta
+    (Anthropic ``content_block_delta`` or OpenAI ``choices[].delta.content``),
+    wraps it in a Connect data frame (flags=0x00), and yields it immediately so
+    Cursor sees tokens arrive incrementally.
+
+    Always terminates with a Connect end-of-stream frame (flags=0x02).  Any error
+    causes a silent fallback: the generator yields only the EOS frame so Cursor
+    sees an empty stream rather than a broken connection.
+    """
+    try:
+        from matrx.cli.cursor_connect import build_connect_frame
+        from matrx.cli.cursor_proto import _PROTOS_AVAILABLE, server_chat_pb2  # type: ignore[import]
+    except Exception:
+        yield _EOS_FRAME
+        return
+
+    if not _PROTOS_AVAILABLE:
+        yield _EOS_FRAME
+        return
+
+    try:
+        async with httpx.AsyncClient(
+            timeout=httpx.Timeout(timeout=90.0, connect=5.0)
+        ) as client:
+            async with client.stream("POST", url, json=payload, headers=headers) as resp:
+                if resp.status_code >= 400:
+                    logger.info(
+                        "cursor_reroute: stream upstream returned %s", resp.status_code
+                    )
+                    yield _EOS_FRAME
+                    return
+                async for raw_line in resp.aiter_lines():
+                    if not raw_line.startswith("data: "):
+                        continue
+                    data_str = raw_line[6:].strip()
+                    if data_str == "[DONE]":
+                        break
+                    try:
+                        chunk = json.loads(data_str)
+                    except json.JSONDecodeError:
+                        continue
+
+                    text = ""
+                    if provider == "anthropic":
+                        # Mirrors extract_from_anthropic_sse_response inner loop
+                        if chunk.get("type") == "content_block_delta":
+                            delta = chunk.get("delta") or {}
+                            if delta.get("type") == "text_delta":
+                                text = delta.get("text") or ""
+                    else:
+                        # Mirrors extract_from_openai_sse_response inner loop
+                        for choice in chunk.get("choices") or []:
+                            delta = choice.get("delta") or {}
+                            text += delta.get("content") or ""
+
+                    if text:
+                        resp_msg = server_chat_pb2.StreamUnifiedChatWithToolsResponse()
+                        resp_msg.content.text = text
+                        yield build_connect_frame(0x00, resp_msg.SerializeToString())
+
+    except Exception:
+        logger.warning("cursor_reroute: streaming reroute error", exc_info=True)
+
+    yield _EOS_FRAME
+
+
 async def try_reroute_to_matrx(
     *,
     path: str,
@@ -121,23 +381,94 @@ async def try_reroute_to_matrx(
     session_id: str | None = None,
     group_id: str | None = None,
     project_id: str | None = None,
-) -> tuple[bool, dict[str, str], bytes | None, bool] | None:
+) -> tuple[bool, dict[str, str], AsyncGenerator[bytes, None] | bytes | None, bool] | None:
     """
     Attempt to reroute a Cursor AI request through MTRX.
 
     Returns:
-        (success, response_headers, response_body, is_streaming) if handled,
+        (success, response_headers, response_body_or_generator, is_streaming) if handled,
         None to fall back to normal forward.
+        response_body_or_generator is an AsyncGenerator[bytes, None] of Connect frames;
+        the proxy must iterate it using chunked transfer encoding.
     """
-    if method != "POST" or not is_ai_path(path):
+    classification = classify_ai_request(method, path, req_headers)
+    if not classification["candidate"]:
+        return None
+    if not classification["reroutable"]:
+        logger.info("cursor_reroute: candidate AI path not yet reroutable: %s", path)
+        return None
+
+    try:
+        from matrx.cli.cursor_connect import is_connect_proto_request, parse_connect_frame
+        from matrx.cli.cursor_extraction import (
+            _PROTOS_AVAILABLE,
+            extract_from_request,
+            parse_request_proto,
+            ship_ai_telemetry,
+        )
+    except Exception:
         return None
 
-    # TODO: Full protobuf parsing. Cursor uses Connect/gRPC with binary frames.
-    # For now we don't have the proto conversion - fall back to forward.
-    # When implemented: parse req_body, extract messages+model, call MTRX,
-    # convert response back to Cursor's gRPC format.
-    logger.debug("cursor_reroute: path=%s would reroute (protobuf conversion not yet implemented)", path)
-    return None
+    if not _PROTOS_AVAILABLE or not is_connect_proto_request(req_headers):
+        logger.info("cursor_reroute: reroutable path lacks compiled proto support: %s", path)
+        return None
+
+    import gzip as _gzip
+
+    body = req_body
+    ce = req_headers.get("content-encoding", "").lower()
+    if ce == "gzip" or (len(body) >= 2 and body[:2] == b"\x1f\x8b"):
+        try:
+            body = _gzip.decompress(body)
+        except Exception:
+            return None
+
+    try:
+        _, proto_bytes = parse_connect_frame(body)
+    except ValueError:
+        proto_bytes = body
+
+    req_proto = parse_request_proto(proto_bytes)
+    if req_proto is None:
+        logger.info("cursor_reroute: parse_request_proto failed for %s", path)
+        return None
+
+    extracted = extract_from_request(req_proto)
+    extracted["session_id"] = extracted.get("session_id") or session_id or ""
+    asyncio.create_task(ship_ai_telemetry(extracted, matrx_base_url, matrx_key))
+
+    upstream_request = _build_matrx_upstream_request(req_proto=req_proto, extracted=extracted)
+    if upstream_request is None:
+        logger.info("cursor_reroute: insufficient prompt/model data for %s", path)
+        return None
+
+    upstream_path, auth_headers, payload = upstream_request
+    headers = {
+        "X-Matrx-Key": matrx_key,
+        "Content-Type": "application/json",
+    }
+    if "x-api-key" in auth_headers:
+        headers["x-api-key"] = matrx_key
+    if "authorization" in auth_headers:
+        headers["Authorization"] = f"Bearer {matrx_key}"
+
+    url = f"{matrx_base_url.rstrip('/')}{upstream_path}"
+    provider = _detect_provider_from_model(str(payload.get("model", "")))
+    gen = _stream_rerouted_frames(
+        url=url,
+        payload={**payload, "stream": True},
+        headers=headers,
+        provider=provider,
+    )
+    return (
+        True,
+        {
+            "content-type": req_headers.get("content-type", "application/connect+proto"),
+            "connect-protocol-version": "1",
+        },
+        gen,
+        True,
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -189,22 +520,50 @@ async def try_inject_context(
         )
         from matrx.cli.cursor_extraction import (
             _PROTOS_AVAILABLE,
+            _raw_extract_request,
+            extract_from_json_request,
             extract_from_request,
             parse_request_proto,
             ship_ai_telemetry,
         )
 
+        # JSON API path (Anthropic Messages API, OpenAI Chat Completions)
+        content_type = req_headers.get("content-type", "").lower()
+        if "application/json" in content_type:
+            extracted = extract_from_json_request(req_body)
+            extracted["session_id"] = extracted.get("session_id") or session_id
+            asyncio.create_task(ship_ai_telemetry(extracted, matrx_base_url, matrx_key))
+            return None
+
         if not is_connect_proto_request(req_headers):
             return None
 
-        # Parse Connect envelope → raw proto bytes
-        flags, proto_bytes = parse_connect_frame(req_body)
+        # Decompress body if gzip-encoded
+        import gzip as _gzip
+        body = req_body
+        ce = req_headers.get("content-encoding", "").lower()
+        if ce == "gzip" or (len(body) >= 2 and body[:2] == b"\x1f\x8b"):
+            try:
+                body = _gzip.decompress(body)
+            except Exception:
+                return None
+
+        # Parse Connect envelope → raw proto bytes.
+        # Fall back to treating the body as raw protobuf if Connect framing is absent
+        # (some Cursor versions send raw proto without the 5-byte envelope).
+        try:
+            flags, proto_bytes = parse_connect_frame(body)
+        except ValueError:
+            flags, proto_bytes = 0, body
 
-        # Deserialize proto
+        # Deserialize proto (compiled path); fall back to raw wire parsing
         req_proto = parse_request_proto(proto_bytes)
+        if req_proto is not None:
+            extracted = extract_from_request(req_proto)
+        else:
+            extracted = _raw_extract_request(proto_bytes)
 
         # Extract structured data and ship telemetry fire-and-forget
-        extracted = extract_from_request(req_proto)
         extracted["session_id"] = extracted.get("session_id") or session_id
         asyncio.create_task(ship_ai_telemetry(extracted, matrx_base_url, matrx_key))
 

package/src/matrx/cli/launcher.py

@@ -582,16 +582,21 @@ def _build_codex_env(
         env_b64 = base64.b64encode(json.dumps(env_snap).encode()).decode() if env_snap else ""
         session_id = str(uuid.uuid4())
         group_id, project_id = _resolve_matrx_context_overrides(state, env)
+        codex_root = ensure_root_url(matrx.get("base_url"))
+        if not group_id:
+            group_id = _auto_resolve_default_group_id(codex_root, mx_key)
         runtime_agent_id = (
             (orchestration or {}).get("agent_id")
             or _runtime_agent_basename("codex")[0]
         )
+        workspace_fp = _compute_workspace_fingerprint(_workspace_cwd(env))
         header_parts = [
             f'"Authorization" = "Bearer {provider_bearer}"',
             f'"X-Matrx-Key" = "{mx_key}"',
             f'"X-Matrx-Agent-Id" = "{runtime_agent_id}"',
             '"X-Matrx-Provider" = "codex"',
             f'"X-Matrx-Session-Id" = "{session_id}"',
+            f'"X-Matrx-Workspace" = "{workspace_fp}"',
         ]
         if group_id:
             header_parts.append(f'"X-Matrx-Group" = "{group_id}"')
@@ -655,24 +660,32 @@ def _build_gemini_env(
             env.pop(key, None)
         env.pop("MTRX_KEY", None)
         group_id, project_id = _resolve_matrx_context_overrides(state, env)
+        if not group_id:
+            group_id = _auto_resolve_default_group_id(proxy_root, mx_key)
         session_id = str(uuid.uuid4())
         runtime_agent_id = (
             (orchestration or {}).get("agent_id")
             or _runtime_agent_basename("gemini")[0]
         )
+        workspace_fp = _compute_workspace_fingerprint(_workspace_cwd(env))
         ctx_params: list[str] = []
         if project_id:
             ctx_params.append(f"mtrx_project={project_id}")
+        if group_id:
+            ctx_params.append(f"mtrx_group={group_id}")
         if session_id:
             ctx_params.append(f"mtrx_session={session_id}")
         if runtime_agent_id:
             ctx_params.append(f"mtrx_agent={runtime_agent_id}")
+        ctx_params.append(f"mtrx_workspace={workspace_fp}")
         git_branch, git_commit = _capture_git_context(_workspace_cwd(env))
         git_repo_url = _capture_git_remote_url(_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}")
+        if git_repo_url:
+            ctx_params.append(f"mtrx_repo_url={git_repo_url}")
 
         query_suffix = f"?{'&'.join(ctx_params)}" if ctx_params else ""
         env_snap = _capture_env_snapshot()
@@ -682,6 +695,7 @@ def _build_gemini_env(
             f"x-matrx-agent-id: {runtime_agent_id}",
             "x-matrx-provider: gemini_code",
             f"x-matrx-session-id: {session_id}",
+            f"x-matrx-workspace: {workspace_fp}",
         ]
         if group_id:
             custom_headers.append(f"x-matrx-group: {group_id}")
@@ -763,6 +777,35 @@ def _build_gemini_env(
     return env, "missing_auth"
 
 
+def _compute_workspace_fingerprint(cwd: str) -> str:
+    return hashlib.sha256(cwd.encode("utf-8")).hexdigest()[:16]
+
+
+def _auto_resolve_default_group_id(base_url: str, mx_key: str) -> str:
+    """Fetch the user's groups; return the sole/default group ID if unambiguous."""
+    if not base_url or not mx_key:
+        return ""
+    try:
+        with httpx.Client(timeout=5) as client:
+            resp = client.get(
+                f"{base_url.rstrip('/')}/v1/groups",
+                headers={"X-Matrx-Key": mx_key},
+            )
+        if resp.status_code != 200:
+            return ""
+        groups = resp.json().get("groups", [])
+        if not groups:
+            return ""
+        if len(groups) == 1:
+            return str(groups[0].get("id", ""))
+        for g in groups:
+            if g.get("is_default"):
+                return str(g.get("id", ""))
+        return ""
+    except (httpx.HTTPError, Exception):
+        return ""
+
+
 def _build_claude_env(
     state: dict,
     route: str,
@@ -789,12 +832,14 @@ def _build_claude_env(
         env["ANTHROPIC_BASE_URL"] = proxy_root
         env.pop("ANTHROPIC_API_KEY", None)
         group_id, project_id = _resolve_matrx_context_overrides(state, env)
+        if not group_id:
+            group_id = _auto_resolve_default_group_id(proxy_root, mx_key)
         session_id = str(uuid.uuid4())
         runtime_agent_id = (
             (orchestration or {}).get("agent_id")
             or _runtime_agent_basename("claude")[0]
         )
-        # Evolutionary scaffolding: env snapshot for AI context injection
+        workspace_fp = _compute_workspace_fingerprint(_workspace_cwd(env))
         env_snap = _capture_env_snapshot()
         env_b64 = base64.b64encode(json.dumps(env_snap).encode()).decode() if env_snap else ""
         custom_headers = "\n".join(
@@ -805,6 +850,7 @@ def _build_claude_env(
                 f"x-matrx-session-id: {session_id}",
             ]
         )
+        custom_headers += f"\nx-matrx-workspace: {workspace_fp}"
         if group_id:
             custom_headers += f"\nx-matrx-group: {group_id}"
         if project_id: