Upload aco/proxy.py

aco/proxy.py

@@ -6,12 +6,13 @@ Start:  aco-proxy --port 8080
 Use:    openai.api_base = "http://localhost:8080/v1"
 
 The proxy intercepts POST /v1/chat/completions and:
-1. Routes to cheapest adequate model (v1 tier-router)
+1. Routes to cheapest adequate model
 2. Gates unnecessary tool calls (v1 tool-gater, F1=0.92)
 3. Lays out prompts for cache reuse (system + tools in prefix)
 4. Compresses verbose error traces and thinking-only turns
 5. Collects telemetry: cost, tokens, latency, cache hits
-6. Reports live dashboard at GET /dashboard
+6. Live dashboard at GET /dashboard
+7. JSON telemetry at GET /telemetry
 
 Zero agent code changes needed.
 """
@@ -45,7 +46,6 @@ MODEL_REGISTRY = {
     "gemini-3-pro":        {"tier": 5, "cost_in": 2.00,  "cost_out": 12.50, "ctx": 1048576},
 }
 
-# Default provider endpoints (override via OPENAI_BASE_URL, ANTHROPIC_BASE_URL, etc.)
 PROVIDER_ENDPOINTS = {
     "openai":    os.environ.get("OPENAI_BASE_URL",    "https://api.openai.com/v1"),
     "anthropic": os.environ.get("ANTHROPIC_BASE_URL", "https://api.anthropic.com/v1"),
@@ -61,216 +61,253 @@ MODEL_PROVIDER = {
     "deepseek-v3.2": "deepseek",
 }
 
-# ── Telemetry Store ──────────────────────────────────────────────────
+# ── Tool-Gater Classifier (v1: DistilBERT, F1=0.92) ─────────────────
+
+_tool_gater = None  # Lazy-loaded singleton
+
+def _get_tool_gater():
+    """Lazy-load the v1 DistilBERT tool-gater."""
+    global _tool_gater
+    if _tool_gater is not None:
+        return _tool_gater
+
+    try:
+        from transformers import AutoTokenizer, AutoModelForSequenceClassification
+        import torch
+
+        model_id = "narcolepticchicken/aco-specialists-tool-gater"
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=2)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model.to(device)
+        model.eval()
+        _tool_gater = (model, tokenizer, device)
+        return _tool_gater
+    except Exception as e:
+        print(f"[ACO] Failed to load tool-gater: {e}. Using heuristic fallback.")
+        return None
+
+
+def should_gate_tools_ml(messages: List[Dict]) -> bool:
+    """
+    ML-based tool gating using the v1 DistilBERT classifier (F1=0.92).
+
+    Falls back to heuristic if classifier is unavailable.
+    """
+    # Find user query
+    user_text = ""
+    system_text = ""
+    for msg in messages:
+        if msg.get("role") == "user":
+            user_text = str(msg.get("content", ""))[:1500]
+            break
+        elif msg.get("role") == "system":
+            system_text = str(msg.get("content", ""))[:500]
+
+    if not user_text:
+        return False
+
+    # Check if tools have already been used in this conversation
+    has_tool_history = any(
+        '<function=' in str(msg.get("content", ""))
+        for msg in messages if msg.get("role") == "assistant"
+    )
+    if has_tool_history:
+        return False
+
+    # Build input text
+    text = f"Query: {user_text}"
+    if system_text:
+        text = f"System: {system_text}\n\n{text}"
+
+    # Try ML classifier
+    gater = _get_tool_gater()
+    if gater:
+        model, tokenizer, device = gater
+        try:
+            import torch
+            inputs = tokenizer(text[:2000], truncation=True, max_length=512,
+                              return_tensors="pt").to(device)
+            with torch.no_grad():
+                logits = model(**inputs).logits
+                probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
+            # Label 0 = "skip_tool", Label 1 = "call_tool"
+            # Gate tools OFF (return True) when prob[no_tool] > prob[tool]
+            return probs[0] >= probs[1]
+        except Exception as e:
+            print(f"[ACO] Tool-gater inference failed: {e}")
+
+    # Heuristic fallback
+    return heuristic_tool_gate(user_text)
+
+
+def heuristic_tool_gate(user_text: str) -> bool:
+    """Heuristic: should we gate (suppress) tools for this query?"""
+    ut = user_text.lower()
+    simple_patterns = [
+        r'\bwhat is\b', r'\bwho (is|was)\b', r'\bwhen (is|was)\b',
+        r'\bdefine\b', r'\bexplain\b', r'\bsummarize\b',
+        r'\bhow (do|does|to)\b', r'\bdifference between\b',
+        r'\bcapital of\b', r'\bmeaning of\b', r'\btranslate\b',
+    ]
+    return any(re.search(p, ut) for p in simple_patterns)
 
-@dataclass
-class TraceRecord:
-    request_id: str
-    timestamp: str
-    model: str
-    provider: str
-    tier: int
-    input_tokens: int
-    output_tokens: int
-    cache_hit_tokens: int
-    latency_ms: float
-    cost: float
-    tool_gated: bool
-    context_compressed: float  # ratio
-    success: bool
-    error: Optional[str] = None
 
 # ── Cache-Aware Prompt Layout ────────────────────────────────────────
 
 def layout_cache_prompt(messages: List[Dict], tools: Optional[List[Dict]] = None) -> List[Dict]:
     """
