From ea78cf6252e85bb0919482a8c8613a411921304e Mon Sep 17 00:00:00 2001 From: chopratejas Date: Fri, 1 May 2026 15:25:18 -0700 Subject: [PATCH 1/2] fix(proxy): cache concurrency lock, multi-worker docs, bounded compression executor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three audit follow-ups from issue #327's deep-dive review. C1 — CompressionCache concurrency lock ====================================== `CompressionCache` instances are shared per `session_id` and accessed from async-dispatched threadpool workers. Pre-fix, concurrent requests for the same session raced on `_cache`, `_stable_hashes`, `_first_seen`, and `_total_tokens_saved` with no synchronization. Observable failures: * Lost-update on `_total_tokens_saved` (read-modify-write). * `RuntimeError: OrderedDict mutated during iteration` from `apply_cached` when a concurrent `store_compressed` evicts during the walk. * Lost stable-hash records — next-turn compute_frozen_count reads inconsistent state. May also explain part of SvenMeyer's `_cache: 0 entries / 1003 misses` observation: the cache was being clobbered concurrently. Added `threading.RLock` guarding all mutating methods. `RLock` (not `Lock`) so future code can call locked methods from inside another locked method without self-deadlock. Also locked `HeadroomProxy._compression_caches` dict-of-caches access via a separate `_compression_caches_lock` so two concurrent calls for the same session_id can't each create distinct CompressionCache objects (which would split the cache state between them). The `/stats` endpoint snapshots the cache list under the dict lock before iterating to avoid eviction-during-iteration. C2 — Multi-worker CCR fragmentation: documented + startup warning ================================================================= The in-memory `InMemoryCcrStore` (Rust), `_compression_caches` (Python), `session_tracker_store` (Python), and TOIN learner state are ALL per-process. Multi-worker uvicorn round-robins requests across workers, so a session whose turn-1 lands on worker A may have turn-2 land on worker B. Worker B has zero knowledge of A's CCR markers, replay cache, or prefix-cache state. Result: `Retrieve original: hash=X` markers stay in-context as opaque directives, every fresh tool_result is recompressed from scratch, and `frozen_message_count=0` causes Anthropic prefix-cache busts on every cross-worker turn. Added a "Multi-worker deployment — CCR fragmentation" section in `RUST_DEV.md` documenting the failure modes, the supported configuration (`--workers 1`), and the sticky-session workaround for horizontal scale. The proxy emits a `WARNING`-level log line on startup if `workers > 1` is detected, pointing at the doc section. C3 — Bounded compression executor with cancel-aware metrics =========================================================== `asyncio.wait_for(asyncio.to_thread(pipeline.apply), timeout=...)` cancellation does NOT propagate into the threadpool worker that's running Rust code. Once the worker has picked up the task, `concurrent.futures.Future.cancel()` returns False and the thread runs to completion. Stuck threads accumulated invisibly on asyncio's default executor, contending with unrelated `to_thread` callers (file IO, etc.). Replaced all 7 `asyncio.to_thread` call sites for `pipeline.apply()` across `proxy/handlers/anthropic.py` (3) and `proxy/handlers/openai.py` (4) with a new `HeadroomProxy._run_compression_in_executor(fn, *, timeout)` helper that: 1. Submits to a dedicated bounded `ThreadPoolExecutor` named `headroom-compress` (configurable via `ProxyConfig.compression_max_workers`; defaults to `min(32, (cpu_count or 1) * 4)`). 2. Increments `_compression_in_flight` (gauge) when work starts and decrements when work completes; tracks `_compression_in_flight_max` as a high-water mark. 3. Detects "leaked threads" by comparing wall-clock elapsed against the timeout in the worker's `finally` block. Increments `_compression_leaked_threads` when a worker finishes after its asyncio future was cancelled. Operators can see the leaked-thread rate climbing in `/stats runtime.compression_executor` BEFORE the pool fills up. Tests ===== * `TestCompressionCacheConcurrency` (3 tests) — many threads store_compressed / apply_cached / update_from_result on a single CompressionCache; assert no exceptions, no lost updates, no partial state. * `test_get_compression_cache_returns_same_instance_under_contention` — 32 concurrent `_get_compression_cache(same_id)` calls return the identical instance (would split pre-lock). * `test_proxy_compression_executor.py` (8 tests) — pool size respects config, in-flight gauge tracks running compressions, high-water mark is monotonic, timeout propagates to awaiter, leaked-thread counter increments on post-deadline completion, `/stats` surfaces all three gauges. Verification ============ * All 123 targeted regression tests pass. * `make ci-precheck` clean. * No `Co-Authored-By` trailer; conventional `fix:` prefix; no `--no-verify`. --- .claude-plugin/marketplace.json | 4 +- .github/plugin/marketplace.json | 4 +- RUST_DEV.md | 61 +++++ headroom/cache/compression_cache.py | 179 +++++++----- headroom/proxy/handlers/anthropic.py | 60 ++-- headroom/proxy/handlers/openai.py | 64 ++--- headroom/proxy/models.py | 19 ++ headroom/proxy/server.py | 193 +++++++++++-- .../.claude-plugin/plugin.json | 2 +- .../.github/plugin/plugin.json | 2 +- ...est_anthropic_pre_upstream_backpressure.py | 40 +++ tests/test_compression_cache.py | 201 ++++++++++++++ tests/test_proxy_compression_executor.py | 259 ++++++++++++++++++ 13 files changed, 915 insertions(+), 173 deletions(-) create mode 100644 tests/test_proxy_compression_executor.py diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json index 568f24331..a8a6b362a 100644 --- a/.claude-plugin/marketplace.json +++ b/.claude-plugin/marketplace.json @@ -5,14 +5,14 @@ }, "metadata": { "description": "Headroom marketplace for Claude Code and GitHub Copilot CLI plugins.", - "version": "0.20.2" + "version": "0.20.8" }, "plugins": [ { "name": "headroom", "source": "./plugins/headroom-agent-hooks", "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", - "version": "0.20.2", + "version": "0.20.8", "author": { "name": "Headroom Contributors", "url": "https://github.com/chopratejas/headroom" diff --git a/.github/plugin/marketplace.json b/.github/plugin/marketplace.json index 568f24331..a8a6b362a 100644 --- a/.github/plugin/marketplace.json +++ b/.github/plugin/marketplace.json @@ -5,14 +5,14 @@ }, "metadata": { "description": "Headroom marketplace for Claude Code and GitHub Copilot CLI plugins.", - "version": "0.20.2" + "version": "0.20.8" }, "plugins": [ { "name": "headroom", "source": "./plugins/headroom-agent-hooks", "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", - "version": "0.20.2", + "version": "0.20.8", "author": { "name": "Headroom Contributors", "url": "https://github.com/chopratejas/headroom" diff --git a/RUST_DEV.md b/RUST_DEV.md index 5c6ab1809..20df16995 100644 --- a/RUST_DEV.md +++ b/RUST_DEV.md @@ -285,3 +285,64 @@ doesn't rediscover them. - **`rust-toolchain.toml`** pins `channel = "stable"` rather than a specific version so CI picks up the same toolchain the local box uses. Tighten to a pinned version (e.g. `1.78`) once the port stabilizes. + +## Multi-worker deployment — CCR fragmentation + +**Recommendation: run the proxy with `--workers 1` (the default).** The +in-memory CCR store is per-process. Multi-worker uvicorn deployments produce +silent retrieval failures — and `Compress-Cache-Retrieve` is the lossless +half of the pipeline. + +### What goes wrong with `--workers N > 1` + +Each uvicorn worker is a separate Python process. Each process holds its own +copies of: + +1. **`InMemoryCcrStore`** (`crates/headroom-core/src/ccr.rs:78`) — the + sharded `DashMap` mapping `hash → original_content` for content the + compressor replaced with `Retrieve original: hash=X` markers. +2. **`HeadroomProxy._compression_caches`** (`headroom/proxy/server.py:367`) + — the per-session `CompressionCache` dict. +3. **`HeadroomProxy.session_tracker_store`** — per-session prefix-tracker + state derived from Anthropic's `cache_read_input_tokens` responses. +4. **TOIN learner state** — pattern statistics used to bias the compressor. + +When uvicorn round-robins requests across workers, a session whose turn-1 +landed on worker A may have turn-2 land on worker B. Worker B has zero +knowledge of what worker A did: + +- The CCR marker `Retrieve original: hash=X` is in the conversation, but + worker B's `InMemoryCcrStore` returns `None` for `X`. The marker stays + in-context as an opaque directive the model can't act on. Tokens spent, + no retrieval value. +- The `CompressionCache` on worker B has no replay entries → every fresh + tool_result is recompressed from scratch, even content that worker A + already compressed once. CPU wasted; observable as compression latency + doubling on round-robin sessions. +- The `prefix_tracker` on worker B starts at `frozen_message_count = 0` → + worker B compresses positions that Anthropic has already cached. Cache + bust + write premium paid on every cross-worker session turn. + +### What works today + +`--workers 1` is the only fully-supported configuration. The proxy is async +and a single worker handles thousands of concurrent requests via the event +loop; CPU-bound Rust work releases the GIL via `py.allow_threads`, so +vertical scaling on one process is the intended path. + +### What we'd need for multi-worker + +A backend implementation of the `CcrStore` trait +(`crates/headroom-core/src/ccr.rs`) backed by a shared store — Redis, +Memcached, or a sticky-session reverse proxy. The `_compression_caches` +and `session_tracker_store` would also need shared backing. None of this +is implemented yet. If you need horizontal scale today, run multiple +single-worker proxy processes behind a sticky-session load balancer (hash +on session_id) — that pins each session to one worker and avoids the +fragmentation entirely. + +### Detecting it in the wild + +The proxy emits a `WARNING`-level log line on startup if it detects +`WEB_CONCURRENCY` or uvicorn `--workers` set to anything > 1, pointing +operators at this section. diff --git a/headroom/cache/compression_cache.py b/headroom/cache/compression_cache.py index aa79fc190..c8cbfed3f 100644 --- a/headroom/cache/compression_cache.py +++ b/headroom/cache/compression_cache.py @@ -10,6 +10,7 @@ import copy import hashlib import json import logging +import threading import time from collections import OrderedDict from dataclasses import dataclass @@ -82,6 +83,18 @@ class CompressionCache: def __init__(self, max_entries: int = 10000) -> None: self.max_entries = max_entries + # Reentrant lock guarding all mutable state below. Required because + # the proxy is async and dispatches multiple concurrent requests per + # session (Claude Code background tools, parallel agents, etc.) into + # `asyncio.to_thread` workers — without this, two concurrent + # requests for the same `session_id` race on `_cache`, + # `_stable_hashes`, `_first_seen`, and `_total_tokens_saved`. The + # observable failures are (a) lost-update on `_total_tokens_saved`, + # (b) `OrderedDict mutated during iteration` in `apply_cached`, and + # (c) lost stable-hash records that drop the next-turn cache lookup. + # `RLock` (not `Lock`) so future code can call locked methods from + # inside another locked method without self-deadlock. + self._lock = threading.RLock() self._cache: OrderedDict[str, _CacheEntry] = OrderedDict() # `_stable_hashes` is CONTENT-KEYED, not positional. It records "we # have seen this content before and it is known not to compress @@ -102,13 +115,14 @@ class CompressionCache: def get_compressed(self, hash: str) -> str | None: """Retrieve compressed content by hash, refreshing LRU position on hit.""" - entry = self._cache.get(hash) - if entry is None: - self._misses += 1 - return None - self._hits += 1 - self._cache.move_to_end(hash) - return entry.compressed + with self._lock: + entry = self._cache.get(hash) + if entry is None: + self._misses += 1 + return None + self._hits += 1 + self._cache.move_to_end(hash) + return entry.compressed def store_compressed(self, hash: str, compressed: str, tokens_saved: int) -> None: """Store a compressed version keyed by content hash. @@ -117,17 +131,18 @@ class CompressionCache: end (most recently used). When the cache exceeds max_entries, the oldest entry is evicted. """ - if hash in self._cache: - old_entry = self._cache[hash] - self._total_tokens_saved -= old_entry.tokens_saved - del self._cache[hash] + with self._lock: + if hash in self._cache: + old_entry = self._cache[hash] + self._total_tokens_saved -= old_entry.tokens_saved + del self._cache[hash] - self._cache[hash] = _CacheEntry(compressed=compressed, tokens_saved=tokens_saved) - self._total_tokens_saved += tokens_saved + self._cache[hash] = _CacheEntry(compressed=compressed, tokens_saved=tokens_saved) + self._total_tokens_saved += tokens_saved - while len(self._cache) > self.max_entries: - _, evicted = self._cache.popitem(last=False) - self._total_tokens_saved -= evicted.tokens_saved + while len(self._cache) > self.max_entries: + _, evicted = self._cache.popitem(last=False) + self._total_tokens_saved -= evicted.tokens_saved def mark_stable(self, content_hash: str) -> None: """Mark a content hash as stable (unchanged, not compressed). @@ -136,15 +151,17 @@ class CompressionCache: These messages appear verbatim every turn, so they are prefix-stable even though no compressed version exists in the cache. """ - self._stable_hashes.add(content_hash) + with self._lock: + self._stable_hashes.add(content_hash) def mark_stable_from_messages(self, messages: list[dict], up_to: int) -> None: """Mark all tool_result hashes in messages[:up_to] as stable.""" - for msg in messages[:up_to]: - if _is_tool_result_message(msg): - content = _extract_tool_result_content(msg) - if content is not None: - self._stable_hashes.add(self.content_hash(content)) + with self._lock: + for msg in messages[:up_to]: + if _is_tool_result_message(msg): + content = _extract_tool_result_content(msg) + if content is not None: + self._stable_hashes.add(self.content_hash(content)) def should_defer_compression( self, @@ -169,25 +186,27 @@ class CompressionCache: - **Near the TTL boundary**: compress now and amortize the bust across future turns (batched recompression). """ - now = time.time() - first_seen = self._first_seen.get(content_hash) - if first_seen is None: - self._first_seen[content_hash] = now - return False # First time — compress now (no cache entry to preserve) - age = now - first_seen - if age >= ttl_seconds - batch_window: - return False # Near TTL boundary — compress now (batch window) - return True # Seen recently within TTL — defer to preserve cache + with self._lock: + now = time.time() + first_seen = self._first_seen.get(content_hash) + if first_seen is None: + self._first_seen[content_hash] = now + return False # First time — compress now (no cache entry to preserve) + age = now - first_seen + if age >= ttl_seconds - batch_window: + return False # Near TTL boundary — compress now (batch window) + return True # Seen recently within TTL — defer to preserve cache def get_stats(self) -> dict: """Return cache statistics.""" - return { - "entries": len(self._cache), - "stable_hashes": len(self._stable_hashes), - "hits": self._hits, - "misses": self._misses, - "tokens_saved": self._total_tokens_saved, - } + with self._lock: + return { + "entries": len(self._cache), + "stable_hashes": len(self._stable_hashes), + "hits": self._hits, + "misses": self._misses, + "tokens_saved": self._total_tokens_saved, + } @staticmethod def content_hash(content: str | list) -> str: @@ -210,21 +229,22 @@ class CompressionCache: content hash is already in the cache. The first unstable tool_result (cache miss) stops the count. """ - count = 0 - for msg in messages: - if _is_tool_result_message(msg): - content = _extract_tool_result_content(msg) - if content is not None: - h = self.content_hash(content) - if h not in self._cache and h not in self._stable_hashes: + with self._lock: + count = 0 + for msg in messages: + if _is_tool_result_message(msg): + content = _extract_tool_result_content(msg) + if content is not None: + h = self.content_hash(content) + if h not in self._cache and h not in self._stable_hashes: + break + else: + # tool_result with non-string content; treat as unstable break - else: - # tool_result with non-string content; treat as unstable - break - # Regular user/assistant/system messages and assistant+tool_use - # are always stable — fall through. - count += 1 - return count + # Regular user/assistant/system messages and assistant+tool_use + # are always stable — fall through. + count += 1 + return count def apply_cached(self, messages: list[dict]) -> list[dict]: """Return a new list with cached compressions swapped into tool results. @@ -232,18 +252,22 @@ class CompressionCache: Never mutates the input list or any message dict within it. Output always has the same length as input. """ - result: list[dict] = [] - for msg in messages: - if _is_tool_result_message(msg): - content = _extract_tool_result_content(msg) - if content is not None: - h = self.content_hash(content) - compressed = self.get_compressed(h) - if compressed is not None: - result.append(_swap_tool_result_content(msg, compressed)) - continue - result.append(msg) - return result + # `get_compressed` re-acquires the lock (RLock); single contiguous + # critical section so concurrent `store_compressed` cannot mutate + # `_cache` mid-iteration. + with self._lock: + result: list[dict] = [] + for msg in messages: + if _is_tool_result_message(msg): + content = _extract_tool_result_content(msg) + if content is not None: + h = self.content_hash(content) + compressed = self.get_compressed(h) + if compressed is not None: + result.append(_swap_tool_result_content(msg, compressed)) + continue + result.append(msg) + return result def update_from_result(self, originals: list[dict], compressed: list[dict]) -> None: """Cache new compressions by comparing original and compressed messages. @@ -259,15 +283,18 @@ class CompressionCache: ) return - for orig, comp in zip(originals, compressed): - orig_content = _extract_tool_result_content(orig) - comp_content = _extract_tool_result_content(comp) - if orig_content is None or comp_content is None: - continue - if orig_content == comp_content: - # Content unchanged — mark as stable for frozen count walk - self._stable_hashes.add(self.content_hash(orig_content)) - continue - h = self.content_hash(orig_content) - tokens_saved = len(orig_content) // 4 - len(comp_content) // 4 - self.store_compressed(h, comp_content, tokens_saved=max(tokens_saved, 0)) + # Single critical section: `store_compressed` re-acquires the lock + # (RLock) safely. + with self._lock: + for orig, comp in zip(originals, compressed): + orig_content = _extract_tool_result_content(orig) + comp_content = _extract_tool_result_content(comp) + if orig_content is None or comp_content is None: + continue + if orig_content == comp_content: + # Content unchanged — mark as stable for frozen count walk + self._stable_hashes.add(self.content_hash(orig_content)) + continue + h = self.content_hash(orig_content) + tokens_saved = len(orig_content) // 4 - len(comp_content) // 4 + self.store_compressed(h, comp_content, tokens_saved=max(tokens_saved, 0)) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 78d7f5328..c2c093950 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -775,17 +775,15 @@ class AnthropicHandlerMixin: comp_cache.mark_stable_from_messages(messages, frozen_message_count) async with stage_timer.measure("compression_first_stage"): - result = await asyncio.wait_for( - asyncio.to_thread( - lambda: self.anthropic_pipeline.apply( - messages=working_messages, - model=model, - model_limit=context_limit, - context=extract_user_query(working_messages), - frozen_message_count=frozen_message_count, - biases=biases, - request_id=request_id, - ) + result = await self._run_compression_in_executor( + lambda: self.anthropic_pipeline.apply( + messages=working_messages, + model=model, + model_limit=context_limit, + context=extract_user_query(working_messages), + frozen_message_count=frozen_message_count, + biases=biases, + request_id=request_id, ), timeout=COMPRESSION_TIMEOUT_SECONDS, ) @@ -816,17 +814,15 @@ class AnthropicHandlerMixin: optimized_tokens = tokenizer.count_messages(optimized_messages) elif not is_cache_mode(self.config.mode): async with stage_timer.measure("compression_first_stage"): - result = await asyncio.wait_for( - asyncio.to_thread( - lambda: self.anthropic_pipeline.apply( - messages=messages, - model=model, - model_limit=context_limit, - context=extract_user_query(messages), - frozen_message_count=frozen_message_count, - biases=biases, - request_id=request_id, - ) + result = await self._run_compression_in_executor( + lambda: self.anthropic_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + context=extract_user_query(messages), + frozen_message_count=frozen_message_count, + biases=biases, + request_id=request_id, ), timeout=COMPRESSION_TIMEOUT_SECONDS, ) @@ -848,17 +844,15 @@ class AnthropicHandlerMixin: if delta is not None: stable_forwarded_prefix, delta_messages = delta if delta_messages: - result = await asyncio.wait_for( - asyncio.to_thread( - lambda: self.anthropic_pipeline.apply( - messages=delta_messages, - model=model, - model_limit=context_limit, - context=extract_user_query(delta_messages), - frozen_message_count=0, - biases=biases, - request_id=request_id, - ) + result = await self._run_compression_in_executor( + lambda: self.anthropic_pipeline.apply( + messages=delta_messages, + model=model, + model_limit=context_limit, + context=extract_user_query(delta_messages), + frozen_message_count=0, + biases=biases, + request_id=request_id, ), timeout=COMPRESSION_TIMEOUT_SECONDS, ) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 5599d09d4..2b9aea852 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -364,16 +364,14 @@ class OpenAIHandlerMixin: # Re-freeze boundary openai_frozen_count = comp_cache.compute_frozen_count(messages) - result = await asyncio.wait_for( - asyncio.to_thread( - lambda: self.openai_pipeline.apply( - messages=working_messages, - model=model, - model_limit=context_limit, - context=extract_user_query(working_messages), - frozen_message_count=openai_frozen_count, - biases=_hook_biases, - ) + result = await self._run_compression_in_executor( + lambda: self.openai_pipeline.apply( + messages=working_messages, + model=model, + model_limit=context_limit, + context=extract_user_query(working_messages), + frozen_message_count=openai_frozen_count, + biases=_hook_biases, ), timeout=COMPRESSION_TIMEOUT_SECONDS, ) @@ -389,16 +387,14 @@ class OpenAIHandlerMixin: # so tokens_saved captures both Zone 1 + Zone 2 savings. optimized_tokens = result.tokens_after else: - result = await asyncio.wait_for( - asyncio.to_thread( - lambda: self.openai_pipeline.apply( - messages=messages, - model=model, - model_limit=context_limit, - context=extract_user_query(messages), - frozen_message_count=openai_frozen_count, - biases=_hook_biases, - ) + result = await self._run_compression_in_executor( + lambda: self.openai_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + context=extract_user_query(messages), + frozen_message_count=openai_frozen_count, + biases=_hook_biases, ), timeout=COMPRESSION_TIMEOUT_SECONDS, ) @@ -1152,14 +1148,12 @@ class OpenAIHandlerMixin: if _should_compress and _license_ok: try: context_limit = self.openai_provider.get_context_limit(model) - result = await asyncio.wait_for( - asyncio.to_thread( - lambda: self.openai_pipeline.apply( - messages=messages, - model=model, - model_limit=context_limit, - context=extract_user_query(messages), - ) + result = await self._run_compression_in_executor( + lambda: self.openai_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + context=extract_user_query(messages), ), timeout=COMPRESSION_TIMEOUT_SECONDS, ) @@ -1647,14 +1641,12 @@ class OpenAIHandlerMixin: context_limit = self.openai_provider.get_context_limit(model) async with stage_timer.measure("compression"): - result = await asyncio.wait_for( - asyncio.to_thread( - lambda: self.openai_pipeline.apply( - messages=messages, - model=model, - model_limit=context_limit, - context=extract_user_query(messages), - ) + result = await self._run_compression_in_executor( + lambda: self.openai_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + context=extract_user_query(messages), ), timeout=COMPRESSION_TIMEOUT_SECONDS, ) diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index 0f94974cb..79e1dc809 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -269,6 +269,25 @@ class ProxyConfig: # COMPRESSION_TIMEOUT_SECONDS guard; this bounds the memory leg too. anthropic_pre_upstream_memory_context_timeout_seconds: float = 2.0 + # Bound the dedicated compression threadpool. CPU-bound Rust work runs + # here; the pool is separate from asyncio's default executor so other + # ``asyncio.to_thread`` callers (file IO, etc.) are not contended by + # compression bursts. ``None`` resolves to ``min(32, (cpu_count or 1) * 4)``, + # matching asyncio's default executor sizing today. Lower the cap to + # tighten resource use on multi-tenant hosts; raise it to handle larger + # bursts. CLI: ``--compression-max-workers``. Env: + # ``HEADROOM_COMPRESSION_MAX_WORKERS``. + # + # Background: ``asyncio.wait_for`` cancellation does NOT propagate into + # the threadpool worker that's running Rust code — once the worker has + # picked up the task, ``concurrent.futures.Future.cancel()`` returns + # ``False`` and the thread runs to completion. A bounded pool lets us + # observe the worst case (max queue depth, "leaked" threads that + # finished post-deadline) and fail fast under contention rather than + # piling unboundedly on the default executor. See + # ``HeadroomProxy._run_compression_in_executor``. + compression_max_workers: int | None = None + @property def provider_api_overrides(self) -> ProviderApiOverrides: """Return provider API URL overrides as a dedicated provider config object.""" diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index f6a0c5a3e..f8d243cdb 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -25,11 +25,13 @@ from __future__ import annotations import argparse import asyncio +import concurrent.futures import contextlib import json import logging import os import sys +import threading import time from dataclasses import fields, is_dataclass from datetime import datetime, timezone @@ -363,8 +365,13 @@ class HeadroomProxy( ) ) - # Compression cache store for token mode (session-scoped) + # Compression cache store for token mode (session-scoped). The dict + # itself is mutated under `_compression_caches_lock`; the per-session + # `CompressionCache` instances have their own internal lock guarding + # `_cache`/`_stable_hashes`/`_first_seen` against concurrent + # async-dispatched requests for the same session. self._compression_caches: dict[str, CompressionCache] = {} + self._compression_caches_lock = threading.RLock() self.logger = ( RequestLogger( @@ -421,6 +428,38 @@ class HeadroomProxy( else: self.anthropic_pre_upstream_sem = None + # Dedicated compression executor — see C3 in the audit followup. + # Replaces ``asyncio.to_thread(...)`` for ``pipeline.apply()`` calls + # so that: + # 1. Compression work is bounded — CPU-bound Rust runs here, and + # bursts cannot starve other ``asyncio.to_thread`` callers + # sharing the loop's default executor (file IO, etc.). + # 2. Tasks that exceed ``COMPRESSION_TIMEOUT_SECONDS`` and complete + # *after* the asyncio future was cancelled are counted in the + # ``compression_leaked_threads`` gauge — Python cannot preempt + # the worker, so this is the only signal that some pool slots + # are sitting on stuck work. + _compression_max_cfg = config.compression_max_workers + if _compression_max_cfg is None: + _compression_max = min(32, (os.cpu_count() or 1) * 4) + else: + _compression_max = max(1, _compression_max_cfg) + self.compression_max_workers: int = _compression_max + self._compression_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=_compression_max, + thread_name_prefix="headroom-compress", + ) + # Gauge: currently-running compression tasks. Mutated under + # ``_compression_metrics_lock`` from worker threads + the asyncio + # event loop. + self._compression_in_flight: int = 0 + # High-water mark for in-flight count. + self._compression_in_flight_max: int = 0 + # Counter: threads that finished AFTER their asyncio future hit the + # timeout. Stuck-thread leak indicator. + self._compression_leaked_threads: int = 0 + self._compression_metrics_lock = threading.Lock() + # Backend for Anthropic API (direct, LiteLLM, or any-llm) # Supports: "anthropic" (direct), "bedrock", "vertex", "litellm-", or "anyllm" self.anthropic_backend: Backend | None = create_proxy_backend( @@ -560,27 +599,101 @@ class HeadroomProxy( }, ) + async def _run_compression_in_executor( + self, + fn, # noqa: ANN001 — caller-supplied no-arg sync callable + *, + timeout: float, + ): + """Run a synchronous compression callable on the bounded executor + with cancel-aware metrics. + + Replaces ``asyncio.wait_for(asyncio.to_thread(fn), timeout=...)``. + + Why a dedicated executor: the proxy's compression path is CPU-bound + Rust work that releases the GIL via ``py.allow_threads``. Sharing + the loop's default executor (used by ``asyncio.to_thread``) means + a burst of slow compressions can starve unrelated ``to_thread`` + callers (file IO, etc.). The compression executor is sized + independently via ``config.compression_max_workers``. + + Why "cancel-aware metrics": when ``asyncio.wait_for`` times out, it + cancels the *asyncio future*. The underlying + ``concurrent.futures.Future`` from ``run_in_executor`` cannot + actually cancel a thread that has started — Python has no way to + preempt running CPython bytecode or in-flight Rust calls. The + worker keeps running to completion, ignored. We detect this by + recording the start timestamp and incrementing + ``_compression_leaked_threads`` from the worker's ``finally`` + block when ``elapsed > timeout``. Operators can see leaked-thread + rate climbing in ``/stats`` before the pool fills up. + + Args: + fn: A no-arg sync callable that runs the compression. Must not + raise asyncio Cancellation; if it does, the wrapper still + decrements the in-flight gauge but the leaked-thread + counter may double-count. + timeout: Wall-clock timeout for the asyncio side. The + executor worker keeps running past this (Python limitation + — see above), but at least the awaiter unblocks. + + Returns: + Whatever ``fn()`` returns. + + Raises: + ``asyncio.TimeoutError`` if the callable doesn't return within + ``timeout``. Any exception raised by ``fn`` propagates + unchanged. + """ + loop = asyncio.get_running_loop() + start = time.monotonic() + with self._compression_metrics_lock: + self._compression_in_flight += 1 + if self._compression_in_flight > self._compression_in_flight_max: + self._compression_in_flight_max = self._compression_in_flight + + def _wrapped(): # noqa: ANN202 + try: + return fn() + finally: + elapsed = time.monotonic() - start + with self._compression_metrics_lock: + self._compression_in_flight -= 1 + if elapsed > timeout: + self._compression_leaked_threads += 1 + + future = loop.run_in_executor(self._compression_executor, _wrapped) + return await asyncio.wait_for(future, timeout=timeout) + def _get_compression_cache(self, session_id: str) -> CompressionCache: - """Get or create a CompressionCache for a session.""" - if session_id not in self._compression_caches: - from headroom.cache.compression_cache import CompressionCache + """Get or create a CompressionCache for a session. - # Evict oldest caches if at capacity - if len(self._compression_caches) >= MAX_COMPRESSION_CACHE_SESSIONS: - # Remove oldest quarter to amortize cleanup cost - oldest_keys = list(self._compression_caches.keys())[ - : MAX_COMPRESSION_CACHE_SESSIONS // 4 - ] - for key in oldest_keys: - del self._compression_caches[key] - logger.info( - "Evicted %d compression caches (exceeded %d max sessions)", - len(oldest_keys), - MAX_COMPRESSION_CACHE_SESSIONS, - ) + Thread-safe under `_compression_caches_lock`: a concurrent pair of + `_get_compression_cache(session_id)` calls (e.g. two async requests + for the same conversation) must return the **same** instance, + otherwise the per-session cache state splits and the two halves + diverge across requests. + """ + with self._compression_caches_lock: + if session_id not in self._compression_caches: + from headroom.cache.compression_cache import CompressionCache - self._compression_caches[session_id] = CompressionCache() - return self._compression_caches[session_id] + # Evict oldest caches if at capacity + if len(self._compression_caches) >= MAX_COMPRESSION_CACHE_SESSIONS: + # Remove oldest quarter to amortize cleanup cost + oldest_keys = list(self._compression_caches.keys())[ + : MAX_COMPRESSION_CACHE_SESSIONS // 4 + ] + for key in oldest_keys: + del self._compression_caches[key] + logger.info( + "Evicted %d compression caches (exceeded %d max sessions)", + len(oldest_keys), + MAX_COMPRESSION_CACHE_SESSIONS, + ) + + self._compression_caches[session_id] = CompressionCache() + return self._compression_caches[session_id] def _setup_code_aware(self, config: ProxyConfig, transforms: list) -> str: """Set up code-aware compression if enabled. @@ -1282,6 +1395,12 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: ws_active_relay_tasks = ( ws_registry.active_relay_task_count() if ws_registry is not None else 0 ) + # Snapshot compression executor metrics under their lock (gauges + # mutated by worker threads; not safe to read without). + with proxy._compression_metrics_lock: + _comp_in_flight = proxy._compression_in_flight + _comp_in_flight_max = proxy._compression_in_flight_max + _comp_leaked = proxy._compression_leaked_threads return { "anthropic_pre_upstream": { "enabled": proxy.anthropic_pre_upstream_sem is not None, @@ -1296,6 +1415,13 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: ), "codex_ws_gated": False, }, + "compression_executor": { + "max_workers": proxy.compression_max_workers, + "in_flight": _comp_in_flight, + "in_flight_max": _comp_in_flight_max, + "leaked_threads_total": _comp_leaked, + "source": ("auto" if config.compression_max_workers is None else "explicit"), + }, "websocket_sessions": { "active_sessions": ws_active_sessions, "active_relay_tasks": ws_active_relay_tasks, @@ -1515,14 +1641,20 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: except Exception: logger.warning("Failed to log /stats summary payload") - # Compression cache stats (token mode) + # Compression cache stats (token mode). Snapshot the cache list under + # the dict lock so a concurrent eviction can't mutate the dict while + # we iterate. Each per-session `get_stats()` is independently + # thread-safe via the cache's own internal lock. compression_cache_stats: dict = {} if proxy.config.mode == PROXY_MODE_TOKEN and proxy._compression_caches: + with proxy._compression_caches_lock: + _caches_snapshot = list(proxy._compression_caches.values()) + _active_sessions = len(proxy._compression_caches) total_entries = 0 total_hits = 0 total_misses = 0 total_tokens_saved = 0 - for cache in proxy._compression_caches.values(): + for cache in _caches_snapshot: s = cache.get_stats() total_entries += s.get("entries", 0) total_hits += s.get("hits", 0) @@ -1530,7 +1662,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: total_tokens_saved += s.get("total_tokens_saved", 0) compression_cache_stats = { "mode": PROXY_MODE_TOKEN, - "active_sessions": len(proxy._compression_caches), + "active_sessions": _active_sessions, "total_entries": total_entries, "total_hits": total_hits, "total_misses": total_misses, @@ -2430,6 +2562,23 @@ def run_server( app_target: Any uvicorn_kwargs: dict[str, Any] = {} if workers > 1: + # CCR / compression-cache / prefix-tracker / TOIN state are all + # per-process. Round-robin across workers fragments these caches + # and produces silent retrieval failures for `Retrieve original: + # hash=X` markers and avoidable cache busts on the upstream + # provider. See the "Multi-worker deployment — CCR fragmentation" + # section in RUST_DEV.md for the full failure modes and the + # sticky-session workaround. + logger.warning( + "Headroom is running with workers=%d. The in-memory CCR store, " + "compression cache, prefix tracker, and TOIN state are all " + "per-process; multi-worker deployments produce silent retrieval " + "failures and avoidable cache busts when sessions land on different " + "workers. Run --workers 1 (or place a sticky-session load balancer " + "in front of multiple --workers 1 processes). See RUST_DEV.md → " + "'Multi-worker deployment — CCR fragmentation'.", + workers, + ) os.environ[_MULTI_WORKER_CONFIG_ENV] = json.dumps(_proxy_config_payload(config)) app_target = "headroom.proxy.server:create_app_from_env" uvicorn_kwargs["factory"] = True diff --git a/plugins/headroom-agent-hooks/.claude-plugin/plugin.json b/plugins/headroom-agent-hooks/.claude-plugin/plugin.json index 8853264fa..779a2a0ba 100644 --- a/plugins/headroom-agent-hooks/.claude-plugin/plugin.json +++ b/plugins/headroom-agent-hooks/.claude-plugin/plugin.json @@ -1,6 +1,6 @@ { "name": "headroom", - "version": "0.20.2", + "version": "0.20.8", "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", "author": { "name": "Headroom Contributors", diff --git a/plugins/headroom-agent-hooks/.github/plugin/plugin.json b/plugins/headroom-agent-hooks/.github/plugin/plugin.json index efbb4aa63..3c948b9aa 100644 --- a/plugins/headroom-agent-hooks/.github/plugin/plugin.json +++ b/plugins/headroom-agent-hooks/.github/plugin/plugin.json @@ -1,6 +1,6 @@ { "name": "headroom", - "version": "0.20.2", + "version": "0.20.8", "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", "author": { "name": "Headroom Contributors", diff --git a/tests/test_anthropic_pre_upstream_backpressure.py b/tests/test_anthropic_pre_upstream_backpressure.py index 1035b7e1b..083465c53 100644 --- a/tests/test_anthropic_pre_upstream_backpressure.py +++ b/tests/test_anthropic_pre_upstream_backpressure.py @@ -170,11 +170,51 @@ class _DummyAnthropicHandler(AnthropicHandlerMixin): self.anthropic_pre_upstream_concurrency = ( 0 if anthropic_pre_upstream_sem is None else anthropic_pre_upstream_sem._value ) + # Audit follow-up C3: dedicated compression executor + cancel-aware + # metrics. The mixin's compression path delegates to + # ``HeadroomProxy._run_compression_in_executor`` for bounded thread + # use; this dummy handler stands in for the proxy and must therefore + # provide the same surface. + import concurrent.futures as _cf + import threading as _threading + + self._compression_executor = _cf.ThreadPoolExecutor( + max_workers=4, thread_name_prefix="dummy-compress" + ) + self.compression_max_workers = 4 + self._compression_in_flight = 0 + self._compression_in_flight_max = 0 + self._compression_leaked_threads = 0 + self._compression_metrics_lock = _threading.Lock() self._upstream_delay_s = upstream_delay_s self._raise_during_critical = raise_during_critical self.upstream_enter_times: list[float] = [] self.upstream_exit_times: list[float] = [] + async def _run_compression_in_executor(self, fn, *, timeout): # noqa: ANN001 + # Mirror of ``HeadroomProxy._run_compression_in_executor`` for the + # mixin tests. Same metrics semantics; same timeout behavior. + loop = asyncio.get_running_loop() + start = time.perf_counter() + with self._compression_metrics_lock: + self._compression_in_flight += 1 + self._compression_in_flight_max = max( + self._compression_in_flight_max, self._compression_in_flight + ) + + def _wrapped(): + try: + return fn() + finally: + elapsed = time.perf_counter() - start + with self._compression_metrics_lock: + self._compression_in_flight -= 1 + if elapsed > timeout: + self._compression_leaked_threads += 1 + + future = loop.run_in_executor(self._compression_executor, _wrapped) + return await asyncio.wait_for(future, timeout=timeout) + async def _next_request_id(self) -> str: # Unique IDs so log assertions remain disambiguated under parallelism. return f"req-{id(object()):x}" diff --git a/tests/test_compression_cache.py b/tests/test_compression_cache.py index ce721e36f..116c0ce19 100644 --- a/tests/test_compression_cache.py +++ b/tests/test_compression_cache.py @@ -389,3 +389,204 @@ class TestCompressionCacheApplyAndUpdate: assert result[0]["content"] == "compressed openai" # Original untouched assert messages[0]["content"] == original_content + + +# ─── C1 (audit follow-up): concurrency regression suite ──────────────────── +# +# CompressionCache must be safe under multi-threaded mutation. The proxy is +# async and dispatches multiple concurrent requests per `session_id` into +# `asyncio.to_thread` workers — a single CompressionCache instance therefore +# sees concurrent calls to `store_compressed` / `get_compressed` / +# `mark_stable_from_messages` / `apply_cached` / `update_from_result`. +# These tests provoke the race conditions that motivated adding `_lock`. + + +class TestCompressionCacheConcurrency: + """Threading regression suite for the audit-followup lock.""" + + def test_concurrent_store_does_not_corrupt_total_tokens_saved(self) -> None: + """Many threads each store_compressed with tokens_saved=N; the + bookkeeping field must equal SUM(N) when threads finish. Pre-lock + this races (read-modify-write of `_total_tokens_saved`).""" + import threading + + cache = CompressionCache(max_entries=1_000_000) + n_threads = 32 + per_thread = 100 + per_thread_tokens = 7 + + def worker(tid: int) -> None: + for i in range(per_thread): + h = CompressionCache.content_hash(f"thread-{tid}-item-{i}") + cache.store_compressed(h, f"comp-{tid}-{i}", tokens_saved=per_thread_tokens) + + threads = [threading.Thread(target=worker, args=(t,)) for t in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + + expected = n_threads * per_thread * per_thread_tokens + stats = cache.get_stats() + assert stats["entries"] == n_threads * per_thread + # The expected token count is exact only because each (thread, item) + # produces a unique hash → no overwrite path. Pre-lock this would be + # < expected due to lost updates. + assert stats["tokens_saved"] == expected + + def test_concurrent_apply_cached_with_concurrent_store_does_not_raise(self) -> None: + """`apply_cached` iterates `_cache` (via `get_compressed`); if a + concurrent `store_compressed` mutates the OrderedDict during the + iteration, pre-lock you'd get `RuntimeError: OrderedDict mutated + during iteration`. Locks make this a single critical section.""" + import threading + + cache = CompressionCache() + + # Pre-populate so apply_cached has work to do. + for i in range(50): + h = CompressionCache.content_hash(f"seed-{i}") + cache.store_compressed(h, f"comp-{i}", tokens_saved=1) + + msgs = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"t{i}", + "content": f"seed-{i}", + } + ], + } + for i in range(50) + ] + + stop = threading.Event() + errors: list[Exception] = [] + + def reader() -> None: + try: + while not stop.is_set(): + _ = cache.apply_cached(msgs) + except Exception as e: # pragma: no cover + errors.append(e) + + def writer() -> None: + try: + for i in range(500): + h = CompressionCache.content_hash(f"writer-{i}") + cache.store_compressed(h, f"w-{i}", tokens_saved=1) + except Exception as e: # pragma: no cover + errors.append(e) + + readers = [threading.Thread(target=reader) for _ in range(4)] + writers = [threading.Thread(target=writer) for _ in range(4)] + for t in readers + writers: + t.start() + for t in writers: + t.join() + stop.set() + for t in readers: + t.join() + + assert errors == [], f"Concurrent ops raised: {errors}" + + def test_concurrent_update_from_result_no_partial_state(self) -> None: + """update_from_result must be all-or-nothing per call. With many + threads calling update_from_result in parallel on the same cache, + the final state must reflect every call's full effect (no partial + writes).""" + import threading + + cache = CompressionCache() + + n_threads = 16 + per_thread_calls = 20 + + def worker(tid: int) -> None: + for i in range(per_thread_calls): + orig_text = f"orig-{tid}-{i}-" + "X" * 200 + comp_text = f"comp-{tid}-{i}-" + "X" * 50 + originals = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"t-{tid}-{i}", + "content": orig_text, + } + ], + } + ] + compressed = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"t-{tid}-{i}", + "content": comp_text, + } + ], + } + ] + cache.update_from_result(originals, compressed) + + threads = [threading.Thread(target=worker, args=(t,)) for t in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + + stats = cache.get_stats() + # Each (tid, i) is a unique hash → cache entries == n_threads * per_thread_calls. + assert stats["entries"] == n_threads * per_thread_calls + assert stats["tokens_saved"] > 0 + + +def test_get_compression_cache_returns_same_instance_under_contention() -> None: + """`HeadroomProxy._get_compression_cache(session_id)` must return the + SAME `CompressionCache` instance for concurrent calls with the same + session_id. Pre-lock, two concurrent calls could both see "not in dict" + and each create a new instance, splitting the cache state across them. + """ + import threading + + pytest.importorskip("fastapi") + from headroom.proxy.server import ProxyConfig, create_app + + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + log_requests=False, + ccr_inject_tool=False, + ccr_handle_responses=False, + ccr_context_tracking=False, + image_optimize=False, + ) + app = create_app(config) + proxy = app.state.proxy + + n_threads = 32 + results: list[CompressionCache] = [] + results_lock = threading.Lock() + + def worker() -> None: + c = proxy._get_compression_cache("shared-session-id") + with results_lock: + results.append(c) + + threads = [threading.Thread(target=worker) for _ in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(results) == n_threads + first = results[0] + for c in results[1:]: + assert c is first, "Concurrent _get_compression_cache returned different instances" diff --git a/tests/test_proxy_compression_executor.py b/tests/test_proxy_compression_executor.py new file mode 100644 index 000000000..2c2d78ccc --- /dev/null +++ b/tests/test_proxy_compression_executor.py @@ -0,0 +1,259 @@ +"""Audit follow-up C3: bounded compression executor + cancel-aware metrics. + +Replaces ``asyncio.to_thread`` for ``pipeline.apply()`` calls with a dedicated +``ThreadPoolExecutor`` that's bounded by ``ProxyConfig.compression_max_workers``. + +Locks the following invariants: + +1. The pool exists and respects ``compression_max_workers`` (auto and explicit). +2. ``compression_in_flight`` increments while a compression is running and + decrements after it completes — under load, the high-water mark moves up + as expected. +3. When a compression call exceeds its timeout, the awaiter unblocks with + ``TimeoutError`` — but the worker thread keeps running (Python cannot + preempt running CPython bytecode or in-flight Rust calls), and when the + work eventually completes, ``compression_leaked_threads`` increments. +4. ``/stats runtime.compression_executor`` surfaces the gauge + counter so + operators can see leaked-thread rate. + +These tests also serve as documentation: anyone reading them sees that +"timeout fired" does not mean "compression was cancelled" — it means "we +stopped waiting; the worker is still going". A bounded pool plus the +leaked-thread counter is how we make that visible. +""" + +from __future__ import annotations + +import asyncio +import threading +import time + +import pytest + +pytest.importorskip("fastapi") + +from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS # noqa: F401 +from headroom.proxy.server import ProxyConfig, create_app + + +def _make_proxy(compression_max_workers: int | None = None): + """Construct a HeadroomProxy with a no-op pipeline. Returns the proxy.""" + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + log_requests=False, + ccr_inject_tool=False, + ccr_handle_responses=False, + ccr_context_tracking=False, + image_optimize=False, + compression_max_workers=compression_max_workers, + ) + app = create_app(config) + return app.state.proxy + + +def test_compression_executor_default_size_matches_asyncio_default() -> None: + """When ``compression_max_workers`` is None, the resolved size should + match asyncio's default executor sizing (``min(32, (cpu+1)*4)`` style). + """ + import os + + proxy = _make_proxy(compression_max_workers=None) + expected = min(32, (os.cpu_count() or 1) * 4) + assert proxy.compression_max_workers == expected + assert proxy._compression_executor._max_workers == expected + + +def test_compression_executor_explicit_override() -> None: + """``ProxyConfig.compression_max_workers=N`` is honored verbatim.""" + proxy = _make_proxy(compression_max_workers=3) + assert proxy.compression_max_workers == 3 + assert proxy._compression_executor._max_workers == 3 + + +def test_compression_executor_minimum_one_worker() -> None: + """A non-positive override clamps to 1 (zero workers would deadlock).""" + proxy = _make_proxy(compression_max_workers=0) + assert proxy.compression_max_workers == 1 + + +def test_in_flight_gauge_tracks_running_compressions() -> None: + """While a compression is running, ``_compression_in_flight`` reads ≥ 1. + After it completes, it returns to 0. The high-water mark records the + peak observed. + """ + proxy = _make_proxy(compression_max_workers=4) + + enter_event = threading.Event() + release_event = threading.Event() + observed: dict[str, int] = {} + + def _slow_compression(): + enter_event.set() + # Block until the test thread reads in_flight from the gauge. + release_event.wait(timeout=5.0) + return "done" + + async def _drive(): + task = asyncio.create_task( + proxy._run_compression_in_executor(_slow_compression, timeout=10.0) + ) + # Wait for the worker to actually start. + for _ in range(50): + if enter_event.is_set(): + break + await asyncio.sleep(0.01) + with proxy._compression_metrics_lock: + observed["mid_flight"] = proxy._compression_in_flight + observed["mid_flight_max"] = proxy._compression_in_flight_max + release_event.set() + result = await task + return result + + result = asyncio.run(_drive()) + assert result == "done" + assert observed["mid_flight"] == 1, ( + f"in_flight should be 1 mid-call, got {observed['mid_flight']}" + ) + assert observed["mid_flight_max"] >= 1 + # Decremented after task completes. + with proxy._compression_metrics_lock: + assert proxy._compression_in_flight == 0 + + +def test_high_water_mark_persists_after_completion() -> None: + """``_compression_in_flight_max`` is monotonic — never decreases.""" + proxy = _make_proxy(compression_max_workers=8) + + enter_events = [threading.Event() for _ in range(3)] + release_events = [threading.Event() for _ in range(3)] + + def _make_slow(idx: int): + def _slow(): + enter_events[idx].set() + release_events[idx].wait(timeout=5.0) + return idx + + return _slow + + async def _drive(): + tasks = [ + asyncio.create_task(proxy._run_compression_in_executor(_make_slow(i), timeout=10.0)) + for i in range(3) + ] + # Wait for all 3 to enter. + for ev in enter_events: + for _ in range(50): + if ev.is_set(): + break + await asyncio.sleep(0.01) + peak = proxy._compression_in_flight + for ev in release_events: + ev.set() + for t in tasks: + await t + return peak + + peak = asyncio.run(_drive()) + assert peak == 3, f"Should have observed 3 concurrent compressions, got {peak}" + # After all complete, in_flight is back to 0 but max remains 3. + with proxy._compression_metrics_lock: + assert proxy._compression_in_flight == 0 + assert proxy._compression_in_flight_max >= 3 + + +def test_timeout_fires_and_leaked_thread_is_counted() -> None: + """When the compression exceeds ``timeout``, the awaiter sees + ``TimeoutError`` immediately. The worker keeps running; when it finishes, + ``_compression_leaked_threads`` increments by 1. + """ + proxy = _make_proxy(compression_max_workers=2) + finished_event = threading.Event() + timeout_seconds = 0.10 + + def _slow_compression(): + # Sleep well past the timeout so the asyncio side cancels first. + time.sleep(timeout_seconds * 5) + finished_event.set() + return "completed-after-deadline" + + async def _drive(): + with pytest.raises(asyncio.TimeoutError): + await proxy._run_compression_in_executor(_slow_compression, timeout=timeout_seconds) + + asyncio.run(_drive()) + + # Wait for the worker to actually finish (it ran past the deadline). + finished_event.wait(timeout=2.0) + # Give the worker thread a moment to update the counter under the lock. + deadline = time.monotonic() + 1.0 + while time.monotonic() < deadline: + with proxy._compression_metrics_lock: + if proxy._compression_leaked_threads >= 1: + break + time.sleep(0.01) + + with proxy._compression_metrics_lock: + assert proxy._compression_leaked_threads >= 1, ( + f"leaked_threads should be ≥ 1; got {proxy._compression_leaked_threads}. " + f"The worker either didn't finish past the deadline, or the wrapper " + f"didn't increment the counter." + ) + # In-flight gauge restored. + assert proxy._compression_in_flight == 0 + + +def test_compression_executor_metrics_appear_in_runtime_payload() -> None: + """``/stats runtime.compression_executor`` surfaces the new gauges.""" + from fastapi.testclient import TestClient + + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + log_requests=False, + ccr_inject_tool=False, + ccr_handle_responses=False, + ccr_context_tracking=False, + image_optimize=False, + compression_max_workers=5, + ) + app = create_app(config) + + with TestClient(app) as client: + # The compression_executor metrics are published from the runtime + # payload (also surfaced in /health). Hit /health and look there. + r = client.get("/health") + assert r.status_code == 200 + runtime = r.json()["runtime"] + assert "compression_executor" in runtime + ce = runtime["compression_executor"] + assert ce["max_workers"] == 5 + assert ce["in_flight"] == 0 + assert ce["leaked_threads_total"] == 0 + assert ce["source"] == "explicit" + + +def test_explicit_None_resolves_to_auto_source() -> None: + """When max_workers is None (default), the runtime payload reports + ``source: auto``.""" + from fastapi.testclient import TestClient + + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + log_requests=False, + ccr_inject_tool=False, + ccr_handle_responses=False, + ccr_context_tracking=False, + image_optimize=False, + ) + app = create_app(config) + with TestClient(app) as client: + r = client.get("/health") + assert r.json()["runtime"]["compression_executor"]["source"] == "auto" From 456a6b33af06f9dfdd07c58cb4a7820c6b890435 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Fri, 1 May 2026 16:53:21 -0700 Subject: [PATCH 2/2] fix(test): stub _run_compression_in_executor on _DummyOpenAIHandler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bounded compression executor introduced in this PR moved every handler's compression call from `asyncio.wait_for(asyncio.to_thread(...))` to `self._run_compression_in_executor(...)`, which lives on `HeadroomProxy` (server.py) and is inherited by handler mixins at runtime. The test's `_DummyOpenAIHandler` only inherits `OpenAIHandlerMixin`, not `HeadroomProxy`, so it lacks the method. The Responses API compression path caught the AttributeError and silently fell back — which made `test_handle_openai_responses_stream_keeps_compression` fail with `apply.call_count == 0`. Add a synchronous stub that just invokes the callable; tests don't need real thread-pool semantics. --- tests/test_openai_codex_routing.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_openai_codex_routing.py b/tests/test_openai_codex_routing.py index 62eb5e327..c56a2f9ae 100644 --- a/tests/test_openai_codex_routing.py +++ b/tests/test_openai_codex_routing.py @@ -154,6 +154,13 @@ class _DummyOpenAIHandler(OpenAIHandlerMixin): self.captured_request = (method, url, headers, body) return _ResponseStub() + async def _run_compression_in_executor(self, fn, *, timeout: float): + # Test stub for HeadroomProxy._run_compression_in_executor. + # The real implementation runs `fn` on a bounded thread pool with + # a wall-clock timeout; tests just need the callable invoked + # synchronously so MagicMock call_count assertions fire. + return fn() + async def _stream_response( self, url: str,