mtrx-cli 0.1.24 → 0.1.26

package/src/matrx/cli/cursor_reroute.py

@@ -13,23 +13,84 @@ Refs: cursor-tap (https://github.com/burpheart/cursor-tap), everestmz/cursor-rpc
 
 from __future__ import annotations
 
+import asyncio
 import json
 import logging
 import re
 from typing import Any
 
+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",
+    r"ChatService",
+    r"StreamUnifiedChat",
+    r"StreamDiff",
+    r"GetCompletion",
     r"RunSSE",
     r"StreamCpp",
     r"BidiAppend",
     r"AgentService",
+    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",
+    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:
     """Return True if this path is an AI/LLM endpoint we should reroute to MTRX."""
     if not path:
@@ -37,6 +98,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
@@ -46,6 +223,72 @@ def _cursor_model_to_openai(cursor_model: str) -> str:
     return cursor_model
 
 
+def _build_search_query(extracted: dict[str, Any]) -> str:
+    files = extracted.get("files", [])
+    query_parts = [f.get("path", "").strip() for f in files[:3] if f.get("path")]
+    if query_parts:
+        return " ".join(query_parts)
+    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:
+    for item in reversed(injected_items):
+        try:
+            context_items.insert(0, item)
+        except Exception:
+            context_items.append(item)
+
+
+def _inject_memory_context_items(
+    *,
+    req_proto: Any,
+    memory_results: list[dict[str, Any]],
+    server_chat_pb2: Any,
+    existing_files: list[dict[str, Any]] | None = None,
+    limit: int = 5,
+) -> int:
+    existing_contents = {
+        (entry.get("content") or "").strip() for entry in (existing_files or []) if entry.get("content")
+    }
+    injected_contents: set[str] = set()
+    injected_items: list[Any] = []
+
+    for mem in memory_results:
+        if len(injected_items) >= limit:
+            break
+        content = (mem.get("content") or "").strip()
+        if not content or content in existing_contents or content in injected_contents:
+            continue
+
+        cached_item = server_chat_pb2.PotentiallyCachedContextItem()
+        ctx_item = server_chat_pb2.ContextItem()
+        file_chunk = server_chat_pb2.FileChunk()
+        file_chunk.file_path = f"[matrx:{(mem.get('id') or '')[:8]}]"
+        file_chunk.content = content[:4096]
+        if hasattr(getattr(ctx_item, "file_chunk", None), "CopyFrom"):
+            ctx_item.file_chunk.CopyFrom(file_chunk)
+        else:
+            ctx_item.file_chunk = file_chunk
+        if hasattr(getattr(cached_item, "context_item", None), "CopyFrom"):
+            cached_item.context_item.CopyFrom(ctx_item)
+        else:
+            cached_item.context_item = ctx_item
+
+        injected_items.append(cached_item)
+        injected_contents.add(content)
+
+    if not injected_items:
+        return 0
+
+    _prepend_context_items(req_proto.context_items, injected_items)
+    return len(injected_items)
+
+
 async def try_reroute_to_matrx(
     *,
     path: str,
@@ -65,12 +308,251 @@ async def try_reroute_to_matrx(
         (success, response_headers, response_body, is_streaming) if handled,
         None to fall back to normal forward.
     """
-    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_anthropic_sse_response,
+            extract_from_openai_sse_response,
+            extract_from_request,
+            parse_request_proto,
+            ship_ai_telemetry,
+        )
+    except Exception:
+        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}"
+    try:
+        async with httpx.AsyncClient(timeout=httpx.Timeout(timeout=90.0, connect=5.0)) as client:
+            resp = await client.post(url, json={**payload, "stream": True}, headers=headers)
+    except Exception:
+        logger.warning("cursor_reroute: upstream request failed for %s", path, exc_info=True)
+        return None
+
+    if resp.status_code >= 400:
+        logger.info(
+            "cursor_reroute: upstream returned %s for %s; forwarding unchanged",
+            resp.status_code,
+            path,
+        )
+        return None
+
+    provider = _detect_provider_from_model(str(payload.get("model", "")))
+    if provider == "anthropic":
+        frame_data = extract_from_anthropic_sse_response(resp.content)
+    else:
+        frame_data = extract_from_openai_sse_response(resp.content)
+
+    text = frame_data.get("text", "")
+    usage = frame_data.get("usage")
+    response_body = _build_cursor_response_bytes(text=text, usage=usage)
+    if response_body is None:
+        return None
+
+    response_telemetry = {
+        "session_id": extracted.get("session_id") or session_id or "",
+        "conversation_id": extracted.get("conversation_id") or "",
+        "model": extracted.get("model") or "",
+        "files": extracted.get("files", []),
+        "edits": extracted.get("edits", []),
+        "response_text": text,
+        "tool_calls": [],
+        "usage": usage,
+    }
+    asyncio.create_task(ship_ai_telemetry(response_telemetry, matrx_base_url, matrx_key))
+
+    return (
+        True,
+        {
+            "content-type": req_headers.get("content-type", "application/connect+proto"),
+            "connect-protocol-version": "1",
+        },
+        response_body,
+        True,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Context injection
+# ---------------------------------------------------------------------------
+
 
-    # 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
+async def _query_memory(
+    *,
+    query: str,
+    matrx_base_url: str,
+    matrx_key: str,
+    limit: int = 5,
+) -> list[dict]:
+    """Query MTRX memory search API. Returns list of memory entry dicts."""
+    try:
+        async with httpx.AsyncClient(timeout=0.1) as client:  # 100 ms budget
+            resp = await client.get(
+                f"{matrx_base_url.rstrip('/')}/v1/memory/search",
+                params={"q": query, "limit": limit},
+                headers={"X-Matrx-Key": matrx_key},
+            )
+            if resp.status_code == 200:
+                return resp.json().get("entries", [])
+    except Exception:
+        logger.debug("cursor_reroute: memory query failed", exc_info=True)
+    return []
+
+
+async def try_inject_context(
+    *,
+    req_body: bytes,
+    req_headers: dict[str, str],
+    matrx_base_url: str,
+    matrx_key: str,
+    session_id: str,
+) -> bytes | None:
+    """Parse the Connect+proto request, extract data, inject MTRX memory context items.
+
+    Returns modified request bytes with injected context items prepended, or
+    ``None`` to signal that the original request should be forwarded unchanged.
+    Wraps everything in try/except — never raises.
+    """
+    try:
+        from matrx.cli.cursor_connect import (
+            build_connect_frame,
+            is_connect_proto_request,
+            parse_connect_frame,
+        )
+        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
+
+        # 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 (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["session_id"] = extracted.get("session_id") or session_id
+        asyncio.create_task(ship_ai_telemetry(extracted, matrx_base_url, matrx_key))
+
+        # Cannot inject without compiled protos or a successfully parsed proto
+        if not _PROTOS_AVAILABLE or req_proto is None:
+            return None
+
+        # Build search query from the open file paths
+        search_query = _build_search_query(extracted)
+        if not search_query:
+            return None
+
+        # Query MTRX memory (100 ms budget — never blocks the forward path)
+        memory_results = await _query_memory(
+            query=search_query,
+            matrx_base_url=matrx_base_url,
+            matrx_key=matrx_key,
+            limit=5,
+        )
+        if not memory_results:
+            return None
+
+        # Inject memory results as PotentiallyCachedContextItem entries
+        from matrx.cli.cursor_proto import server_chat_pb2  # type: ignore[import]
+
+        inserted = _inject_memory_context_items(
+            req_proto=req_proto,
+            memory_results=memory_results,
+            server_chat_pb2=server_chat_pb2,
+            existing_files=extracted.get("files", []),
+        )
+        if inserted == 0:
+            return None
+
+        # Serialize modified proto and re-wrap in Connect frame
+        new_proto_bytes = req_proto.SerializeToString()
+        return build_connect_frame(flags, new_proto_bytes)
+
+    except Exception:
+        logger.warning("cursor_reroute: try_inject_context failed", exc_info=True)
+        return None

package/src/matrx/cli/launcher.py

@@ -663,6 +663,8 @@ def _build_gemini_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:
@@ -673,6 +675,8 @@ def _build_gemini_env(
             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()