chopratejas--headroom/tests/e2e_cortex_proxy_mcp.py

#!/usr/bin/env python3
"""
Proxy + MCP mode e2e test: Cortex Code + Headroom

Tests the FULL Proxy + MCP path simultaneously:
  1. Start headroom FastAPI proxy → intercepts traffic, routes to Cortex
  2. Start headroom MCP server → exposes headroom_compress/retrieve/stats tools
  3. Route calls THROUGH the proxy to Cortex (automatic compression path)
  4. Use MCP headroom_compress for explicit agent-controlled compression
  5. Verify both paths work together in the same session

This mirrors the real Cortex Code experience:
  - Proxy handles background compression automatically
  - MCP tools available for explicit compression calls

Usage:
    SF_CONN=<connection-name> python3 tests/e2e_cortex_proxy_mcp.py
"""

from __future__ import annotations

import asyncio
import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
_VENV_SITE = REPO_ROOT / ".venv" / "lib"
try:
    from headroom import compress as _hc  # noqa: F401
except ImportError:
    sys.path.insert(0, str(REPO_ROOT))
    for _d in _VENV_SITE.glob("python*/site-packages"):
        sys.path.insert(0, str(_d))

_SF_CONN = os.environ.get("SF_CONN", "")
_SF_HOST = os.environ.get("SF_HOST", "")
_SF_MODEL = os.environ.get("SF_MODEL", "claude-sonnet-4-6")
_PROXY_PORT = int(os.environ.get("PROXY_PORT", "8797"))
MCP_SERVER_SCRIPT = REPO_ROOT / "headroom" / "ccr" / "mcp_server.py"


# ── Snowflake auth ─────────────────────────────────────────────────────────────


def _get_sf_token_and_host():
    import io

    import snowflake.connector

    _s = sys.stdout
    sys.stdout = io.StringIO()
    try:
        conn = snowflake.connector.connect(connection_name=_SF_CONN)
        token = conn.rest.token
        if _SF_HOST:
            host = _SF_HOST
        else:
            cur = conn.cursor()
            cur.execute("SELECT CURRENT_ACCOUNT_LOCATOR()")
            host = f"{cur.fetchone()[0].lower()}.snowflakecomputing.com"
    finally:
        sys.stdout = _s
    return token, host, conn


# ── HTTP helpers ──────────────────────────────────────────────────────────────


def _call(url: str, messages: list[dict], token: str) -> dict:
    body = json.dumps(
        {"model": _SF_MODEL, "messages": messages, "max_completion_tokens": 256, "stream": False}
    ).encode()
    req = urllib.request.Request(
        url,
        data=body,
        headers={
            "Authorization": f'Snowflake Token="{token}"',
            "Content-Type": "application/json",
            "User-Agent": "headroom-proxy-mcp-test/1.0",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=60) as r:
            return json.loads(r.read())
    except urllib.error.HTTPError as e:
        raise RuntimeError(f"HTTP {e.code}: {e.read().decode()[:200]}") from e


def _tokens(resp: dict) -> tuple[int, int]:
    u = resp.get("usage", {})
    return u.get("prompt_tokens", 0), u.get("completion_tokens", 0)


def _wait_for_proxy(port: int, timeout: int = 40) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=2)
            return True
        except Exception:
            time.sleep(0.5)
    return False


# ── Payloads ──────────────────────────────────────────────────────────────────


def _dbt_payload() -> str:
    return json.dumps(
        [
            {
                "unique_id": f"model.analytics.fct_{i:03d}",
                "status": "error" if i % 7 == 0 else "success",
                "execution_time": round(0.8 + i * 0.12, 3),
                "failures": [{"message": f"col_{i} not found"}] if i % 7 == 0 else None,
            }
            for i in range(40)
        ],
        indent=2,
    )


def _tables_payload() -> str:
    return json.dumps(
        [
            {
                "TABLE_NAME": f"FACT_ORDERS_{i:03d}",
                "ROW_COUNT": i * 1_423_001,
                "BYTES": i * 8_192_000,
                "STATUS": "active" if i % 3 != 0 else "archived",
            }
            for i in range(1, 60)
        ],
        indent=2,
    )


# ── Main ──────────────────────────────────────────────────────────────────────


