chopratejas--headroom/examples/strands_via_proxy_demo.py

#!/usr/bin/env python3
"""End-to-end demo: Strands -> Headroom proxy -> Bedrock.

Proves the four Path-B fixes work together against live AWS Bedrock:

  Fix #1  PrefixCacheTracker.update_from_response on the backend path
  Fix #2  CCR response intercept for the OpenAI-shape proxy
  Fix #3  LiteLLM's native cache_control -> cachePoint translation
  Fix #4  Strands harness label in CLIENT_UA_MAP

What this script does
---------------------
1. Spawns the Headroom proxy as a subprocess (backend=bedrock).
2. Waits for /readyz.
3. Sends two requests in the same session via Strands' OpenAIModel
   pointed at the proxy. The system prompt is intentionally large
   (>1024 tokens; Bedrock's minimum cacheable block) and tagged with
   ``cache_control: {type: "ephemeral"}`` via Headroom's CacheAligner.
4. Reports:
     * compression numbers (tokens before/after, on each turn)
     * Bedrock cache hits (cache_read_input_tokens on turn 2 -- proves
       LiteLLM translated cache_control to cachePoint AND Bedrock
       served from the cache)
     * harness label (proves X-Client: strands flows through)
5. Tears the proxy back down.

Requirements
------------
- AWS credentials in ~/.aws/credentials or environment.
- ``pip install -e .[strands,bedrock]`` from the repo root (already
  done if you've been running the existing Strands demo).
- A free local TCP port (default 8765; override via ``--port``).

Run
---
    AWS_REGION=us-west-2 python examples/strands_via_proxy_demo.py
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.request
from contextlib import suppress
from pathlib import Path
from typing import Any

# Defer Strands / openai imports until after we've validated the proxy
# starts -- that way the error message for a missing dep doesn't bury
# a more useful "proxy refused to start" trace.


# ----------------------------------------------------------------------------
# Constants
# ----------------------------------------------------------------------------

DEFAULT_PORT = 8765
DEFAULT_REGION = "us-west-2"
DEFAULT_MODEL = "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
SESSION_ID = "strands-via-proxy-demo-1"

# A ~2.5K-token block. Sonnet 4.5 caches empirically at this size on
# Bedrock (verified: cache_write=2206 with the same prompt below).
# The CacheAligner (Anthropic-style ephemeral cache_control) marks
# this as cacheable; LiteLLM translates the marker to Bedrock
# cachePoint; Bedrock serves it from the read cache on turn 2.
LARGE_SYSTEM_PROMPT = (
    "You are a precise technical assistant. "
    "Treat the following as authoritative reference context for every "
    "question in this conversation. Quote it accurately, do not "
    "fabricate. Reference context:\n\n"
    + "Headroom is an open-source context compression layer for LLM "
    "applications. It sits in front of provider APIs (Anthropic, "
    "OpenAI, Bedrock, Vertex) and shrinks the prompt without losing "
    "semantically important information. " * 200
)


# ----------------------------------------------------------------------------
# Proxy lifecycle
# ----------------------------------------------------------------------------


def start_proxy(port: int, region: str) -> subprocess.Popen[bytes]:
    """Spawn `headroom proxy --backend bedrock` as a subprocess."""
    env = os.environ.copy()
    env.setdefault("AWS_REGION", region)
    env.setdefault("AWS_DEFAULT_REGION", region)
    # Crank logging up so we can read pipeline decisions live.
    env.setdefault("HEADROOM_LOG", "INFO")
    cmd = [
        sys.executable,
        "-m",
        "headroom.cli",
        "proxy",
        "--backend",
        "bedrock",
        "--region",
        region,
        "--port",
        str(port),
    ]
    print(f"  $ {' '.join(cmd)}", file=sys.stderr)
    log_path = Path("/tmp") / f"strands_via_proxy_demo_{port}.log"
    log_file = log_path.open("wb")
    proc = subprocess.Popen(  # noqa: S603 — argv is fixed above
        cmd,
        env=env,
        stdout=log_file,
        stderr=subprocess.STDOUT,
    )
    print(f"  proxy logs -> {log_path}", file=sys.stderr)
    return proc


def wait_for_proxy_ready(port: int, timeout_s: float = 30.0) -> None:
    """Poll /readyz until the proxy answers or timeout."""
    url = f"http://127.0.0.1:{port}/readyz"
    deadline = time.time() + timeout_s
    last_err: Exception | None = None
    while time.time() < deadline:
        try:
            with urllib.request.urlopen(url, timeout=1) as resp:  # noqa: S310
                if resp.status == 200:
                    return
        except (urllib.error.URLError, ConnectionError, TimeoutError) as e:
            last_err = e
        time.sleep(0.5)
    raise RuntimeError(
        f"Proxy on port {port} did not become ready within {timeout_s}s; last error: {last_err!r}"
    )


def stop_proxy(proc: subprocess.Popen[bytes]) -> None:
    """Politely shut the proxy down."""
    with suppress(ProcessLookupError):
        proc.terminate()
    try:
        proc.wait(timeout=5)
    except subprocess.TimeoutExpired:
        proc.kill()
        proc.wait(timeout=5)


# ----------------------------------------------------------------------------
# Strands wiring
# ----------------------------------------------------------------------------


def build_agent(port: int, model_id: str) -> Any:
    """Construct a Strands Agent pointed at the proxy.

    Uses OpenAIModel + base_url because that's the proxy-friendly path
    (Bedrock's native auth would bypass the proxy entirely).
    """
    from strands import Agent
    from strands.models.openai import OpenAIModel

    model = OpenAIModel(
        model_id=model_id,
        client_args={
            "api_key": "dummy-bedrock-uses-aws-creds-at-proxy",
            "base_url": f"http://127.0.0.1:{port}/v1",
            "default_headers": {
                # Stable session key so the proxy's PrefixCacheTracker
                # treats both turns as the same conversation.
                "x-headroom-session-id": SESSION_ID,
                # Harness identification (Fix #4) — the proxy labels
                # this request as 'strands' in metrics + outcomes.
                "X-Client": "strands",
            },
        },
        # Bedrock-Claude rejects the OpenAI default of temperature=1.0
        # for some Opus versions; pinning a Bedrock-compatible value.
        params={"max_tokens": 200, "temperature": 0.2},
    )
    return Agent(model=model, system_prompt=LARGE_SYSTEM_PROMPT)


# ----------------------------------------------------------------------------
# Cache-stat probes
# ----------------------------------------------------------------------------


def fetch_proxy_stats(port: int) -> dict[str, Any]:
    """Fetch overall proxy stats so we can correlate per-turn behaviour."""
    url = f"http://127.0.0.1:{port}/stats"
    try:
        with urllib.request.urlopen(url, timeout=2) as resp:  # noqa: S310
            return json.loads(resp.read().decode())
    except Exception as e:
        return {"_error": str(e)}


# ----------------------------------------------------------------------------
# Direct HTTP smoke test (no Strands) -- proves the proxy alone
# ----------------------------------------------------------------------------


def direct_smoke_test(
    port: int, model_id: str, session_id: str, with_cache_control: bool
) -> dict[str, Any]:
    """Issue a single chat.completion via raw HTTP -- proves the wiring
    end-to-end without the Strands layer in the way.

    When ``with_cache_control=True`` the system message carries an
    explicit Anthropic-style ``cache_control: ephemeral`` block; this
    isolates "does the proxy correctly forward cache_control + extract
    response cache stats" from "does CacheAligner insert cache_control
    on its own". Both questions must answer "yes" for the Path-B claim
    to hold end-to-end.
    """
    url = f"http://127.0.0.1:{port}/v1/chat/completions"
    if with_cache_control:
        system_content: Any = [
            {
                "type": "text",
                "text": LARGE_SYSTEM_PROMPT,
                "cache_control": {"type": "ephemeral"},
            }
        ]
    else:
        system_content = LARGE_SYSTEM_PROMPT
    body = json.dumps(
        {
            "model": model_id,
            "messages": [
                {"role": "system", "content": system_content},
                {"role": "user", "content": "In one short sentence, what is Headroom?"},
            ],
            "max_tokens": 60,
            "temperature": 0.2,
        }
    ).encode()
    req = urllib.request.Request(  # noqa: S310
        url,
        data=body,
        headers={
            "Content-Type": "application/json",
            "Authorization": "Bearer dummy-bedrock-uses-aws-creds-at-proxy",
            "x-headroom-session-id": session_id,
            "X-Client": "strands",
        },
        method="POST",
    )
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=60) as resp:  # noqa: S310
        body_bytes = resp.read()
    elapsed_ms = (time.time() - t0) * 1000
    parsed = json.loads(body_bytes)
    return {"elapsed_ms": elapsed_ms, "body": parsed}


# ----------------------------------------------------------------------------
# Main demo
# ----------------------------------------------------------------------------


def usage_summary(usage: dict[str, Any]) -> str:
    """One-line digest of the usage block returned by the proxy."""
    return (
        f"prompt={usage.get('prompt_tokens', 0)} "
        f"completion={usage.get('completion_tokens', 0)} "
        f"cache_read={usage.get('cache_read_input_tokens', 0)} "
        f"cache_write={usage.get('cache_creation_input_tokens', 0)}"
    )


async def run_demo(port: int, region: str, model_id: str) -> int:
    print("=" * 76)
    print(" Headroom Path-B E2E: Strands -> Headroom proxy -> Bedrock")
    print("=" * 76)
    print(f" port={port} region={region} model={model_id}")
    print(f" session_id={SESSION_ID}")
    print()

    print("[1/4] Spawning Headroom proxy ...")
    proxy = start_proxy(port=port, region=region)
    try:
        try:
            wait_for_proxy_ready(port=port, timeout_s=45.0)
        except Exception as e:
            print(f"  ! Proxy failed to start: {e}", file=sys.stderr)
            return 2
        print("  proxy ready.")

        # ----------------------------------------------------------------
        # 2. Direct HTTP smoke test WITH explicit cache_control.
        # This isolates "proxy forwards cache_control + extracts stats"
        # from "CacheAligner inserts cache_control on its own".
        # ----------------------------------------------------------------
        print("\n[2/4] Direct HTTP probe -- explicit cache_control (turn A, turn B same session)")
        smoke_session = "cc-smoke-1"
        try:
            smoke_a = direct_smoke_test(
                port=port, model_id=model_id, session_id=smoke_session, with_cache_control=True
            )
            smoke_b = direct_smoke_test(
                port=port, model_id=model_id, session_id=smoke_session, with_cache_control=True
            )
        except urllib.error.HTTPError as e:
            err_body = e.read().decode("utf-8", errors="replace")[:500]
            print(f"  ! smoke test failed: HTTP {e.code}: {err_body}", file=sys.stderr)
            return 3
        ua = smoke_a["body"].get("usage", {})
        ub = smoke_b["body"].get("usage", {})
        print(f"  turn A: {usage_summary(ua)}  ({smoke_a['elapsed_ms']:.0f}ms)")
        print(f"  turn B: {usage_summary(ub)}  ({smoke_b['elapsed_ms']:.0f}ms)")
        cache_works = ub.get("cache_read_input_tokens", 0) > 0
        cache_write_a = ua.get("cache_creation_input_tokens", 0) > 0
        if cache_works:
            print(
                f"  ✓ cache hit on turn B (read={ub['cache_read_input_tokens']}) -- proxy chain OK."
            )
        elif cache_write_a:
            print("  ! turn A wrote cache but turn B didn't read -- session keying may be off.")
        else:
            print(
                "  ! no cache write on turn A -- LiteLLM cache_control translation OR proxy did not forward it."
            )

        # ----------------------------------------------------------------
        # 2b. Streaming probe: same chain but stream=True. The non-
        # streaming smoke proved the synchronous path; this proves the
        # streaming path also (a) forwards cache_control and (b) parses
        # cache stats from the SSE usage frame and (c) updates the
        # prefix tracker on stream end (Fix #1 streaming half).
        # ----------------------------------------------------------------
        print("\n[2b/4] Streaming probe -- same explicit cache_control payload")
        try:
            stream_url = f"http://127.0.0.1:{port}/v1/chat/completions"
            stream_body = json.dumps(
                {
                    "model": model_id,
                    "messages": [
                        {
                            "role": "system",
                            "content": [
                                {
                                    "type": "text",
                                    "text": LARGE_SYSTEM_PROMPT,
                                    "cache_control": {"type": "ephemeral"},
                                }
                            ],
                        },
                        {"role": "user", "content": "Reply in 5 words."},
                    ],
                    "max_tokens": 30,
                    "temperature": 0.2,
                    "stream": True,
                    "stream_options": {"include_usage": True},
                }
            ).encode()
            stream_req = urllib.request.Request(  # noqa: S310
                stream_url,
                data=stream_body,
                headers={
                    "Content-Type": "application/json",
                    "Authorization": "Bearer dummy",
                    "x-headroom-session-id": smoke_session,
                    "X-Client": "strands",
                },
                method="POST",
            )
            t0 = time.time()
            last_usage_frame: dict[str, Any] | None = None
            chunk_count = 0
            with urllib.request.urlopen(stream_req, timeout=60) as stream_resp:  # noqa: S310
                for raw_line in stream_resp:
                    line = raw_line.decode("utf-8", errors="replace").strip()
                    if not line.startswith("data: "):
                        continue
                    payload = line[6:]
                    if payload == "[DONE]":
                        continue
                    try:
                        event = json.loads(payload)
                    except json.JSONDecodeError:
                        continue
                    chunk_count += 1
                    if event.get("usage"):
                        last_usage_frame = event["usage"]
            stream_elapsed_ms = (time.time() - t0) * 1000
            print(
                f"  streamed chunks={chunk_count} elapsed={stream_elapsed_ms:.0f}ms "
                f"final_usage={last_usage_frame}"
            )
            if last_usage_frame and last_usage_frame.get("cache_read_input_tokens", 0) > 0:
                print("  ✓ streaming path also returned cache_read_input_tokens > 0.")
            elif last_usage_frame is None:
                print("  ! no final usage frame surfaced -- check include_usage wiring.")
            else:
                print("  - no cache hit on streaming probe (may be a 3rd-call eviction edge).")
        except urllib.error.HTTPError as e:
            err_body = e.read().decode("utf-8", errors="replace")[:500]
            print(f"  ! streaming probe failed: HTTP {e.code}: {err_body}", file=sys.stderr)

        # ----------------------------------------------------------------
        # 2c. Compression probe.
        # ContentRouter SKIPS user + system messages by design
        # (skip_user_messages=True, skip_system=True at content_router.py:456
        # and :2294). The bulk savings in real agent loops come from
        # compressing TOOL RESULTS (and assistant turns), not from
        # rewriting the user's question or paraphrasing the system
        # prompt. To exercise the compression pipeline we send a fake
        # assistant turn that just returned a verbose JSON tool result.
        # SmartCrusher (Rust-backed, always available) targets exactly
        # this shape.
        # ----------------------------------------------------------------
        print("\n[2c/4] Compression probe -- tool_result with verbose JSON")
        # ContentRouter defaults (headroom/transforms/content_router.py):
        #   skip_user_messages: True   (line 456) -- "subject of conversation"
        #   skip_system: True          (line 2294) -- system prompt is sacred
        #   compress_assistant_text_blocks: False (line 472) -- conservative
        # The ONE shape that compresses by default is the tool_result. This
        # matches the real-world AWS agent-loop pattern: tool calls return
        # large JSON/log/diff blobs that accumulate across turns and dominate
        # the prompt. ContentRouter classifies the tool_result content,
        # dispatches via the magika/unidiff detection chain to a per-type
        # compressor (SmartCrusher for JSON arrays here), records % saved.
        big_tool_result = json.dumps(
            [
                {
                    "id": f"order-{i}",
                    "customer_id": f"cust-{i % 100}",
                    "status": "completed",
                    "total_usd": 100 + i,
                    "items": [
                        {"sku": f"sku-{j}", "qty": 1, "name": f"Product {j}"} for j in range(5)
                    ],
                    "created_at": f"2026-05-{(i % 28) + 1:02d}T10:00:00Z",
                    "notes": "Standard processing, no exceptions",
                }
                for i in range(250)
            ]
        )
        print(
            f"  tool_result size: {len(big_tool_result)} chars (~{len(big_tool_result) // 4} tokens)"
        )
        tool_probe_body = json.dumps(
            {
                "model": model_id,
                "messages": [
                    {"role": "user", "content": "List recent completed orders."},
                    {
                        "role": "assistant",
                        "content": None,
                        "tool_calls": [
                            {
                                "id": "call_1",
                                "type": "function",
                                "function": {
                                    "name": "list_orders",
                                    "arguments": "{}",
                                },
                            }
                        ],
                    },
                    {
                        "role": "tool",
                        "tool_call_id": "call_1",
                        "content": big_tool_result,
                    },
                    {
                        "role": "user",
                        "content": "How many orders are in 'completed' status? Reply with the number only.",
                    },
                ],
                "tools": [
                    {
                        "type": "function",
                        "function": {
                            "name": "list_orders",
                            "description": "List recent orders.",
                            "parameters": {
                                "type": "object",
                                "properties": {},
                                "required": [],
                            },
                        },
                    }
                ],
                "max_tokens": 16,
                "temperature": 0.0,
            }
        ).encode()
        tool_req = urllib.request.Request(  # noqa: S310
            f"http://127.0.0.1:{port}/v1/chat/completions",
            data=tool_probe_body,
            headers={
                "Content-Type": "application/json",
                "Authorization": "Bearer dummy",
                "x-headroom-session-id": "compression-probe-session",
                "X-Client": "strands",
            },
            method="POST",
        )
        t0 = time.time()
        try:
            with urllib.request.urlopen(tool_req, timeout=60) as tr:  # noqa: S310
                tool_resp = json.loads(tr.read())
            print(f"  elapsed={time.time() - t0:.1f}s usage={tool_resp.get('usage', {})}")
            # Check proxy /stats AFTER this call so we can see the compression delta.
            post_stats = fetch_proxy_stats(port=port)
            comp = post_stats.get("summary", {}).get("compression", {})
            uncomp = post_stats.get("summary", {}).get("uncompressed_requests", {})
            print(
                f"  cumulative compression: requests_compressed={comp.get('requests_compressed', 0)} "
                f"tokens_removed={comp.get('total_tokens_removed', 0)} "
                f"best_pct={comp.get('best_compression_pct', 0.0):.1f}%"
            )
            print(f"  uncompressed reasons: {uncomp}")
        except urllib.error.HTTPError as e:
            err_body = e.read().decode("utf-8", errors="replace")[:500]
            print(f"  ! compression probe failed: HTTP {e.code}: {err_body}")

        # ----------------------------------------------------------------
        # 3. Strands agent: two turns, same session.
        # No explicit cache_control here -- this tests whether
        # CacheAligner (inside the proxy) inserts the marker itself.
        # ----------------------------------------------------------------
        print("\n[3/4] Building Strands agent ...")
        agent = build_agent(port=port, model_id=model_id)

        print(
            "\n[4/4] Two-turn cache test via Strands (cache_control inserted by CacheAligner) ..."
        )
        print("  turn 1: priming the cache with the large system prompt")
        r1 = agent("In one short sentence, what is Headroom?")
        print(f"  turn 1 response: {str(r1)[:160]}")

        print("\n  turn 2: same session -> should hit Bedrock prompt cache")
        r2 = agent("In one short sentence, what providers does it support?")
        print(f"  turn 2 response: {str(r2)[:160]}")

        # Verdict: pull the proxy stats (cumulative) so we can SEE cache stats
        stats = fetch_proxy_stats(port=port)
        print("\n  proxy /stats snapshot:")
        print(f"    {json.dumps(stats, indent=2, default=str)[:1200]}")

        # Tail the proxy log for cache_read mentions on turn 2 -- this is
        # the load-bearing assertion: Fix #1 + Fix #3 worked iff the proxy
        # logged a non-zero cache_read_input_tokens on the second call.
        log_path = Path("/tmp") / f"strands_via_proxy_demo_{port}.log"
        log_tail = log_path.read_text(errors="replace").splitlines()[-200:]
        cache_lines = [
            line
            for line in log_tail
            if "cache_read" in line.lower() or "cache stats" in line.lower()
        ]
        ccr_lines = [line for line in log_tail if "ccr" in line.lower()]
        print("\n  proxy log cache lines (last 200 lines):")
        if cache_lines:
            for line in cache_lines[-10:]:
                print(f"    {line[:240]}")
        else:
            print("    (no cache_read events surfaced — possible miss on this run)")
        if ccr_lines:
            print("\n  proxy log CCR lines (last 200 lines):")
            for line in ccr_lines[-10:]:
                print(f"    {line[:240]}")

        print("\n" + "=" * 76)
        print(" PATH-B E2E COMPLETE.")
        print(" If you see cache_read_input_tokens > 0 on the second call,")
        print(" the prefix-cache + cachePoint chain is working end-to-end.")
        print("=" * 76)
        return 0
    finally:
        print("\n  shutting down proxy ...")
        stop_proxy(proxy)


def main() -> int:
    ap = argparse.ArgumentParser(description="Strands -> Headroom proxy -> Bedrock E2E")
    ap.add_argument("--port", type=int, default=DEFAULT_PORT)
    ap.add_argument("--region", default=DEFAULT_REGION)
    ap.add_argument("--model", default=DEFAULT_MODEL)
    args = ap.parse_args()
    return asyncio.run(run_demo(port=args.port, region=args.region, model_id=args.model))


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