mtrx-cli 0.1.24 → 0.1.25

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mtrx-cli",
-  "version": "0.1.24",
+  "version": "0.1.25",
   "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.24"
+__version__ = "0.1.25"

package/src/matrx/cli/cursor_proxy.py

@@ -33,7 +33,7 @@ import httpx
 from matrx.cli.cursor_ca import CertCache, load_ca
 
 try:
-    from matrx.cli.cursor_reroute import is_ai_path, try_reroute_to_matrx
+    from matrx.cli.cursor_reroute import is_ai_path, try_inject_context, try_reroute_to_matrx
 except ImportError:
     # Stubs when cursor_reroute not available (e.g. npm package omit).
     def is_ai_path(path: str) -> bool:
@@ -42,6 +42,9 @@ except ImportError:
     async def try_reroute_to_matrx(*, path: str, method: str, **kwargs: Any) -> None:
         return None
 
+    async def try_inject_context(**kwargs: Any) -> None:
+        return None
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_PORT = 8842
@@ -287,8 +290,10 @@ class MITMProxy:
             path = parts[1] if len(parts) > 1 else "/"
 
             req_body_size = 0
+            _is_ai_req = method == "POST" and is_ai_path(path)
+            _req_session_id = str(uuid.uuid4()) if _is_ai_req else ""
             # For AI paths: buffer request and try rerouting through MTRX (live injection)
-            if method == "POST" and is_ai_path(path):
+            if _is_ai_req:
                 req_headers, req_cl, req_chunked = await self._read_headers_only(
                     client_reader
                 )
@@ -303,7 +308,7 @@ class MITMProxy:
                     req_body=req_body,
                     matrx_base_url=self.matrx_base_url,
                     matrx_key=self.matrx_key,
-                    session_id=str(uuid.uuid4()),
+                    session_id=_req_session_id,
                 )
                 if result is not None:
                     success, resp_headers, resp_body, is_streaming = result
@@ -327,10 +332,20 @@ class MITMProxy:
                         )
                         continue
                     # Reroute returned but failed — fall through to forward
-                # Reroute not implemented or failed — forward to upstream
+                # Inject MTRX memory context into request before forwarding
+                injected_body = await try_inject_context(
+                    req_body=req_body,
+                    req_headers=req_headers,
+                    matrx_base_url=self.matrx_base_url,
+                    matrx_key=self.matrx_key,
+                    session_id=_req_session_id,
+                )
+                body_to_forward = injected_body if injected_body is not None else req_body
+                fwd_headers = dict(req_headers)
+                fwd_headers["content-length"] = str(len(body_to_forward))
                 up_writer.write(req_line)
-                await self._write_headers(up_writer, req_headers)
-                up_writer.write(req_body)
+                await self._write_headers(up_writer, fwd_headers)
+                up_writer.write(body_to_forward)
                 await up_writer.drain()
             else:
                 up_writer.write(req_line)
@@ -369,9 +384,20 @@ class MITMProxy:
                 for t in ("text/event-stream", "grpc", "proto", "connect")
             )
 
-            resp_body_size = await self._forward_body(
-                up_reader, client_writer, resp_cl, resp_chunked
-            )
+            if _is_ai_req:
+                resp_body_size, resp_captured = await self._forward_body_with_capture(
+                    up_reader, client_writer, resp_cl, resp_chunked
+                )
+                if resp_captured:
+                    asyncio.create_task(
+                        self._extract_ai_response(
+                            resp_captured, _req_session_id, hostname
+                        )
+                    )
+            else:
+                resp_body_size = await self._forward_body(
+                    up_reader, client_writer, resp_cl, resp_chunked
+                )
 
             elapsed_ms = int((time.monotonic() - started) * 1000)
             self._request_count += 1
@@ -397,6 +423,114 @@ class MITMProxy:
             if "close" in conn_h:
                 break
 