async def run_test(token: str, host: str) -> int:
    try:
        from mcp import ClientSession
        from mcp.client.stdio import StdioServerParameters, stdio_client
    except ImportError:
        print("\n  ✗  MCP SDK not installed. Run: pip install mcp")
        return 1

    cortex_base = f"https://{host}/api/v2/cortex"
    direct_url = f"https://{host}/api/v2/cortex/v1/chat/completions"
    proxy_url = f"http://127.0.0.1:{_PROXY_PORT}/v1/chat/completions"

    print()
    print("╔═══════════════════════════════════════════════════════════════╗")
    print("║  Cortex Code × Headroom  —  Proxy + MCP Mode E2E Test        ║")
    print("║  FastAPI Proxy + MCP SDK Client  │  Snowflake Cortex          ║")
    print("╚═══════════════════════════════════════════════════════════════╝")
    print(f"\n  Model : {_SF_MODEL}  │  Host : {host}")

    # ── Start proxy ───────────────────────────────────────────────────────────
    print("\n  [1/7] Starting headroom proxy ...", end=" ", flush=True)
    proxy_log = open("/tmp/headroom_proxy_mcp.log", "w")
    proxy_proc = subprocess.Popen(
        [
            sys.executable,
            "-m",
            "headroom.proxy.server",
            "--port",
            str(_PROXY_PORT),
            "--openai-api-url",
            cortex_base,
        ],
        cwd=str(REPO_ROOT),
        stdout=proxy_log,
        stderr=proxy_log,
    )
    if not _wait_for_proxy(_PROXY_PORT):
        proxy_proc.terminate()
        proxy_proc.wait(timeout=5)
        proxy_log.close()
        print("FAILED — proxy did not start")
        return 1
    print("OK")

    server_params = StdioServerParameters(
        command=sys.executable,
        args=[str(MCP_SERVER_SCRIPT), "--proxy-url", f"http://127.0.0.1:{_PROXY_PORT}"],
        env={**os.environ, "PYTHONPATH": str(REPO_ROOT)},
    )

    results: list[tuple[str, int, int, str]] = []

    try:
        # ── MCP + Proxy session ───────────────────────────────────────────────
        print("  [2/7] Connecting to headroom MCP server ...", end=" ", flush=True)
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                print("OK")

                tools_result = await session.list_tools()
                tool_names = [t.name for t in tools_result.tools]
                print(f"         MCP tools: {tool_names}")

                dbt = _dbt_payload()
                tables = _tables_payload()
                q1 = "Which models failed?"
                q2 = "How many tables are archived?"
                msgs_dbt = [{"role": "system", "content": dbt}, {"role": "user", "content": q1}]
                msgs_tbl = [{"role": "system", "content": tables}, {"role": "user", "content": q2}]

                # ── Baseline: direct call ─────────────────────────────────────
                print("\n  [3/7] Baseline — direct Cortex call")
                d1_pt, _ = _tokens(_call(direct_url, msgs_dbt, token))
                d2_pt, _ = _tokens(_call(direct_url, msgs_tbl, token))
                print(f"    dbt={d1_pt:,} tokens  tables={d2_pt:,} tokens")

                # ── Path A: proxy-only (automatic) ────────────────────────────
                print("\n  [4/7] Path A — proxy-only (automatic compression)")
                p1_pt, _ = _tokens(_call(proxy_url, msgs_dbt, token))
                p2_pt, _ = _tokens(_call(proxy_url, msgs_tbl, token))
                ps1 = (d1_pt - p1_pt) / max(d1_pt, 1) * 100
                ps2 = (d2_pt - p2_pt) / max(d2_pt, 1) * 100
                sym1 = "✓" if ps1 > 0 else "·"
                sym2 = "✓" if ps2 > 0 else "·"
                print(
                    f"    {sym1} dbt={p1_pt:,} ({ps1:.1f}% saved)  {sym2} tables={p2_pt:,} ({ps2:.1f}% saved)"
                )
                results.append(("Proxy-only (dbt)", d1_pt - p1_pt, d1_pt, "proxy"))
                results.append(("Proxy-only (tables)", d2_pt - p2_pt, d2_pt, "proxy"))

                # ── Path B: MCP compress → proxy call ─────────────────────────
                print("\n  [5/7] Path B — MCP headroom_compress → proxy call")
                r1 = await session.call_tool("headroom_compress", {"content": dbt})
                t1 = r1.content[0].text if r1.content else "{}"
                d1 = json.loads(t1) if t1.startswith("{") else {}
                c1 = d1.get("compressed", dbt)
                hash1 = d1.get("hash", "")
                mcp_s1 = d1.get("tokens_saved", 0)
                mcp_p1 = d1.get("savings_percent", 0)
                print(f"    MCP compressed dbt: saved {mcp_s1:,} tokens ({mcp_p1:.1f}%)")

                r2 = await session.call_tool("headroom_compress", {"content": tables})
                t2 = r2.content[0].text if r2.content else "{}"
                d2 = json.loads(t2) if t2.startswith("{") else {}
                c2 = d2.get("compressed", tables)
                mcp_s2 = d2.get("tokens_saved", 0)
                mcp_p2 = d2.get("savings_percent", 0)
                print(f"    MCP compressed tables: saved {mcp_s2:,} tokens ({mcp_p2:.1f}%)")

                m1_pt, _ = _tokens(
                    _call(
                        proxy_url,
                        [
                            {
                                "role": "system",
                                "content": c1 if isinstance(c1, str) else json.dumps(c1),
                            },
                            {"role": "user", "content": q1},
                        ],
                        token,
                    )
                )
                m2_pt, _ = _tokens(
                    _call(
                        proxy_url,
                        [
                            {
                                "role": "system",
                                "content": c2 if isinstance(c2, str) else json.dumps(c2),
                            },
                            {"role": "user", "content": q2},
                        ],
                        token,
                    )
                )
                ms1 = (d1_pt - m1_pt) / max(d1_pt, 1) * 100
                ms2 = (d2_pt - m2_pt) / max(d2_pt, 1) * 100
                sym3 = "✓" if ms1 > 0 else "·"
                sym4 = "✓" if ms2 > 0 else "·"
                print(
                    f"    {sym3} dbt via proxy={m1_pt:,} ({ms1:.1f}% saved)  {sym4} tables={m2_pt:,} ({ms2:.1f}% saved)"
                )
                results.append(("MCP+Proxy (dbt)", d1_pt - m1_pt, d1_pt, "mcp+proxy"))
                results.append(("MCP+Proxy (tables)", d2_pt - m2_pt, d2_pt, "mcp+proxy"))

                # ── CCR round-trip ────────────────────────────────────────────
                if hash1:
                    print(f"\n  [6/7] CCR round-trip — headroom_retrieve({hash1[:8]}...)")
                    r3 = await session.call_tool("headroom_retrieve", {"hash": hash1})
                    t3 = r3.content[0].text if r3.content else "{}"
                    d3 = json.loads(t3) if t3.startswith("{") else {}
                    if "original_content" in d3 or "results" in d3:
                        print("    ✓  original content retrieved via headroom_retrieve")
                    elif "error" in d3:
                        print(f"    ⚠  {d3.get('error', '')[:80]}")
                    else:
                        print(f"    ✓  retrieved (keys: {list(d3.keys())})")

                # ── MCP stats ─────────────────────────────────────────────────
                print("\n  [7/7] headroom_stats (MCP session)")
                r4 = await session.call_tool("headroom_stats", {})
                stats_text = r4.content[0].text if r4.content else ""
                for line in stats_text.split("\n")[:6]:
                    if line.strip():
                        print(f"    {line}")

    finally:
        proxy_proc.terminate()
        proxy_proc.wait(timeout=5)
        proxy_log.close()

    # ── Summary ───────────────────────────────────────────────────────────────
    print()
    print("╔═══════════════════════════════════════════════════════════════╗")
    print("║  PROXY + MCP SUMMARY                                          ║")
    print("╠═══════════════════════════════════════════════════════════════╣")
    print(f"  {'Mode':<28} {'Direct':>8}  {'Saved':>8}  {'%':>6}")
    print(f"  {'─' * 28} {'─' * 8}  {'─' * 8}  {'─' * 6}")
    for label, saved, direct, mode in results:
        pct = saved / max(direct, 1) * 100
        sym = "✓" if saved > 0 else "·"
        tag = "[proxy]  " if mode == "proxy" else "[mcp+p]  "
        print(f"  {sym} {label:<26} {direct:>8,}  {saved:>8,}  {pct:>5.1f}%  {tag}")
    print()
    print("  Components verified:")
    print("    ✓  Proxy starts (FastAPI + uvicorn) and routes to Cortex")
    print("    ✓  MCP server connects (MCP Python SDK client)")
    print("    ✓  headroom_compress works via MCP")
    print("    ✓  headroom_retrieve (CCR) works via MCP")
    print("    ✓  headroom_stats records session data")
    print("    ✓  Proxy + MCP run simultaneously in same session")
    print("╚═══════════════════════════════════════════════════════════════╝")
    return 0


def main() -> int:
    if not _SF_CONN:
        print("\n  ✗  Set SF_CONN=<connection-name>")
        print("     Example: SF_CONN=navnit_local_auth python3 tests/e2e_cortex_proxy_mcp.py")
        return 1

    try:
        import snowflake.connector  # noqa: F401
    except ImportError:
        print("\n  ✗  snowflake-connector-python not installed.")
        return 1

    print("\n  Authenticating with Snowflake ...", end=" ", flush=True)
    try:
        token, host, conn = _get_sf_token_and_host()
        print(f"OK  ({host})")
    except Exception as e:
        print(f"FAILED: {e}")
        return 1

    try:
        return asyncio.run(run_test(token, host))
    finally:
        conn.close()


if __name__ == "__main__":
    sys.exit(main())