@miller-tech/uap 1.133.2 → 1.134.0

package/tools/agents/scripts/anthropic_proxy.py

@@ -184,6 +184,27 @@ PROXY_CONTEXT_WINDOW = int(os.environ.get("PROXY_CONTEXT_WINDOW", "0"))
 PROXY_CONTEXT_PRUNE_THRESHOLD = float(
     os.environ.get("PROXY_CONTEXT_PRUNE_THRESHOLD", "0.85")
 )
+# Compaction forcing (Option A, 2026-07-10): Claude Code decides when to
+# auto-compact against ITS believed model window (~200k) using
+# /v1/messages/count_tokens, so an HONEST count on a smaller local rail lets
+# every session grow into the "thrash band" — above the rail (per-request
+# critical prunes, findings evaporate, read-only doom loops; observed live:
+# 87 reads / 0 writes in 30 min at 122% utilization) but below the client's
+# compact trigger (~92.5% of 200k). Scaling the reported count makes the
+# client compact BEFORE the rail. "auto" (default) derives the scale from the
+# LIVE rail at call time — assumed_window / (rail * target_fraction) — so it
+# tracks rail resizes; a number forces that scale; "1"/"off" disables.
+PROXY_COUNT_TOKENS_SCALE = os.environ.get("PROXY_COUNT_TOKENS_SCALE", "auto")
+PROXY_CLIENT_ASSUMED_WINDOW = int(
+    os.environ.get("PROXY_CLIENT_ASSUMED_WINDOW", "200000")
+)
+# Where compaction should land sessions, as a fraction of the rail. Unset
+# (0) = auto: just under the pruner threshold (x0.95), so the client's own
+# compaction fires BEFORE the proxy ever needs to prune — the pruner stays on
+# purely as a backstop. An explicit value overrides (clamped to <1).
+PROXY_COMPACT_TARGET_FRACTION = float(
+    os.environ.get("PROXY_COMPACT_TARGET_FRACTION", "0")
+)
 PROXY_CONTEXT_PRUNE_TARGET_FRACTION = float(
     os.environ.get("PROXY_CONTEXT_PRUNE_TARGET_FRACTION", "0.50")
 )
