#!/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())