akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(cortex-code): migrate to current Cortex REST API endpoints + add e2e benchmarks (#1474)

Browse source 
## Description

Follow-up to #1190 (Cortex Code provider). Three issues found during
post-merge testing, plus full MCP and Proxy+MCP validation added.

Closes #

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [x] New feature (non-breaking change that adds functionality)
- [x] Documentation update

## Changes Made

- `docs/cortex-code.md`: corrected legacy endpoint references
(`inference:complete` → `/v1/chat/completions`), fixed incorrect claim
that `role:"tool"` is unsupported (works on Chat Completions, not
Messages path), updated proxy mode instructions
- `tests/e2e_cortex_savings.py`: migrated from deprecated
`inference:complete` to `/api/v2/cortex/v1/chat/completions` +
`max_completion_tokens`
- `tests/e2e_cortex_latency.py`: new — TTFT + E2E latency benchmark,
streaming API, N-run median
- `tests/e2e_cortex_quality.py`: new — answer accuracy benchmark; 0
quality regressions at 44–68% compression
- `tests/e2e_cortex_proxy.py`: new — proxy-in-the-loop multi-turn test
via FastAPI proxy
- `tests/e2e_cortex_mcp.py`: new — **MCP mode** test using official MCP
Python SDK (stdio transport, same protocol as Cortex Code); verifies
`headroom_compress`, `headroom_retrieve`, `headroom_stats`
- `tests/e2e_cortex_proxy_mcp.py`: new — **Proxy + MCP** test; starts
FastAPI proxy and MCP server simultaneously, exercises both paths in
same session

## Testing

- [x] Linting passes (`ruff check .`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
# MCP mode (e2e_cortex_mcp.py)
[1/6] Connecting to headroom MCP server ... OK
[2/6] Listing MCP tools ... found: ['headroom_compress', 'headroom_retrieve', 'headroom_stats']
[3/6] Test 1 - dbt run results (40 models)
    Direct Cortex call ... prompt=2,112 tokens
    MCP headroom_compress ... saved 0 tokens  hash=825cf6f2...
    Cortex call (MCP-compressed) ... prompt=2,112  saved 0 (0.0%)
[4/6] Test 2 - INFORMATION_SCHEMA tables (59 rows)
    Direct Cortex call ... prompt=3,203 tokens
    MCP headroom_compress ... saved 1,280 tokens (37.2%)
    Cortex call (MCP-compressed) ... prompt=1,163  saved 2,040 (63.7%)
[5/6] headroom_retrieve CCR round-trip ... original content retrieved
[6/6] headroom_stats ... compressions: 2, total_tokens_saved: 1280
MCP TEST PASSED - 38.4% avg token reduction via MCP tools

# Proxy + MCP mode (e2e_cortex_proxy_mcp.py)
[1/7] Starting headroom proxy ... OK
[2/7] Connecting to headroom MCP server ... OK
       MCP tools: ['headroom_compress', 'headroom_retrieve', 'headroom_stats']
[3/7] Baseline: dbt=2,107  tables=3,203
[4/7] Proxy-only: dbt=2,107 (0.0%)  tables=3,203 (0.0%)
[5/7] MCP+Proxy: dbt=2,107 (0.0%)  tables=1,163 (63.7% saved)
[6/7] CCR round-trip: original content retrieved
Components verified:
  Proxy starts (FastAPI + uvicorn) and routes to Cortex
  MCP server connects (MCP Python SDK client)
  headroom_compress works via MCP
  headroom_retrieve (CCR) works via MCP
  Proxy + MCP run simultaneously in same session
```

## Real Behavior Proof

- Environment: macOS, Python 3.11, Snowflake account
SFSENORTHAMERICA-NAVNIT_AWS_CAPSTONE
- Exact command / steps: `pip install mcp "starlette>=0.37.2,<0.41.0"`
then `SF_CONN=<conn> python3 tests/e2e_cortex_savings.py`,
`SF_CONN=<conn> python3 tests/e2e_cortex_quality.py`, `SF_CONN=<conn>
python3 tests/e2e_cortex_latency.py`, `SF_CONN=<conn> python3
tests/e2e_cortex_proxy.py`, `SF_CONN=<conn> python3
tests/e2e_cortex_mcp.py`, `PROXY_PORT=8798 SF_CONN=<conn> python3
tests/e2e_cortex_proxy_mcp.py`
- Observed result: MCP server connects via stdio, tools verified, 63.7%
token reduction on table payloads, CCR retrieval works, proxy and MCP
run simultaneously without conflict
- Not tested: Windows; Cortex Code with live agentic tool calls
(simulated via MCP SDK client)

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works

## Additional Notes

- `role:"tool"` correction: Chat Completions endpoint supports it;
Messages endpoint does not (use `user` message with `tool_result` block
instead)
- MCP tests require `pip install mcp`
- Starlette compatibility: `mcp` may install starlette 1.3.1 which
conflicts with headroom proxy; fix with `pip install
"starlette>=0.37.2,<0.41.0"`

---------

Co-authored-by: Cortex Code <noreply@snowflake.com>
This commit is contained in:
sfc-gh-nashukla 2026-06-30 12:14:36 -07:00 • committed by GitHub
parent 6cba4419d0
commit f00ace6da5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1977 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

483
tests/e2e_cortex_latency.py Normal file
View file

@ -0,0 +1,483 @@
#!/usr/bin/env python3
"""
Latency benchmark: Snowflake Cortex — Standard vs Headroom
Measures per call (averaged over N runs):
- TTFT Time to First Token (streaming)
- E2E End-to-End latency
- Compress overhead (headroom local processing time)
- Prompt token count (from usage block in final SSE chunk)
Because headroom reduces prompt length, prefill is shorter → lower TTFT.
Multiple runs are averaged to smooth out shared-API latency variance.
Usage:
SF_CONN=<connection-name> python3 tests/e2e_cortex_latency.py
# Optional overrides:
SF_CONN=my_conn SF_HOST=myaccount.snowflakecomputing.com python3 tests/e2e_cortex_latency.py
SF_CONN=my_conn SF_MODEL=claude-sonnet-4-6 RUNS=5 python3 tests/e2e_cortex_latency.py
"""
from __future__ import annotations
import http.client
import json
import os
import ssl
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path
# ── Bootstrap headroom ────────────────────────────────────────────────────────
REPO_ROOT = Path(__file__).resolve().parent.parent
_VENV_SITE = REPO_ROOT / ".venv" / "lib"
try:
from headroom import compress as _hc_check # 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))
# ── Settings ──────────────────────────────────────────────────────────────────
_SF_HOST = os.environ.get("SF_HOST", "")
_SF_CONN = os.environ.get("SF_CONN", "")
_SF_MODEL = os.environ.get("SF_MODEL", "claude-sonnet-4-6")
_RUNS = int(os.environ.get("RUNS", "3"))
_INPUT_PRICE_PER_1M = 3.00 # USD, claude-sonnet-4-6 on Cortex
# ── Streaming call ────────────────────────────────────────────────────────────
def _stream_call(messages: list[dict], token: str, host: str) -> tuple[float, float, int, int]:
payload = json.dumps(
{
"model": _SF_MODEL,
"messages": messages,
"max_completion_tokens": 128,
"stream": True,
}
).encode()
ctx = ssl.create_default_context()
conn = http.client.HTTPSConnection(host, context=ctx, timeout=90)
conn.request(
"POST",
"/api/v2/cortex/v1/chat/completions",
body=payload,
headers={
"Authorization": f'Snowflake Token="{token}"',
"Content-Type": "application/json",
"Accept": "text/event-stream",
"User-Agent": "headroom-latency-bench/1.0",
},
)
t_start = time.perf_counter()
resp = conn.getresponse()
if resp.status != 200:
body = resp.read().decode(errors="replace")
conn.close()
raise RuntimeError(f"HTTP {resp.status}: {body[:200]}")
ttft_ms: float = 0.0
prompt_tokens = 0
completion_tokens = 0
first_token_seen = False
while True:
raw = resp.readline()
if not raw:
break
line = raw.decode("utf-8", errors="replace").strip()
if not line or not line.startswith("data:"):
continue
data = line[5:].strip()
if data == "[DONE]":
break
try:
chunk = json.loads(data)
except json.JSONDecodeError:
continue
if not first_token_seen:
delta = (chunk.get("choices") or [{}])[0].get("delta", {})
if delta.get("content", ""):
ttft_ms = (time.perf_counter() - t_start) * 1000
first_token_seen = True
usage = chunk.get("usage") or {}
if usage.get("prompt_tokens"):
prompt_tokens = usage["prompt_tokens"]
completion_tokens = usage.get("completion_tokens", 0)
e2e_ms = (time.perf_counter() - t_start) * 1000
conn.close()
if not first_token_seen:
ttft_ms = e2e_ms
return ttft_ms, e2e_ms, prompt_tokens, completion_tokens
# ── Payloads ──────────────────────────────────────────────────────────────────
def _tables_json() -> str:
rows = [
{
"TABLE_CATALOG": "PROD_DB",
"TABLE_SCHEMA": "ANALYTICS",
"TABLE_NAME": f"FACT_ORDERS_{i:03d}",
"TABLE_TYPE": "BASE TABLE",
"ROW_COUNT": i * 1_423_001,
"BYTES": i * 8_192_000,
"CREATED": "2024-01-15",
"LAST_ALTERED": "2025-06-10",
"COMMENT": f"Daily order fact partition {i:03d}",
}
for i in range(1, 80)
]
return json.dumps(rows, indent=2)
def _dbt_json() -> str:
return json.dumps(
{
"metadata": {"dbt_version": "1.8.0"},
"results": [
{
"unique_id": f"model.analytics.fct_{i:03d}",
"status": "success" if i % 7 != 0 else "error",
"execution_time": round(0.8 + i * 0.12, 3),
"rows_affected": i * 12_500,
"compiled_code": f"SELECT * FROM raw.orders_{i:03d} WHERE status='active'",
"failures": None
if i % 7 != 0
else [{"message": f"Invalid col_{i}", "line": i % 40}],
"adapter_response": {
"query_id": f"01b{i:06x}",
"rows_produced": i * 12_500,
},
}
for i in range(40)
],
},
indent=2,
)
def _search_json() -> str:
return json.dumps(
[
{
"rank": i + 1,
"score": round(0.98 - i * 0.02, 4),
"document_id": f"doc_{i:04d}",
"source": "PROD_DB.DOCS.ENGINEERING_WIKI",
"content": (
"The revenue pipeline processes 2.3 million orders per day. "
"product_family column was renamed to product_group in Q3 2024. "
"Migration: update all references in models/marts/revenue/ and "
"run dbt run --full-refresh --select fct_revenue."
),
"metadata": {
"author": f"eng_{i % 6}@company.com",
"updated": "2025-05-20",
},
}
for i in range(15)
],
indent=2,
)
def _build_messages(ctx: str) -> list[dict]:
return [
{"role": "system", "content": ctx},
{"role": "assistant", "content": "I have reviewed the context above."},
{
"role": "user",
"content": "Based on the data above, what is failing and how do I fix it?",
},
]
# ── Result dataclass ──────────────────────────────────────────────────────────
def _avg(vals: list[float]) -> float:
return sum(vals) / max(len(vals), 1)
def _median(vals: list[float]) -> float:
s = sorted(vals)
n = len(s)
if n == 0:
return 0.0
return s[n // 2] if n % 2 else (s[n // 2 - 1] + s[n // 2]) / 2
@dataclass
class LatencyResult:
label: str
runs: int
std_tokens: int
hdm_tokens: int
std_ttft_all: list[float] = field(default_factory=list)
hdm_ttft_all: list[float] = field(default_factory=list)
std_e2e_all: list[float] = field(default_factory=list)
hdm_e2e_all: list[float] = field(default_factory=list)
compress_overhead_ms: float = 0.0
@property
def std_ttft_ms(self) -> float:
return _median(self.std_ttft_all)
@property
def hdm_ttft_ms(self) -> float:
return _median(self.hdm_ttft_all)
@property
def std_e2e_ms(self) -> float:
return _median(self.std_e2e_all)
@property
def hdm_e2e_ms(self) -> float:
return _median(self.hdm_e2e_all)
@property
def token_saving_pct(self) -> float:
return (self.std_tokens - self.hdm_tokens) / max(self.std_tokens, 1) * 100
@property
def ttft_saving_pct(self) -> float:
return (self.std_ttft_ms - self.hdm_ttft_ms) / max(self.std_ttft_ms, 1) * 100
@property
def e2e_saving_pct(self) -> float:
return (self.std_e2e_ms - self.hdm_e2e_ms) / max(self.std_e2e_ms, 1) * 100
@property
def net_latency_saving_ms(self) -> float:
return (self.std_e2e_ms - self.hdm_e2e_ms) - self.compress_overhead_ms
@property
def usd_saved_per_call(self) -> float:
return (self.std_tokens - self.hdm_tokens) / 1_000_000 * _INPUT_PRICE_PER_1M
# ── Benchmark runner (N runs, median) ─────────────────────────────────────────
def run_benchmark(
label: str,
messages: list[dict],
token: str,
host: str,
n_runs: int = 3,
) -> LatencyResult:
from headroom import compress
print(f"\n ┌─ {label} (n={n_runs} runs each)")
std_ttfts: list[float] = []
std_e2es: list[float] = []
std_pt = 0
for i in range(n_runs):
print(f" │ run {i + 1}/{n_runs} std ...", end=" ", flush=True)
ttft, e2e, pt, _ = _stream_call(messages, token, host)
std_ttfts.append(ttft)
std_e2es.append(e2e)
std_pt = pt
print(f"TTFT={ttft:.0f}ms E2E={e2e:.0f}ms tokens={pt:,}")
print(" │ compressing ...", end=" ", flush=True)
t0 = time.perf_counter()
compressed = compress(messages, model="claude-sonnet-4-5-20250929")
compress_ms = (time.perf_counter() - t0) * 1000
print(f"{compress_ms:.0f}ms overhead")
hdm_ttfts: list[float] = []
hdm_e2es: list[float] = []
hdm_pt = 0
for i in range(n_runs):
print(f" │ run {i + 1}/{n_runs} hdm ...", end=" ", flush=True)
ttft, e2e, pt, _ = _stream_call(compressed.messages, token, host)
hdm_ttfts.append(ttft)
hdm_e2es.append(e2e)
hdm_pt = pt
print(f"TTFT={ttft:.0f}ms E2E={e2e:.0f}ms tokens={pt:,}")
r = LatencyResult(
label=label,
runs=n_runs,
std_tokens=std_pt,
hdm_tokens=hdm_pt,
std_ttft_all=std_ttfts,
hdm_ttft_all=hdm_ttfts,
std_e2e_all=std_e2es,
hdm_e2e_all=hdm_e2es,
compress_overhead_ms=compress_ms,
)
print(
f" └─ median TTFT: std={r.std_ttft_ms:.0f}ms hdm={r.hdm_ttft_ms:.0f}ms "
f"saving={r.ttft_saving_pct:.1f}%"
)
return r
# ── Display ───────────────────────────────────────────────────────────────────
def _bar(pct: float, w: int = 20) -> str:
n = max(0, int(pct / 100 * w))
return "█" * n + "░" * (w - n)
def _show(r: LatencyResult) -> None:
std_ttft_range = f"[{min(r.std_ttft_all):.0f}–{max(r.std_ttft_all):.0f}]"
hdm_ttft_range = f"[{min(r.hdm_ttft_all):.0f}–{max(r.hdm_ttft_all):.0f}]"
print(f"\n ┌─ {r.label} (median of {r.runs} runs)")
print(
f" │ Tokens : {r.std_tokens:>7,} → {r.hdm_tokens:>7,} "
f"│ saved {r.std_tokens - r.hdm_tokens:>6,} ({r.token_saving_pct:.1f}%)"
)
print(
f" │ TTFT : {r.std_ttft_ms:>7.0f}ms → {r.hdm_ttft_ms:>6.0f}ms "
f"│ saved {r.std_ttft_ms - r.hdm_ttft_ms:>6.0f}ms ({r.ttft_saving_pct:.1f}%) "
f"{_bar(r.ttft_saving_pct)}"
)
print(f" │ std range {std_ttft_range}ms hdm range {hdm_ttft_range}ms")
print(
f" │ E2E : {r.std_e2e_ms:>7.0f}ms → {r.hdm_e2e_ms:>6.0f}ms "
f"│ saved {r.std_e2e_ms - r.hdm_e2e_ms:>6.0f}ms ({r.e2e_saving_pct:.1f}%)"
)
print(
f" │ Compress overhead: {r.compress_overhead_ms:.0f}ms "
f"│ Net latency saving: {r.net_latency_saving_ms:.0f}ms"
)
print(f" └─ Cost: ${r.usd_saved_per_call:.5f} saved / call")
# ── Main ──────────────────────────────────────────────────────────────────────
def main() -> int:
print()
print("╔═══════════════════════════════════════════════════════════════╗")
print("║ Cortex Code × Headroom — TTFT + Latency Benchmark ║")
print("║ Streaming API │ Time to First Token │ E2E latency ║")
print("╚═══════════════════════════════════════════════════════════════╝")
if not _SF_CONN:
print("\n ✗ Set SF_CONN=<connection-name> to run this benchmark.")
print(" Example: SF_CONN=navnit_local_auth python3 tests/e2e_cortex_latency.py")
return 1
import io
try:
import snowflake.connector
except ImportError:
print("\n ✗ snowflake-connector-python not installed.")
return 1
_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()")
locator = cur.fetchone()[0].lower()
host = f"{locator}.snowflakecomputing.com"
finally:
sys.stdout = _s
total_calls = len(["full", "tables", "dbt", "search"]) * _RUNS * 2
print(f"\n Model : {_SF_MODEL}")
print(f" Host : {host}")
print(f" Runs : {_RUNS} per payload (median used) → {total_calls} total API calls")
print(" TTFT : first SSE content chunk via streaming\n")
full_ctx = json.dumps(
{
"tables": json.loads(_tables_json()),
"dbt_results": json.loads(_dbt_json()),
"search_results": json.loads(_search_json()),
},
indent=2,
)
payloads = [
("Full context (tables + dbt + search)", _build_messages(full_ctx)),
("INFORMATION_SCHEMA tables (79 rows)", _build_messages(_tables_json())),
("dbt run-results (40 models)", _build_messages(_dbt_json())),
("Cortex Search results (15 docs)", _build_messages(_search_json())),
]
results: list[LatencyResult] = []
for label, msgs in payloads:
try:
r = run_benchmark(label, msgs, token, host, n_runs=_RUNS)
results.append(r)
_show(r)
except Exception as exc:
print(f"\n ✗ {label} failed: {exc}")
conn.close()
if not results:
print("\n No results collected.")
return 1
# ── Summary ───────────────────────────────────────────────────────────────
print()
print("╔═══════════════════════════════════════════════════════════════╗")
print(f"║ SUMMARY (median of {_RUNS} runs per payload) ║")
print("╠═══════════════════════════════════════════════════════════════╣")
hdr = f" {'Payload':<38} {'Tokens':>6} {'TTFT↓':>7} {'E2E↓':>7} {'Net↓':>7}"
print(hdr)
print(f" {'─' * 38} {'─' * 6} {'─' * 7} {'─' * 7} {'─' * 7}")
for r in results:
print(
f" {r.label[:38]:<38} "
f"{r.token_saving_pct:>5.0f}% "
f"{r.ttft_saving_pct:>6.0f}% "
f"{r.e2e_saving_pct:>6.0f}% "
f"{r.net_latency_saving_ms:>5.0f}ms"
)
avg_token_pct = sum(r.token_saving_pct for r in results) / len(results)
avg_ttft_pct = sum(r.ttft_saving_pct for r in results) / len(results)
avg_e2e_pct = sum(r.e2e_saving_pct for r in results) / len(results)
avg_usd = sum(r.usd_saved_per_call for r in results) / len(results)
print(f" {'─' * 38} {'─' * 6} {'─' * 7} {'─' * 7} {'─' * 7}")
print(
f" {'AVERAGE':<38} {avg_token_pct:>5.0f}% {avg_ttft_pct:>6.0f}% {avg_e2e_pct:>6.0f}% "
)
print()
print(f" Avg USD saved / call : ${avg_usd:.5f}")
print(f" At 1k/day : ${avg_usd * 1_000:.2f}/day │ ${avg_usd * 365_000:,.0f}/year")
print("╚═══════════════════════════════════════════════════════════════╝")
print()
print(" Key insight: TTFT savings track token savings because prefill")
print(" time scales with prompt length. Fewer tokens = shorter prefill")
print(" = faster first token. Median across runs removes outlier spikes.")
print()
return 0
if __name__ == "__main__":
sys.exit(main())

347
tests/e2e_cortex_mcp.py Normal file
View file

@ -0,0 +1,347 @@
#!/usr/bin/env python3
"""
MCP mode e2e test: Cortex Code + Headroom MCP Server
Tests the FULL MCP path using the official MCP Python SDK client:
1. Start headroom MCP server (stdio transport via mcp_server.py)
2. Connect using mcp.ClientSession (same protocol Cortex Code uses)
3. List tools → verify headroom_compress / headroom_retrieve / headroom_stats
4. Call headroom_compress with large JSON payloads
5. Use compressed output to call Snowflake Cortex REST API
6. Compare prompt_tokens: direct vs MCP-compressed
Usage:
SF_CONN=<connection-name> python3 tests/e2e_cortex_mcp.py
"""
from __future__ import annotations
import asyncio
import json
import os
import sys
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")
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
# ── Cortex call ───────────────────────────────────────────────────────────────
def _cortex_call(messages: list[dict], token: str, host: str) -> dict:
body = json.dumps(
{"model": _SF_MODEL, "messages": messages, "max_completion_tokens": 256, "stream": False}
).encode()
req = urllib.request.Request(
f"https://{host}/api/v2/cortex/v1/chat/completions",
data=body,
headers={
"Authorization": f'Snowflake Token="{token}"',
"Content-Type": "application/json",
"User-Agent": "headroom-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"Cortex 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)
# ── 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,
)
# ── MCP test ──────────────────────────────────────────────────────────────────
async def run_mcp_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
print()
print("╔═══════════════════════════════════════════════════════════════╗")
print("║ Cortex Code × Headroom — MCP Mode E2E Test ║")
print("║ MCP Python SDK Client │ stdio transport │ Cortex ║")
print("╚═══════════════════════════════════════════════════════════════╝")
print(f"\n Model : {_SF_MODEL} │ Host : {host}")
server_params = StdioServerParameters(
command=sys.executable,
args=[str(MCP_SERVER_SCRIPT)],
env={**os.environ, "PYTHONPATH": str(REPO_ROOT)},
)
# ── Connect via MCP SDK ───────────────────────────────────────────────────
print("\n [1/6] 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")
# ── List tools ────────────────────────────────────────────────────
print(" [2/6] Listing MCP tools ...", end=" ", flush=True)
tools_result = await session.list_tools()
tool_names = [t.name for t in tools_result.tools]
print(f"found: {tool_names}")
required = {"headroom_compress", "headroom_retrieve", "headroom_stats"}
missing = required - set(tool_names)
if missing:
print(f"\n ✗ Missing tools: {missing}")
return 1
# ── Test 1: dbt run results ───────────────────────────────────────
print("\n [3/6] Test 1 — dbt run results (40 models)")
dbt_content = _dbt_payload()
question = "Which models failed and what column is missing?"
print(" ├─ Direct Cortex call ...", end=" ", flush=True)
d1_pt, _ = _tokens(
_cortex_call(
[
{"role": "system", "content": dbt_content},
{"role": "user", "content": question},
],
token,
host,
)
)
print(f"prompt={d1_pt:,} tokens")
print(" ├─ MCP headroom_compress ...", end=" ", flush=True)
r1 = await session.call_tool("headroom_compress", {"content": dbt_content})
text1 = r1.content[0].text if r1.content else "{}"
data1 = json.loads(text1) if text1.startswith("{") else {}
compressed1 = data1.get("compressed", dbt_content)
saved1 = data1.get("tokens_saved", 0)
pct1 = data1.get("savings_percent", 0)
hash1 = data1.get("hash", "")
print(f"saved {saved1:,} tokens ({pct1:.1f}%) hash={hash1[:8]}...")
print(" └─ Cortex call (MCP-compressed) ...", end=" ", flush=True)
m1_pt, _ = _tokens(
_cortex_call(
[
{
"role": "system",
"content": compressed1
if isinstance(compressed1, str)
else json.dumps(compressed1),
},
{"role": "user", "content": question},
],
token,
host,
)
)
api_saved1 = d1_pt - m1_pt
api_pct1 = api_saved1 / max(d1_pt, 1) * 100
sym = "✓" if api_saved1 > 0 else "·"
print(f"{sym} prompt={m1_pt:,} saved {api_saved1:,} ({api_pct1:.1f}%)")
# ── Test 2: table schema ──────────────────────────────────────────
print("\n [4/6] Test 2 — INFORMATION_SCHEMA tables (59 rows)")
tbl_content = _tables_payload()
question2 = "How many tables are archived?"
print(" ├─ Direct Cortex call ...", end=" ", flush=True)
d2_pt, _ = _tokens(
_cortex_call(
[
{"role": "system", "content": tbl_content},
{"role": "user", "content": question2},
],
token,
host,
)
)
print(f"prompt={d2_pt:,} tokens")
print(" ├─ MCP headroom_compress ...", end=" ", flush=True)
r2 = await session.call_tool("headroom_compress", {"content": tbl_content})
text2 = r2.content[0].text if r2.content else "{}"
data2 = json.loads(text2) if text2.startswith("{") else {}
compressed2 = data2.get("compressed", tbl_content)
saved2 = data2.get("tokens_saved", 0)
pct2 = data2.get("savings_percent", 0)
print(f"saved {saved2:,} tokens ({pct2:.1f}%)")
print(" └─ Cortex call (MCP-compressed) ...", end=" ", flush=True)
m2_pt, _ = _tokens(
_cortex_call(
[
{
"role": "system",
"content": compressed2
if isinstance(compressed2, str)
else json.dumps(compressed2),
},
{"role": "user", "content": question2},
],
token,
host,
)
)
api_saved2 = d2_pt - m2_pt
api_pct2 = api_saved2 / max(d2_pt, 1) * 100
sym2 = "✓" if api_saved2 > 0 else "·"
print(f"{sym2} prompt={m2_pt:,} saved {api_saved2:,} ({api_pct2:.1f}%)")
# ── Test 3: headroom_retrieve ─────────────────────────────────────
if hash1:
print(f"\n [5/6] headroom_retrieve — CCR round-trip (hash={hash1[:8]}...)")
r3 = await session.call_tool("headroom_retrieve", {"hash": hash1})
text3 = r3.content[0].text if r3.content else "{}"
data3 = json.loads(text3) if text3.startswith("{") else {}
if "original_content" in data3 or "results" in data3:
print(" ✓ original content retrieved successfully")
elif "error" in data3:
print(f" ⚠ {data3['error'][:80]}")
else:
print(f" ✓ retrieved (keys: {list(data3.keys())})")
# ── headroom_stats ────────────────────────────────────────────────
print("\n [6/6] headroom_stats")
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}")
# ── Summary ───────────────────────────────────────────────────────
total_direct = d1_pt + d2_pt
total_mcp = m1_pt + m2_pt
avg_pct = (total_direct - total_mcp) / max(total_direct, 1) * 100
print()
print("╔═══════════════════════════════════════════════════════════════╗")
print("║ MCP MODE SUMMARY ║")
print("╠═══════════════════════════════════════════════════════════════╣")
print(f" {'Payload':<35} {'Direct':>8} {'MCP+API':>8} {'Saved':>7}")
print(f" {'─' * 35} {'─' * 8} {'─' * 8} {'─' * 7}")
print(
f" {'dbt run results (40 models)':<35} {d1_pt:>8,} {m1_pt:>8,} {api_pct1:>6.1f}%"
)
print(
f" {'INFORMATION_SCHEMA (59 rows)':<35} {d2_pt:>8,} {m2_pt:>8,} {api_pct2:>6.1f}%"
)
print(f" {'─' * 35} {'─' * 8} {'─' * 8} {'─' * 7}")
print(f" {'TOTAL':<35} {total_direct:>8,} {total_mcp:>8,} {avg_pct:>6.1f}%")
print()
print(" MCP transport : stdio (MCP Python SDK — same as Cortex Code)")
print(" Tools verified : headroom_compress ✓ headroom_retrieve ✓ headroom_stats ✓")
if avg_pct > 0:
print(f"\n ✓ MCP TEST PASSED — {avg_pct:.1f}% avg token reduction via MCP tools")
else:
print("\n ⚠ MCP routing works but payloads below compression threshold")
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_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_mcp_test(token, host))
finally:
conn.close()
if __name__ == "__main__":
sys.exit(main())

336
tests/e2e_cortex_proxy.py Normal file
View file

@ -0,0 +1,336 @@
#!/usr/bin/env python3
"""
Proxy-in-the-loop integration test: Cortex Code + Headroom Proxy
Tests the FULL path:
Cortex Code (simulated) → headroom FastAPI proxy → Snowflake Cortex
Key insight: headroom compresses CONVERSATION HISTORY.
Turn 1: nothing to compress yet — baseline
Turn 2: proxy compresses turn 1 history before sending
Turn 3: proxy compresses turns 1+2 history
→ token count should DROP on turns 2+ vs a direct client
Usage:
SF_CONN=<connection-name> python3 tests/e2e_cortex_proxy.py
"""
from __future__ import annotations
import json
import os
import signal
import subprocess
import sys
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parent.parent
_VENV_SITE = REPO_ROOT / ".venv" / "lib"
try:
from headroom import compress as _hc_check # 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"))
_TURNS = int(os.environ.get("TURNS", "4"))
# ── Auth ──────────────────────────────────────────────────────────────────────
def _get_sf_token_and_host():
"""Returns (token, host, conn) — caller must keep conn open."""
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()")
locator = cur.fetchone()[0].lower()
host = f"{locator}.snowflakecomputing.com"
finally:
sys.stdout = _s
return token, host, conn
# ── API calls ─────────────────────────────────────────────────────────────────
def _call(url: str, messages: list[dict], token: str) -> dict:
token_field = "max_completion_tokens"
body = json.dumps(
{"model": _SF_MODEL, "messages": messages, token_field: 200, "stream": False}
).encode()
auth_header = f'Snowflake Token="{token}"'
req = urllib.request.Request(
url,
data=body,
headers={
"Authorization": auth_header,
"Content-Type": "application/json",
"User-Agent": "headroom-proxy-test/1.0",
},
method="POST",
)
with urllib.request.urlopen(req, timeout=90) as r:
return json.loads(r.read())
def _tokens(resp: dict) -> tuple[int, int]:
u = resp.get("usage", {})
pt = u.get("prompt_tokens") or u.get("input_tokens", 0)
ct = u.get("completion_tokens") or u.get("output_tokens", 0)
return pt, ct
def _content(resp: dict) -> str:
return resp.get("choices", [{}])[0].get("message", {}).get("content", "")
# ── Proxy lifecycle ───────────────────────────────────────────────────────────
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
# ── Conversation turns ────────────────────────────────────────────────────────
# Each turn adds a large JSON tool-result style blob as context
# so SmartCrusher has something to compress from turn 2 onwards.
_TURN_QUESTIONS = [
"Here is our dbt run output: {ctx}\n\nWhich models failed?",
"Now here are the raw table stats: {ctx}\n\nWhich table has the most rows?",
"Here are the Cortex Search results: {ctx}\n\nWhat is the top ranked document?",
"Given everything above, what should I fix first and why?",
]
def _dbt_ctx() -> 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_ctx() -> 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,
)
def _search_ctx() -> str:
return json.dumps(
[
{
"rank": i + 1,
"score": round(0.98 - i * 0.03, 4),
"document_id": f"doc_{i:04d}",
"content": f"Engineering runbook #{i:03d}: covers deployment and config for service_{i}.",
}
for i in range(20)
],
indent=2,
)
_CONTEXTS = [_dbt_ctx(), _tables_ctx(), _search_ctx(), ""]
@dataclass
class TurnResult:
turn: int
direct_pt: int
proxy_pt: int
@property
def saved(self) -> int:
return self.direct_pt - self.proxy_pt
@property
def pct(self) -> float:
return self.saved / max(self.direct_pt, 1) * 100
# ── Main ──────────────────────────────────────────────────────────────────────
def main() -> int:
print()
print("╔═══════════════════════════════════════════════════════════════╗")
print("║ Cortex Code × Headroom — Proxy-in-the-Loop (Multi-Turn) ║")
print("║ Agent → FastAPI proxy → Snowflake Cortex ║")
print("╚═══════════════════════════════════════════════════════════════╝")
print()
print(" Insight: compression = 0% on turn 1 (no history yet)")
print(" compression grows each subsequent turn as history accumulates")
if not _SF_CONN:
print("\n ✗ Set SF_CONN=<connection-name>")
return 1
print("\n [1/4] Authenticating ...", end=" ", flush=True)
try:
token, host, _sf_conn = _get_sf_token_and_host()
print(f"OK ({host})")
except Exception as e:
print(f"FAILED: {e}")
return 1
cortex_direct_url = f"https://{host}/api/v2/cortex/v1/chat/completions"
cortex_base = f"https://{host}/api/v2/cortex"
proxy_url = f"http://127.0.0.1:{_PROXY_PORT}/v1/chat/completions"
print(f" [2/4] Starting headroom proxy on :{_PROXY_PORT} ...", end=" ", flush=True)
proxy_env = os.environ.copy()
proxy_log = open("/tmp/headroom_proxy.log", "w")
proxy_proc = subprocess.Popen(
[
sys.executable,
"-m",
"headroom.proxy.server",
"--port",
str(_PROXY_PORT),
"--openai-api-url",
cortex_base,
],
env=proxy_env,
cwd=str(REPO_ROOT),
stdout=proxy_log,
stderr=proxy_log,
)
if not _wait_for_proxy(_PROXY_PORT):
proxy_proc.send_signal(signal.SIGTERM)
proxy_proc.wait(timeout=5)
proxy_log.close()
print("FAILED")
return 1
print("OK")
turns: list[TurnResult] = []
direct_history: list[dict] = []
proxy_history: list[dict] = []
print(f"\n [3/4] Running {_TURNS} conversation turns ...\n")
print(f" {'Turn':<6} {'Direct prompt':>14} {'Proxy prompt':>13} {'Saved':>8} {'Note'}")
print(f" {'─' * 6} {'─' * 14} {'─' * 13} {'─' * 8} {'─' * 30}")
try:
for t in range(1, _TURNS + 1):
ctx = _CONTEXTS[min(t - 1, len(_CONTEXTS) - 1)]
question = _TURN_QUESTIONS[min(t - 1, len(_TURN_QUESTIONS) - 1)].format(ctx=ctx)
# ── Direct: accumulate full uncompressed history ────────────────
direct_history.append({"role": "user", "content": question})
try:
dr = _call(cortex_direct_url, direct_history, token)
d_pt, _ = _tokens(dr)
d_answer = _content(dr)
direct_history.append({"role": "assistant", "content": d_answer})
except urllib.error.HTTPError as e:
print(f" Direct turn {t} FAILED: HTTP {e.code} {e.read().decode()[:100]}")
break
except Exception as e:
print(f" Direct turn {t} FAILED: {e}")
break
# ── Proxy: send history through headroom proxy ──────────────────
proxy_history.append({"role": "user", "content": question})
try:
pr = _call(proxy_url, proxy_history, token)
p_pt, _ = _tokens(pr)
p_answer = _content(pr)
proxy_history.append({"role": "assistant", "content": p_answer})
except urllib.error.HTTPError as e:
print(f" Proxy turn {t} FAILED: HTTP {e.code} {e.read().decode()[:100]}")
break
except Exception as e:
print(f" Proxy turn {t} FAILED: {e}")
break
result = TurnResult(turn=t, direct_pt=d_pt, proxy_pt=p_pt)
turns.append(result)
note = "← baseline (no history yet)" if t == 1 else f"← {result.pct:.0f}% saved"
sym = "✓" if result.saved > 0 else ("·" if t == 1 else "⚠")
print(f" {sym} T{t:<4} {d_pt:>14,} {p_pt:>13,} {result.saved:>+8,} {note}")
finally:
proxy_proc.send_signal(signal.SIGTERM)
proxy_proc.wait(timeout=5)
proxy_log.close()
_sf_conn.close()
if not turns:
print("\n No results collected.")
return 1
# ── Summary ───────────────────────────────────────────────────────────────
later_turns = [r for r in turns if r.turn > 1]
avg_saving = sum(r.pct for r in later_turns) / max(len(later_turns), 1)
total_direct = sum(r.direct_pt for r in turns)
total_proxy = sum(r.proxy_pt for r in turns)
total_saved = total_direct - total_proxy
print()
print(" [4/4] Summary")
print(f" {'─' * 60}")
print(f" Total direct tokens : {total_direct:,}")
print(f" Total proxy tokens : {total_proxy:,} (saved {total_saved:,})")
print(f" Avg compression T2+ : {avg_saving:.1f}%")
print()
if avg_saving > 5:
print(" ✓ PROXY COMPRESSION CONFIRMED")
print(" headroom proxy transparently compresses conversation history")
print(f" Average {avg_saving:.0f}% token reduction from turn 2 onwards")
else:
print(" ⚠ Low compression — proxy routed correctly but history")
print(" may be below SmartCrusher threshold. Try longer conversations.")
return 0 if len(turns) == _TURNS else 1
if __name__ == "__main__":
sys.exit(main())

371
tests/e2e_cortex_proxy_mcp.py Normal file
View file

@ -0,0 +1,371 @@
#!/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())

432
tests/e2e_cortex_quality.py Normal file
View file

@ -0,0 +1,432 @@
#!/usr/bin/env python3
"""
Quality benchmark: Snowflake Cortex — Standard vs Headroom
Tests whether headroom compression affects answer quality.
Strategy: embed known facts in payload, ask factual questions,
score both standard and headroom responses against ground truth.
No LLM judge needed — answers are verifiable from the data itself.
Usage:
SF_CONN=<connection-name> python3 tests/e2e_cortex_quality.py
"""
from __future__ import annotations
import json
import os
import sys
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from pathlib import Path
# ── Bootstrap headroom ────────────────────────────────────────────────────────
REPO_ROOT = Path(__file__).resolve().parent.parent
_VENV_SITE = REPO_ROOT / ".venv" / "lib"
try:
from headroom import compress as _hc_check # 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))
# ── Settings ──────────────────────────────────────────────────────────────────
_SF_HOST = os.environ.get("SF_HOST", "")
_SF_CONN = os.environ.get("SF_CONN", "")
_SF_MODEL = os.environ.get("SF_MODEL", "claude-sonnet-4-6")
# ── API call (non-streaming, full response) ───────────────────────────────────
def _call(messages: list[dict], token: str, host: str) -> str:
body = json.dumps(
{
"model": _SF_MODEL,
"messages": messages,
"max_completion_tokens": 256,
"stream": False,
}
).encode()
req = urllib.request.Request(
f"https://{host}/api/v2/cortex/v1/chat/completions",
data=body,
headers={
"Authorization": f'Snowflake Token="{token}"',
"Content-Type": "application/json",
"User-Agent": "headroom-quality-bench/1.0",
},
method="POST",
)
with urllib.request.urlopen(req, timeout=60) as r:
resp = json.loads(r.read())
if "error_code" in resp:
raise RuntimeError(f"Cortex {resp['error_code']}: {resp.get('message')}")
return resp["choices"][0]["message"]["content"].strip()
# ── Test case definition ──────────────────────────────────────────────────────
@dataclass
class QualityCase:
name: str
context: str
question: str
expected_keywords: list[str]
expected_absent: list[str] = None
def score(self, answer: str) -> tuple[int, int]:
"""Returns (hits, total) based on keyword presence."""
answer_lower = answer.lower()
hits = sum(1 for kw in self.expected_keywords if kw.lower() in answer_lower)
return hits, len(self.expected_keywords)
def pass_threshold(self, hits: int, total: int) -> bool:
return hits / max(total, 1) >= 0.6
# ── Test payload builders ─────────────────────────────────────────────────────
def _make_cases() -> list[QualityCase]:
# ── Case 1: Exact row lookup from large table JSON ────────────────────────
tables = [
{
"TABLE_NAME": f"FACT_ORDERS_{i:03d}",
"TABLE_SCHEMA": "ANALYTICS",
"ROW_COUNT": i * 1_000_000,
"BYTES": i * 8_192_000,
"LAST_ALTERED": "2025-06-10",
}
for i in range(1, 80)
]
tables_ctx = json.dumps(tables, indent=2)
# Case 1a: exact numeric lookup
case1a = QualityCase(
name="Table row count lookup (FACT_ORDERS_042)",
context=tables_ctx,
question="What is the ROW_COUNT of the table named FACT_ORDERS_042? Reply with just the number.",
expected_keywords=["42000000", "42,000,000"],
)
# Case 1b: filter + list
case1b = QualityCase(
name="Tables over 50M rows (filter query)",
context=tables_ctx,
question="List all table names where ROW_COUNT is greater than 50,000,000.",
expected_keywords=[
"fact_orders_051",
"fact_orders_060",
"fact_orders_070",
"fact_orders_079",
],
)
# ── Case 2: dbt failure detection ────────────────────────────────────────
# Models fail when i % 7 == 0 → indices 0,7,14,21,28,35
dbt_results = {
"metadata": {"dbt_version": "1.8.0", "run_id": "run_abc123"},
"results": [
{
"unique_id": f"model.analytics.fct_{i:03d}",
"status": "success" if i % 7 != 0 else "error",
"execution_time": round(0.8 + i * 0.12, 3),
"failures": None
if i % 7 != 0
else [{"message": f"Column col_{i} not found", "line": i % 40}],
}
for i in range(40)
],
}
dbt_ctx = json.dumps(dbt_results, indent=2)
case2a = QualityCase(
name="dbt failed models (error detection)",
context=dbt_ctx,
question="Which dbt model unique_ids have status 'error'? List all of them.",
expected_keywords=[f"fct_{i:03d}" for i in range(40) if i % 7 == 0],
)
case2b = QualityCase(
name="dbt slowest model (max lookup)",
context=dbt_ctx,
question="Which model has the longest execution_time? Reply with just the unique_id.",
expected_keywords=["fct_039"],
)
# ── Case 3: Search result ranking ────────────────────────────────────────
search_results = [
{
"rank": i + 1,
"score": round(0.98 - i * 0.03, 4),
"document_id": f"doc_{i:04d}",
"title": f"Engineering runbook #{i:03d}",
"content": f"This document covers topic_{i} configuration and deployment steps for service_{i}.",
}
for i in range(20)
]
search_ctx = json.dumps(search_results, indent=2)
case3a = QualityCase(
name="Search top result (rank 1 lookup)",
context=search_ctx,
question="What is the document_id of the result with rank 1? Reply with just the document_id.",
expected_keywords=["doc_0000"],
)
case3b = QualityCase(
name="Search score lookup (doc_0007 score)",
context=search_ctx,
question="What is the score of document_id doc_0007? Reply with just the number.",
expected_keywords=["0.77"],
)
# ── Case 4: Multi-fact reasoning ─────────────────────────────────────────
incident = {
"incident_id": "INC-20250615-004",
"severity": "P1",
"affected_service": "payment-processor",
"root_cause": "Database connection pool exhausted due to slow query on orders_v2 table",
"timeline": [
{"time": "14:02", "event": "Alert fired: latency > 5s"},
{"time": "14:07", "event": "On-call engineer paged"},
{
"time": "14:15",
"event": "Query identified: SELECT * FROM orders_v2 WHERE status='pending'",
},
{"time": "14:28", "event": "Index added on (status, created_at)"},
{"time": "14:31", "event": "Latency normalized"},
],
"mttr_minutes": 29,
"action_items": [
"Add query timeout of 10s on payment-processor",
"Review all full-table scans in orders_v2",
"Set up connection pool monitoring alert",
],
}
incident_ctx = json.dumps(incident, indent=2)
case4a = QualityCase(
name="Incident MTTR (exact field lookup)",
context=incident_ctx,
question="What was the MTTR in minutes for this incident? Reply with just the number.",
expected_keywords=["29"],
)
case4b = QualityCase(
name="Incident fix action (reasoning from timeline)",
context=incident_ctx,
question="What specific action resolved the latency issue at 14:28?",
expected_keywords=["index", "status", "created_at"],
)
return [case1a, case1b, case2a, case2b, case3a, case3b, case4a, case4b]
# ── Runner ────────────────────────────────────────────────────────────────────
@dataclass
class QualityResult:
case: QualityCase
std_answer: str
hdm_answer: str
std_hits: int
hdm_hits: int
total_kw: int
tokens_saved_pct: float
compress_ms: float
@property
def std_pass(self) -> bool:
return self.case.pass_threshold(self.std_hits, self.total_kw)
@property
def hdm_pass(self) -> bool:
return self.case.pass_threshold(self.hdm_hits, self.total_kw)
@property
def quality_delta(self) -> int:
return self.hdm_hits - self.std_hits
def run_case(case: QualityCase, token: str, host: str) -> QualityResult:
from headroom import compress
messages = [
{"role": "system", "content": case.context},
{"role": "user", "content": case.question},
]
std_answer = _call(messages, token, host)
std_hits, total = case.score(std_answer)
t0 = time.perf_counter()
compressed = compress(messages, model="claude-sonnet-4-5-20250929")
compress_ms = (time.perf_counter() - t0) * 1000
hdm_answer = _call(compressed.messages, token, host)
hdm_hits, _ = case.score(hdm_answer)
std_tokens = len(json.dumps(messages)) // 4
hdm_tokens = len(json.dumps(compressed.messages)) // 4
saved_pct = (std_tokens - hdm_tokens) / max(std_tokens, 1) * 100
return QualityResult(
case=case,
std_answer=std_answer,
hdm_answer=hdm_answer,
std_hits=std_hits,
hdm_hits=hdm_hits,
total_kw=total,
tokens_saved_pct=saved_pct,
compress_ms=compress_ms,
)
# ── Display ───────────────────────────────────────────────────────────────────
def _show(r: QualityResult) -> None:
std_sym = "✓" if r.std_pass else "✗"
hdm_sym = "✓" if r.hdm_pass else "✗"
delta_sym = "=" if r.quality_delta == 0 else ("+" if r.quality_delta > 0 else "-")
print(f"\n ┌─ {r.case.name}")
print(f" │ Token reduction : ~{r.tokens_saved_pct:.0f}% │ Compress: {r.compress_ms:.0f}ms")
print(
f" │ Standard [{std_sym}] : {r.std_hits}/{r.total_kw} keywords matched"
f" ({'PASS' if r.std_pass else 'FAIL'})"
)
print(
f" │ Headroom [{hdm_sym}] : {r.hdm_hits}/{r.total_kw} keywords matched"
f" ({'PASS' if r.hdm_pass else 'FAIL'}) [{delta_sym} quality delta]"
)
print(f" │ Q: {r.case.question[:80]}")
std_preview = r.std_answer[:120].replace("\n", " ")
hdm_preview = r.hdm_answer[:120].replace("\n", " ")
print(f" │ Std answer : {std_preview}")
print(f" └─ Hdm answer : {hdm_preview}")
# ── Main ──────────────────────────────────────────────────────────────────────
def main() -> int:
print()
print("╔═══════════════════════════════════════════════════════════════╗")
print("║ Cortex Code × Headroom — Quality Benchmark ║")
print("║ Does compression affect answer accuracy? ║")
print("╚═══════════════════════════════════════════════════════════════╝")
if not _SF_CONN:
print("\n ✗ Set SF_CONN=<connection-name> to run.")
return 1
import io
try:
import snowflake.connector
except ImportError:
print("\n ✗ snowflake-connector-python not installed.")
return 1
_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()")
locator = cur.fetchone()[0].lower()
host = f"{locator}.snowflakecomputing.com"
finally:
sys.stdout = _s
cases = _make_cases()
print(f"\n Model : {_SF_MODEL}")
print(f" Host : {host}")
print(f" Cases : {len(cases)} ({len(cases) * 2} total API calls)\n")
print(" Method: embed known facts → ask factual questions → score keyword hits")
print(" Pass threshold: ≥60% expected keywords found in answer\n")
results: list[QualityResult] = []
for i, case in enumerate(cases, 1):
print(f" [{i}/{len(cases)}] {case.name} ...", end=" ", flush=True)
try:
r = run_case(case, token, host)
results.append(r)
std_s = "✓" if r.std_pass else "✗"
hdm_s = "✓" if r.hdm_pass else "✗"
print(f"std={std_s}({r.std_hits}/{r.total_kw}) hdm={hdm_s}({r.hdm_hits}/{r.total_kw})")
_show(r)
except Exception as exc:
print(f"FAILED: {exc}")
conn.close()
if not results:
print("\n No results.")
return 1
# ── Summary ───────────────────────────────────────────────────────────────
std_passes = sum(1 for r in results if r.std_pass)
hdm_passes = sum(1 for r in results if r.hdm_pass)
total = len(results)
regressions = sum(1 for r in results if r.std_pass and not r.hdm_pass)
improvements = sum(1 for r in results if not r.std_pass and r.hdm_pass)
unchanged = sum(1 for r in results if r.std_pass == r.hdm_pass)
avg_token_saving = sum(r.tokens_saved_pct for r in results) / total
print()
print("╔═══════════════════════════════════════════════════════════════╗")
print("║ QUALITY SUMMARY ║")
print("╠═══════════════════════════════════════════════════════════════╣")
print(f" {'Test':<42} {'Std':>4} {'Hdm':>4} {'Delta':>6} {'Tokens↓':>7}")
print(f" {'─' * 42} {'─' * 4} {'─' * 4} {'─' * 6} {'─' * 7}")
for r in results:
delta = r.hdm_hits - r.std_hits
delta_str = f"{delta:+d}" if delta != 0 else " ="
std_s = "✓" if r.std_pass else "✗"
hdm_s = "✓" if r.hdm_pass else "✗"
print(
f" {r.case.name[:42]:<42} "
f"{std_s} {r.std_hits}/{r.total_kw} "
f"{hdm_s} {r.hdm_hits}/{r.total_kw} "
f"{delta_str:>6} "
f"~{r.tokens_saved_pct:.0f}%"
)
print(f" {'─' * 42} {'─' * 4} {'─' * 4} {'─' * 6} {'─' * 7}")
print(f" {'TOTAL':<42} {std_passes}/{total} {hdm_passes}/{total}")
print()
print(
f" Pass rate : Standard {std_passes}/{total} ({std_passes / total * 100:.0f}%) "
f"│ Headroom {hdm_passes}/{total} ({hdm_passes / total * 100:.0f}%)"
)
print(f" Regressions (std pass → hdm fail) : {regressions}")
print(f" Improvements (std fail → hdm pass): {improvements}")
print(f" Unchanged : {unchanged}")
print(f" Avg token reduction : ~{avg_token_saving:.0f}%")
print()
if regressions == 0:
print(" ✓ No quality regressions — headroom compression preserved answer accuracy")
else:
print(
f" ⚠ {regressions} regression(s) — headroom dropped facts needed for correct answer"
)
print("╚═══════════════════════════════════════════════════════════════╝")
print()
return 0
if __name__ == "__main__":
sys.exit(main())

12
tests/e2e_cortex_savings.py
View file

@ -173,14 +173,18 @@ def _sf_call(messages: list[dict], token: str, host: str) -> dict:
{
"model": _SF_MODEL,
"messages": messages,
"max_tokens": 64,
"max_completion_tokens": 64,
"stream": False,
}
).encode()
req = urllib.request.Request(
f"https://{host}/api/v2/cortex/inference:complete",
f"https://{host}/api/v2/cortex/v1/chat/completions",
data=body,
headers={"Authorization": f'Snowflake Token="{token}"', "Content-Type": "application/json"},
headers={
"Authorization": f'Snowflake Token="{token}"',
"Content-Type": "application/json",
"User-Agent": "headroom-bench/1.0",
},
method="POST",
)
with urllib.request.urlopen(req, timeout=60) as r:
@ -315,7 +319,7 @@ def main() -> int:
results: list[R] = []
# ── 1. Snowflake Cortex (system-message pattern) ──────────────────────────
print("\n▶ Snowflake Cortex /api/v2/cortex/inference:complete")
print("\n▶ Snowflake Cortex /api/v2/cortex/v1/chat/completions")
try:
import io