@@ -8936,7 +8957,7 @@ async def _maybe_apply_recipe(anthropic_resp, anthropic_body, openai_body, clien
         return anthropic_resp
 
 
-async def _heartbeat_then_buffered(produce_coro, model: str):
+async def _heartbeat_then_buffered(produce_coro, model: str, input_tokens: int = 0):
     """SSE generator: keep-alive heartbeat wrapper for the guarded-non-stream path.
 
     Emits an immediate ``message_start`` so the client registers an active
@@ -8953,7 +8974,7 @@ async def _heartbeat_then_buffered(produce_coro, model: str):
     msg_id = f"msg_{uuid.uuid4().hex[:24]}"
     yield (
         f"event: message_start\n"
-        f"data: {json.dumps({'type': 'message_start', 'message': {'id': msg_id, 'type': 'message', 'role': 'assistant', 'content': [], 'model': model, 'stop_reason': None, 'stop_sequence': None, 'usage': {'input_tokens': 0, 'output_tokens': 0}}})}\n\n"
+        f"data: {json.dumps({'type': 'message_start', 'message': {'id': msg_id, 'type': 'message', 'role': 'assistant', 'content': [], 'model': model, 'stop_reason': None, 'stop_sequence': None, 'usage': {'input_tokens': input_tokens, 'output_tokens': 0}}})}\n\n"
     )
     interval = PROXY_STREAM_HEARTBEAT_SECS if PROXY_STREAM_HEARTBEAT_SECS > 0 else 15.0
     task = asyncio.ensure_future(produce_coro)
@@ -9011,7 +9032,13 @@ async def stream_anthropic_message(anthropic_resp: dict, emit_message_start: boo
             "model": anthropic_resp.get("model", "unknown"),
             "stop_reason": None,
             "stop_sequence": None,
-            "usage": {"input_tokens": 0, "output_tokens": 0},
+            "usage": {
+                "input_tokens": int(
+                    (anthropic_resp.get("usage", {}) or {}).get("input_tokens", 0)
+                    * _count_tokens_scale()
+                ),
+                "output_tokens": 0,
+            },
         }
         yield f"event: message_start\ndata: {json.dumps({'type': 'message_start', 'message': message})}\n\n"
 
@@ -9053,11 +9080,13 @@ async def stream_anthropic_message(anthropic_resp: dict, emit_message_start: boo
             )
         block_index += 1
 
-    output_tokens = anthropic_resp.get("usage", {}).get("output_tokens", 0)
+    resp_usage = anthropic_resp.get("usage", {}) or {}
+    output_tokens = resp_usage.get("output_tokens", 0)
+    scaled_input = int(resp_usage.get("input_tokens", 0) * _count_tokens_scale())
     stop_reason = anthropic_resp.get("stop_reason", "end_turn")
     yield (
         "event: message_delta\n"
-        f"data: {json.dumps({'type': 'message_delta', 'delta': {'stop_reason': stop_reason, 'stop_sequence': None}, 'usage': {'output_tokens': output_tokens}})}\n\n"
+        f"data: {json.dumps({'type': 'message_delta', 'delta': {'stop_reason': stop_reason, 'stop_sequence': None}, 'usage': {'input_tokens': scaled_input, 'output_tokens': output_tokens}})}\n\n"
     )
     yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n"
 
@@ -9084,10 +9113,11 @@ async def stream_anthropic_response(
     msg_id = f"msg_{uuid.uuid4().hex[:24]}"
     stream_started = time.monotonic()
 
-    # message_start
+    # message_start — carries the client's (scaled) input size; a hardcoded 0
+    # here blinded clients to their own context usage (compaction never fired).
     yield (
         f"event: message_start\n"
-        f"data: {json.dumps({'type': 'message_start', 'message': {'id': msg_id, 'type': 'message', 'role': 'assistant', 'content': [], 'model': model, 'stop_reason': None, 'stop_sequence': None, 'usage': {'input_tokens': 0, 'output_tokens': 0}}})}\n\n"
+        f"data: {json.dumps({'type': 'message_start', 'message': {'id': msg_id, 'type': 'message', 'role': 'assistant', 'content': [], 'model': model, 'stop_reason': None, 'stop_sequence': None, 'usage': {'input_tokens': _client_input_tokens(monitor), 'output_tokens': 0}}})}\n\n"
     )
 
     # content_block_start for text (index 0)
@@ -9392,10 +9422,12 @@ async def stream_anthropic_response(
     if _is_unexpected_end_turn(synthetic_openai_resp, anthropic_body):
         monitor.unexpected_end_turn_count += 1
 
-    # message_delta with final stop reason
+    # message_delta with final stop reason. Repeats the (scaled) input size so
+    # clients that read usage from the delta rather than message_start still
+    # see their context usage — the compaction-forcing signal.
     yield (
         f"event: message_delta\n"
-        f"data: {json.dumps({'type': 'message_delta', 'delta': {'stop_reason': finish_reason, 'stop_sequence': None}, 'usage': {'output_tokens': output_tokens}})}\n\n"
+        f"data: {json.dumps({'type': 'message_delta', 'delta': {'stop_reason': finish_reason, 'stop_sequence': None}, 'usage': {'input_tokens': _client_input_tokens(monitor), 'output_tokens': output_tokens}})}\n\n"
     )
 
     # message_stop
@@ -9563,7 +9595,98 @@ async def count_tokens(request: Request):
             status_code=400,
             media_type="application/json",
         )
-    return {"input_tokens": estimate_total_tokens(body)}
+    est = estimate_total_tokens(body)
+    scale = _count_tokens_scale()
+    if scale > 1.0:
+        scaled = int(est * scale)
+        # Once per scale value, explain the discrepancy an operator would
+        # otherwise chase between this endpoint and the pruner's numbers.
+        global _count_scale_logged
+        if _count_scale_logged != scale:
+            _count_scale_logged = scale
+            logger.info(
+                "COMPACT FORCING: scaling count_tokens by %.2f "
+                "(client assumed window %d, rail %d) — client auto-compact "
+                "fires at ~%d real tokens, before the pruner",
+                scale,
+                PROXY_CLIENT_ASSUMED_WINDOW,
+                default_context_window if default_context_window > 0 else PROXY_CONTEXT_WINDOW,
+                int(PROXY_CLIENT_ASSUMED_WINDOW * 0.925 / scale),
+            )
+        return {"input_tokens": scaled}
+    return {"input_tokens": est}
+
+
+_count_scale_logged: float = 0.0
+
+
+def _count_tokens_scale() -> float:
+    """Resolve the count_tokens compaction-forcing scale (>= 1.0; 1.0 = off).
+
+    "auto" derives it from the LIVE rail each call so a rail resize (server
+    restart with a different --ctx-size) re-tunes the client's compact point
+    without a proxy restart.
+    """
+    raw = PROXY_COUNT_TOKENS_SCALE.strip().lower()
+    if raw in ("", "off", "none", "0", "1", "1.0"):
+        return 1.0
+    if raw != "auto":
+        try:
+            return max(1.0, float(raw))
+        except ValueError:
+            return 1.0
+    window = default_context_window if default_context_window > 0 else PROXY_CONTEXT_WINDOW
+    if window <= 0:
+        return 1.0
+    frac = (
+        PROXY_COMPACT_TARGET_FRACTION
+        if 0 < PROXY_COMPACT_TARGET_FRACTION < 1
+        # Auto: land compaction just under the pruner trigger so the client
+        # compacts before the proxy ever prunes (pruner = backstop only).
+        else min(0.9, PROXY_CONTEXT_PRUNE_THRESHOLD * 0.95)
+    )
+    target = window * frac
+    if target <= 0 or PROXY_CLIENT_ASSUMED_WINDOW <= target:
+        # Rail already exceeds the client's own compact point — honest counts
+        # are fine, the client compacts before the rail unaided.
+        return 1.0
+    return PROXY_CLIENT_ASSUMED_WINDOW / target
+
+
+def _client_input_tokens(monitor) -> int:
+    """Input size to REPORT to the client, scaled for compaction forcing.
+
+    Streaming responses historically hardcoded ``input_tokens: 0`` in
+    message_start — the client could never see its own context usage, so its
+    auto-compaction NEVER fired and sessions ground against the pruner forever
+    (the actual root cause of the interactive thrash band; count_tokens
+    scaling alone was insufficient because the dominant clients don't call
+    it). Reports the PRE-prune estimate — the client's own conversation size,
+    which is what its compaction decision is about — times the forcing scale.
+    """
+    est = (
+        getattr(monitor, "pre_prune_input_tokens", 0)
+        or getattr(monitor, "last_input_tokens", 0)
+        or 0
+    )
+    return int(est * _count_tokens_scale())
+
+
+def _scale_client_usage(anthropic_resp: dict) -> dict:
+    """Shallow-copy a finalized response with usage.input_tokens scaled for the
+    CLIENT. Internal consumers (pruner accounting, telemetry) read the original
+    dict's honest numbers — only the wire copy lies, and it lies consistently
+    with count_tokens and the streaming frames."""
+    scale = _count_tokens_scale()
+    usage = anthropic_resp.get("usage")
+    if scale <= 1.0 or not isinstance(usage, dict):
+        return anthropic_resp
+    out = dict(anthropic_resp)
+    out["usage"] = {
+        **usage,
+        "input_tokens": int(usage.get("input_tokens", 0) * scale),
+    }
+    return out
 
 
 @app.post("/v1/messages")
@@ -10037,7 +10160,7 @@ async def messages(request: Request):
 
         if PROXY_STREAM_HEARTBEAT_SECS > 0:
             return StreamingResponse(
-                _heartbeat_then_buffered(_produce_guarded(), model),
+                _heartbeat_then_buffered(_produce_guarded(), model, _client_input_tokens(monitor)),
                 media_type="text/event-stream",
                 headers={
                     "Cache-Control": "no-cache",
@@ -10466,7 +10589,9 @@ async def messages(request: Request):
         await asyncio.to_thread(
             _record_project_telemetry, body, model, anthropic_resp.get("usage", {})
         )
-        return anthropic_resp
+        # Wire copy only: telemetry above and monitor accounting keep the
+        # honest numbers; the client sees the compaction-forcing scale.
+        return _scale_client_usage(anthropic_resp)
 
 
 @app.post("/anthropic/v1/messages")