-    Reorder messages so cache-stable content is in prefix:
-    - system prompt (stable) → first
-    - tool definitions (stable) → next (converted to system-like)
-    - user messages → after
-    - dynamic content (timestamps, request IDs) → moved to suffix or stripped
+    Reorder messages for maximum prefix-cache reuse.
+    Stable content (system, tool defs) first; dynamic content last.
+    Strips timestamps/request IDs from user messages to improve cache hits.
     """
     laid_out = []
-    tool_block = None
+    has_tool_block = False
 
-    # Find and remove system message, move to front
     for msg in messages:
         if msg.get("role") == "system":
-            laid_out.insert(0, {"role": "system", "content": msg["content"]})
-        elif msg.get("role") == "tool" and tool_block is None:
-            # First tool message: convert to system for cache
-            laid_out.insert(1, {"role": "system", 
-                               "content": f"[TOOL_DEFS]\n{msg.get('content', '')}"})
-            tool_block = len(laid_out)
+            # System prompt first (most stable — best cache target)
+            content = str(msg.get("content", ""))
+            # Normalize: strip timestamps
+            content = re.sub(r'\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}[^\s]*', '[TIME]', content)
+            content = re.sub(r'run_[a-f0-9]{8,}', 'run_xxx', content)
+            laid_out.insert(0, {"role": "system", "content": content})
+
+        elif msg.get("role") == "tool" and not has_tool_block:
+            # Tool definitions: convert to system block for cache
+            laid_out.insert(1, {"role": "system",
+                                "content": f"[TOOL_DEFS]\n{str(msg.get('content', ''))}"})
+            has_tool_block = True
+
+        elif msg.get("role") == "user":
+            content = str(msg.get("content", ""))
+            # Normalize dynamic markers
+            content = re.sub(r'(?:req|trace|run)_[a-f0-9]{8,32}', 'xxx', content)
+            content = re.sub(r'\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}[^\s]*', '[TIME]', content)
+            laid_out.append({"role": "user", "content": content})
+
         else:
             laid_out.append(msg)
 
-    # Append tools as last system-like block (Anthropic: suffix, OpenAI: doesn't matter)
-    if tools and not tool_block:
-        tool_defs = json.dumps([{"name": t.get("function", {}).get("name", t.get("name", "")),
-                                 "desc": t.get("function", {}).get("description", 
-                                          t.get("description", ""))[:200]}
-                                for t in tools])
-        laid_out.insert(1, {"role": "system", "content": f"[TOOL_DEFS]\n{tool_defs}"})
-
-    # Strip dynamic metadata from user messages
-    for msg in laid_out:
-        if msg.get("role") == "user" and isinstance(msg.get("content"), str):
-            # Remove timestamps, request IDs, run IDs
-            content = msg["content"]
-            content = re.sub(r'\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}[^\s]*', '[TIME]', content)
-            content = re.sub(r'run_[a-f0-9]{8,}', 'run_xxx', content)
-            content = re.sub(r'req_[a-f0-9]{8,}', 'req_xxx', content)
-            content = re.sub(r'trace_[a-f0-9]{8,}', 'trace_xxx', content)
-            msg["content"] = content
+    # Append tools as stable suffix block if provided externally
+    if tools and not has_tool_block:
+        tool_names = [t.get("function", {}).get("name", t.get("name", "?"))
+                      for t in tools][:20]
+        laid_out.insert(1, {"role": "system",
+                            "content": f"[TOOL_DEFS]\n{json.dumps(tool_names)}"})
 
     return laid_out
 
 
 # ── Context Compression ──────────────────────────────────────────────
 
-def compress_context(messages: List[Dict]) -> List[Dict]:
-    """Compress verbose messages while preserving signal."""
+def compress_context(messages: List[Dict]) -> tuple:
+    """Compress verbose agent messages while preserving signal."""
     compressed = []
     total_orig = 0
     total_comp = 0
 
     for msg in messages:
-        content = msg.get("content", "")
+        content = str(msg.get("content", ""))
         role = msg.get("role", "")
-        total_orig += len(str(content))
+        total_orig += len(content)
 
         if role == "user":
-            cl = str(content).lower()
-            # Trim long error outputs
-            if len(str(content)) > 2000 and any(k in cl for k in ['error', 'traceback', 'exception']):
-                content = str(content)[:300] + '\n... [error trace trimmed] ...\n' + str(content)[-300:]
-            # Trim huge file listings
-            elif len(str(content)) > 3000 and 'ls -' not in cl[:50]:
-                content = str(content)[:2000] + '\n... [output trimmed] ...'
+            cl = content.lower()
+            # Trim stack traces to head + tail
+            if len(content) > 2000 and any(k in cl for k in ['traceback', 'error:', 'exception']):
+                lines = content.split('\n')
+                head = '\n'.join(lines[:8])
+                tail = '\n'.join(lines[-5:])
+                content = f"{head}\n... [{len(lines)-13} lines trimmed] ...\n{tail}"
+            elif len(content) > 3000:
+                content = content[:2000] + '\n... [output trimmed] ...'
 
         elif role == "assistant":
-            # Strip thinking-only turns (no function calls, no code blocks)
-            s = str(content)
+            s = content
+            # Drop pure-thinking turns (no function calls, no code blocks)
             if len(s) > 800 and '```' not in s and '<function=' not in s:
-                # Check if it's mostly reasoning without action
-                if not re.search(r'\b(?:execute|run|apply|create|delete|modify|write|patch|fix)\b', s, re.IGNORECASE):
+                if not re.search(r'\b(?:execute|run|apply|create|delete|modify|write|patch|fix|submit)\b',
+                                 s, re.IGNORECASE):
                     content = s[:200] + '\n... [thinking trimmed] ...'
-            # Trim long code blocks
+            # Trim large code blocks
             if len(s) > 4000:
-                content = s[:3000] + '\n... [code trimmed] ...'
+                content = s[:3000] + '\n... [truncated] ...'
 
-        total_comp += len(str(content))
+        total_comp += len(content)
         compressed.append({**msg, "content": content})
 
     ratio = total_comp / max(total_orig, 1)
     return compressed, ratio
 
 
-# ── Tool Gate ────────────────────────────────────────────────────────
-
-def should_gate_tools(messages: List[Dict]) -> bool:
-    """Quick heuristic: does this request really need tools?
-    
-    If the agent is asking a simple question with no tool-calling
-    history, suppress tools to save input tokens.
-    
-    For production, replace with v1 tool-gater classifier (F1=0.92).
-    """
-    # Find the last user message
-    user_text = ""
-    for msg in reversed(messages):
-        if msg.get("role") == "user":
-            user_text = str(msg.get("content", ""))[:500].lower()
-            break
-
-    # Check if any prior assistant message already used tools
-    has_tool_history = any(
-        '<function=' in str(msg.get("content", ""))
-        for msg in messages if msg.get("role") == "assistant"
-    )
-
-    # If there's tool history, don't gate
-    if has_tool_history:
-        return False
-
-    # Quick-answer patterns that rarely need tools
-    simple_patterns = [
-        r'\bwhat is\b', r'\bwho (is|was)\b', r'\bwhen (is|was)\b',
-        r'\bdefine\b', r'\bexplain\b', r'\bsummarize\b',
-        r'\bhow (do|does|to)\b', r'\bdifference between\b',
-        r'\bcapital of\b', r'\bmeaning of\b', r'\btranslate\b',
-    ]
-    return any(re.search(p, user_text) for p in simple_patterns)
-
-
 # ── Model Router ─────────────────────────────────────────────────────
 
 def route_model(requested_model: str, messages: List[Dict]) -> str:
-    """Route to cheapest model that can handle this request.
-    
-    Strategy: trust the agent's model choice but downgrade when safe.
-    """
-    # Get the last user message for difficulty estimation
+    """Route to cheapest model that can handle this request."""
+    info = MODEL_REGISTRY.get(requested_model)
+    if not info:
+        return requested_model
+
+    tier = info["tier"]
+
+    # Get last user message
     user_text = ""
     for msg in reversed(messages):
         if msg.get("role") == "user":
             user_text = str(msg.get("content", ""))
             break
 
-    info = MODEL_REGISTRY.get(requested_model)
-    if not info:
-        return requested_model  # Unknown model, pass through
+    # Downgrade: if using tier 3+ for short simple text, use tier 1
+    if tier >= 3 and len(user_text) < 300:
+        return "deepseek-v4-flash"
 
-    tier = info["tier"]
-    text_len = len(user_text)
-
-    # Downgrade rules: if using tier 3+ for short, simple text
-    if tier >= 3 and text_len < 300:
-        # Try tier 1
-        t1_models = [m for m, d in MODEL_REGISTRY.items() if d["tier"] == 1]
-        if t1_models:
-            return t1_models[0]  # deepseek-v4-flash
-
-    # For coding tasks, keep tier 2 minimum
-    code_indicators = ['def ', 'class ', 'function', 'import ', '```', 'fix ', 'bug',
-                       'implement', 'refactor', 'test_', 'pytest']
-    if any(c in user_text for c in code_indicators) and tier < 2:
-        t2_models = [m for m, d in MODEL_REGISTRY.items() if d["tier"] == 2]
-        if t2_models:
-            return t2_models[0]
+    # Coding floor: keep tier 2 minimum
+    code_words = ['def ', 'class ', 'function', 'import ', '```', 'fix ', 'bug',
+                  'implement', 'refactor', 'test_', 'pytest', 'traceback', 'error:']
+    if any(c in user_text for c in code_words) and tier < 2:
+        return "gpt-5-mini"
 
     return requested_model
 
 
 # ── Cost Calculator ──────────────────────────────────────────────────
 
-def estimate_cost(model: str, input_tokens: int, output_tokens: int,
-                  cache_hit_tokens: int = 0) -> float:
-    """Estimate cost in USD."""
+def compute_cost(model: str, input_tokens: int, output_tokens: int,
+                 cache_hit_tokens: int = 0) -> float:
+    """Estimate cost in USD per current provider pricing."""
     info = MODEL_REGISTRY.get(model)
     if not info:
         return 0.0
+    chargeable_input = max(0, input_tokens - cache_hit_tokens)
+    return round(
+        (chargeable_input / 1_000_000) * info["cost_in"] +
+        (output_tokens / 1_000_000) * info["cost_out"],
+        6)
 
-    # Cache hit: input tokens that hit cache are free or heavily discounted
-    chargeable_input = input_tokens - cache_hit_tokens
-    cost = (chargeable_input / 1_000_000) * info["cost_in"]
-    cost += (output_tokens / 1_000_000) * info["cost_out"]
-    return round(cost, 6)
 
+# ── Telemetry Store ──────────────────────────────────────────────────
 
-# ── Response Rewriter ────────────────────────────────────────────────
-
-def add_usage_headers(response: dict, model: str, cost: float, 
-                      cache_hit: int = 0, compressed: float = 1.0) -> dict:
-    """Inject ACO telemetry into response headers and usage."""
-    if "usage" not in response:
-        response["usage"] = {}
-    response["usage"]["aco_cost_usd"] = cost
-    response["usage"]["aco_model"] = model
-    response["usage"]["aco_cache_hit_tokens"] = cache_hit
-    response["usage"]["aco_compression_ratio"] = compressed
-    return response
+@dataclass
+class TraceRecord:
+    request_id: str
+    timestamp: str
+    model: str
+    provider: str
+    tier: int
+    input_tokens: int
+    output_tokens: int
+    cache_hit_tokens: int
+    latency_ms: float
+    cost: float
+    tool_gated: bool
+    gated_by: str   # "ml" | "heuristic" | "none"
+    context_compressed: float     # ratio
+    cache_layout_applied: bool
+    model_routed: bool            # was the model changed?
+    original_model: str
+    success: bool
+    error: Optional[str] = None
 
 
 # ══════════════════════════════════════════════════════════════════════
@@ -278,148 +315,153 @@ def add_usage_headers(response: dict, model: str, cost: float,
 # ══════════════════════════════════════════════════════════════════════
 
 if FASTAPI_AVAILABLE:
-    app = FastAPI(title="ACO Proxy", version="1.0.0")
+    app = FastAPI(title="ACO Proxy", version="1.1.0")
     telemetry_store: List[TraceRecord] = []
     telemetry_lock = threading.Lock()
     start_time = datetime.utcnow()
 
     @app.get("/health")
     async def health():
-        return {"status": "ok", "uptime": str(datetime.utcnow() - start_time)}
+        return {"status": "ok", "uptime_seconds": (datetime.utcnow() - start_time).total_seconds()}
 
     @app.get("/v1/models")
     async def list_models():
-        """List all known models."""
         return {
             "object": "list",
-            "data": [{"id": m, "object": "model", "owned_by": PROVIDER_ENDPOINTS.get(
-                MODEL_PROVIDER.get(m, "unknown"), "unknown")} 
+            "data": [{"id": m, "object": "model",
+                       "owned_by": MODEL_PROVIDER.get(m, "unknown")}
                       for m in MODEL_REGISTRY]
         }
 
     @app.post("/v1/chat/completions")
     async def chat_completions(request: Request):
-        """
-        OpenAI-compatible chat completions endpoint.
-        Applies ACO optimizations transparently.
-        """
+        """OpenAI-compatible endpoint. ACO optimizations applied transparently."""
         body = await request.json()
-        request_id = body.get("user", hashlib.md5(str(time.time()).encode()).hexdigest()[:8])
+        import uuid
+        request_id = body.get("user", str(uuid.uuid4())[:8])
 
-        # ── Step 1: Extract params ──
+        # ── Extract params ──
         messages = body.get("messages", [])
         tools = body.get("tools")
         requested_model = body.get("model", "gpt-5-mini")
         stream = body.get("stream", False)
+        original_model = requested_model
 
-        # ── Step 2: Route model ──
+        # ── Route model ──
         routed_model = route_model(requested_model, messages)
+        model_routed = routed_model != requested_model
         provider = MODEL_PROVIDER.get(routed_model, "openai")
 
-        # ── Step 3: Gate tools ──
+        # ── Gate tools (ML classifier with heuristic fallback) ──
         tools_gated = False
-        if tools and should_gate_tools(messages):
+        gated_by = "none"
+        if tools and should_gate_tools_ml(messages):
             tools = None
             tools_gated = True
+            gated_by = "ml" if _get_tool_gater() else "heuristic"
 
-        # ── Step 4: Layout for cache ──
+        # ── Layout for cache ──
         laid_out = layout_cache_prompt(messages, tools)
+        cache_applied = laid_out != messages
 
-        # ── Step 5: Compress context ──
+        # ── Compress context ──
         compressed_messages, compression_ratio = compress_context(laid_out)
 
-        # ── Step 6: Build forwarding request ──
+        # ── Forward request ──
         forward_body = {**body}
         forward_body["model"] = routed_model
         forward_body["messages"] = compressed_messages
         if tools is None and "tools" in forward_body:
             del forward_body["tools"]
 
-        # ── Step 7: Determine provider endpoint ──
         endpoint = PROVIDER_ENDPOINTS.get(provider, PROVIDER_ENDPOINTS["openai"])
         target_url = f"{endpoint}/chat/completions"
 
-        # ── Step 8: Forward with auth ──
-        auth_header = request.headers.get("authorization", "")
-        headers = {"content-type": "application/json"}
-        if auth_header:
-            headers["authorization"] = auth_header
-
-        # API key overrides
         api_key_map = {
+            "openai": os.environ.get("OPENAI_API_KEY"),
             "anthropic": os.environ.get("ANTHROPIC_API_KEY"),
             "google": os.environ.get("GOOGLE_API_KEY"),
             "deepseek": os.environ.get("DEEPSEEK_API_KEY"),
-            "openai": os.environ.get("OPENAI_API_KEY"),
         }
-        provider_key = api_key_map.get(provider)
-        if provider_key and not auth_header:
-            headers["authorization"] = f"Bearer {provider_key}"
+        auth = request.headers.get("authorization", "")
+        headers = {"content-type": "application/json"}
+        if not auth:
+            key = api_key_map.get(provider)
+            if key:
+                headers["authorization"] = f"Bearer {key}"
+        else:
+            headers["authorization"] = auth
 
-        # ── Step 9: Make upstream call ──
+        # ── Make upstream call ──
         t_start = time.time()
         error = None
         success = True
+        response_data = {}
 
         try:
             async with httpx.AsyncClient(timeout=300.0) as client:
                 upstream = await client.post(target_url, json=forward_body, headers=headers)
                 latency = (time.time() - t_start) * 1000
 
-                if upstream.status_code != 200:
-                    error = f"Upstream {upstream.status_code}: {upstream.text[:200]}"
-                    success = False
-                    # Fall back to original model on failure
-                    if routed_model != requested_model:
-                        forward_body["model"] = requested_model
-                        t_start2 = time.time()
-                        upstream2 = await client.post(target_url, json=forward_body, headers=headers)
-                        latency = (time.time() - t_start) * 1000
-                        upstream = upstream2
-                        routed_model = requested_model
-                        success = upstream.status_code == 200
-                        if not success:
-                            error = f"Fallback also failed: {upstream.status_code}"
+                if upstream.status_code != 200 and model_routed:
+                    # Fall back to original model
+                    forward_body["model"] = original_model
+                    upstream2 = await client.post(target_url, json=forward_body, headers=headers)
+                    latency = (time.time() - t_start) * 1000
+                    upstream = upstream2
+                    routed_model = original_model
+                    model_routed = False
 
                 if stream:
                     return StreamingResponse(
                         upstream.aiter_bytes(),
                         media_type="text/event-stream",
-                        headers={"x-aco-model": routed_model,
-                                 "x-aco-cost": "streaming"}
+                        headers={"x-aco-model": routed_model, "x-aco-tier": str(
+                            MODEL_REGISTRY.get(routed_model, {}).get("tier", "?"))}
                     )
 
-                response_data = upstream.json()
+                if upstream.status_code == 200:
+                    response_data = upstream.json()
+                else:
+                    error = f"Upstream {upstream.status_code}: {upstream.text[:200]}"
+                    success = False
 
         except Exception as e:
             latency = (time.time() - t_start) * 1000
             error = str(e)
             success = False
+
+        if not success:
             response_data = {
                 "id": f"aco-err-{request_id}",
                 "object": "chat.completion",
                 "created": int(time.time()),
                 "model": requested_model,
                 "choices": [{"index": 0, "message": {"role": "assistant",
-                              "content": f"ACO proxy error: {error}"},
+                              "content": f"[ACO proxy error: {error}]"},
                               "finish_reason": "error"}],
                 "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
             }
 
-        # ── Step 10: Compute cost ──
+        # ── Compute cost ──
         usage = response_data.get("usage", {})
         input_tokens = usage.get("prompt_tokens", 0)
         output_tokens = usage.get("completion_tokens", 0)
-        # Some providers report cache hits
-        cache_hit = usage.get("cache_read_input_tokens", 0) or \
-                    usage.get("prompt_tokens_details", {}).get("cached_tokens", 0)
-        cost = estimate_cost(routed_model, input_tokens, output_tokens, cache_hit)
-
-        # ── Step 11: Rewrite response ──
-        response_data = add_usage_headers(response_data, routed_model, cost, 
-                                          cache_hit, compression_ratio)
-
-        # ── Step 12: Record telemetry ──
+        cache_hit = (usage.get("cache_read_input_tokens", 0) or
+                     usage.get("prompt_tokens_details", {}).get("cached_tokens", 0))
+        cost = compute_cost(routed_model, input_tokens, output_tokens, cache_hit)
+
+        # ── Rewrite response ──
+        response_data.setdefault("usage", {})
+        response_data["usage"]["aco_cost_usd"] = cost
+        response_data["usage"]["aco_model"] = routed_model
+        response_data["usage"]["aco_tier"] = MODEL_REGISTRY.get(routed_model, {}).get("tier", 0)
+        response_data["usage"]["aco_cache_hit_tokens"] = cache_hit
+        response_data["usage"]["aco_compression_ratio"] = round(compression_ratio, 2)
+        response_data["usage"]["aco_tool_gated"] = tools_gated
+        response_data["model"] = requested_model  # Agent sees original model
+
+        # ── Record telemetry ──
         trace = TraceRecord(
             request_id=request_id,
             timestamp=datetime.utcnow().isoformat(),
@@ -432,7 +474,11 @@ if FASTAPI_AVAILABLE:
             latency_ms=round(latency, 1),
             cost=cost,
             tool_gated=tools_gated,
+            gated_by=gated_by,
             context_compressed=round(compression_ratio, 3),
+            cache_layout_applied=cache_applied,
+            model_routed=model_routed,
+            original_model=original_model,
             success=success,
             error=error,
         )
@@ -443,193 +489,118 @@ if FASTAPI_AVAILABLE:
 
     @app.get("/dashboard")
     async def dashboard():
-        """Live cost dashboard as HTML."""
+        """Live HTML cost dashboard."""
         with telemetry_lock:
             traces = list(telemetry_store)
 
+        n = len(traces)
+        if n == 0:
+            return HTMLResponse("<h2>No traffic yet. Send requests to /v1/chat/completions</h2>")
+
         total_cost = sum(t.cost for t in traces)
-        total_calls = len(traces)
         successful = sum(1 for t in traces if t.success)
-        total_input = sum(t.input_tokens for t in traces)
-        total_output = sum(t.output_tokens for t in traces)
+        total_in = sum(t.input_tokens for t in traces)
+        total_out = sum(t.output_tokens for t in traces)
         total_cache = sum(t.cache_hit_tokens for t in traces)
-        avg_latency = sum(t.latency_ms for t in traces) / max(total_calls, 1)
-
-        # Tier distribution
-        tier_counts = defaultdict(int)
-        tier_costs = defaultdict(float)
-        for t in traces:
-            tier_counts[t.tier] += 1
-            tier_costs[t.tier] += t.cost
-
-        # Tool gating stats
+        avg_lat = sum(t.latency_ms for t in traces) / n
         gated = sum(1 for t in traces if t.tool_gated)
-        gated_success = sum(1 for t in traces if t.tool_gated and t.success)
+        routed = sum(1 for t in traces if t.model_routed)
 
-        # Model distribution
-        model_counts = defaultdict(int)
+        tier_calls = defaultdict(int)
+        tier_cost = defaultdict(float)
+        model_calls = defaultdict(int)
         for t in traces:
-            model_counts[t.model] += 1
+            tier_calls[t.tier] += 1
+            tier_cost[t.tier] += t.cost
+            model_calls[t.model] += 1
 
-        html = f"""
-        <!DOCTYPE html>
-        <html><head>
-        <title>ACO Proxy Dashboard</title>
-        <meta charset="utf-8">
+        html = f"""<!DOCTYPE html><html><head>
+        <title>ACO Proxy</title><meta charset="utf-8"><meta http-equiv="refresh" content="3">
         <style>
-            body {{ font-family: system-ui; margin: 2rem; background: #111; color: #eee; }}
-            .card {{ background: #1a1a2e; border-radius: 8px; padding: 1rem; margin: 0.5rem; }}
-            .grid {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 1rem; }}
-            .metric {{ font-size: 2rem; font-weight: bold; color: #00d4aa; }}
-            .label {{ font-size: 0.8rem; color: #888; text-transform: uppercase; }}
-            table {{ width: 100%; border-collapse: collapse; margin: 1rem 0; }}
-            th, td {{ padding: 0.5rem; text-align: left; border-bottom: 1px solid #333; }}
-            th {{ color: #888; font-size: 0.8rem; text-transform: uppercase; }}
-            .error {{ color: #ff5555; }}
-            .success {{ color: #00d4aa; }}
-        </style>
-        <meta http-equiv="refresh" content="5">
-        </head><body>
-        <h1>🤖 ACO Proxy Dashboard</h1>
-        <p>Uptime: {datetime.utcnow() - start_time}</p>
-
+            * {{ margin:0; padding:0; box-sizing:border-box; }}
+            body {{ font-family: system-ui; background: #0d1117; color: #c9d1d9; padding: 1.5rem; }}
+            h1 {{ font-size: 1.2rem; margin-bottom: 0.5rem; }}
+            .grid {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(180px,1fr)); gap: 0.5rem; margin-bottom: 1rem; }}
+            .card {{ background: #161b22; border: 1px solid #30363d; border-radius: 6px; padding: 0.8rem; }}
+            .card .val {{ font-size: 1.6rem; font-weight: 700; color: #58a6ff; }}
+            .card .lbl {{ font-size: 0.7rem; color: #8b949e; text-transform: uppercase; letter-spacing: 0.5px; }}
+            table {{ width: 100%; border-collapse: collapse; font-size: 0.8rem; margin-bottom: 1rem; }}
+            th, td {{ padding: 0.4rem 0.5rem; text-align: right; border-bottom: 1px solid #21262d; }}
+            th {{ color: #8b949e; font-weight: 500; text-transform: uppercase; font-size: 0.65rem; }}
+            td:first-child, th:first-child {{ text-align: left; }}
+            .good {{ color: #3fb950; }} .bad {{ color: #f85149; }} .dim {{ color: #8b949e; }}
+        </style></head><body>
+        <h1>🤖 ACO Proxy <span class="dim">— {n} calls, ${total_cost:.4f} total</span></h1>
         <div class="grid">
-            <div class="card">
-                <div class="label">Total Calls</div>
-                <div class="metric">{total_calls}</div>
-            </div>
-            <div class="card">
-                <div class="label">Total Cost</div>
-                <div class="metric">${total_cost:.4f}</div>
-            </div>
-            <div class="card">
-                <div class="label">Success Rate</div>
-                <div class="metric">{successful/max(total_calls,1)*100:.1f}%</div>
-            </div>
-            <div class="card">
-                <div class="label">Avg Latency</div>
-                <div class="metric">{avg_latency:.0f}ms</div>
-            </div>
-            <div class="card">
-                <div class="label">Total Tokens In/Out</div>
-                <div class="metric">{total_input//1000}k / {total_output//1000}k</div>
-            </div>
-            <div class="card">
-                <div class="label">Cache Hit Tokens</div>
-                <div class="metric">{total_cache//1000}k</div>
-            </div>
-            <div class="card">
-                <div class="label">Tools Gated</div>
-                <div class="metric">{gated} calls</div>
-            </div>
-            <div class="card">
-                <div class="label">Cost/Call</div>
-                <div class="metric">${total_cost/max(total_calls,1):.4f}</div>
-            </div>
+            <div class="card"><div class="val">{successful/n*100:.0f}%</div><div class="lbl">Success Rate</div></div>
+            <div class="card"><div class="val">${total_cost:.4f}</div><div class="lbl">Total Cost</div></div>
+            <div class="card"><div class="val">${total_cost/max(n,1):.5f}</div><div class="lbl">Avg Cost/Call</div></div>
+            <div class="card"><div class="val">{avg_lat:.0f}ms</div><div class="lbl">Avg Latency</div></div>
+            <div class="card"><div class="val">{total_in//1000}k</div><div class="lbl">Tokens In</div></div>
+            <div class="card"><div class="val">{total_out//1000}k</div><div class="lbl">Tokens Out</div></div>
+            <div class="card"><div class="val">{total_cache//1000}k</div><div class="lbl">Cache Hits</div></div>
+            <div class="card"><div class="val">{gated}</div><div class="lbl">Tools Gated</div></div>
+            <div class="card"><div class="val">{routed}</div><div class="lbl">Models Rerouted</div></div>
         </div>
-
-        <h2>By Tier</h2>
-        <table>
-            <tr><th>Tier</th><th>Calls</th><th>Cost</th><th>Cost/Call</th></tr>
-        """
-        for tier in sorted(tier_counts):
-            html += f"<tr><td>{tier}</td><td>{tier_counts[tier]}</td>"
-            html += f"<td>${tier_costs[tier]:.4f}</td>"
-            html += f"<td>${tier_costs[tier]/max(tier_counts[tier],1):.4f}</td></tr>"
-
-        html += """
-        </table>
-        <h2>By Model</h2>
-        <table>
-            <tr><th>Model</th><th>Calls</th></tr>
-        """
-        for model, count in sorted(model_counts.items(), key=lambda x: -x[1]):
-            html += f"<tr><td>{model}</td><td>{count}</td></tr>"
-
-        html += """
-        </table>
-        <h2>Recent Calls</h2>
-        <table>
-            <tr><th>Time</th><th>Model</th><th>Tier</th><th>Tokens In</th><th>Tokens Out</th>
-            <th>Cache</th><th>Latency</th><th>Cost</th><th>Gated</th><th>Compress</th><th>Status</th></tr>
-        """
-        for t in reversed(traces[-50:]):
-            status = '<span class="success">✓</span>' if t.success else '<span class="error">✗</span>'
-            html += f"<tr><td>{t.timestamp[-12:]}</td><td>{t.model}</td><td>{t.tier}</td>"
-            html += f"<td>{t.input_tokens}</td><td>{t.output_tokens}</td>"
-            html += f"<td>{t.cache_hit_tokens}</td><td>{t.latency_ms:.0f}ms</td>"
-            html += f"<td>${t.cost:.6f}</td><td>{'✓' if t.tool_gated else ''}</td>"
-            html += f"<td>{t.context_compressed:.2f}</td><td>{status}</td></tr>"
-
+        <table><tr><th>Model</th><th>Calls</th></tr>"""
+        for m, c in sorted(model_calls.items(), key=lambda x: -x[1]):
+            html += f"<tr><td>{m}</td><td>{c}</td></tr>"
+        html += """</table>
+        <table><tr><th>Time</th><th>Model</th><th>Tier</th><th>Tokens</th><th>Cost</th><th>Lat</th><th>Gated</th><th>Routed</th></tr>"""
+        for t in reversed(traces[-30:]):
+            s = 'good' if t.success else 'bad'
+            html += f"<tr><td class='dim'>{t.timestamp[-8:]}</td>"
+            html += f"<td>{t.model[:20]}</td><td>{t.tier}</td>"
+            html += f"<td>{t.input_tokens}+{t.output_tokens}</td>"
+            html += f"<td>${t.cost:.5f}</td><td>{t.latency_ms:.0f}ms</td>"
+            html += f"<td>{'✓' if t.tool_gated else ''}</td>"
+            html += f"<td>{'✓' if t.model_routed else ''}</td></tr>"
         html += "</table></body></html>"
         return HTMLResponse(html)
 
     @app.get("/telemetry")
-    async def telemetry():
-        """JSON telemetry endpoint."""
+    async def telemetry_json():
+        """JSON telemetry for programmatic consumption."""
         with telemetry_lock:
-            traces = [{
-                "request_id": t.request_id,
-                "timestamp": t.timestamp,
-                "model": t.model,
-                "provider": t.provider,
-                "tier": t.tier,
-                "input_tokens": t.input_tokens,
-                "output_tokens": t.output_tokens,
-                "cache_hit_tokens": t.cache_hit_tokens,
-                "latency_ms": t.latency_ms,
-                "cost": t.cost,
-                "tool_gated": t.tool_gated,
-                "context_compressed": t.context_compressed,
-                "success": t.success,
-                "error": t.error,
-            } for t in telemetry_store]
-
-        total_cost = sum(t.cost for t in telemetry_store)
+            traces = [{"model": t.model, "tier": t.tier, "cost": t.cost,
+                       "input_tokens": t.input_tokens, "output_tokens": t.output_tokens,
+                       "cache_hit_tokens": t.cache_hit_tokens, "latency_ms": t.latency_ms,
+                       "tool_gated": t.tool_gated, "gated_by": t.gated_by,
+                       "model_routed": t.model_routed, "original_model": t.original_model,
+                       "context_compressed": t.context_compressed,
+                       "success": t.success, "error": t.error,
+                       "timestamp": t.timestamp}
+                      for t in telemetry_store]
         return {
             "total_calls": len(traces),
-            "total_cost": total_cost,
+            "total_cost": round(sum(t["cost"] for t in traces), 6),
             "calls": traces,
-            "summary": {
-                "success_rate": sum(1 for t in telemetry_store if t.success) / max(len(traces), 1),
-                "avg_latency_ms": sum(t.latency_ms for t in telemetry_store) / max(len(traces), 1),
-                "total_input_tokens": sum(t.input_tokens for t in telemetry_store),
-                "total_output_tokens": sum(t.output_tokens for t in telemetry_store),
-                "total_cache_hit_tokens": sum(t.cache_hit_tokens for t in telemetry_store),
-                "tools_gated": sum(1 for t in telemetry_store if t.tool_gated),
-            }
         }
 
     @app.get("/telemetry/reset")
     async def reset_telemetry():
-        """Reset telemetry store."""
         with telemetry_lock:
             telemetry_store.clear()
-        return {"status": "ok", "message": "Telemetry reset"}
+        return {"status": "ok"}
 
 
 def serve(host: str = "0.0.0.0", port: int = 8080):
-    """Start the ACO proxy server."""
     if not FASTAPI_AVAILABLE:
-        print("ERROR: fastapi and uvicorn required. pip install fastapi uvicorn httpx")
+        print("ERROR: pip install fastapi uvicorn httpx")
         return
-    print(f"🚀 ACO Proxy starting on http://{host}:{port}")
-    print(f"   Dashboard: http://localhost:{port}/dashboard")
-    print(f"   Telemetry: http://localhost:{port}/telemetry")
-    print(f"   Set openai.api_base = 'http://localhost:{port}/v1'")
-    uvicorn.run(app, host=host, port=port, log_level="info")
+    print(f"🚀 ACO Proxy → http://{host}:{port}")
+    print(f"   Dashboard:   http://localhost:{port}/dashboard")
+    print(f"   Telemetry:   http://localhost:{port}/telemetry")
+    print(f"   Agent usage: openai.api_base = 'http://localhost:{port}/v1'")
+    uvicorn.run(app, host=host, port=port, log_level="warning")
 
 
-# ── CLI entry point ──────────────────────────────────────────────────
-
 def main():
     import argparse
-    parser = argparse.ArgumentParser(description="ACO Proxy Server")
-    parser.add_argument("--host", default="0.0.0.0")
-    parser.add_argument("--port", type=int, default=8080)
-    args = parser.parse_args()
-    serve(host=args.host, port=args.port)
-
+    p = argparse.ArgumentParser(description="ACO Proxy Server")
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=8080)
+    serve(**vars(p.parse_args()))
 
 if __name__ == "__main__":
     main()