+    async def _forward_body_with_capture(
+        self,
+        reader: asyncio.StreamReader,
+        writer: asyncio.StreamWriter,
+        content_length: int,
+        chunked: bool,
+    ) -> tuple[int, bytes]:
+        """Forward body like ``_forward_body`` while also capturing a copy.
+
+        Returns ``(bytes_forwarded, captured_bytes)``.  The capture enables
+        background response extraction without blocking the forward path.
+        """
+        parts: list[bytes] = []
+
+        if content_length > 0:
+            total = 0
+            remaining = content_length
+            while remaining > 0:
+                chunk = await reader.read(min(remaining, 65536))
+                if not chunk:
+                    break
+                writer.write(chunk)
+                await writer.drain()
+                parts.append(chunk)
+                total += len(chunk)
+                remaining -= len(chunk)
+            return total, b"".join(parts)
+
+        if chunked:
+            total = 0
+            while True:
+                size_line = await reader.readline()
+                if not size_line:
+                    break
+                writer.write(size_line)
+                await writer.drain()
+                size_str = size_line.decode("utf-8", errors="replace").strip()
+                try:
+                    chunk_size = int(size_str.split(";")[0], 16)
+                except ValueError:
+                    break
+                if chunk_size == 0:
+                    trailer = await reader.readline()
+                    writer.write(trailer)
+                    await writer.drain()
+                    break
+                remaining = chunk_size
+                chunk_parts: list[bytes] = []
+                while remaining > 0:
+                    data = await reader.read(min(remaining, 65536))
+                    if not data:
+                        return total, b"".join(parts)
+                    writer.write(data)
+                    await writer.drain()
+                    chunk_parts.append(data)
+                    total += len(data)
+                    remaining -= len(data)
+                chunk_data = b"".join(chunk_parts)
+                parts.append(chunk_data)
+                crlf = await reader.readline()
+                writer.write(crlf)
+                await writer.drain()
+            return total, b"".join(parts)
+
+        return 0, b""
+
+    async def _extract_ai_response(
+        self,
+        resp_bytes: bytes,
+        session_id: str,
+        hostname: str,
+    ) -> None:
+        """Parse Connect frames from *resp_bytes* and ship response telemetry.
+
+        Fire-and-forget — never raises, never blocks the forward path.
+        """
+        try:
+            from matrx.cli.cursor_connect import parse_all_frames
+            from matrx.cli.cursor_extraction import (
+                extract_from_response_frame,
+                parse_response_proto,
+                ship_ai_telemetry,
+            )
+
+            frames = parse_all_frames(resp_bytes)
+            accumulated: dict = {
+                "session_id": session_id,
+                "response_text": "",
+                "tool_calls": [],
+                "usage": None,
+            }
+            for flags, payload in frames:
+                if flags == 0x02:  # end-of-stream trailer — stop
+                    break
+                resp_proto = parse_response_proto(payload)
+                frame_data = extract_from_response_frame(resp_proto)
+                if frame_data:
+                    accumulated["response_text"] = (
+                        accumulated.get("response_text", "") + frame_data.get("text", "")
+                    )
+                    accumulated["tool_calls"].extend(frame_data.get("tool_calls", []))
+                    if frame_data.get("usage"):
+                        accumulated["usage"] = frame_data["usage"]
+
+            await ship_ai_telemetry(accumulated, self.matrx_base_url, self.matrx_key)
+        except Exception:
+            logger.debug("proxy: _extract_ai_response failed", exc_info=True)
+
     async def _read_headers_only(
         self, reader: asyncio.StreamReader
     ) -> tuple[dict[str, str], int, bool]:

package/src/matrx/cli/cursor_reroute.py

@@ -13,11 +13,14 @@ 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.
@@ -46,6 +49,67 @@ 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)
+    return (extracted.get("conversation_summary") or "").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,
@@ -74,3 +138,111 @@ async def try_reroute_to_matrx(
     # convert response back to Cursor's gRPC format.
     logger.debug("cursor_reroute: path=%s would reroute (protobuf conversion not yet implemented)", path)
     return None
+
+
+# ---------------------------------------------------------------------------
+# Context injection
+# ---------------------------------------------------------------------------
+
+
+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,
+            extract_from_request,
+            parse_request_proto,
+            ship_ai_telemetry,
+        )
+
+        if not is_connect_proto_request(req_headers):
+            return None
+
+        # Parse Connect envelope → raw proto bytes
+        flags, proto_bytes = parse_connect_frame(req_body)
+
+        # Deserialize proto
+        req_proto = parse_request_proto(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))
+
+        # 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